Update app.py

app.py

@@ -159,25 +159,25 @@ pipe.to("cuda")
 # pipe.load_lora_weights("TODO/TODO", adapter_name="ltx-lora")
 # pipe.set_adapters(["lrx-lora"], adapter_weights=[1.0])
 
+INTERRUPT_PIPELINE = False
+
+def interrupt_inference():
+    INTERRUPT_PIPELINE = True
+
+def interrupt_callback(pipeline, i, t, callback_kwargs):
+    stop_idx = 19
+    if i >= stop_idx:
+        pipeline._interrupt = False
+        return callback_kwargs
+    
+    pipeline._interrupt = INTERRUPT_PIPELINE
+    return callback_kwargs
+
 @spaces.GPU(duration=120)
 @torch.inference_mode()
 def generate_video(prompt, negative_prompt, height, width, num_frames, num_inference_steps, fps, seed, progress=gr.Progress(track_tqdm=True)):
+    INTERRUPT_PIPELINE = False
     progress_steps = []
-
-    def setup_progressbar_length(_num_steps=num_inference_steps):
-        for _step in bytes(range(_num_steps)):
-            progress_steps.append(_step) # one step one byte - fq the logic
-    
-    def progress_step():
-        if len(progress_steps) == 0:
-            return
-        for done_step in tqdm(enumerate(range(len(progress_steps)))):
-            progress_steps.pop()
-            if len(progress_steps) == 0:        
-                tqdm.close()
-            break
-
-    setup_progressbar_length()
     
     # Randomize seed if seed is 0
     if seed == 0:
@@ -194,7 +194,7 @@ def generate_video(prompt, negative_prompt, height, width, num_frames, num_infer
         num_frames=num_frames,
         num_inference_steps=num_inference_steps,
         generator=torch.Generator(device='cuda').manual_seed(seed),
-        callback_on_step_end=progress_step
+        callback_on_step_end=interrupt_callback
     ).frames[0]
 
     # Create output filename based on prompt and timestamp
@@ -233,6 +233,7 @@ with gr.Blocks() as demo:
 
             output_video = gr.Video(label="Generated Video", show_label=True)
             generate_button = gr.Button("Generate Video")
+            cancel_button = gr.Button("Cancel")
             save_state_button = gr.Button("Save State")
 
             random_seed_button.click(randomize_seed, outputs=seed)
@@ -241,6 +242,10 @@ with gr.Blocks() as demo:
                 inputs=[prompt, negative_prompt, height, width, num_frames, num_inference_steps, fps, seed],
                 outputs=output_video
             )
+            cancel_button.click(
+                interrupt_inference,
+                outputs=gr.Text(label="Interrupted.")
+            )
             save_state_button.click(
                 save_ui_state,
                 inputs=[prompt, negative_prompt, height, width, num_frames, num_inference_steps, fps, seed],