app.py

import os
import requests
import json
import time
import threading
import uuid
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
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:
            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), "#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_scale,
    negative_prompt,
    inference_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

    # API调用
    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 = {
            "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
        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"]
            
            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():
        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")
                
        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="示例输入"
    )

    # 事件处理
    random_seed_btn.click(
        fn=lambda: random.randint(0, 999999),
        outputs=seed
    )
    
    generate_btn.click(
        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__":
    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
    )