app.py

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
    )