# 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 # 尝试导入 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 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 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 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 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() state = ensure_enhancer(state, model_path, device_map, torch_dtype) enhancer = state["enhancer"] try: out = enhancer.predict( prompt_cot=prompt, sys_prompt=sys_prompt, temperature=temperature, max_new_tokens=max_new_tokens, device=device ) dt = time.time() - t0 return out, 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="/apdcephfs_jn3/share_302243908/aladdinwang/model_weight/cot_taurus_v6_50/global_step0", ) 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( # 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, model_path, device_map, torch_dtype, state], outputs=[out_text, out_info, state] ) gr.Markdown( "提示：如有任何问题可email联系：linqing1995@buaa.edu.cn" ) # 为避免多并发导致显存爆，限制并发 # 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)