import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces
import torch
import gradio as gr
from transformers import AutoProcessor, AutoModelForMultimodalLM

MODEL_ID = "microsoft/GELab-Zero-4B-preview-Sico-Evolution"

processor = AutoProcessor.from_pretrained(MODEL_ID)
model = AutoModelForMultimodalLM.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda").eval()


@spaces.GPU(duration=120)
def predict(image, text_input, max_new_tokens, history):
    """Run inference on the image+text input."""
    # Build the messages list for the chat template
    messages = []
    # Include conversation history
    for h in history:
        messages.append(h)
    # Add the current user message
    user_content = []
    if image is not None:
        user_content.append({"type": "image", "image": image})
    user_content.append({"type": "text", "text": text_input})
    messages.append({"role": "user", "content": user_content})

    # Apply chat template
    inputs = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    ).to("cuda")

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

    # Decode only the generated tokens
    generated_ids = output_ids[0][inputs["input_ids"].shape[-1]:]
    generated_text = processor.decode(
        generated_ids,
        skip_special_tokens=True,
    )

    # Update history
    new_history = history + [
        {"role": "user", "content": [{"type": "text", "text": text_input}] + ([{"type": "image"}] if image is not None else [])},
        {"role": "assistant", "content": [{"type": "text", "text": generated_text}]},
    ]

    return generated_text, new_history


@spaces.GPU(duration=120)
def predict_single(image, text_input, max_new_tokens):
    """Single-turn prediction for the simple interface."""
    messages = [
        {
            "role": "user",
            "content": [],
        }
    ]
    if image is not None:
        messages[0]["content"].append({"type": "image", "image": image})
    messages[0]["content"].append({"type": "text", "text": text_input})

    inputs = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    ).to("cuda")

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

    generated_ids = output_ids[0][inputs["input_ids"].shape[-1]:]
    generated_text = processor.decode(
        generated_ids,
        skip_special_tokens=True,
    )

    return generated_text


# --- UI ---

title = "GELab-Zero-4B Sico Evolution"
description = (
    "Demo for [microsoft/GELab-Zero-4B-preview-Sico-Evolution](https://huggingface.co/microsoft/GELab-Zero-4B-preview-Sico-Evolution), "
    "a 4B Qwen3-VL based vision-language model fine-tuned as a GUI agent. "
    "Upload an image and ask a question — the model will describe or analyze the content."
)

with gr.Blocks(title=title) as demo:
    gr.Markdown(f"# {title}\n\n{description}")

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(type="pil", label="Input Image")
            text_input = gr.Textbox(
                label="Question / Instruction",
                placeholder="e.g., Describe this image in detail.",
                lines=3,
            )
            max_tokens = gr.Slider(
                minimum=64, maximum=2048, value=512, step=64,
                label="Max New Tokens",
            )
            submit_btn = gr.Button("Submit", variant="primary")

        with gr.Column(scale=1):
            output_text = gr.Textbox(
                label="Model Response",
                lines=15,
                interactive=False,
            )

    gr.Examples(
        examples=[
            [None, "Hello! What can you help me with?"],
            [None, "Explain quantum computing in simple terms."],
        ],
        inputs=[image_input, text_input],
        outputs=output_text,
        fn=predict_single,
        cache_examples=False,
    )

    submit_btn.click(
        fn=predict_single,
        inputs=[image_input, text_input, max_tokens],
        outputs=output_text,
    )


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