Update app_t2v.py

app_t2v.py

@@ -1,55 +1,58 @@
+# PyTorch nightly for CUDA compatibility
 import os
 os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces')
 
+# Imports
+import spaces
 import torch
+from diffusers import WanPipeline, AutoencoderKLWan
+from diffusers.utils import export_to_video
 import gradio as gr
 import tempfile
 import random
 import numpy as np
-import spaces
-from diffusers import WanPipeline, AutoencoderKLWan
-from diffusers.utils import export_to_video
 
 # Constants
 MODEL_ID = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
-MAX_SEED = np.iinfo(np.int32).max
 FIXED_FPS = 16
-DEFAULT_NEGATIVE_PROMPT = (
+MAX_SEED = np.iinfo(np.int32).max
+DEFAULT_HEIGHT = 720
+DEFAULT_WIDTH = 1280
+MAX_FRAMES = 81
+
+# Prompts
+default_prompt_t2v = "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
+default_negative_prompt = (
     "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，"
     "最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，"
     "画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
 )
 
-# Setup
-dtype = torch.float16  # using float16 for broader compatibility
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+# Load pipeline
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 
-# Load model components on correct device
-vae = AutoencoderKLWan.from_pretrained(
-    MODEL_ID, subfolder="vae", torch_dtype=torch.float32
-).to(device)
+vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=torch.float32).to(device)
 
-pipe = WanPipeline.from_pretrained(
-    MODEL_ID, vae=vae, torch_dtype=dtype
-).to(device)
+pipe = WanPipeline.from_pretrained(MODEL_ID, vae=vae, torch_dtype=dtype).to(device)
 
-# Warm-up call to reduce cold-start latency
+# Optional: warm-up
 _ = pipe(
     prompt="warmup",
-    negative_prompt=DEFAULT_NEGATIVE_PROMPT,
+    negative_prompt=default_negative_prompt,
     height=512,
     width=768,
     num_frames=8,
     num_inference_steps=2,
-    generator=torch.Generator(device=device).manual_seed(0),
+    generator=torch.Generator(device=device).manual_seed(0)
 ).frames[0]
 
-# Estimate duration for Hugging Face Spaces GPU usage
-def get_duration(prompt, negative_prompt, height, width, num_frames, guidance_scale, guidance_scale_2, num_steps, seed, randomize_seed):
-    return int(num_steps * 15)
+# Space-aware duration helper
+def get_duration(prompt, negative_prompt, height, width, num_frames, guidance_scale, guidance_scale_2, steps, seed, randomize_seed, progress):
+    return int(steps * 15)
 
 @spaces.GPU(duration=get_duration)
-def generate_video(
+def generate_t2v(
     prompt,
     negative_prompt,
     height,
@@ -57,57 +60,64 @@ def generate_video(
     num_frames,
     guidance_scale,
     guidance_scale_2,
-    num_steps,
+    steps,
     seed,
-    randomize_seed
+    randomize_seed,
+    progress=gr.Progress(track_tqdm=True),
 ):
     current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
     generator = torch.Generator(device=device).manual_seed(current_seed)
 
-    output = pipe(
+    output_frames = pipe(
         prompt=prompt,
         negative_prompt=negative_prompt,
-        height=height,
-        width=width,
-        num_frames=num_frames,
-        guidance_scale=guidance_scale,
-        guidance_scale_2=guidance_scale_2,
-        num_inference_steps=num_steps,
+        height=int(height),
+        width=int(width),
+        num_frames=int(num_frames),
+        guidance_scale=float(guidance_scale),
+        guidance_scale_2=float(guidance_scale_2),
+        num_inference_steps=int(steps),
         generator=generator,
     ).frames[0]
 
     with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
-        export_to_video(output, tmpfile.name, fps=FIXED_FPS)
+        export_to_video(output_frames, tmpfile.name, fps=FIXED_FPS)
         return tmpfile.name, current_seed
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## 🎬 Wan2.2 Text-to-Video Generator with Hugging Face Spaces GPU")
+    gr.Markdown("## 🎬 Wan 2.2 T2V: Text-to-Video via Wan-AI")
 
     with gr.Row():
         with gr.Column():
-            prompt = gr.Textbox(label="Prompt", value="Two anthropomorphic cats in comfy boxing gear fight intensely.")
-            negative_prompt = gr.Textbox(label="Negative Prompt", value=DEFAULT_NEGATIVE_PROMPT, lines=3)
-            height = gr.Slider(360, 1024, value=720, step=16, label="Height")
-            width = gr.Slider(360, 1920, value=1280, step=16, label="Width")
-            num_frames = gr.Slider(8, 81, value=81, step=1, label="Number of Frames")
-            num_steps = gr.Slider(10, 60, value=40, step=1, label="Inference Steps")
-            guidance_scale = gr.Slider(1.0, 10.0, value=4.0, step=0.5, label="Guidance Scale")
-            guidance_scale_2 = gr.Slider(1.0, 10.0, value=3.0, step=0.5, label="Guidance Scale 2")
-            seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
-            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-
-            generate_button = gr.Button("🎥 Generate Video")
+            prompt_input = gr.Textbox(label="Prompt", value=default_prompt_t2v)
+            negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
+            height_slider = gr.Slider(360, 1024, step=16, value=DEFAULT_HEIGHT, label="Height")
+            width_slider = gr.Slider(360, 1920, step=16, value=DEFAULT_WIDTH, label="Width")
+            frames_slider = gr.Slider(8, MAX_FRAMES, value=MAX_FRAMES, step=1, label="Frames")
+
+            with gr.Accordion("Advanced Settings", open=False):
+                guidance_slider = gr.Slider(0.0, 20.0, step=0.5, value=4.0, label="Guidance Scale")
+                guidance2_slider = gr.Slider(0.0, 20.0, step=0.5, value=3.0, label="Guidance Scale 2")
+                steps_slider = gr.Slider(1, 60, step=1, value=40, label="Inference Steps")
+                seed_slider = gr.Slider(0, MAX_SEED, step=1, value=42, label="Seed", interactive=True)
+                randomize_seed_check = gr.Checkbox(label="Randomize Seed", value=True)
+
+            generate_button = gr.Button("🎥 Generate Video", variant="primary")
 
         with gr.Column():
             video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False)
-            final_seed_display = gr.Number(label="Used Seed", interactive=False)
+            used_seed = gr.Number(label="Used Seed", interactive=False)
+
+    inputs = [
+        prompt_input, negative_prompt_input,
+        height_slider, width_slider,
+        frames_slider,
+        guidance_slider, guidance2_slider,
+        steps_slider, seed_slider, randomize_seed_check
+    ]
 
-    generate_button.click(
-        fn=generate_video,
-        inputs=[prompt, negative_prompt, height, width, num_frames, guidance_scale, guidance_scale_2, num_steps, seed, randomize_seed],
-        outputs=[video_output, final_seed_display],
-    )
+    generate_button.click(fn=generate_t2v, inputs=inputs, outputs=[video_output, used_seed])
 
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.queue().launch(mcp_server=True)