Update app.py

app.py

@@ -1,16 +1,12 @@
-# app.py
-# Gradio UI for PromptEnhancerV2
-
 import os
-from threading import Thread
-from transformers import TextIteratorStreamer, AutoTokenizer
 import time
 import logging
 import re
-import torch
 import gradio as gr
-import spaces
+from spaces import zero  # 关键：引入 zero 装饰器
 
+# 不要在这里 import torch 或加载模型
+# from transformers import TextIteratorStreamer, AutoTokenizer  # 不再需要
 
 # 尝试导入 qwen_vl_utils，若失败则提供降级实现（返回空的图像/视频输入）
 try:
@@ -25,120 +21,114 @@ def replace_single_quotes(text):
     replaced_text = replaced_text.replace("’", "”").replace("‘", "“")
     return replaced_text
 
-class PromptEnhancerV2:
-    @spaces.GPU 
-    def __init__(self, models_root_path, device_map="auto", torch_dtype="bfloat16"):#auto
-        from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
-        if not logging.getLogger(__name__).handlers:
-            logging.basicConfig(level=logging.INFO)
-        self.logger = logging.getLogger(__name__)
-
-        # dtype 兼容处理
-        if torch_dtype == "bfloat16":
-            dtype = torch.bfloat16
-        elif torch_dtype == "float16":
-            dtype = torch.float16
-        else:
-            dtype = torch.float32
-
-        self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-            models_root_path,
-            torch_dtype=dtype,
-            device_map=device_map,
-        )
-        self.processor = AutoProcessor.from_pretrained(models_root_path)
-
-    # @torch.inference_mode()
-    @spaces.GPU
-    def predict(
-        self,
-        prompt_cot,
-        sys_prompt="请根据用户的输入，生成思考过程的思维链并改写提示词：",
-        temperature=0.1,
-        top_p=1.0,
-        max_new_tokens=2048,
-        device="cuda",
-    ):
-        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 = self.processor.apply_chat_template(
-                messages, tokenize=False, add_generation_prompt=True
-            )
-            image_inputs, video_inputs = process_vision_info(messages)
-            inputs = self.processor(
-                text=[text],
-                images=image_inputs,
-                videos=video_inputs,
-                padding=True,
-                return_tensors="pt",
-            )
-            inputs = inputs.to(device)
-
-            # 注意：原始代码固定 do_sample=False，top_k=5, top_p=0.9，这里保持一致
-            generated_ids = self.model.generate(
-                **inputs,
-                max_new_tokens=2048,  # 与原始代码保持一致（未使用 max_new_tokens 参数）
-                temperature=float(temperature),
-                do_sample=False,
-                top_k=5,
-                top_p=0.9
-            )
-            generated_ids_trimmed = [
-                out_ids[len(in_ids):]
-                for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
-            ]
-            output_text = self.processor.batch_decode(
-                generated_ids_trimmed,
-                skip_special_tokens=True,
-                clean_up_tokenization_spaces=False,
-            )
-            output_res = output_text[0]
-            assert output_res.count("think>") == 2
-            prompt_cot = output_res.split("think>")[-1]
-            if prompt_cot.startswith("\n"):
-                prompt_cot = prompt_cot[1:]
-            prompt_cot = replace_single_quotes(prompt_cot)
-        except Exception as e:
-            prompt_cot = org_prompt_cot
-            print(f"✗ Re-prompting failed, so we are using the original prompt. Error: {e}")
-
-        return prompt_cot
-# -------------------------
-# Gradio app helpers
-# -------------------------
-
 DEFAULT_MODEL_PATH = os.environ.get("MODEL_OUTPUT_PATH", "PromptEnhancer/PromptEnhancer-32B")
 
-def ensure_enhancer(state, model_path, device_map, torch_dtype):
-    """
-    state: dict or None
-    Returns: (state_dict)
-    """
-    need_reload = False
-    if state is None or not isinstance(state, dict):
-        need_reload = True
+def _str_to_dtype(dtype_str):
+    # 在子进程中再真正用 torch；这里仅返回字符串用于传参
+    if dtype_str in ("bfloat16", "float16", "float32"):
+        return dtype_str
+    return "float32"
+
+@zero.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:
-        prev_path = state.get("model_path")
-        prev_map = state.get("device_map")
-        prev_dtype = state.get("torch_dtype")
-        if prev_path != model_path or prev_map != device_map or prev_dtype != torch_dtype:
-            need_reload = True
+        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)
 
-    if need_reload:
-        enhancer = PromptEnhancerV2(model_path, device_map=device_map, torch_dtype=torch_dtype)
-        return {"enhancer": enhancer, "model_path": model_path, "device_map": device_map, "torch_dtype": torch_dtype}
-    return state
+    # 组装消息
+    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):
@@ -146,21 +136,24 @@ def run_single(prompt, sys_prompt, temperature, max_new_tokens, device,
         return "", "请先输入提示词。", state
 
     t0 = time.time()
-    state = ensure_enhancer(state, model_path, device_map, torch_dtype)
-    enhancer = state["enhancer"]
     try:
-        out = enhancer.predict(
+        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
+            device=device,
         )
         dt = time.time() - t0
-        return out, f"耗时：{dt:.2f}s", state
+        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
-        
+        return "", f"调用失败：{e}", state
+
 # 示例数据
 test_list_zh = [
     "第三人称视角，赛车在城市赛道上飞驰，左上角是小地图，地图下面是当前名次，右下角仪表盘显示当前速度。",
@@ -183,13 +176,13 @@ with gr.Blocks(title="Prompt Enhancer_V2") as demo:
             model_path = gr.Textbox(
                 label="模型路径（本地或HF地址）",
                 value=DEFAULT_MODEL_PATH,
-                placeholder="/apdcephfs_jn3/share_302243908/aladdinwang/model_weight/cot_taurus_v6_50/global_step0",
+                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",
@@ -204,7 +197,7 @@ with gr.Blocks(title="Prompt Enhancer_V2") as demo:
             )
             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（原代码未使用该参数）")
+                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)
@@ -223,12 +216,6 @@ with gr.Blocks(title="Prompt Enhancer_V2") as demo:
                 out_text = gr.Textbox(label="重写结果", lines=10)
                 out_info = gr.Markdown("准备就绪。")
 
-        # run_btn.click(
-        #     stream_single,
-        #     inputs=[prompt, sys_prompt, temperature, max_new_tokens, device,
-        #     model_path, device_map, torch_dtype, state],
-        #     outputs=[out_text, out_info, state]
-        #     )
         run_btn.click(
             run_single,
             inputs=[prompt, sys_prompt, temperature, max_new_tokens, device,
@@ -236,12 +223,10 @@ with gr.Blocks(title="Prompt Enhancer_V2") as demo:
             outputs=[out_text, out_info, state]
         )
 
-    gr.Markdown(
-        "提示：如有任何问题可email联系：[email protected]"
-    )
+    gr.Markdown("提示：如有任何问题可 email 联系：[email protected]")
 
-# 为避免多并发导致显存爆，限制并发
+# 为避免多并发导致显存爆，可限制并发（ZeroGPU 本身是无状态，仍建议限制）
 # demo.queue(concurrency_count=1, max_size=10)
+
 if __name__ == "__main__":
-    # demo.launch(server_name="0.0.0.0", server_port=8080, show_error=True)
     demo.launch(ssr_mode=False, show_error=True, share=True)