import gradio as gr
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info
import torch

# ============== 官方INT4 AWQ量化模型，免费硬件专属 ==============
# 先跑通7B版本，免费A10G硬件100%能跑，后续想换32B直接改这里即可
MODEL_NAME = "Qwen/Qwen2.5-VL-7B-Instruct-AWQ"

# 加载处理器（官方原生写法，无导入错误）
processor = AutoProcessor.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True
)

# 加载模型（免费A10G硬件优化配置，避免显存卡死）
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True,
    low_cpu_mem_usage=True
).eval()

# ============== 推理函数（官方原生逻辑，稳定不报错） ==============
def chat_with_image(image, prompt):
    # 空输入校验，避免程序崩溃
    if image is None:
        return "❌ 请先上传一张图片！"
    if not prompt.strip():
        prompt = "请详细描述这张图片的内容"

    # 构造官方标准输入格式，完全对齐模型要求
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": prompt}
            ]
        }
    ]

    # 官方预处理逻辑，兼容性100%
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt"
    ).to("cuda")

    # 生成优化：免费硬件稳定不超时，响应速度更快
    with torch.no_grad():
        generated_ids = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.7,
            do_sample=True
        )

    generated_ids_trimmed = [
        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    output_text = processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )[0]

    return output_text

# ============== Gradio界面，小白友好版 ==============
with gr.Blocks(title="Qwen2.5-VL 免费在线演示") as demo:
    gr.Markdown("# 🤖 Qwen2.5-VL-7B-Instruct-AWQ 免费在线演示")
    gr.Markdown("✅ 上传图片 → 输入问题 → 一键获取AI识图回答")

    with gr.Row():
        image_input = gr.Image(type="pil", label="📸 上传图片", height=300)
        prompt_input = gr.Textbox(label="💬 输入你的问题", placeholder="例如：图片里有什么？", value="请详细描述这张图片的内容")

    output_text = gr.Textbox(label="🤖 AI回答", lines=10)
    submit_btn = gr.Button("🚀 开始推理", variant="primary")

    submit_btn.click(
        fn=chat_with_image,
        inputs=[image_input, prompt_input],
        outputs=output_text
    )

# 启动配置，优化队列，避免多用户请求卡死
if __name__ == "__main__":
    demo.queue(max_size=5).launch(show_error=True)