import gradio as gr
import torch
from PIL import Image
from transformers import AutoModel, AutoTokenizer

# =========================
# 1️⃣ LOAD MODEL
# =========================
model_name = "MathLLMs/MathCoder-VL-8B"

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(
    model_name,
    torch_dtype=torch.float32,
    low_cpu_mem_usage=True,
    device_map="cpu",
    trust_remote_code=True
)
model.eval()

# =========================
# 2️⃣ XỬ LÝ ẢNH + VĂN BẢN
# =========================
def solve_math(image, question):
    if image is None and (not question or question.strip() == ""):
        return "⚠️ Hãy tải ảnh hoặc nhập câu hỏi."

    # Model multimodal — gọi forward tự động xử lý cả 2 modality
    with torch.no_grad():
        response, _ = model.chat(tokenizer, image=image, question=question)
    return response


# =========================
# 3️⃣ GIAO DIỆN GRADIO
# =========================
with gr.Blocks(theme="soft") as demo:
    gr.Markdown("## 🧮 MathCoder-VL-8B — Multimodal Reasoning (CPU Demo)")
    gr.Markdown(
        "Model có thể nhận **ảnh + câu hỏi toán học**, sau đó sinh lời giải bằng ngôn ngữ tự nhiên."
    )

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(type="pil", label="Tải ảnh toán học (optional)")
            text_input = gr.Textbox(
                label="Câu hỏi hoặc mô tả", 
                placeholder="Ví dụ: Hãy giải bài toán trong hình, hoặc nhập đề bài...",
                lines=3
            )
            btn = gr.Button("💡 Giải toán")

        with gr.Column(scale=1):
            output = gr.Textbox(label="Lời giải", lines=15)

    btn.click(fn=solve_math, inputs=[image_input, text_input], outputs=output)

demo.launch()