test_caption.py

"""
test_caption.py — verify STAGE 1 (object identification) in isolation.

Calls the caption Space and checks we get back a short object phrase (the dedup
key). Run `inspect_api.py` first and make sure the /chat_joycaption arg order
below matches; if it doesn't, fix it here AND in app.py's caption_object().

Consumes a little ZeroGPU quota (one call). Set HF_TOKEN so it's billed to you
and gets real quota; without a token you get the tiny anonymous allowance.

Usage:
    HF_TOKEN=hf_xxx python test_caption.py --image photo.jpg
    HF_TOKEN=hf_xxx python test_caption.py            # uses a synthetic test image
"""

import os
import sys
import argparse
import tempfile

from gradio_client import Client, handle_file

SPACE = os.getenv("CAPTION_SPACE", "fancyfeast/joy-caption-beta-one")
TOKEN = os.getenv("HF_TOKEN")

INSTRUCTION = (
    "Identify the single main physical object in this image as a short, generic "
    "noun phrase of 1 to 4 words (for example 'ceramic coffee mug' or 'wooden "
    "chair'). Ignore the background. Reply with ONLY the object name."
)


def make_test_image() -> str:
    """A trivial image so the test runs with zero setup. A real photo gives a more
    meaningful caption, but this is enough to prove the API call works."""
    from PIL import Image, ImageDraw
    img = Image.new("RGB", (512, 512), (232, 232, 236))
    ImageDraw.Draw(img).ellipse([150, 150, 362, 362], fill=(198, 92, 70))
    path = os.path.join(tempfile.gettempdir(), "piclets_test.png")
    img.save(path)
    return path


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--image", help="path to a test photo (optional)")
    args = ap.parse_args()
    image_path = args.image or make_test_image()

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

    client = Client(SPACE, hf_token=TOKEN)

    # /chat_joycaption positional args (verify with inspect_api.py):
    #   (input_image, prompt, temperature, top_p, max_new_tokens, log_prompt)
    #     -> caption (a single string)
    # temperature=0 for deterministic, terse object names. Mirror in app.py
    # caption_object().
    result = client.predict(
        handle_file(image_path),
        INSTRUCTION,    # prompt
        0,              # temperature (deterministic)
        api_name="/chat_joycaption",
    )

    print("\n[caption] RAW RESULT (inspect this if anything looks off):")
    print(repr(result)[:1500])

    # beta-one returns the caption as a bare string (not a tuple).
    caption = result if isinstance(result, str) else (
        result[1] if isinstance(result, (list, tuple)) and len(result) > 1 else str(result)
    )
    phrase = " ".join(caption.strip().split()[:5])
    print(f"\n[caption] extracted object phrase: {phrase!r}")

    ok = bool(phrase) and len(phrase) < 80
    print("[caption] PASS" if ok else
          "[caption] FAIL — check RAW RESULT, then fix the arg order / result index "
          "here and in app.py caption_object().")
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()