Update app.py

app.py

@@ -1,241 +1,297 @@
-import os
-import requests
-import json
-import time
-import threading
-import shutil
-from datetime import datetime
-from pathlib import Path
-from http.server import HTTPServer, SimpleHTTPRequestHandler
-import base64
-from dotenv import load_dotenv
-import gradio as gr
-import random
-
-load_dotenv()
-
-def image_to_base64(file_path):
-    try:
-        with open(file_path, "rb") as image_file:
-            # 处理特殊MIME类型
-            ext = Path(file_path).suffix.lower().lstrip('.')
-            mime_map = {
-                'jpg': 'jpeg',
-                'jpeg': 'jpeg',
-                'png': 'png',
-                'webp': 'webp',
-                'gif': 'gif'
-            }
-            mime_type = mime_map.get(ext, 'jpeg')
-            
-            # 读取并编码
-            raw_data = image_file.read()
-            encoded = base64.b64encode(raw_data)
-            missing_padding = len(encoded) % 4
-            if missing_padding:
-                encoded += b'=' * (4 - missing_padding)
-                
-            return f"data:image/{mime_type};base64,{encoded.decode('utf-8')}"
-            
-    except Exception as e:
-        raise ValueError(f"Base64编码失败: {str(e)}")
-
-def generate_random_seed():
-    return random.randint(0, 999999)
-
-def generate_video(
-    image,
-    prompt,
-    duration,
-    enable_safety,
-    flow_shift,
-    guidance_scale,
-    negative_prompt,
-    inference_steps,
-    seed,
-    size
-):
-    API_KEY = os.getenv("WAVESPEED_API_KEY")
-    if not API_KEY:
-        yield "❌ Error: Missing API Key", None
-        return
-
-    try:
-        base64_image = image_to_base64(image)
-    except Exception as e:
-        yield f"❌ File upload failed: {str(e)}", None
-        return
-
-    payload = {
-        "duration": duration,
-        "enable_safety_checker": enable_safety,
-        "flow_shift": flow_shift,
-        "guidance_scale": guidance_scale,
-        "image": base64_image,
-        "negative_prompt": negative_prompt,
-        "num_inference_steps": inference_steps,
-        "prompt": prompt,
-        "seed": seed,
-        "size": size
-    }
-
-    headers = {
-        "Content-Type": "application/json",
-        "Authorization": f"Bearer {API_KEY}",
-    }
-
-    try:
-        response = requests.post(
-            "https://api.wavespeed.ai/api/v2/wavespeed-ai/wan-2.1/i2v-480p-ultra-fast",
-            headers=headers,
-            data=json.dumps(payload)
-        )
-        
-        if response.status_code != 200:
-            yield f"❌ API Error ({response.status_code}): {response.text}", None
-            return
-            
-        request_id = response.json()["data"]["id"]
-        yield f"✅ Task submitted (ID: {request_id})", None
-    except Exception as e:
-        yield f"❌ Connection Error: {str(e)}", None
-        return
-
-    # 轮询结果
-    result_url = f"https://api.wavespeed.ai/api/v2/predictions/{request_id}/result"
-    start_time = time.time()
-    video_url = None
-    
-    while True:
-        time.sleep(1)
-        try:
-            response = requests.get(result_url, headers={"Authorization": f"Bearer {API_KEY}"})
-            if response.status_code != 200:
-                yield f"❌ Polling Error ({response.status_code}): {response.text}", None
-                return
-
-            data = response.json()["data"]
-            status = data["status"]
-            
-            if status == "completed":
-                elapsed = time.time() - start_time
-                video_url = data['outputs'][0]
-                yield (f"🎉 Completed in {elapsed:.1f}s!\n"
-                       f"Download URL: {video_url}"), video_url
-                return
-                
-            elif status == "failed":
-                yield f"❌ Failed: {data.get('error', 'Unknown error')}", None
-                return
-                
-            else:
-                yield f"⏳ Status: {status.capitalize()}...", None
-                
-        except Exception as e:
-            yield f"❌ Polling Failed: {str(e)}", None
-            return
-
-# Gradio UI
-with gr.Blocks(
-    theme=gr.themes.Soft(), 
-    css="""
-    .video-preview {
-        max-width: 600px !important
-        }
-    .example-preview {
-        border: 1px solid #e0e0e0;
-        border-radius: 8px;
-        padding: 10px;
-        margin: 5px;
-    }
-    .example-preview img {
-        max-width: 200px;
-        max-height: 150px;
-    }
-    """
-    ) as app:
-
-    session_id = gr.State(None)
-
-    gr.Markdown("# 🌊 Wan-2.1-i2v-480p-Ultra-Fast Run On WaveSpeedAI")
-
-    gr.Markdown("""
-        [WaveSpeedAI](https://wavespeed.ai/) is the global pioneer in accelerating AI-powered video and image generation.
-        Our in-house inference accelerator provides lossless speedup on image & video generation based on our rich inference optimization software stack, including our in-house inference compiler, CUDA kernel libraries and parallel computing libraries.
-        """)
-    gr.Markdown("""
-        The Wan2.1 14B model is an advanced image-to-video model that offers accelerated inference capabilities, enabling high-res video generation with high visual quality and motion diversity.
-        """)
-    
-    with gr.Row():        
-        # 右侧控制面板
-        with gr.Column(scale=1):
-            with gr.Row():
-                with gr.Column(scale=1):
-                    img_input = gr.Image(type="filepath", label="Upload Image")
-                    prompt = gr.Textbox(label="Prompt", lines=5, placeholder="Describe your scene...")
-                    negative_prompt = gr.Textbox(label="Negative Prompt", lines=2)
-                    size = gr.Dropdown(["832*480"], value="832*480", label="Resolution")
-                    steps = gr.Slider(1, 50, value=30, step=1, label="Inference Steps")
-                    duration = gr.Slider(0, 10, value=5, step=5, label="Duration (seconds)")
-                    guidance = gr.Slider(1, 30, value=5, step=0.1, label="Guidance Scale") 
-                    seed = gr.Number(-1, label="Seed")
-                    random_seed_btn = gr.Button("🎲random seed", variant="secondary")
-                    flow_shift = gr.Number(3, label="Flow Shift",interactive=False)
-                    enable_safety = gr.Checkbox(True, label="Safety Checker",interactive=False)
-        # 左侧视频展示区域
-        with gr.Column(scale=1):
-                    video_output = gr.Video(label="Generated Video",format="mp4",interactive=False,elem_classes=["video-preview"]
-                    )
-                    generate_btn = gr.Button("Generate Video", variant="primary")
-                    output = gr.Textbox(label="Status", interactive=False, lines=4)
-                    gr.Examples(
-                        examples=[
-                            [
-                                "Victorian era, 19th-century gentleman wearing a black top hat and tuxedo, standing on a cobblestone street, dim gaslight lamps, passersby in vintage clothing, gentle breeze moving his coat, slow cinematic pan around him, nostalgic retro film style, realistic textures",
-                                "https://d2g64w682n9w0w.cloudfront.net/media/images/1745725874603980753_95mFCAxu.jpg"
-                            ],
-                            [
-                                "A cyberpunk female warrior with short silver hair and glowing green eyes, wearing a futuristic armored suit, standing in a neon-lit rainy city street, camera slowly circling around her, raindrops falling in slow motion, neon reflections on wet pavement, cinematic atmosphere, highly detailed, ultra realistic, 4K",
-                                "https://d2g64w682n9w0w.cloudfront.net/media/images/1745726299175719855_pFO0WSRM.jpg"
-                            ],
-                            [
-                                "Wide shot of a brave medieval female knight in shining silver armor and a red cape, standing on a castle rooftop at sunset, slowly drawing a large ornate sword from its scabbard, seen from a distance with the vast castle and surrounding landscape in the background, golden light bathing the scene, hair and cape flowing gently in the wind, cinematic epic atmosphere, dynamic motion, majestic clouds drifting, ultra realistic, high fantasy world, 4K ultra-detailed",
-                                "https://d2g64w682n9w0w.cloudfront.net/media/images/1745727436576834405_rtsokheb.jpg"
-                            ],
-                            [
-                                "A girl stands in a lively 17th-century market. She holds a red tomato, looks gently into the camera and smiles briefly. Then, she glances at the tomato in her hand, slowly sets it back into the basket, turns around gracefully, and walks away with her back to the camera. The market around her is rich with colorful vegetables, meats hanging above, and bustling townsfolk. Golden-hour painterly lighting, subtle facial expressions, smooth cinematic motion, ultra-realistic detail, Vermeer-inspired style",
-                                "https://d2g64w682n9w0w.cloudfront.net/media/images/1745079024013078406_QT6jKNPZ.png"
-                            ],
-                            [
-                                "A calming video explaining diabetes management and prevention tips to reduce anxiety.",
-                                "https://d2g64w682n9w0w.cloudfront.net/predictions/517d518c28ef49ed9464610af48528f5/1.jpg"
-                            ],
-                            [
-                                "Girl dancing and spinning with friends.",
-                                "https://d2g64w682n9w0w.cloudfront.net/media/d45e0d4893d44712b359f3ad0b3c2795/images/1745449961409630099_KISOKGEB.jpg"
-                            ]
-                        ],
-                        inputs=[prompt, img_input],  # 同时绑定到���片和提示输入框
-                        label="Example Inputs",
-                        examples_per_page=3
-                    )
-                
-    random_seed_btn.click(
-        fn=lambda: random.randint(0, 999999),
-        outputs=seed
-    )   
-
-    generate_btn.click(
-        generate_video,
-        inputs=[img_input, prompt, duration, enable_safety, flow_shift, 
-               guidance, negative_prompt, steps, seed, size],
-        outputs=[output, video_output]
-    )
-
-if __name__ == "__main__":
-    app.queue(max_size=4).launch(
-        server_name="0.0.0.0", 
-        max_threads=16,
-        debug=True
+import os
+import requests
+import json
+import time
+import random
+import base64
+import uuid
+import threading
+from pathlib import Path
+from dotenv import load_dotenv
+import gradio as gr
+import torch
+from PIL import Image, ImageDraw, ImageFont
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+load_dotenv()
+
+MODEL_URL = "TostAI/nsfw-text-detection-large"
+CLASS_NAMES = {0: "✅ SAFE", 1: "⚠️ QUESTIONABLE", 2: "🚫 UNSAFE"}
+tokenizer = AutoTokenizer.from_pretrained(MODEL_URL)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_URL)
+
+class SessionManager:
+    _instances = {}
+    _lock = threading.Lock()
+
+    @classmethod
+    def get_session(cls, session_id):
+        with cls._lock:
+            if session_id not in cls._instances:
+                cls._instances[session_id] = {
+                    'count': 0,
+                    'history': [],
+                    'last_active': time.time()
+                }
+            return cls._instances[session_id]
+
+    @classmethod
+    def cleanup_sessions(cls):
+        with cls._lock:
+            now = time.time()
+            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.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 create_error_image(message):
+    img = Image.new("RGB", (832, 480), "#ffdddd")
+    try:
+        font = ImageFont.truetype("arial.ttf", 24)
+    except:
+        font = ImageFont.load_default()
+    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 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 image_to_base64(file_path):
+    try:
+        with open(file_path, "rb") as image_file:
+            ext = Path(file_path).suffix.lower().lstrip('.')
+            mime_map = {
+                'jpg': 'jpeg',
+                'jpeg': 'jpeg',
+                'png': 'png',
+                'webp': 'webp',
+                'gif': 'gif'
+            }
+            mime_type = mime_map.get(ext, 'jpeg')
+            
+            raw_data = image_file.read()
+            encoded = base64.b64encode(raw_data)
+            missing_padding = len(encoded) % 4
+            if missing_padding:
+                encoded += b'=' * (4 - missing_padding)
+                
+            return f"data:image/{mime_type};base64,{encoded.decode('utf-8')}"
+    except Exception as e:
+        raise ValueError(f"Base64编码失败: {str(e)}")
+
+def generate_video(
+    image,
+    prompt,
+    enable_safety,
+    flow_shift,
+    guidance_scale,
+    negative_prompt,
+    seed,
+    size,
+    session_id
+):
+
+    safety_level = classify_prompt(prompt)
+    if safety_level != 0:
+        error_img = create_error_image(CLASS_NAMES[safety_level])
+        yield f"❌ Blocked: {CLASS_NAMES[safety_level]}", error_img
+        return
+
+    if not rate_limiter.check(session_id):
+        error_img = create_error_image("每小时限制20次请求")
+        yield "❌ 请求过于频繁，请稍后再试", error_img
+        return
+
+    session = session_manager.get_session(session_id)
+    session['last_active'] = time.time()
+    session['count'] += 1
+
+    API_KEY = os.getenv("WAVESPEED_API_KEY")
+    if not API_KEY:
+        error_img = create_error_image("API密钥缺失")
+        yield "❌ Error: Missing API Key", error_img
+        return
+
+    try:
+        base64_image = image_to_base64(image)
+    except Exception as e:
+        error_img = create_error_image(str(e))
+        yield f"❌ 文件上传失败: {str(e)}", error_img
+        return
+
+    payload = {
+        "enable_safety_checker": enable_safety,
+        "flow_shift": flow_shift,
+        "guidance_scale": guidance_scale,
+        "image": base64_image,
+        "negative_prompt": negative_prompt,
+        "prompt": prompt,
+        "seed": seed if seed != -1 else random.randint(0, 999999),
+        "size": size
+    }
+
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {API_KEY}",
+    }
+
+    try:
+        response = requests.post(
+            "https://api.wavespeed.ai/api/v2/wavespeed-ai/hunyuan-custom-ref2v-480p",
+            headers=headers,
+            data=json.dumps(payload)
+        )
+        
+        if response.status_code != 200:
+            error_img = create_error_image(response.text)
+            yield f"❌ API错误 ({response.status_code}): {response.text}", error_img
+            return
+            
+        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
+        return
+
+    result_url = f"https://api.wavespeed.ai/api/v2/predictions/{request_id}/result"
+    start_time = time.time()
+    
+    while True:
+        time.sleep(0.5)
+        try:
+            response = requests.get(result_url, headers=headers)
+            if response.status_code != 200:
+                error_img = create_error_image(response.text)
+                yield f"❌ 轮询错误 ({response.status_code}): {response.text}", error_img
+                return
+
+            data = response.json()["data"]
+            status = data["status"]
+            
+            if status == "completed":
+                elapsed = time.time() - start_time
+                video_url = data['outputs'][0]
+                session["history"].append(video_url)
+                yield (f"🎉 完成! 耗时 {elapsed:.1f}秒\n"
+                       f"下载链接: {video_url}"), video_url
+                return
+                
+            elif status == "failed":
+                error_img = create_error_image(data.get('error', '未知错误'))
+                yield f"❌ 任务失败: {data.get('error', '未知错误')}", error_img
+                return
+                
+            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)
+
+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("# 🌊Hunyuan-Custom-Ref2v Run On [WaveSpeedAI](https://wavespeed.ai/)")
+    gr.Markdown("""HunyuanCustom, a multi-modal, conditional, and controllable generation model centered on subject consistency, built upon the Hunyuan Video generation framework. It enables the generation of subject-consistent videos conditioned on text, images, audio, and video inputs.""")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            img_input = gr.Image(type="filepath", label="Input Image")
+            prompt = gr.Textbox(label="Prompt", lines=5, placeholder="Prompt...")
+            negative_prompt = gr.Textbox(label="Negative Prompt", lines=2)
+            size = gr.Dropdown(["832*480", "480*832"], value="832*480", label="Size")
+            guidance = gr.Slider(1, 30, value=5, step=0.1, label="Guidance")
+            seed = gr.Number(-1, label="Seed")
+            random_seed_btn = gr.Button("Random🎲Seed", variant="secondary")
+            flow_shift = gr.Number(3, label="Shift", interactive=False)
+            enable_safety = gr.Checkbox(True, label="Enable Safety Checker", interactive=False)
+
+        with gr.Column(scale=1):
+            video_output = gr.Video(label="Video Output", format="mp4", interactive=False,  elem_classes=["video-preview"])
+            generate_btn = gr.Button("Generate", variant="primary")
+            status_output = gr.Textbox(label="status", interactive=False, lines=4)
+                
+            gr.Examples(
+                examples=[
+                    [
+                        "Create a dynamic and intense scene depicting a warrior fighting a fearsome fire dragon. The setting should be grand and immersive, capturing the scale and drama of the confrontation.",
+                        "https://d2g64w682n9w0w.cloudfront.net/media/images/1745337597019716935_RS1XUQMI.jpg"
+                    ]
+                ],
+                inputs=[prompt, img_input],
+                label="Examples Prompt",
+                examples_per_page=3
+            )
+
+    random_seed_btn.click(
+        fn=lambda: random.randint(0, 999999),
+        outputs=seed
+    )
+
+    generate_btn.click(
+        generate_video,
+        inputs=[
+            img_input,
+            prompt,
+            enable_safety,
+            flow_shift,
+            guidance,
+            negative_prompt,
+            seed,
+            size,
+            session_id
+        ],
+        outputs=[status_output, video_output]
+    )
+
+if __name__ == "__main__":
+    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
     )