#!/usr/bin/env python3
"""Run one image-description turn with the published Vitrus LoRA adapter."""

from __future__ import annotations

import sys

import torch
from huggingface_hub import snapshot_download
from peft import PeftModel
from PIL import Image
from transformers import AutoModelForImageTextToText, AutoProcessor


REPO_ID = "lucas-vitrus/liquid-crow"
BASE_MODEL = "LiquidAI/LFM2.5-VL-450M-Extract"
PROMPT = "Describe the world you see in details."


def main() -> None:
    if len(sys.argv) != 2:
        raise SystemExit("usage: python examples/load_lora.py path/to/image.jpg")

    image = Image.open(sys.argv[1]).convert("RGB")
    adapter_dir = snapshot_download(REPO_ID, allow_patterns=["lora/*"])
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if device == "cuda" else torch.float32

    processor = AutoProcessor.from_pretrained(
        BASE_MODEL,
        min_image_tokens=64,
        max_image_tokens=256,
        do_image_splitting=True,
    )
    base = AutoModelForImageTextToText.from_pretrained(
        BASE_MODEL,
        dtype=dtype,
        low_cpu_mem_usage=True,
    ).to(device)
    model = PeftModel.from_pretrained(base, f"{adapter_dir}/lora").eval()

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": PROMPT},
                {"type": "image", "image": image},
            ],
        }
    ]
    inputs = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_tensors="pt",
        return_dict=True,
    ).to(device)

    with torch.inference_mode():
        output_ids = model.generate(**inputs, max_new_tokens=256, do_sample=False)

    generated_ids = output_ids[:, inputs["input_ids"].shape[1] :]
    print(processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip())


if __name__ == "__main__":
    main()