math

app.py

@@ -1,4 +1,3 @@
-import spaces
 import gradio as gr
 import torch
 from diffusers import DiffusionPipeline
@@ -21,38 +20,31 @@ quant_config = PipelineQuantizationConfig(
     components_to_quantize=["transformer", "text_encoder"]
 )
 
-# Load pipeline with quantization, without immediately moving to CUDA to avoid ConstantVariable error
+# Load pipeline with quantization
 pipe = DiffusionPipeline.from_pretrained(
     ckpt_id,
     quantization_config=quant_config,
     torch_dtype=torch.bfloat16
 )
-
-# Enable CPU offload and compile after offload
 pipe.enable_model_cpu_offload()
 torch._dynamo.config.recompile_limit = 1000
 torch._dynamo.config.capture_dynamic_output_shape_ops = True
 
-# Duration function
-def get_duration(prompt, height, width, 
-                 negative_prompt, duration_seconds,
-                 guidance_scale, steps,
-                 seed, randomize_seed):
-    if steps > 4 and duration_seconds > 2:
-        return 90
-    elif steps > 4 or duration_seconds > 2:
-        return 75
+# Smart duration function using all UI params
+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 60
+        return steps * 25
 
 # Gradio inference function with spaces GPU decorator
-@spaces.GPU(duration=60)
-def generate_video(prompt, seed, steps, duration_seconds,progress=gr.Progress(track_tqdm=True)):
+@spaces.GPU(duration=get_duration)
+def generate_video(prompt, height, width, negative_prompt, duration_seconds, guidance_scale, steps, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
     generator = torch.manual_seed(seed) if seed else None
     fps = 8
-    num_frames = duration_seconds * fps if duration_seconds else 16
+    num_frames = int(duration_seconds * fps) if duration_seconds else 16
 
-    # Run pipeline on default device with automatic offload
     video_frames = pipe(
         prompt=prompt,
         num_frames=num_frames,
@@ -60,7 +52,6 @@ def generate_video(prompt, seed, steps, duration_seconds,progress=gr.Progress(tr
         num_inference_steps=steps
     ).frames[0]
 
-    # Ensure frames are uint8 numpy arrays for imageio
     processed_frames = [
         (np.clip(frame * 255, 0, 255).astype(np.uint8) if frame.dtype in [np.float32, np.float64] else frame)
         for frame in video_frames
@@ -70,20 +61,26 @@ def generate_video(prompt, seed, steps, duration_seconds,progress=gr.Progress(tr
     imageio.mimsave(out_path, processed_frames, fps=fps)
     return out_path
 
-# Build Gradio UI
-with gr.Blocks() as demo:
-    gr.Markdown("## 🚀 Wan2.1 T2V - Text to Video Generator (Quantized, Dynamic Duration)")
+# Build Gradio UI with all parameters
+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():
         with gr.Column():
             prompt_input = gr.Textbox(label="Prompt", lines=3, value="A futuristic cityscape with flying cars and neon lights.")
-            seed_input = gr.Number(value=42, label="Seed (optional)")
+            negative_prompt_input = gr.Textbox(label="Negative Prompt", lines=3, value="")
+            height_input = gr.Slider(256, 1024, step=8, value=512, label="Height")
+            width_input = gr.Slider(256, 1024, step=8, value=512, label="Width")
+            duration_input = gr.Slider(1, 10, value=2, step=0.1, label="Duration (seconds)")
             steps_input = gr.Slider(1, 50, value=20, step=1, label="Inference Steps")
-            duration_input = gr.Slider(1, 10, value=2, step=1, label="Video Duration (seconds)")
+            guidance_scale_input = gr.Slider(0.0, 20.0, step=0.5, value=7.5, label="Guidance Scale")
+            seed_input = gr.Number(value=42, label="Seed (optional)")
+            randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True)
             run_btn = gr.Button("Generate Video")
         with gr.Column():
             output_video = gr.Video(label="Generated Video")
 
-    run_btn.click(fn=generate_video, inputs=[prompt_input, seed_input, steps_input, duration_input], outputs=output_video)
+    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
 demo.launch()