dynamic gpu duration

app.py

@@ -463,8 +463,38 @@ def use_output_as_input(output_images):
         return []
     return output_images
 
+
+def get_edit_duration(
+    images,
+    prompt,
+    seed=42,
+    randomize_seed=False,
+    true_guidance_scale=1.0,
+    num_inference_steps=4,
+    height=None,
+    width=None,
+    rewrite_prompt=True,
+    zerogpu_budget=0,
+    progress=None,
+):
+    if zerogpu_budget and int(zerogpu_budget) > 0:
+        return max(20, min(120, int(zerogpu_budget)))
+    h = int(height) if height and int(height) > 256 else 1024
+    w = int(width) if width and int(width) > 256 else 1024
+    n_inputs = 0
+    if images:
+        try:
+            n_inputs = len(images)
+        except Exception:
+            n_inputs = 1
+    steps = max(1, int(num_inference_steps))
+    res_scale = ((h * w) / (1024 * 1024)) ** 1.3
+    estimate = int(8 + n_inputs * 1.0 + steps * 3.0 * res_scale)
+    return max(20, min(120, estimate))
+
+
 # --- Main Inference Function (with hardcoded negative prompt) ---
-@spaces.GPU()
+@spaces.GPU(duration=get_edit_duration)
 def infer(
     images,
     prompt,
@@ -475,6 +505,7 @@ def infer(
     height=None,
     width=None,
     rewrite_prompt=True,
+    zerogpu_budget=0,
     num_images_per_prompt=1,
     progress=gr.Progress(track_tqdm=True),
 ):
@@ -659,6 +690,14 @@ with gr.Blocks(css=css) as demo:
                 
                 rewrite_prompt = gr.Checkbox(label="Rewrite prompt", value=True)
 
+            zerogpu_budget = gr.Slider(
+                label="ZeroGPU max duration (0 = auto)",
+                minimum=0,
+                maximum=120,
+                step=5,
+                value=0,
+            )
+
         # gr.Examples(examples=examples, inputs=[prompt], outputs=[result, seed], fn=infer, cache_examples=False)
 
     gr.on(
@@ -674,6 +713,7 @@ with gr.Blocks(css=css) as demo:
             height,
             width,
             rewrite_prompt,
+            zerogpu_budget,
         ],
         outputs=[result, seed, use_output_btn],  # Added use_output_btn to outputs
     )