Create app.py

app.py

@@ -0,0 +1,332 @@
+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
+import logging
+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:
+            raw_data = image_file.read()
+            encoded = base64.b64encode(raw_data)
+            missing_padding = len(encoded) % 4
+            if missing_padding:
+                encoded += b'=' * (4 - missing_padding)
+            return encoded.decode('utf-8')
+    except Exception as e:
+        raise ValueError(f"Base64编码失败: {str(e)}")
+
+def generate_video(
+    context_scale,
+    enable_safety_checker,
+    flow_shift,
+    guidance_scale,
+    images,
+    negative_prompt,
+    num_inference_steps,
+    prompt,
+    seed,
+    size,
+    task,
+    video,
+    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:
+        if not images or len(images) < 2:
+            raise ValueError("需要上传至少两张图片")
+            
+        base64_images = []
+        for img_path in images[:2]:
+            base64_img = image_to_base64(img_path)
+            base64_images.append(base64_img)
+            
+    except Exception as e:
+        error_img = create_error_image(str(e))
+        yield f"❌ 文件处理失败: {str(e)}", error_img
+        return
+
+    video_payload = ""
+    if video is not None:
+        if isinstance(video, (list, tuple)):
+            video_payload = video[0] if video else ""
+        else:
+            video_payload = video
+
+    payload = {
+    "context_scale": context_scale,
+    "enable_fast_mode": False,
+    "enable_safety_checker": enable_safety_checker,
+    "flow_shift": flow_shift,
+    "guidance_scale": guidance_scale,
+    "images": base64_images,
+    "negative_prompt": negative_prompt,
+    "num_inference_steps": num_inference_steps,
+    "prompt": prompt,
+    "seed": seed if seed != -1 else random.randint(0, 999999),
+    "size": size,
+    "task": task,
+    "video": str(video_payload) if video_payload else "",
+    }
+
+    logging.debug(f"API请求payload: {json.dumps(payload, indent=2)}")
+
+    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-14b-vace",
+            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("# 🌊Wan-2.1-14B-Vace Run On [WaveSpeedAI](https://wavespeed.ai/)")
+    gr.Markdown("""VACE is an all-in-one model designed for video creation and editing. It encompasses various tasks, including reference-to-video generation (R2V), video-to-video editing (V2V), and masked video-to-video editing (MV2V), allowing users to compose these tasks freely. This functionality enables users to explore diverse possibilities and streamlines their workflows effectively, offering a range of capabilities, such as Move-Anything, Swap-Anything, Reference-Anything, Expand-Anything, Animate-Anything, and more.""")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            images = gr.File(label="upload image", file_count="multiple",  file_types=["image"], type="filepath", elem_id="image-uploader")
+            video = gr.Video(label="Input Video", format="mp4", sources=["upload"])
+            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")
+            context_scale = gr.Slider(0, 2, value=1, step=0.1, label="Context Scale")
+            num_inference_steps = gr.Slider(1, 100, value=20, step=1, label="Inference Steps")
+            task = gr.Dropdown(["depth", "pose"], value="depth", label="Task") 
+            seed = gr.Number(-1, label="Seed")
+            random_seed_btn = gr.Button("Random🎲Seed", variant="secondary")
+            guidance = gr.Slider(1, 20, value=7.5, step=0.1, label="Guidance_Scale")
+            flow_shift = gr.Slider(1, 20, value=16, step=1, label="Shift")
+            enable_safety_checker = gr.Checkbox(True, label="Enable Safety Checker", interactive=True)
+        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=[
+                    [
+                        "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, as well as her pretty face. 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",
+                            "https://d2g64w682n9w0w.cloudfront.net/media/ec44bbf6abac4c25998dd2c4af1a46a7/images/1747413586520964413_7bkgc9ol.png"
+                        ]
+                    ]
+                ],
+                inputs=[prompt, images],
+                )
+
+    random_seed_btn.click(
+        fn=lambda: random.randint(0, 999999),
+        outputs=seed
+    )
+
+    generate_btn.click(
+    generate_video,
+    inputs=[
+        context_scale,
+        enable_safety_checker, 
+        flow_shift,
+        guidance,
+        images,
+        negative_prompt,
+        num_inference_steps,
+        prompt,
+        seed,
+        size,
+        task,
+        video,
+        session_id,
+    ],
+    outputs=[status_output, video_output]
+)
+
+logging.basicConfig(
+    level=logging.DEBUG,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.FileHandler("gradio_app.log"),
+        logging.StreamHandler()
+    ]
+)
+
+gradio_logger = logging.getLogger("gradio")
+gradio_logger.setLevel(logging.INFO)
+
+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
+    )