Update app.py

app.py

@@ -14,6 +14,12 @@ from videosys.models.cogvideo.pipeline import CogVideoPABConfig
 import psutil
 import GPUtil
 
+import threading
+
+# tracking the task status
+task_running = threading.Event()
+
+
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -312,19 +318,47 @@ with gr.Blocks(css=css) as demo:
             status['gpu_memory']
         )
 
+    def check_task_status():
+        if task_running.is_set():
+            return gr.update(value="Please wait for the previous running end", visible=True)
+        return gr.update(value="", visible=False)
 
+    def run_task(func, *args):
+        if task_running.is_set():
+            return None, gr.update(value="Please wait for the previous running end", visible=True), gr.update(visible=False)
+        
+        task_running.set()
+        try:
+            result = func(*args)
+            return result, gr.update(visible=False), gr.update(visible=True)
+        finally:
+            task_running.clear()
+
+   # click events
     generate_button.click(
-        generate_vanilla,
-        inputs=[prompt, num_inference_steps, guidance_scale],
-        outputs=[video_output, download_video_button, elapsed_time],
+        run_task,
+        inputs=[lambda: generate_vanilla, prompt, num_inference_steps, guidance_scale],
+        outputs=[
+            gr.Group([video_output, download_video_button, elapsed_time]), 
+            task_status,
+            gr.Group([generate_button, generate_button_vs])
+        ]
     )
 
     generate_button_vs.click(
-        generate_vs,
-        inputs=[prompt, num_inference_steps, guidance_scale, pab_threshold, pab_gap],
-        outputs=[video_output_vs, download_video_button_vs, elapsed_time_vs],
+        run_task,
+        inputs=[lambda: generate_vs, prompt, num_inference_steps, guidance_scale, pab_threshold, pab_gap],
+        outputs=[
+            gr.Group([video_output_vs, download_video_button_vs, elapsed_time_vs]), 
+            task_status,
+            gr.Group([generate_button, generate_button_vs])
+        ]
     )
 
+    # check click status
+    generate_button.click(check_task_status, outputs=task_status)
+    generate_button_vs.click(check_task_status, outputs=task_status)
+
     enhance_button.click(enhance_prompt_func, inputs=[prompt], outputs=[prompt])