Update app.py

app.py

@@ -4,11 +4,8 @@ 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
@@ -84,7 +81,6 @@ def image_to_base64(file_path):
             if len(img_data) == 0:
                 raise ValueError("空文件")
             
-            # 使用URL安全编码并自动填充
             encoded = base64.urlsafe_b64encode(img_data)
             missing_padding = len(encoded) % 4
             if missing_padding:
@@ -118,7 +114,7 @@ def classify_prompt(prompt):
     return torch.argmax(outputs.logits).item()
 
 def generate_video(
-    image,
+    image_files,
     prompt,
     duration,
     enable_safety,
@@ -131,6 +127,11 @@ def generate_video(
     session_id
 ):
 
+    if len(image_files) != 2:
+        error_img = create_error_image("upload 2 images")
+        yield "❌ error: upload 2 images", error_img
+        return
+
     safety_level = classify_prompt(prompt)
     if safety_level != 0:
         error_img = create_error_image(CLASS_NAMES[safety_level])
@@ -150,24 +151,26 @@ def generate_video(
         api_key = os.getenv("WAVESPEED_API_KEY")
         if not api_key:
             raise ValueError("API key missing")
-            
-        base64_img = image_to_base64(image)
+
+        base64_images = [image_to_base64(img) for img in image_files]
+        
         headers = {
             "Authorization": f"Bearer {api_key}", 
             "Content-Type": "application/json"
-            }
+        }
         
         payload = {
-            "context_scale": 1,
-            "enable_safety_checker": True,
+            "seed": seed if seed != -1 else random.randint(0, 999999),
+            "size": size.replace(" ", ""),
+            "images": base64_images,
+            "prompt": prompt,
             "flow_shift": flow_shift,
+            "context_scale": 1,
             "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"
+            "enable_safety_checker": enable_safety,
+            "model_id": "wavespeed-ai/wan-2.1-14b-vace"
         }
 
         response = requests.post(
@@ -236,44 +239,56 @@ with gr.Blocks(
     
     session_id = gr.State(str(uuid.uuid4()))
     
-    gr.Markdown("# 🌊 Wan-2.1-i2v-480p-Ultra-Fast Run On WaveSpeedAI")
+    gr.Markdown("# 🌊 Wan-2.1-14B-VACE")
+    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."
+    )
     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.
+        [WaveSpeedAI](https://wavespeed.ai/) 提供先进的AI视频生成加速技术
         """)
     
     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)
+            img_input = gr.File(
+                file_count="multiple",
+                file_types=["image"],
+                label="upload 2 images"
+            )
+            prompt = gr.Textbox(label="prompt", lines=3, placeholder="请输入描述...")
+            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")
+                size = gr.Dropdown(
+                    ["480*832", "832*480"], 
+                    value="480*832",
+                    label="resolution"
+                )
+                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="Guidance Scale")
+                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="Seed")
-                random_seed_btn = gr.Button("Random🎲Seed", variant="secondary")
+                seed = gr.Number(-1, label="随机种子")
+                random_seed_btn = gr.Button("随机种子🎲", 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")
+                enable_safety = gr.Checkbox(label="🔒 安全检测", value=True)
+                flow_shift = gr.Slider(1, 50, value=16, label="运动幅度")
                 
         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")
+            video_output = gr.Video(label="生成结果", format="mp4")
+            status_output = gr.Textbox(label="系统状态", interactive=False, lines=4)
+            generate_btn = gr.Button("开始生成", 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"
+        examples=[[
+            "The elegant lady carefully selects bags in the boutique...",
+            [
+                "https://d2g64w682n9w0w.cloudfront.net/media/ec44bbf6abac4c25998dd2c4af1a46a7/images/1747413751234102420_md9ywspl.png",
+                "https://d2g64w682n9w0w.cloudfront.net/media/ec44bbf6abac4c25998dd2c4af1a46a7/images/1747413586520964413_7bkgc9ol.png"
             ]
-        ],
+        ]],
         inputs=[prompt, img_input],
-        label="Example Inputs",
+        label="示例输入",
         examples_per_page=3
     )
 
@@ -297,10 +312,7 @@ with gr.Blocks(
             size,
             session_id 
         ],
-        outputs=[
-            status_output, 
-            video_output
-        ]
+        outputs=[status_output, video_output]
     )
 
 if __name__ == "__main__":