Update app.py

app.py

@@ -0,0 +1,312 @@
+import os
+import requests
+import json
+import time
+import threading
+import uuid
+import shutil
+import base64
+from datetime import datetime
+from pathlib import Path
+from http.server import HTTPServer, SimpleHTTPRequestHandler
+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",
+    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 image_to_base64(file_path):
+    try:
+        with open(file_path, "rb") as f:
+            img_data = f.read()
+            if len(img_data) == 0:
+                raise ValueError("空文件")
+            
+            # 使用URL安全编码并自动填充
+            encoded = base64.urlsafe_b64encode(img_data)
+            missing_padding = len(encoded) % 4
+            if missing_padding:
+                encoded += b'=' * (4 - missing_padding)
+                
+            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')
+            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), "#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"
+
+@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,
+    negative_prompt,
+    steps,
+    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("Hourly limit exceeded (20 requests)")
+        yield "❌ 请求过于频繁，请稍后再试", error_img
+        return
+
+    session = session_manager.get_session(session_id)
+    session['last_active'] = time.time()
+    session['count'] += 1
+
+    try:
+        api_key = os.getenv("WAVESPEED_API_KEY")
+        if not api_key:
+            raise ValueError("API key missing")
+            
+        base64_img = image_to_base64(image)
+        headers = {
+            "Authorization": f"Bearer {api_key}", 
+            "Content-Type": "application/json"
+            }
+        
+        payload = {
+            "context_scale": 1,
+            "enable_safety_checker": True,
+            "flow_shift": flow_shift,
+            "guidance_scale": guidance,
+            "images": [base64_img],
+            "negative_prompt": negative_prompt,
+            "num_inference_steps": steps,
+            "prompt": prompt,
+            "seed": seed if seed != -1 else random.randint(0, 999999),
+            "size": "480*832"
+        }
+
+        response = requests.post(
+            "https://api.wavespeed.ai/api/v3/wavespeed-ai/wan-2.1-14b-vace",
+            headers=headers,
+            json=payload
+        )
+        
+        if response.status_code != 200:
+            raise Exception(f"API Error {response.status_code}: {response.text}")
+            
+        requestId = response.json()["data"]["id"]
+        yield f"✅ 任务已提交 (ID: {requestId})", 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/v3/predictions/{requestId}/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"]
+            
+            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)
+
+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("# 🌊 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.
+        """)
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            img_input = gr.Image(type="filepath", label="Upload Image")
+            prompt = gr.Textbox(label="Prompt", lines=3, placeholder="Prompt...")
+            negative_prompt = gr.Textbox(label="Negative Prompt", lines=2)
+            
+            with gr.Row():
+                size = gr.Dropdown(["832*480", "480*832"], value="832 * 480", interactive=True, label="Resolution")
+                steps = gr.Slider(1, 50, value=30, label="Inference Steps")
+            with gr.Row():
+                duration = gr.Slider(1, 10, value=5, step=1, label="时长(秒)")
+                guidance = gr.Slider(1, 20, value=7, label="Guidance Scale")
+            with gr.Row():
+                seed = gr.Number(-1, label="Seed")
+                random_seed_btn = gr.Button("Random🎲Seed", variant="secondary")
+            with gr.Row():
+                enable_safety = gr.Checkbox(label="🔒 Enable Safety Checker",value=True, interactive=False)
+                flow_shift = gr.Slider(1, 50, value=16, label="flow_shift")
+                
+        with gr.Column(scale=1):
+            video_output = gr.Video(label="Generated Video", format="mp4", elem_classes=["video-preview"])
+            status_output = gr.Textbox(label="System Status", interactive=False, lines=4)
+            generate_btn = gr.Button("Generate Video", variant="primary")
+                
+    gr.Examples(
+        examples=[
+            ["The elegant lady carefully selects bags in the boutique, and she shows the charm of a mature woman in a black slim dress with a pearl necklace. Holding a vintage-inspired blue leather half-moon handbag, she is carefully observing its craftsmanship and texture. The interior of the store is a haven of sophistication and luxury. Soft, ambient lighting casts a warm glow over the polished wooden floors", 
+            "https://d2g64w682n9w0w.cloudfront.net/media/ec44bbf6abac4c25998dd2c4af1a46a7/images/1747413751234102420_md9ywspl.png"
+            ]
+        ],
+        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,
+            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
+    )