import os
from threading import Thread

import gradio as gr
import spaces
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from transformers import TextIteratorStreamer

from model import VLM, VLMConfig, transform

# ── Load the model once at startup ───────────────────────────────────────────
# The ~10GB checkpoint lives in a (private) model repo; download it with the
# Space's HF_TOKEN secret, then load the weights into our custom VLM.
CKPT_REPO = "ndrugov/encoder-free-vlm-densefusion-sharegpt4v"
CKPT_FILE = "vlm_best.pt"

ckpt_path = hf_hub_download(
    repo_id=CKPT_REPO,
    filename=CKPT_FILE,
    token=os.environ.get("HF_TOKEN"),
)

cfg = VLMConfig()
vlm = VLM(cfg)
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
state = ckpt.get("model_state_dict", ckpt)
vlm.load_state_dict(state)
vlm.eval()
vlm.to("cuda")
print(f"Loaded {CKPT_FILE} (step={ckpt.get('step')}, "
      f"best_val_loss={ckpt.get('best_val_loss')})")

N_IMG = cfg.vision.num_patches


@spaces.GPU
def model_inference(input_dict, history):
    text = input_dict["text"]
    files = input_dict.get("files", [])

    if not files:
        raise gr.Error("Please upload an image along with your question.")
    if text == "":
        raise gr.Error("Please type a question about the image.")

    # This encoder-free VLM is trained on a single 512x512 image per turn.
    image = files[-1]
    if isinstance(image, str):
        image = Image.open(image)
    img = transform(image.convert("RGB")).unsqueeze(0).to("cuda")  # (1, 3, 512, 512)

    # Prompt = one <|image|> per patch, then the question, as an open chat turn.
    content = vlm.tokenizer.image_token * N_IMG + text
    prompt = vlm.tokenizer.apply_chat_template(
        [{"role": "user", "content": content}],
        tokenize=False, add_generation_prompt=True,
    )
    input_ids = torch.tensor([vlm.tokenizer.encode(prompt)], device="cuda")

    # Splice projected patch embeddings into the <|image|> slots.
    image_embd = vlm.connector(vlm.vision_embedder(img))
    token_embd = vlm.decoder.get_input_embeddings()(input_ids)
    combined = vlm._replace_img_tokens_with_embd(input_ids, token_embd, image_embd)
    attn = torch.ones(combined.shape[:2], dtype=torch.long, device="cuda")

    streamer = TextIteratorStreamer(
        vlm.tokenizer, skip_prompt=True, skip_special_tokens=True
    )
    gen_kwargs = dict(
        inputs_embeds=combined,
        attention_mask=attn,
        max_new_tokens=256,
        do_sample=False,
        pad_token_id=vlm.tokenizer.pad_token_id,
        streamer=streamer,
    )

    def _run():
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16), torch.no_grad():
            vlm.decoder.generate(**gen_kwargs)

    thread = Thread(target=_run)
    thread.start()

    buffer = ""
    yield "…"
    for new_text in streamer:
        buffer += new_text
        yield buffer


examples = [
    [{"text": "Describe this image.", "files": ["example_images/cat.jpeg"]}],
    [{"text": "What is in this image?", "files": ["example_images/sf.jpeg"]}],
    [{"text": "Describe this image in detail.", "files": ["example_images/tree.jpeg"]}],
    [{"text": "Describe the scene.", "files": ["example_images/ski.jpeg"]}],
]

demo = gr.ChatInterface(
    fn=model_inference,
    title="Demo of Encoder-Free VLM Trained for $100",
    description=(
        "Play with this encoder-free vision-language model, inspired by the "
        "architecture of Gemma 4 12B Unified. Our model was trained for about "
        "$100 (43 hours on a single H100). It used Qwen 3 1.7B as a decoder and "
        "a subset of FineVision as training data.\n\n"
        "To get started, upload an image and text or try one of the examples. "
        "This demo doesn't use history for the chat, so every chat you start is "
        "a new conversation.\n\n"
        "Read more about how this model was trained in our blogpost: "
        "[Train Your Own Encoder-Free VLM in $100]"
        "(https://huggingface.co/spaces/ndrugov/encoder-free-vlm)."
    ),
    examples=examples,
    textbox=gr.MultimodalTextbox(
        label="Query Input", file_types=["image"], file_count="single"
    ),
    stop_btn="Stop Generation",
    multimodal=True,
    cache_examples=False,
)

if __name__ == "__main__":
    demo.launch()