Update app.py

app.py

@@ -1,10 +1,9 @@
-
-import spaces
-
 import gradio as gr
 import torch
+import spaces
 from diffusers import DiffusionPipeline
 from diffusers.quantizers import PipelineQuantizationConfig
+import imageio
 
 # Checkpoint ID
 ckpt_id = "Wan-AI/Wan2.1-T2V-14B-Diffusers"
@@ -27,45 +26,61 @@ pipe = DiffusionPipeline.from_pretrained(
     torch_dtype=torch.bfloat16
 ).to("cuda")
 
-# Optimize memory
+# Optimize memory and performance
 pipe.enable_model_cpu_offload()
 torch._dynamo.config.recompile_limit = 1000
 torch._dynamo.config.capture_dynamic_output_shape_ops = True
 pipe.transformer.compile()
 
-# Gradio inference function
-@spaces.GPU(duration=20)
-def generate_video(prompt, seed):
+# Duration function
+
+def get_duration(prompt, height, width, 
+                   negative_prompt, duration_seconds,
+                   guidance_scale, steps,
+                   seed, randomize_seed, 
+                   progress):
+    if steps > 4 and duration_seconds > 2:
+        return 90
+    elif steps > 4 or duration_seconds > 2:
+        return 75
+    else:
+        return 60
+
+# Gradio inference function with GPU duration control
+@spaces.GPU(duration=get_duration)
+def generate_video(prompt, seed, steps, duration_seconds):
     generator = torch.manual_seed(seed) if seed else None
     
-    # Force ~2 second video (e.g., fps=8, frames=16)
-    num_frames = 16
+    # Force duration-based frames
     fps = 8
+    num_frames = duration_seconds * fps if duration_seconds else 16
 
     video_frames = pipe(
         prompt=prompt,
         num_frames=num_frames,
-        generator=generator
+        generator=generator,
+        num_inference_steps=steps
     ).frames[0]  # Take first video
 
     # Save as GIF for Gradio preview
-    import imageio
     out_path = "output.gif"
     imageio.mimsave(out_path, video_frames, fps=fps)
     return out_path
 
 # Build Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## 🚀 Wan2.1 T2V - Text to Video Generator (2 sec duration, 4-bit quantized)")
+    gr.Markdown("## 🚀 Wan2.1 T2V - Text to Video Generator (Quantized, Dynamic Duration)")
     with gr.Row():
         with gr.Column():
             prompt = gr.Textbox(label="Prompt", lines=3, value="A futuristic cityscape with flying cars and neon lights.")
             seed = gr.Number(value=42, label="Seed (optional)")
+            steps = gr.Slider(1, 50, value=20, step=1, label="Inference Steps")
+            duration_seconds = gr.Slider(1, 10, value=2, step=1, label="Video Duration (seconds)")
             run_btn = gr.Button("Generate Video")
         with gr.Column():
             output_video = gr.Video(label="Generated Video")
 
-    run_btn.click(fn=generate_video, inputs=[prompt, seed], outputs=output_video)
+    run_btn.click(fn=generate_video, inputs=[prompt, seed, steps, duration_seconds], outputs=output_video)
 
 # Launch demo
 demo.launch()