app.py

import os
import time
import logging
import re
import gradio as gr
from spaces import zero  # 关键：引入 zero 装饰器
import spaces

# 不要在这里 import torch 或加载模型
# from transformers import TextIteratorStreamer, AutoTokenizer  # 不再需要

# 尝试导入 qwen_vl_utils，若失败则提供降级实现（返回空的图像/视频输入）
try:
    from qwen_vl_utils import process_vision_info
except Exception:
    def process_vision_info(messages):
        return None, None

def replace_single_quotes(text):
    pattern = r"\B'([^']*)'\B"
    replaced_text = re.sub(pattern, r'"\1"', text)
    replaced_text = replaced_text.replace("’", "”").replace("‘", "“")
    return replaced_text

DEFAULT_MODEL_PATH = os.environ.get("MODEL_OUTPUT_PATH", "PromptEnhancer/PromptEnhancer-32B")

def _str_to_dtype(dtype_str):
    # 在子进程中再真正用 torch；这里仅返回字符串用于传参
    if dtype_str in ("bfloat16", "float16", "float32"):
        return dtype_str
    return "float32"

@spaces.GPU  # 在子进程（拥有 GPU）中执行：包含模型加载与推理
def gpu_predict(model_path, device_map, torch_dtype,
                prompt_cot, sys_prompt, temperature, max_new_tokens, device):
    # 注意：所有 CUDA 相关 import 放在子进程函数内部
    import torch
    from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor

    # logger（可选）
    if not logging.getLogger(__name__).handlers:
        logging.basicConfig(level=logging.INFO)
    logger = logging.getLogger(__name__)

    # dtype
    if torch_dtype == "bfloat16":
        dtype = torch.bfloat16
    elif torch_dtype == "float16":
        dtype = torch.float16
    else:
        dtype = torch.float32

    # 设备映射：根据 UI 的 device / device_map 决定
    # ZeroGPU 建议 GPU 推理时用 "cuda"
    target_device = "cuda" if device == "cuda" else "cpu"
    load_device_map = "cuda" if device_map == "cuda" else "cpu"

    # 加载模型与处理器（在子进程）
    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        model_path,
        torch_dtype=dtype,
        device_map=load_device_map,
        attn_implementation="sdpa",  # 禁用 flash-attn，兼容性更好
    )
    processor = AutoProcessor.from_pretrained(model_path)

    # 组装消息
    org_prompt_cot = prompt_cot
    try:
        user_prompt_format = sys_prompt + "\n" + org_prompt_cot
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": user_prompt_format},
                ],
            }
        ]

        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",
        )
        # 把输入移动到目标设备
        inputs = inputs.to(target_device)

        # 生成
        generated_ids = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            temperature=float(temperature),
            do_sample=False,
            top_k=5,
            top_p=0.9,
        )
        # 仅解码新增 token
        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,
        )
        output_res = output_text[0]
        # 兼容原逻辑：提取 think> 之后的内容
        try:
            assert output_res.count("think>") == 2
            new_prompt = output_res.split("think>")[-1]
            if new_prompt.startswith("\n"):
                new_prompt = new_prompt[1:]
            new_prompt = replace_single_quotes(new_prompt)
        except Exception:
            # 如果格式不符合预期，则直接回退为原始输入
            new_prompt = org_prompt_cot
        return new_prompt, ""
    except Exception as e:
        # 失败则返回原始提示词和错误信息
        return org_prompt_cot, f"推理失败：{e}"

# -------------------------
# Gradio app
# -------------------------

def run_single(prompt, sys_prompt, temperature, max_new_tokens, device,
               model_path, device_map, torch_dtype, state):
    if not prompt or not str(prompt).strip():
        return "", "请先输入提示词。", state

    t0 = time.time()
    try:
        new_prompt, err = gpu_predict(
            model_path=model_path,
            device_map=device_map,
            torch_dtype=_str_to_dtype(torch_dtype),
            prompt_cot=prompt,
            sys_prompt=sys_prompt,
            temperature=temperature,
            max_new_tokens=max_new_tokens,
            device=device,
        )
        dt = time.time() - t0
        if err:
            return new_prompt, f"{err}（耗时 {dt:.2f}s）", state
        return new_prompt, f"耗时：{dt:.2f}s", state
    except Exception as e:
        return "", f"调用失败：{e}", state

# 示例数据
test_list_zh = [
    "第三人称视角，赛车在城市赛道上飞驰，左上角是小地图，地图下面是当前名次，右下角仪表盘显示当前速度。",
    "韩系插画风女生头像，粉紫色短发+透明感腮红，侧光渲染。",
    "点彩派，盛夏海滨，两位渔夫正在搬运木箱，三艘帆船停在岸边，对角线构图。",
    "一幅由梵高绘制的梦境麦田，旋转的蓝色星云与燃烧的向日葵相纠缠。",
]
test_list_en = [
    "Create a painting depicting a 30-year-old white female white-collar worker on a business trip by plane.",
    "Depicted in the anime style of Studio Ghibli, a girl stands quietly at the deck with a gentle smile.",
    "Blue background, a lone girl gazes into the distant sea; her expression is sorrowful.",
    "A blend of expressionist and vintage styles, drawing a building with colorful walls.",
    "Paint a winter scene with crystalline ice hangings from an Antarctic research station.",
]

with gr.Blocks(title="Prompt Enhancer_V2") as demo:
    gr.Markdown("## 提示词重写器")
    with gr.Row():
        with gr.Column(scale=2):
            model_path = gr.Textbox(
                label="模型路径（本地或HF地址）",
                value=DEFAULT_MODEL_PATH,
                placeholder="例如：Qwen/Qwen2.5-VL-7B-Instruct",
            )
            device_map = gr.Dropdown(
                choices=["cuda", "cpu"],
                value="cuda",
                label="device_map（模型加载映射）"
            )
            torch_dtype = gr.Dropdown(
                choices=["bfloat16", "float16", "float32"],
                value="bfloat16",
                label="torch_dtype"
            )

        with gr.Column(scale=3):
            sys_prompt = gr.Textbox(
                label="系统提示词(默认无需修改)",
                value="请根据用户的输入，生成思考过程的思维链并改写提示词：",
                lines=3
            )
            with gr.Row():
                temperature = gr.Slider(0, 1, value=0.1, step=0.05, label="Temperature")
                max_new_tokens = gr.Slider(16, 4096, value=2048, step=16, label="Max New Tokens")
                device = gr.Dropdown(choices=["cuda", "cpu"], value="cuda", label="推理device")

    state = gr.State(value=None)

    with gr.Tab("推理"):
        with gr.Row():
            with gr.Column(scale=2):
                prompt = gr.Textbox(label="输入提示词", lines=6, placeholder="在此粘贴要改写的提示词...")
                run_btn = gr.Button("生成重写", variant="primary")
                gr.Examples(
                    examples=test_list_zh + test_list_en,
                    inputs=prompt,
                    label="示例"
                )
            with gr.Column(scale=3):
                out_text = gr.Textbox(label="重写结果", lines=10)
                out_info = gr.Markdown("准备就绪。")

        run_btn.click(
            run_single,
            inputs=[prompt, sys_prompt, temperature, max_new_tokens, device,
                    model_path, device_map, torch_dtype, state],
            outputs=[out_text, out_info, state]
        )

    gr.Markdown("提示：如有任何问题可 email 联系：linqing1995@buaa.edu.cn")

# 为避免多并发导致显存爆，可限制并发（ZeroGPU 本身是无状态，仍建议限制）
# demo.queue(concurrency_count=1, max_size=10)

if __name__ == "__main__":
    demo.launch(ssr_mode=False, show_error=True, share=True)