Update app.py

app.py

@@ -1,73 +1,51 @@
 import spaces
 import gradio as gr
 import torch
+import numpy as np
+import os
+import tempfile
 from diffusers import DiffusionPipeline
 from diffusers.quantizers import PipelineQuantizationConfig
-import imageio
 from diffusers.utils.export_utils import export_to_video
-import gradio as gr
-import tempfile
 
-import os
-import re
-import json
-import random
-import tempfile
-import traceback
-from functools import partial
-import numpy as np
-from PIL import Image
-import random
-import numpy as np
-import random
-import gradio as gr
-import tempfile
-import numpy as np
-from PIL import Image
-import random
+# Constants
 LANDSCAPE_WIDTH = 832
 LANDSCAPE_HEIGHT = 480
 MAX_SEED = np.iinfo(np.int32).max
-
 FIXED_FPS = 16
 MIN_FRAMES_MODEL = 8
 MAX_FRAMES_MODEL = 81
 T2V_FIXED_FPS = 16
-MIN_DURATION = round(MIN_FRAMES_MODEL/FIXED_FPS,1)
-MAX_DURATION = round(MAX_FRAMES_MODEL/FIXED_FPS,1)
+MIN_DURATION = round(MIN_FRAMES_MODEL / FIXED_FPS, 1)
+MAX_DURATION = round(MAX_FRAMES_MODEL / FIXED_FPS, 1)
+
 # Checkpoint ID
 ckpt_id = "Wan-AI/Wan2.1-T2V-14B-Diffusers"
 
-# Configure quantization (bitsandbytes 4-bit)
+# Quantization config
 quant_config = PipelineQuantizationConfig(
     quant_backend="bitsandbytes_4bit",
     quant_kwargs={
         "load_in_4bit": True,
         "bnb_4bit_quant_type": "nf4",
-        "bnb_4bit_compute_dtype": torch.bfloat16
+        "bnb_4bit_compute_dtype": torch.bfloat16,
     },
-    components_to_quantize=["transformer", "text_encoder"]
+    components_to_quantize=["transformer", "text_encoder"],
 )
 
-# Load pipeline with quantization
+# Load pipeline
 pipe = DiffusionPipeline.from_pretrained(
     ckpt_id,
     quantization_config=quant_config,
-    torch_dtype=torch.bfloat16
+    torch_dtype=torch.bfloat16,
 )
 pipe.enable_model_cpu_offload()
-torch._dynamo.config.recompile_limit = 1000
-torch._dynamo.config.capture_dynamic_output_shape_ops = True
 
-# Smart duration function using all UI params
+# Duration estimator
 def get_duration(prompt, height, width, negative_prompt, duration_seconds, guidance_scale, steps, seed, randomize_seed, progress):
-    # Calculate dynamic duration based on steps and requested duration
-    if duration_seconds <= 2.5:
-        return steps * 18
-    else:
-        return steps * 25
+    return steps * 18 if duration_seconds <= 2.5 else steps * 25
 
-# Gradio inference function with spaces GPU decorator
+# Inference function
 @spaces.GPU(duration=get_duration)
 def generate_video(prompt, height, width, negative_prompt, duration_seconds,
                    guidance_scale, steps, seed, randomize_seed,
@@ -75,7 +53,7 @@ def generate_video(prompt, height, width, negative_prompt, duration_seconds,
     
     num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)),
                          MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
-    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
+    current_seed = np.random.randint(0, MAX_SEED) if randomize_seed else int(seed)
 
     output_frames_list = pipe(
         prompt=prompt,
@@ -85,20 +63,17 @@ def generate_video(prompt, height, width, negative_prompt, duration_seconds,
         num_frames=num_frames,
         guidance_scale=float(guidance_scale),
         num_inference_steps=int(steps),
-        generator=torch.Generator(device="cuda").manual_seed(current_seed),
+        generator=torch.manual_seed(current_seed),
     ).frames[0]
 
-    filename = f"t2v_aaa.mp4"
     temp_dir = tempfile.mkdtemp()
-    video_path = os.path.join(temp_dir, filename)
+    video_path = os.path.join(temp_dir, "t2v_output.mp4")
     export_to_video(output_frames_list, video_path, fps=T2V_FIXED_FPS)
 
     print(f"✅ Video saved to: {video_path}")
-    download_label = f"📥 Download: {filename}"
-    return video_path, current_seed, gr.File(value=video_path, visible=True, label=download_label)
-  
+    return video_path  # Only return video
 
-# Build Gradio UI with all parameters
+# Gradio UI
 with gr.Blocks(css="body { max-width: 100vw; overflow-x: hidden; }") as demo:
     gr.Markdown("## 🚀 Wan2.1 T2V - Text to Video Generator (Quantized, Smart Duration)")
     with gr.Row():
@@ -116,8 +91,12 @@ with gr.Blocks(css="body { max-width: 100vw; overflow-x: hidden; }") as demo:
         with gr.Column():
             output_video = gr.Video(label="Generated Video")
 
-    ui_inputs = [prompt_input, height_input, width_input, negative_prompt_input, duration_input, guidance_scale_input, steps_input, seed_input, randomize_seed_checkbox]
+    ui_inputs = [
+        prompt_input, height_input, width_input, negative_prompt_input,
+        duration_input, guidance_scale_input, steps_input, seed_input,
+        randomize_seed_checkbox
+    ]
     run_btn.click(fn=generate_video, inputs=ui_inputs, outputs=output_video)
 
-# Launch demo
+# Launch
 demo.launch()