#!/usr/bin/env python
"""Complete image-understanding example for vision-gpt-oss-120B.

Usage:
    python example_inference.py --image photo.jpg
    python example_inference.py --image photo.jpg --prompt "What is written on the sign?" \
        --reasoning_effort medium --max_new_tokens 768

Run `AutoModelForImageTextToText` + `AutoProcessor` with trust_remote_code=True.
The model speaks gpt-oss "harmony" format: it first thinks in an *analysis* channel
and then emits the user-facing answer in a *final* channel. We parse the final
channel below.
"""
import argparse
import re
import torch
from PIL import Image
from transformers import AutoProcessor, AutoModelForImageTextToText


def extract_final(text: str) -> str:
    """Pull the user-facing answer out of the harmony `final` channel."""
    m = re.search(r"final<\|message\|>(.*?)(?=<\|return\||<\|end\||$)", text, re.DOTALL)
    return m.group(1).strip() if m else text.strip()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default=".", help="path or HF repo id of this model")
    ap.add_argument("--image", default="test_images/photo.jpg",
                    help="defaults to a bundled sample image")
    ap.add_argument("--prompt", default="Describe this image in detail.")
    ap.add_argument("--reasoning_effort", default="low", choices=["low", "medium", "high"])
    ap.add_argument("--max_new_tokens", type=int, default=512)
    ap.add_argument("--max_image_size", type=int, default=1536)
    args = ap.parse_args()

    processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True)
    model = AutoModelForImageTextToText.from_pretrained(
        args.model, trust_remote_code=True, dtype=torch.bfloat16, device_map="cuda"
    ).eval()

    # GOTCHA: very large images blow up the vision sequence / latency and can OOM.
    # Downscale the long side (aspect ratio preserved).
    image = Image.open(args.image).convert("RGB")
    if max(image.size) > args.max_image_size:
        image.thumbnail((args.max_image_size, args.max_image_size))

    batch = processor(
        images=image,
        text=args.prompt,
        reasoning_effort=args.reasoning_effort,  # controls analysis-channel length
    )
    batch = {k: (v.cuda() if torch.is_tensor(v) else v) for k, v in batch.items()}

    with torch.no_grad():
        out = model.generate(**batch, max_new_tokens=args.max_new_tokens, do_sample=False)

    # GOTCHA: keep special tokens so the channel markers survive, then parse `final`.
    decoded = processor.decode(out[0], skip_special_tokens=False)
    answer = extract_final(decoded)

    print("=" * 70)
    print("IMAGE :", args.image)
    print("PROMPT:", args.prompt, f"(reasoning_effort={args.reasoning_effort})")
    print("-" * 70)
    print(answer)
    print("=" * 70)


if __name__ == "__main__":
    main()