Dynamic duration

app.py

@@ -41,8 +41,26 @@ aspect_ratios = {
     "3:4": (1140, 1472)
 }
 
+def get_duration(
+    prompt: str,
+    negative_prompt: str,
+    use_negative_prompt: bool,
+    seed: int,
+    width: int,
+    height: int,
+    guidance_scale: float,
+    randomize_seed: bool,
+    num_inference_steps: int,
+    num_images: int,
+    zip_images: bool,
+    progress=gr.Progress(track_tqdm=True),
+):
+    step_duration = 0.6
+    step_duration *= (width * height) / (1024 * 1024)
+    per_image_duration = 5 + num_inference_steps * step_duration
+    return 10 + num_images * per_image_duration
+
 # Generation function for Qwen/Qwen-Image
-@spaces.GPU
 def generate_qwen(
     prompt: str,
     negative_prompt: str = "",
@@ -89,7 +107,7 @@ def generate_qwen(
     return image_paths, seed, f"{duration:.2f}", zip_path
 
 # Wrapper function to handle UI logic
-@spaces.GPU
+@spaces.GPU(duration=get_duration)
 def generate(
     prompt: str,
     negative_prompt: str,