gen-image.py

#!/usr/bin/env python3
"""Generate images via BlueAI Relay (Nano Banana / Gemini image preview models).

Usage:
    python gen-image.py <out_path> <model> <prompt>

Env (read from ~/.openclaw/.env or process env):
    BLUEAI_RELAY_BASE   default https://bmc-llm-relay.bluemediagroup.cn
    BLUEAI_RELAY_KEY    REQUIRED, contact 金明 (@Dr-xiaoming) if 401/403

Outputs PNG to out_path. Exits non-zero on failure with diagnostic preview.
"""
import base64
import json
import os
import re
import sys
import urllib.request
import urllib.error


DEFAULT_RELAY_BASE = "https://bmc-llm-relay.bluemediagroup.cn"


def _resolve_creds() -> tuple[str, str]:
    base = os.environ.get("BLUEAI_RELAY_BASE", DEFAULT_RELAY_BASE)
    key = os.environ.get("BLUEAI_RELAY_KEY")
    if not key:
        raise RuntimeError(
            "BLUEAI_RELAY_KEY not set. Run: set -a; source ~/.openclaw/.env; set +a\n"
            "If the key is missing or rejected (401/403), contact 金明 (@Dr-xiaoming)."
        )
    return base, key


def gen(model: str, prompt: str, out_path: str) -> None:
    relay_base, relay_key = _resolve_creds()

    payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
    }
    req = urllib.request.Request(
        f"{relay_base}/v1/chat/completions",
        data=json.dumps(payload).encode("utf-8"),
        headers={
            "Authorization": f"Bearer {relay_key}",
            "Content-Type": "application/json",
        },
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=180) as resp:
            data = json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        body = e.read().decode("utf-8", errors="replace")[:500]
        if e.code in (401, 403):
            print(
                f"[ERROR] HTTP {e.code} from BlueAI Relay. Key may be invalid/expired.\n"
                f"        Contact 金明 (@Dr-xiaoming). Body preview: {body}",
                file=sys.stderr,
            )
        else:
            print(f"[ERROR] HTTP {e.code}: {body}", file=sys.stderr)
        sys.exit(3)

    content = data["choices"][0]["message"]["content"]

    # 三种返回格式兼容：
    # 1. markdown data URL: ![](data:image/png;base64,...)
    # 2. 裸 base64 (PNG/JPEG/GIF magic header)
    # 3. JSON 包裹 {"b64_json": "..."} / {"image": "..."} / {"data": "..."}
    b64 = None
    m = re.search(r"data:image/[a-zA-Z]+;base64,([A-Za-z0-9+/=]+)", content)
    if m:
        b64 = m.group(1)
    elif content.strip().startswith(("iVBOR", "/9j/", "R0lGOD")):
        b64 = re.sub(r"\s+", "", content.strip())
    else:
        try:
            j = json.loads(content)
            for k in ("b64_json", "image", "data"):
                if k in j and isinstance(j[k], str):
                    b64 = j[k]
                    break
        except Exception:
            pass

    if not b64:
        print(
            f"[ERROR] could not extract base64. Likely the model returned text "
            f"instead of an image (prompt may have triggered safety policy).\n"
            f"Preview: {content[:500]}",
            file=sys.stderr,
        )
        sys.exit(2)

    raw = base64.b64decode(b64)
    out_dir = os.path.dirname(out_path) or "."
    os.makedirs(out_dir, exist_ok=True)
    with open(out_path, "wb") as f:
        f.write(raw)

    usage = data.get("usage", {}).get("total_tokens", "?")
    print(
        f"OK {out_path} ({len(raw):,} bytes, model={model}, usage={usage} tokens)"
    )


if __name__ == "__main__":
    if len(sys.argv) < 4:
        print("Usage: gen-image.py <out_path> <model> <prompt>", file=sys.stderr)
        sys.exit(1)
    gen(sys.argv[2], sys.argv[3], sys.argv[1])