Update app.py

app.py

@@ -14,6 +14,9 @@ from videosys.models.cogvideo.pipeline import CogVideoPABConfig
 import psutil
 import GPUtil
 
+import threading
+task_running = threading.Event()
+
 
 
 logging.basicConfig(level=logging.INFO)
@@ -234,14 +237,29 @@ with gr.Blocks(css=css) as demo:
                 with gr.Row():
                     generate_button = gr.Button("🎬 Generate Video")
                     generate_button_vs = gr.Button("⚡️ Generate Video with VideoSys (Faster)")
+                # with gr.Row(elem_classes="server-status"):
+                #     gr.Markdown("#### Server Status")
+                #     with gr.Row():
+                #         cpu_status = gr.Textbox(label="CPU", scale=1)
+                #         memory_status = gr.Textbox(label="Memory", scale=1)
+                #         disk_status = gr.Textbox(label="Disk", scale=1)
+                #     gpu_status = gr.Textbox(label="GPU Memory", scale=1)
+                #     refresh_button = gr.Button("Refresh", scale=1, size="sm")
                 with gr.Row(elem_classes="server-status"):
                     gr.Markdown("#### Server Status")
+                    with gr.Column():
+                        with gr.Row():
+                            with gr.Column(scale=1):
+                                cpu_status = gr.Textbox(label="CPU")
+                            with gr.Column(scale=1):
+                                memory_status = gr.Textbox(label="Memory")
+                        with gr.Row():
+                            with gr.Column(scale=1):
+                                disk_status = gr.Textbox(label="Disk")
+                            with gr.Column(scale=1):
+                                gpu_status = gr.Textbox(label="GPU Memory")
                     with gr.Row():
-                        cpu_status = gr.Textbox(label="CPU", scale=1)
-                        memory_status = gr.Textbox(label="Memory", scale=1)
-                        disk_status = gr.Textbox(label="Disk", scale=1)
-                    gpu_status = gr.Textbox(label="GPU Memory", scale=1)
-                    refresh_button = gr.Button("Refresh", scale=1, size="sm")
+                        refresh_button = gr.Button("Refresh", scale=1, size="sm")
 
         with gr.Column():
             with gr.Row():
@@ -258,7 +276,18 @@ with gr.Blocks(css=css) as demo:
             task_status = gr.Textbox(label="任务状态", visible=False)
 
 
-         
+
+
+    def run_task(func, *args):
+        if task_running.is_set():
+            return None, gr.update(value="Please wait for another 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()
 
     def generate_vanilla(prompt, num_inference_steps, guidance_scale, progress=gr.Progress(track_tqdm=True)):
         engine = load_model()
@@ -267,7 +296,7 @@ with gr.Blocks(css=css) as demo:
         elapsed_time = time() - t
         video_update = gr.update(visible=True, value=video_path)
         elapsed_time = gr.update(visible=True, value=f"{elapsed_time:.2f}s")
-
+    
         return video_path, video_update, elapsed_time
 
     def generate_vs(prompt, num_inference_steps, guidance_scale, threshold, gap, progress=gr.Progress(track_tqdm=True)):
@@ -279,7 +308,7 @@ with gr.Blocks(css=css) as demo:
         elapsed_time = time() - t
         video_update = gr.update(visible=True, value=video_path)
         elapsed_time = gr.update(visible=True, value=f"{elapsed_time:.2f}s")
-
+    
         return video_path, video_update, elapsed_time
 
     def enhance_prompt_func(prompt):
@@ -315,17 +344,26 @@ with gr.Blocks(css=css) as demo:
             status['disk'],
             status['gpu_memory']
         )
-    
+
+
     generate_button.click(
-        generate_vanilla,
-        inputs=[prompt, num_inference_steps, guidance_scale],
-        outputs=[video_output, download_video_button, elapsed_time],
+        run_task,
+        inputs=[gr.Textbox(value=str(generate_vanilla), visible=False), 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=[gr.Textbox(value=str(generate_vs), visible=False), 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])
+        ]
     )
 
     enhance_button.click(enhance_prompt_func, inputs=[prompt], outputs=[prompt])