Update app.py

app.py

@@ -12,13 +12,14 @@ from dotenv import load_dotenv
 import gradio as gr
 import random
 import torch
+from PIL import Image, ImageDraw, ImageFont
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 from functools import lru_cache
 
-# 加载环境变量
+# 初始化环境
 load_dotenv()
 
-# ==== 新增安全检测模块 ====
+# 安全检测配置
 MODEL_URL = "TostAI/nsfw-text-detection-large"
 CLASS_NAMES = {
     0: "✅ SAFE",
@@ -26,18 +27,11 @@ CLASS_NAMES = {
     2: "🚫 UNSAFE"
 }
 
-# 加载模型和tokenizer
+# 加载模型
 tokenizer = AutoTokenizer.from_pretrained(MODEL_URL)
 model = AutoModelForSequenceClassification.from_pretrained(MODEL_URL)
 
-@lru_cache(maxsize=128)
-def classify_text(text):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
-    with torch.no_grad():
-        outputs = model(**inputs)
-    return torch.argmax(outputs.logits, dim=1).item()
-
-# ==== 会话管理模块 ====
+# 会话管理
 class SessionManager:
     _instances = {}
     _lock = threading.Lock()
@@ -47,9 +41,9 @@ class SessionManager:
         with cls._lock:
             if session_id not in cls._instances:
                 cls._instances[session_id] = {
-                    'request_count': 0,
-                    'last_request': time.time(),
-                    'history': []
+                    'count': 0,
+                    'history': [],
+                    'last_active': time.time()
                 }
             return cls._instances[session_id]
 
@@ -57,133 +51,256 @@ class SessionManager:
     def cleanup_sessions(cls):
         with cls._lock:
             now = time.time()
-            expired = [k for k, v in cls._instances.items() if now - v['last_request'] > 3600]
+            expired = [k for k, v in cls._instances.items() if now - v['last_active'] > 3600]
             for k in expired:
                 del cls._instances[k]
 
-# ==== 频率限制模块 ====
+# 频率限制
 class RateLimiter:
     def __init__(self):
-        self.client_data = {}
-
-    def check_limit(self, client_id):
-        if client_id not in self.client_data:
-            self.client_data[client_id] = {
-                'count': 0,
-                'reset_time': time.time() + 3600
-            }
-        
-        if time.time() > self.client_data[client_id]['reset_time']:
-            self.client_data[client_id] = {
-                'count': 0,
-                'reset_time': time.time() + 3600
-            }
-        
-        if self.client_data[client_id]['count'] >= 20:
-            return False
-        self.client_data[client_id]['count'] += 1
-        return True
+        self.clients = {}
+        self.lock = threading.Lock()
+
+    def check(self, client_id):
+        with self.lock:
+            now = time.time()
+            if client_id not in self.clients:
+                self.clients[client_id] = {'count': 1, 'reset': now + 3600}
+                return True
+                
+            if now > self.clients[client_id]['reset']:
+                self.clients[client_id] = {'count': 1, 'reset': now + 3600}
+                return True
+                
+            if self.clients[client_id]['count'] >= 20:
+                return False
+                
+            self.clients[client_id]['count'] += 1
+            return True
+
+# 初始化模块
+session_manager = SessionManager()
+rate_limiter = RateLimiter()
+
+# 图像处理函数
+def image_to_base64(file_path):
+    try:
+        with open(file_path, "rb") as f:
+            ext = Path(file_path).suffix.lower()[1:]
+            mime_map = {'jpg':'jpeg','jpeg':'jpeg','png':'png','webp':'webp','gif':'gif'}
+            mime = mime_map.get(ext, 'jpeg')
+            
+            encoded = base64.b64encode(f.read())
+            if len(encoded) % 4:
+                encoded += b'=' * (4 - len(encoded) % 4)
+                
+            return f"data:image/{mime};base64,{encoded.decode()}"
+    except Exception as e:
+        raise ValueError(f"Base64 Error: {str(e)}")
 
-# ==== 错误处理模块 ====
 def create_error_image(message):
-    img = Image.new("RGB", (832, 480), color="#ffdddd")
-    draw = ImageDraw.Draw(img)
-    
+    img = Image.new("RGB", (832, 480), "#ffdddd")
     try:
         font = ImageFont.truetype("arial.ttf", 24)
     except:
         font = ImageFont.load_default()
-    
-    text_width, text_height = draw.textsize(message, font=font)
-    x = (832 - text_width) / 2
-    y = (480 - text_height) / 2
-    
-    draw.text((x, y), message, fill="#ff4444", font=font)
+        
+    draw = ImageDraw.Draw(img)
+    text = f"Error: {message[:60]}..." if len(message) > 60 else message
+    draw.text((50, 200), text, fill="#ff0000", font=font)
     img.save("error.jpg")
     return "error.jpg"
 
-# ==== 核心生成逻辑 ====
-def generate_video(/* 保持原有参数 */):
-    # 新增安全检测
-    safety_level = classify_text(prompt)
+# 核心生成逻辑
+@lru_cache(maxsize=100)
+def classify_prompt(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    return torch.argmax(outputs.logits).item()
+
+def generate_video(
+    image,
+    prompt,
+    duration,
+    enable_safety,
+    flow_shift,
+    guidance_scale,
+    negative_prompt,
+    inference_steps,
+    seed,
+    size,
+    session_id
+):
+    # 安全检查
+    safety_level = classify_prompt(prompt)
     if safety_level != 0:
-        error_msg = f"Content blocked: {CLASS_NAMES[safety_level]}"
-        error_img = create_error_image(error_msg)
-        yield f"❌ {error_msg}", error_img
+        error_img = create_error_image(CLASS_NAMES[safety_level])
+        yield f"❌ Blocked: {CLASS_NAMES[safety_level]}", error_img
         return
 
-    # 新增频率检查
-    session = SessionManager.get_session(session_id)
-    if session['request_count'] >= 20:
-        yield "❌ Hourly limit exceeded (20 requests)", None
+    # 频率检查
+    if not rate_limiter.check(session_id):
+        error_img = create_error_image("Hourly limit exceeded (20 requests)")
+        yield "❌ 请求过于频繁，请稍后再试", error_img
         return
-    session['request_count'] += 1
 
-    # 原有生成逻辑保持不变，增加状态跟踪
+    # 会话更新
+    session = session_manager.get_session(session_id)
+    session['last_active'] = time.time()
+    session['count'] += 1
+
+    # API调用
     try:
-        # API调用部分
-        response = requests.post(/* 保持原有参数 */)
-        
-        # 轮询部分增加进度跟踪
-        while True:
-            # 获取状态
-            status = get_status(request_id)
+        # 准备请求
+        api_key = os.getenv("WAVESPEED_API_KEY")
+        if not api_key:
+            raise ValueError("API key missing")
             
-            # 更新会话最后活动时间
-            session['last_request'] = time.time()
-            
-            # 处理不同状态
-            if status == 'processing':
-                yield f"⏳ 生成进度: {progress}%", None
-            elif status == 'completed':
-                session['history'].append(video_url)
-                yield f"✅ 生成完成", video_url
-                return
+        base64_img = image_to_base64(image)
+        headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
+        
+        payload = {
+            "image": base64_img,
+            "prompt": prompt,
+            "duration": duration,
+            "guidance_scale": guidance_scale,
+            "negative_prompt": negative_prompt,
+            "num_inference_steps": inference_steps,
+            "seed": seed if seed != -1 else random.randint(0, 999999),
+            "size": size
+        }
 
+        # 提交任务
+        response = requests.post(
+            "https://api.wavespeed.ai/api/v2/wavespeed-ai/wan-2.1/i2v-480p-ultra-fast",
+            headers=headers,
+            json=payload
+        )
+        
+        if response.status_code != 200:
+            raise Exception(f"API Error {response.status_code}: {response.text}")
+            
+        request_id = response.json()["data"]["id"]
+        yield f"✅ 任务已提交 (ID: {request_id})", None
+        
     except Exception as e:
         error_img = create_error_image(str(e))
-        yield f"❌ 生成失败: {str(e)}", error_img
-
-# ==== 新增定时清理任务 ====
-def start_cleanup_task():
-    def cleanup():
-        while True:
-            SessionManager.cleanup_sessions()
-            time.sleep(3600)
+        yield f"❌ 提交失败: {str(e)}", error_img
+        return
+
+    # 轮询结果
+    result_url = f"https://api.wavespeed.ai/api/v2/predictions/{request_id}/result"
+    start_time = time.time()
+    
+    while True:
+        time.sleep(1)
+        try:
+            resp = requests.get(result_url, headers=headers)
+            if resp.status_code != 200:
+                raise Exception(f"状态查询失败: {resp.text}")
+                
+            data = resp.json()["data"]
+            status = data["status"]
             
-    thread = threading.Thread(target=cleanup)
-    thread.daemon = True
-    thread.start()
-
-# ==== 界面增强 ====
-with gr.Blocks(/* 保持原有参数 */) as app:
-    # 新增状态组件
-    status_bars = {}
+            if status == "completed":
+                elapsed = time.time() - start_time
+                video_url = data["outputs"][0]
+                session["history"].append(video_url)
+                yield f"🎉 生成成功! 耗时 {elapsed:.1f}s", video_url
+                return
+                
+            elif status == "failed":
+                raise Exception(data.get("error", "Unknown error"))
+                
+            else:
+                yield f"⏳ 当前状态: {status.capitalize()}...", None
+                
+        except Exception as e:
+            error_img = create_error_image(str(e))
+            yield f"❌ 生成失败: {str(e)}", error_img
+            return
+
+# 后台清理线程
+def cleanup_task():
+    while True:
+        session_manager.cleanup_sessions()
+        time.sleep(3600)
+
+# Gradio界面
+with gr.Blocks(
+    theme=gr.themes.Soft(),
+    css="""
+    .video-preview { max-width: 600px !important; }
+    .status-box { padding: 10px; border-radius: 5px; margin: 5px; }
+    .safe { background: #e8f5e9; border: 1px solid #a5d6a7; }
+    .warning { background: #fff3e0; border: 1px solid #ffcc80; }
+    .error { background: #ffebee; border: 1px solid #ef9a9a; }
+    """
+) as app:
+    
+    session_id = gr.State(str(uuid.uuid4()))
+    
+    gr.Markdown("# 🌊 视频生成系统 - WaveSpeedAI")
     
     with gr.Row():
-        for backend in ["WAN-2.1", "FLUX", "TURBO"]:  # 示例后端
-            with gr.Column():
-                gr.Markdown(f"**{backend}**")
-                status_bars[backend] = gr.Textbox(label="状态", value="🟢 空闲")
+        with gr.Column(scale=1):
+            img_input = gr.Image(type="filepath", label="上传图片")
+            prompt = gr.Textbox(label="描述文本", lines=3, placeholder="请输入画面描述...")
+            negative_prompt = gr.Textbox(label="排除内容", lines=2)
+            with gr.Row():
+                size = gr.Dropdown(["832 * 480"], label="分辨率")
+                steps = gr.Slider(1, 50, value=30, label="推理步数")
+            with gr.Row():
+                duration = gr.Slider(1, 10, value=5, step=1, label="时长(秒)")
+                guidance = gr.Slider(1, 20, value=7, label="引导强度")
+            with gr.Row():
+                seed = gr.Number(-1, label="随机种子")
+                random_seed_btn = gr.Button("随机生成", variant="secondary")
                 
-    # 在生成按钮点击时更新状态
-    generate_btn.click(
-        fn=update_status,
-        inputs=[...],
-        outputs=[status_bars]
+        with gr.Column(scale=1):
+            video_output = gr.Video(label="生成结果", format="mp4", elem_classes=["video-preview"])
+            status_output = gr.Textbox(label="系统状态", interactive=False, lines=4)
+            generate_btn = gr.Button("开始生成", variant="primary")
+            
+            with gr.Accordion("生成历史", open=False):
+                history_gallery = gr.Gallery(label="历史记录", columns=3)
+                
+            with gr.Accordion("安全状态", open=True):
+                gr.Markdown("""
+                <div class="status-box safe">
+                    ✅ 当前内容安全检测通过
+                </div>
+                """)
+                
+    # 示例区
+    gr.Examples(
+        examples=[
+            ["19世纪绅士在石板街", "example1.jpg"],
+            ["赛博朋克女战士在雨夜", "example2.jpg"]
+        ],
+        inputs=[prompt, img_input],
+        label="示例输入"
     )
 
-    # 新增历史记录模块
-    history = gr.Gallery(label="生成历史")
+    # 事件处理
+    random_seed_btn.click(
+        fn=lambda: random.randint(0, 999999),
+        outputs=seed
+    )
     
-    # 在生成完成后更新历史
     generate_btn.click(
-        fn=update_history,
-        inputs=[video_output],
-        outputs=[history]
+        generate_video,
+        inputs=[img_input, prompt, duration, gr.State(True), gr.State(3),
+               guidance, negative_prompt, steps, seed, size, session_id],
+        outputs=[status_output, video_output]
     )
 
-# ==== 启动时初始化 ====
+# 启动系统
 if __name__ == "__main__":
-    start_cleanup_task()
-    app.queue(/* 保持原有参数 */)
+    threading.Thread(target=cleanup_task, daemon=True).start()
+    app.queue(max_size=4).launch(
+        server_name="0.0.0.0",
+        max_threads=16,
+        share=False
+    )