"""Single-image inference with UHR-BAT2's native Qwen3-VL architecture."""

import argparse
from pathlib import Path

import torch
from PIL import Image
from transformers import AutoProcessor, Qwen3VLForConditionalGeneration


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", default="RL-MIND/UHR-BAT2")
    parser.add_argument("--image", required=True)
    parser.add_argument("--question", default="Describe this remote-sensing image.")
    parser.add_argument("--max-new-tokens", type=int, default=128)
    args = parser.parse_args()

    processor = AutoProcessor.from_pretrained(args.model)
    model = Qwen3VLForConditionalGeneration.from_pretrained(
        args.model,
        dtype="auto",
        device_map="auto",
        attn_implementation="sdpa",
    ).eval()
    image = Image.open(Path(args.image)).convert("RGB")
    messages = [{"role": "user", "content": [
        {"type": "image"}, {"type": "text", "text": args.question},
    ]}]
    prompt = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True,
    )
    inputs = processor(text=[prompt], images=[image], return_tensors="pt")
    inputs.pop("token_type_ids", None)
    inputs = inputs.to(model.device)
    with torch.inference_mode():
        generated = model.generate(
            **inputs,
            max_new_tokens=args.max_new_tokens,
            do_sample=False,
            temperature=None,
            top_p=None,
            top_k=None,
            use_cache=True,
        )
    answer = generated[:, inputs["input_ids"].shape[1]:]
    print(processor.batch_decode(answer, skip_special_tokens=True)[0].strip())


if __name__ == "__main__":
    main()