import spaces
import torch
from transformers import AutoModel, AutoProcessor, PreTrainedModel

# ============================================================
# PATCH KOMPATIBILITAS
# transformers >= 5.9 mengubah _keys_to_ignore_on_load_unexpected
# dari list-append jadi set-union. Remote code ZDTaichu5.0 masih
# pakai list, sehingga muncul: TypeError: unsupported operand
# type(s) for |: 'list' and 'set'
# Patch ini paksa konversi ke set sebelum method asli dijalankan.
# ============================================================
_orig_adjust = PreTrainedModel._adjust_missing_and_unexpected_keys

def _patched_adjust(self, loading_info):
    attr = getattr(self, "_keys_to_ignore_on_load_unexpected", None)
    if isinstance(attr, (list, tuple)):
        self._keys_to_ignore_on_load_unexpected = set(attr)
    attr2 = getattr(self, "_keys_to_ignore_on_load_missing", None)
    if isinstance(attr2, (list, tuple)):
        self._keys_to_ignore_on_load_missing = set(attr2)
    return _orig_adjust(self, loading_info)

PreTrainedModel._adjust_missing_and_unexpected_keys = _patched_adjust

# ============================================================
# LOAD MODEL DI MODULE LEVEL
# ZeroGPU butuh load di luar @spaces.GPU.
# device_map="cuda" (bukan "auto") supaya ZeroGPU bisa emulate.
# ============================================================
MODEL_ID = "TaichuAI/ZDTaichu5.0-9B"

processor = AutoProcessor.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
)

model = AutoModel.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    device_map="cuda",
    trust_remote_code=True,
    attn_implementation="sdpa",
).eval()

# ============================================================
# INFERENCE FUNCTION
# ============================================================
@spaces.GPU(duration=60)
def generate(image, text):
    if image is None:
        content = [{"type": "text", "text": text}]
    else:
        content = [
            {"type": "image", "image": image},
            {"type": "text", "text": text},
        ]

    messages = [{"role": "user", "content": content}]

    inputs = processor.from_messages(messages, return_tensors="pt").to("cuda")

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

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

# ============================================================
# GRADIO UI
# ============================================================
import gradio as gr

demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Image(type="pil", label="Gambar (opsional)"),
        gr.Textbox(label="Pertanyaan", lines=3),
    ],
    outputs=gr.Textbox(label="Jawaban", lines=10),
    title="ZDTaichu5.0-9B Demo",
    description="Multimodal VLM — teks + gambar. Jalan di ZeroGPU.",
)

demo.launch()