app.py

# filename: app.py
import os
import torch
import gradio as gr
from PIL import Image
import spaces

from transformers import (
    AutoProcessor,
    Qwen2_5_VLForConditionalGeneration,  # ✅ 正确的多模态类
)

# ========================
# 基本设置
# ========================
torch.manual_seed(100)
MODEL_NAME = "Eric3200/OralGPT-7B-Preview"

# ========================
# 模型与处理器
# ========================
processor = AutoProcessor.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
)

model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
    torch_dtype="auto",         # 或 torch.bfloat16 / torch.float16
    device_map="auto",
    # attn_implementation="flash_attention_2",  # 环境支持再打开（速度/显存更优）
).eval()

# pad_token 兜底
try:
    PAD_ID = processor.tokenizer.pad_token_id or processor.tokenizer.eos_token_id
except Exception:
    PAD_ID = None

# ========================
# 输入转 Qwen/OralGPT 风格 content
# ========================
def to_qwen_content(user_msg):
    """
    支持：
      - 纯文本
      - (img_path, text)
      - list 混合（PIL.Image / 文本 / 本地路径字符串）
      - dict {files:[...], text:"..."}
    """
    contents = []
    if isinstance(user_msg, tuple):
        img_path, text = user_msg
        if img_path:
            contents.append({"type": "image", "image": img_path})
        if text:
            contents.append({"type": "text", "text": str(text)})
    elif isinstance(user_msg, list):
        for item in user_msg:
            if isinstance(item, Image.Image):
                contents.append({"type": "image", "image": item})
            elif isinstance(item, str) and os.path.isfile(item):
                contents.append({"type": "image", "image": item})
            else:
                contents.append({"type": "text", "text": str(item)})
    elif isinstance(user_msg, dict):
        if user_msg.get("files"):
            for fp in user_msg["files"]:
                contents.append({"type": "image", "image": fp})
        if user_msg.get("text"):
            contents.append({"type": "text", "text": str(user_msg["text"])})
    else:
        contents.append({"type": "text", "text": str(user_msg)})
    return contents

def build_messages(history, latest_user_msg):
    messages = []
    for u, a in history:
        messages.append({"role": "user", "content": to_qwen_content(u)})
        if a:
            messages.append({"role": "assistant", "content": [{"type": "text", "text": str(a)}]})
    messages.append({"role": "user", "content": to_qwen_content(latest_user_msg)})
    return messages

# ========================
# ZeroGPU 推理（关键：@spaces.GPU）
# ========================
@spaces.GPU(duration=120)
def qwen_infer(messages, max_new_tokens=512, do_sample=False, temperature=0.7, top_p=0.9, pad_token_id=None):
    """
    流程：
      1) apply_chat_template → prompt
      2) 收集 messages 中的图片
      3) processor(...) 打包
      4) model.generate(...)
      5) 取新增 token 解码
    """
    # 1) 文本模板
    text_prompt = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )

    # 2) 收集图片（demo 不处理视频）
    images = []
    for m in messages:
        for c in m.get("content", []):
            if c.get("type") == "image":
                img_obj = c.get("image")
                if isinstance(img_obj, Image.Image):
                    images.append(img_obj)
                elif isinstance(img_obj, str) and os.path.isfile(img_obj):
                    images.append(Image.open(img_obj).convert("RGB"))

    # 3) 打包输入
    inputs = processor(
        text=[text_prompt],
        images=images if images else None,
        padding=True,
        return_tensors="pt",
    )

    # 放到 GPU/对应设备
    if torch.cuda.is_available():
        inputs = inputs.to("cuda")
    else:
        inputs = inputs.to(model.device)

    gen_kwargs = dict(
        max_new_tokens=int(max_new_tokens),
        do_sample=bool(do_sample),
    )
    if do_sample:
        gen_kwargs.update(dict(temperature=float(temperature), top_p=float(top_p)))
    if pad_token_id is not None:
        gen_kwargs.update(dict(pad_token_id=int(pad_token_id)))

    # 4) 生成
    generated_ids = model.generate(**inputs, **gen_kwargs)

    # 5) 只取新增 token 并解码
    trimmed = [
        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    outputs = processor.batch_decode(
        trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )
    return outputs[0].strip() if outputs else "(空响应)"

# ========================
# Gradio 界面
# ========================
with gr.Blocks(title="OralGPT-7B-Preview — Gradio Demo (ZeroGPU)") as demo:
    gr.Markdown(
        """
        # 🦷 OralGPT-7B-Preview — 多模态对话 Demo (ZeroGPU)
        上传牙科相关图片并提问，或进行文字对话。右侧可开启 “Thinking（采样）模式” 拓展输出。
        """
    )

    with gr.Row():
        with gr.Column(scale=4):
            chatbot = gr.Chatbot(
                height=500,
                show_label=False,
                container=True,
                type="tuples"
            )

            with gr.Row():
                msg = gr.MultimodalTextbox(
                    interactive=True,
                    file_types=["image"],
                    placeholder="输入消息或上传图片（可多张）…",
                    show_label=False,
                    container=False
                )

            with gr.Row():
                clear = gr.Button("🗑️ 清空", size="sm")
                submit = gr.Button("📤 发送", variant="primary", size="sm")

        with gr.Column(scale=1):
            enable_thinking = gr.Checkbox(label="启用 Thinking（采样）模式", value=False)
            max_tokens = gr.Slider(64, 2048, value=512, step=32, label="max_new_tokens")

            gr.Markdown(
                """
                ### 示例
                - “请问这张图片中是否存在龋病现象？”
                - “是否存在牙周病现象？”
                - “描述图像中的可疑病灶区域。”
                """
            )
            gr.Markdown("**注意：此模型仅供研究参考，不用于临床诊断或治疗。**")

    # 事件逻辑
    def user_submit(message, history, enable_thinking, max_new_tokens):
        # 组织本轮用户消息
        if isinstance(message, dict) and message.get("files"):
            user_msg = []
            for fp in message["files"]:
                user_msg.append(fp)
            if message.get("text", ""):
                user_msg.append(message["text"])
        else:
            user_msg = message.get("text", "") if isinstance(message, dict) else message

        history = history + [(user_msg, None)]
        messages = build_messages(history[:-1], user_msg)

        do_sample = bool(enable_thinking)
        # generate 需要数值，这里给默认
        temperature = 0.7 if do_sample else 0.7
        top_p = 0.9 if do_sample else 0.9

        try:
            answer = qwen_infer(
                messages=messages,
                max_new_tokens=int(max_new_tokens),
                do_sample=do_sample,
                temperature=temperature,
                top_p=top_p,
                pad_token_id=PAD_ID,
            )
        except Exception as e:
            answer = f"Error: {str(e)}"

        history[-1] = (history[-1][0], answer if answer else "(空响应)")
        return "", history

    msg.submit(user_submit, inputs=[msg, chatbot, enable_thinking, max_tokens], outputs=[msg, chatbot])
    submit.click(user_submit, inputs=[msg, chatbot, enable_thinking, max_tokens], outputs=[msg, chatbot])
    clear.click(lambda: (None, []), outputs=[msg, chatbot])

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