import os
# 1. 环境配置
os.environ['CUDA_VISIBLE_DEVICES'] = '' # 强制 CPU
os.environ['TOKENIZERS_PARALLELISM'] = 'false' # 防止 Gradio 多线程警告

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

# 2. 加载模型 (低内存模式)
model_id = "openbmb/MiniCPM-V-4.6"

print("正在加载模型，请耐心等待...")
model = AutoModel.from_pretrained(
    model_id, 
    trust_remote_code=True, 
    low_cpu_mem_usage=True
).eval() # 确保是推理模式

tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
print("模型加载完成！")

# 3. 推理函数
def analyze_image(image, prompt_text):
    if image is None:
        return "请上传一张图片。"
    
    # 构建符合 MiniCPM-V 4.6 要求的输入格式
    msgs = [
        {
            'role': 'user', 
            'content': [
                {"type": "image", "image": image}, # 包装图片
                {"type": "text", "text": prompt_text} # 包装文本
            ]
        }
    ]
    
    # CPU 上的数据类型处理
    with torch.no_grad():
        # 1. 构建输入
        inputs = tokenizer.apply_chat_template(
            msgs, 
            tokenize=True, 
            add_generation_prompt=True, 
            return_tensors="pt"
        )
        
        # 2. 移动到设备并转换数据类型 (CPU 强制 float32)
        inputs = inputs.to(model.device, dtype=torch.float32)

        # 3. 生成回复
        outputs = model.generate(
            inputs, 
            max_new_tokens=1024,
            temperature=0.7,
            do_sample=True
        )
    
    # 解码输出
    response = tokenizer.decode(outputs, skip_special_tokens=True)
    
    # 提取 Assistant 的回答部分
    if 'assistant' in response:
        response = response.split('assistant')[-1].strip()
    
    return response

# 4. 构建界面 (注意：这里去掉了 theme 参数)
with gr.Blocks() as demo:
    gr.Markdown("## 🖼️ MiniCPM-V 4.6 视觉理解 (CPU 优化版)")
    gr.Markdown("上传图片并输入问题，AI 将为你解答。")

    with gr.Row():
        with gr.Column():
            img_input = gr.Image(type="pil", label="上传图片")
            txt_input = gr.Textbox(label="输入问题", placeholder="例如：描述这张图片", value="请描述这张图片")
            btn = gr.Button("🚀 生成回答", variant="primary")
        with gr.Column():
            output = gr.Textbox(label="AI 回答", lines=15)

    btn.click(fn=analyze_image, inputs=[img_input, txt_input], outputs=output)

# 5. 启动 (theme 移到这里，并删除了 show_api)
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0", 
        server_port=7860,
        ssr_mode=False,      # 禁用服务端渲染，避免异步报错
        theme=gr.themes.Soft() # Gradio 6.0 要求 theme 必须放在 launch 里
    )