Update app.py

app.py

@@ -262,13 +262,15 @@ def encode_prompt(
         print(error_msg)
         return None, prompt, error_msg
 
-def get_duration(input_image, prompt, lora_adapter, duration, enhance_prompt, seed, randomize_seed, height, width, progress):
-    if duration <= 5:
-        return 80
-    else:
+# Function to calculate timeout based on the arguments passed to the GPU decorated function
+def calc_timeout_duration(input_image, prompt, lora_adapter, duration, enhance_prompt, seed, randomize_seed, height, width, gpu_timeout, progress=None):
+    # Determine timeout based on the last argument (gpu_timeout) passed to generate_video
+    try:
+        return int(gpu_timeout)
+    except:
         return 120
 
-@spaces.GPU(duration=get_duration)
+@spaces.GPU(duration=calc_timeout_duration)
 def generate_video(
     input_image,
     prompt: str,
@@ -279,6 +281,7 @@ def generate_video(
     randomize_seed: bool = True,
     height: int = DEFAULT_1_STAGE_HEIGHT,
     width: int = DEFAULT_1_STAGE_WIDTH,
+    gpu_timeout: int = 90,
     progress=gr.Progress(track_tqdm=True),
 ):
     gc.collect()
@@ -346,7 +349,6 @@ def generate_video(
         # ---------------------------
         # We instantiate the pipeline inside the GPU function to ensure LoRAs are applied correctly
         # for this specific run without global state pollution.
-        # Since 'model_ledger' and checkpoints are already downloaded/cached, this is fast.
         pipeline = DistilledPipeline(
             device=torch.device("cuda"),
             checkpoint_path=checkpoint_path,
@@ -389,6 +391,7 @@ def generate_video(
         return None, current_seed
 
 def generate_video_example(input_image, prompt, lora_adapter, duration):
+    # We pass a default 90s timeout for examples
     output, seed = generate_video(
         input_image=input_image, 
         prompt=prompt, 
@@ -398,7 +401,8 @@ def generate_video_example(input_image, prompt, lora_adapter, duration):
         seed=42, 
         randomize_seed=True, 
         height=DEFAULT_1_STAGE_HEIGHT, 
-        width=DEFAULT_1_STAGE_WIDTH
+        width=DEFAULT_1_STAGE_WIDTH,
+        gpu_timeout=90
     )
     return output
 
@@ -414,6 +418,9 @@ def apply_duration(duration: str):
     duration_s = int(duration[:-1])
     return duration_s
 
+def apply_gpu_duration(val: str):
+    return int(val)
+
 class RadioAnimated(gr.HTML):
     def __init__(self, choices, value=None, **kwargs):
         if not choices or len(choices) < 2:
@@ -531,11 +538,13 @@ css = """
     .dark .ra-label { color: var(--neutral-400); }
     .dark .ra-highlight { background: var(--neutral-600); }
     .dark .ra-input:checked + .ra-label { color: white; }
+
+    #main-title h1 { font-size: 2.2em !important; }
 """
 
 with gr.Blocks() as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown("# **LTX-2-LoRAs-Camera-Control-Dolly**")
+        gr.Markdown("# **LTX-2-LoRAs-Camera-Control-Dolly**", elem_id="main-title")
         gr.Markdown("Create cinematic video from text or image using [LTX-2 Distilled](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-19b-distilled-lora-384.safetensors) model. Select LoRA adapters for specific camera movements or styles.")
         
         with gr.Row():
@@ -589,6 +598,17 @@ with gr.Blocks() as demo:
                         width = gr.Number(value=DEFAULT_1_STAGE_WIDTH, visible=False)
                         height = gr.Number(value=DEFAULT_1_STAGE_HEIGHT, visible=False)
 
+                # New GPU Duration Row below Resolution
+                with gr.Row():
+                     with gr.Column():
+                        gr.Markdown("**GPU Duration**")
+                        radioanimated_gpu_duration = RadioAnimated(
+                            choices=["120", "240", "300"],
+                            value="120",
+                            elem_id="radioanimated_gpu_duration"
+                        )
+                        gpu_duration_state = gr.Number(value=90, visible=False)
+
                 with gr.Accordion("Advanced Settings", open=False):
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, value=DEFAULT_SEED, step=1)
                     randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
@@ -596,10 +616,11 @@ with gr.Blocks() as demo:
     # Wire up events
     radioanimated_duration.change(fn=apply_duration, inputs=radioanimated_duration, outputs=[duration], api_visibility="private")
     radioanimated_resolution.change(fn=apply_resolution, inputs=radioanimated_resolution, outputs=[width, height], api_visibility="private")
+    radioanimated_gpu_duration.change(fn=apply_gpu_duration, inputs=radioanimated_gpu_duration, outputs=[gpu_duration_state], api_visibility="private")
 
     generate_btn.click(
         fn=generate_video,
-        inputs=[input_image, prompt, lora_adapter, duration, enhance_prompt, seed, randomize_seed, height, width],
+        inputs=[input_image, prompt, lora_adapter, duration, enhance_prompt, seed, randomize_seed, height, width, gpu_duration_state],
         outputs=[output_video, seed]
     )