test_image.py

"""
test_image.py — verify STAGE 3 (image generation) in isolation.

Run `inspect_api.py` FIRST: the exact positional args vary by Space. krea/Krea-2
exposes `/generate(prompt, negative_prompt, model, steps, guidance, width, height,
seed, randomize)`; prompt and negative_prompt are both required, the rest have
defaults. This script passes the two prompts and relies on those defaults — if
inspect shows otherwise, update here AND in app.py generate_image().

Saves the result to test_image_out.webp so you can open it and judge quality.
Consumes a little ZeroGPU quota (one call).

Usage:
    HF_TOKEN=hf_xxx python test_image.py
    HF_TOKEN=hf_xxx python test_image.py --prompt "a small round ceramic creature"
"""

import os
import io
import sys
import argparse

import requests
from PIL import Image
from gradio_client import Client

SPACE = os.getenv("IMAGE_SPACE", "krea/Krea-2")
TOKEN = os.getenv("HF_TOKEN")

DEFAULT_PROMPT = (
    "a small round ceramic-bodied creature with a looping handle-tail, steam curling "
    "from its head, full-body, centered, plain flat background, cute anime creature design"
)
DEFAULT_NEGATIVE = (
    "text, watermark, signature, logo, blurry, low quality, deformed, "
    "extra limbs, extra faces, background clutter"
)


def read_image_result(result) -> bytes:
    """Mirrors app.py _read_image_result: normalize path/URL/dict -> bytes."""
    ref = result[0] if isinstance(result, (list, tuple)) and result else result
    if isinstance(ref, dict):
        ref = ref.get("url") or ref.get("path") or ref.get("image") or ref.get("name")
    if not isinstance(ref, str):
        raise ValueError(f"unexpected result type: {type(result)}")
    if ref.startswith("http"):
        return requests.get(ref, timeout=60).content
    with open(ref, "rb") as f:
        return f.read()


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--prompt", default=DEFAULT_PROMPT)
    ap.add_argument("--negative", default=DEFAULT_NEGATIVE,
                    help="negative prompt (required by Krea-2 /generate)")
    args = ap.parse_args()

    print(f"[image] space = {SPACE}")
    if not TOKEN:
        print("[image] WARNING: no HF_TOKEN — tiny anonymous quota; may fail.")

    client = Client(SPACE, hf_token=TOKEN)

    # krea/Krea-2 `/generate` positional args (verify with inspect_api.py):
    #   (prompt, negative_prompt, model, steps, guidance, width, height, seed, randomize)
    #     -> (result, seed)
    # prompt + negative_prompt are both required; the rest use the Space's defaults
    # (Turbo, 8 steps, 1024x1024). Mirror any changes here in app.py generate_image().
    result = client.predict(args.prompt, args.negative, api_name="/generate")

    print("\n[image] RAW RESULT:")
    print(repr(result)[:1000])

    try:
        raw = read_image_result(result)
        img = Image.open(io.BytesIO(raw))
        img.load()
    except Exception as exc:
        print(f"\n[image] FAIL — could not read an image from the result: {exc}")
        print("        Inspect RAW RESULT and update read_image_result / the /generate args.")
        sys.exit(1)

    out = "test_image_out.webp"
    img.convert("RGB").save(out, "WEBP", quality=85)
    print(f"\n[image] PASS — got {img.size[0]}x{img.size[1]}, saved {out} (open it to check quality).")
    sys.exit(0)


if __name__ == "__main__":
    main()