Update app.py

app.py

@@ -1,10 +1,10 @@
-import spaces
 import gradio as gr
 import torch
 from diffusers import DiffusionPipeline
 from diffusers.quantizers import PipelineQuantizationConfig
 import imageio
-
+import numpy as np
+import spaces
 
 # Checkpoint ID
 ckpt_id = "Wan-AI/Wan2.1-T2V-14B-Diffusers"
@@ -33,9 +33,20 @@ 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
+    else:
+        return 60
 
-@spaces.GPU(duration=100)
-def generate_video(prompt, seed, steps, duration_seconds,progress=gr.Progress(track_tqdm=True)):
+# Gradio inference function with spaces GPU decorator
+@spaces.GPU(duration=90)
+def generate_video(prompt, seed, steps, duration_seconds):
     generator = torch.manual_seed(seed) if seed else None
     fps = 8
     num_frames = duration_seconds * fps if duration_seconds else 16
@@ -48,8 +59,14 @@ 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
+    ]
+
     out_path = "output.gif"
-    imageio.mimsave(out_path, video_frames, fps=fps)
+    imageio.mimsave(out_path, processed_frames, fps=fps)
     return out_path
 
 # Build Gradio UI