"""
Orvion-vl-3b — Gradio ZeroGPU Space
Model: sanaX3065/Orvion-vl-3b (Qwen2.5-VL-3B-Instruct fine-tune)

API: POST /gradio_api/call/generate  {"data": ["prompt"]}
     POST /gradio_api/call/generate  {"data": ["prompt", null]}
"""

import sys
import traceback

import spaces
import gradio as gr
import torch
from PIL import Image
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor, BitsAndBytesConfig

MODEL_ID = "sanaX3065/aegis-qwen2vl-3b"

def _fix_lm_head(m):
    try:
        m.tie_weights()
        if m.lm_head.weight.data_ptr() == m.model.embed_tokens.weight.data_ptr():
            print("[fix] lm_head tied via tie_weights() ✓")
            return
    except Exception as e:
        print(f"[fix] tie_weights() failed: {e}")
    for attr_path in ["model.embed_tokens.weight", "model.language_model.embed_tokens.weight"]:
        try:
            parts = attr_path.split(".")
            obj = m
            for p in parts:
                obj = getattr(obj, p)
            m.lm_head.weight = obj
            print(f"[fix] lm_head tied via {attr_path} ✓")
            return
        except AttributeError:
            continue
    print("[fix] WARNING: could not tie lm_head.weight")


print("[init] Loading processor...")
processor = AutoProcessor.from_pretrained(
    MODEL_ID,
    trust_remote_code=True
)

print("[init] Preparing 4-bit config...")
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
)

print("[init] Loading model in 4-bit...")
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
)

model.eval()
_fix_lm_head(model)


@spaces.GPU(duration=60)   # ← ZeroGPU free tier max is 60s per call
def generate(prompt_text: str, image_file=None) -> str:
    global model, processor
    try:
        device = next(model.parameters()).device

        # Normalise image
        if image_file is not None and not isinstance(image_file, Image.Image):
            try:
                image_file = Image.open(image_file).convert("RGB")
            except Exception:
                image_file = None

        content = []
        if image_file is not None:
            content.append({"type": "image", "image": image_file})
        content.append({"type": "text", "text": prompt_text})
        messages = [{"role": "user", "content": content}]

        text = processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )

        if image_file is not None:
            try:
                from qwen_vl_utils import process_vision_info
                image_inputs, video_inputs = process_vision_info(messages)
            except ImportError:
                image_inputs = [image_file]
                video_inputs = None
        else:
            image_inputs = None
            video_inputs = None

        inputs = processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            return_tensors="pt",
            padding=True,
        ).to(device)

        with torch.inference_mode():
            out_ids = model.generate(
                **inputs,
                max_new_tokens=256,
                do_sample=False,
                temperature=0,
                top_p=0.9,
                repetition_penalty=1.1,
            )

        new_tokens = out_ids[:, inputs["input_ids"].shape[1]:]
        reply = processor.batch_decode(new_tokens, skip_special_tokens=True)[0].strip()
        print(f"[generate] reply={reply[:80]}")
        return reply

    except Exception as e:
        tb = traceback.format_exc()
        print(f"[generate] CRASHED:\n{tb}")
        return f"[Server Error] {type(e).__name__}: {e}"


demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Prompt", lines=5, placeholder="Type your message…"),
        gr.Image(label="Image (optional)", type="pil", value=None),
    ],
    outputs=gr.Textbox(label="Response", lines=8),
    title="Orvion · Qwen2.5-VL-3B",
    api_name="generate",
    flagging_mode="never",
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)