Update app.py

app.py

@@ -18,6 +18,7 @@ MAX_SEED = np.iinfo(np.int32).max
 
 # Optional HF login (works in Spaces if you set HF token as secret env var "hf")
 from huggingface_hub import login
+
 login(token=os.environ.get("hf"))
 
 dtype = torch.bfloat16
@@ -67,7 +68,34 @@ def imagelist_to_pptx(img_files):
         return tmp.name
 
 
-@spaces.GPU(duration=1000)
+def _clamp_int(x, default: int, lo: int, hi: int) -> int:
+    try:
+        v = int(x)
+    except Exception:
+        v = default
+    return max(lo, min(hi, v))
+
+
+# Dynamic duration callable: must accept the same args as infer(). It returns seconds.
+def get_duration(
+    input_image,
+    seed=777,
+    randomize_seed=False,
+    prompt=None,
+    neg_prompt=" ",
+    true_guidance_scale=4.0,
+    num_inference_steps=50,
+    layer=4,
+    cfg_norm=True,
+    use_en_prompt=True,
+    resolution=640,
+    gpu_duration=1000,  # <-- NEW
+):
+    # Allow user override via UI (text field), but keep it sane
+    return _clamp_int(gpu_duration, default=1000, lo=20, hi=1500)
+
+
+@spaces.GPU(duration=get_duration)
 def infer(
     input_image,
     seed=777,
@@ -80,16 +108,14 @@ def infer(
     cfg_norm=True,
     use_en_prompt=True,
     resolution=640,
+    gpu_duration=1000,  # <-- NEW (must match get_duration signature)
 ):
     # Seed
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
     # Normalize resolution input
-    try:
-        resolution = int(resolution)
-    except Exception:
-        resolution = 640
+    resolution = _clamp_int(resolution, default=640, lo=640, hi=1024)
     if resolution not in (640, 1024):
         resolution = 640
 
@@ -117,12 +143,13 @@ def infer(
         "num_inference_steps": num_inference_steps,
         "num_images_per_prompt": 1,
         "layers": layer,
-        "resolution": resolution,  # <-- 640 or 1024
+        "resolution": resolution,  # 640 or 1024
         "cfg_normalize": cfg_norm,
         "use_en_prompt": use_en_prompt,
     }
 
     print("INFER INPUTS:", inputs)
+    print("REQUESTED GPU DURATION:", gpu_duration)
 
     with torch.inference_mode():
         out = pipeline(**inputs)
@@ -244,6 +271,14 @@ The text prompt is intended to describe the overall content of the input image
                         value=True,
                     )
 
+                    # NEW: text field for GPU duration override (seconds)
+                    gpu_duration = gr.Textbox(
+                        label="GPU duration override (seconds, 20..1500)",
+                        value="1000",
+                        lines=1,
+                        placeholder="e.g. 60, 120, 300, 1000, 1500",
+                    )
+
                 run_button = gr.Button("Decompose!", variant="primary")
 
             with gr.Column(scale=2):
@@ -275,7 +310,8 @@ The text prompt is intended to describe the overall content of the input image
             layer,
             cfg_norm,
             use_en_prompt,
-            resolution,  # <-- NEW
+            resolution,
+            gpu_duration,  # <-- NEW
         ],
         outputs=[gallery, export_file, export_zip_file],
     )