Update app.py

app.py

@@ -1,7 +1,4 @@
 import os
-
-os.environ["GRADIO_TEMP_DIR"] = os.path.join(os.getcwd(), ".tmp_outputs")
-
 import torch
 from openai import OpenAI
 from time import time
@@ -13,6 +10,11 @@ from videosys import CogVideoConfig, VideoSysEngine
 from videosys.models.cogvideo.pipeline import CogVideoPABConfig
 import psutil
 import GPUtil
+import queue
+import threading
+import pandas as pd
+
+os.environ["GRADIO_TEMP_DIR"] = os.path.join(os.getcwd(), ".tmp_outputs")
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -31,6 +33,15 @@ Other times the user will not want modifications , but instead want a new image
 Video descriptions must have the same num of words as examples below. Extra words will be ignored.
 """
 
+# 创建一个全局任务队列
+task_queue = queue.Queue()
+
+# 创建一个锁来保护共享资源
+lock = threading.Lock()
+
+# 创建一个列表来存储所有任务的状态
+tasks = []
+
 def convert_prompt(prompt: str, retry_times: int = 3) -> str:
     if not os.environ.get("OPENAI_API_KEY"):
         return prompt
@@ -101,30 +112,68 @@ def generate(engine, prompt, num_inference_steps=50, guidance_scale=6.0):
         logger.error(f"An error occurred: {str(e)}")
         return None
 
-
 def get_server_status():
     cpu_percent = psutil.cpu_percent()
     memory = psutil.virtual_memory()
     disk = psutil.disk_usage('/')
-    gpus = GPUtil.getGPUs()
-    gpu_info = []
-    for gpu in gpus:
-        gpu_info.append({
-            'id': gpu.id,
-            'name': gpu.name,
-            'load': f"{gpu.load*100:.1f}%",
-            'memory_used': f"{gpu.memoryUsed}MB",
-            'memory_total': f"{gpu.memoryTotal}MB"
-        })
+    try:
+        gpus = GPUtil.getGPUs()
+        if gpus:
+            gpu = gpus[0]  # 只获取第一个GPU的信息
+            gpu_memory = f"{gpu.memoryUsed}/{gpu.memoryTotal}MB ({gpu.memoryUtil*100:.1f}%)"
+        else:
+            gpu_memory = "No GPU found"
+    except:
+        gpu_memory = "GPU information unavailable"
     
     return {
         'cpu': f"{cpu_percent}%",
         'memory': f"{memory.percent}%",
         'disk': f"{disk.percent}%",
-        'gpu': gpu_info
+        'gpu_memory': gpu_memory
     }
 
+def task_processor():
+    while True:
+        task = task_queue.get()
+        if task is None:
+            break
+        
+        # 更新任务状态为"运行中"
+        with lock:
+            task['status'] = 'running'
+        
+        # 执行任务
+        result = task['function'](*task['args'])
+        
+        # 更新任务状态为"完成"
+        with lock:
+            task['status'] = 'completed'
+            task['result'] = result
+        
+        task_queue.task_done()
+
+# 启动任务处理器线程
+processor_thread = threading.Thread(target=task_processor)
+processor_thread.start()
+
+def add_task(function, args, task_name):
+    task = {
+        'id': len(tasks),
+        'name': task_name,
+        'status': 'waiting',
+        'function': function,
+        'args': args,
+        'result': None
+    }
+    with lock:
+        tasks.append(task)
+    task_queue.put(task)
+    return task['id']
 
+def get_task_status():
+    with lock:
+        return pd.DataFrame([{'ID': task['id'], 'Task': task['name'], 'Status': task['status']} for task in tasks])
 
 css = """
 body {
@@ -253,56 +302,27 @@ with gr.Blocks(css=css) as demo:
                 download_video_button_vs = gr.File(label="📥 Download Video", visible=False)
                 elapsed_time_vs = gr.Textbox(label="Elapsed Time", value="0s", visible=False)
 
-
-         
+    with gr.Row():
+        task_status = gr.Dataframe(
+            headers=["ID", "Task", "Status"],
+            label="Task Queue",
+            interactive=False
+        )
+        refresh_tasks_button = gr.Button("Refresh Tasks")
 
     def generate_vanilla(prompt, num_inference_steps, guidance_scale, progress=gr.Progress(track_tqdm=True)):
-        engine = load_model()
-        t = time()
-        video_path = generate(engine, prompt, num_inference_steps, guidance_scale)
-        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
+        task_id = add_task(generate, (load_model(), prompt, num_inference_steps, guidance_scale), f"Generate: {prompt[:20]}...")
+        return get_task_status()
 
     def generate_vs(prompt, num_inference_steps, guidance_scale, threshold, gap, progress=gr.Progress(track_tqdm=True)):
         threshold = [int(i) for i in threshold.split(",")]
         gap = int(gap)
-        engine = load_model(enable_video_sys=True, pab_threshold=threshold, pab_gap=gap)
-        t = time()
-        video_path = generate(engine, prompt, num_inference_steps, guidance_scale)
-        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
+        task_id = add_task(generate, (load_model(enable_video_sys=True, pab_threshold=threshold, pab_gap=gap), prompt, num_inference_steps, guidance_scale), f"Generate VS: {prompt[:20]}...")
+        return get_task_status()
 
     def enhance_prompt_func(prompt):
         return convert_prompt(prompt, retry_times=1)
 
-    def get_server_status():
-        cpu_percent = psutil.cpu_percent()
-        memory = psutil.virtual_memory()
-        disk = psutil.disk_usage('/')
-        try:
-            gpus = GPUtil.getGPUs()
-            if gpus:
-                gpu = gpus[0]  # 只获取第一个GPU的信息
-                gpu_memory = f"{gpu.memoryUsed}/{gpu.memoryTotal}MB ({gpu.memoryUtil*100:.1f}%)"
-            else:
-                gpu_memory = "No GPU found"
-        except:
-            gpu_memory = "GPU information unavailable"
-        
-        return {
-            'cpu': f"{cpu_percent}%",
-            'memory': f"{memory.percent}%",
-            'disk': f"{disk.percent}%",
-            'gpu_memory': gpu_memory
-        }
-    
-
     def update_server_status():
         status = get_server_status()
         return (
@@ -312,25 +332,29 @@ with gr.Blocks(css=css) as demo:
             status['gpu_memory']
         )
 
+    def update_task_status():
+        return get_task_status()
 
     generate_button.click(
         generate_vanilla,
         inputs=[prompt, num_inference_steps, guidance_scale],
-        outputs=[video_output, download_video_button, elapsed_time],
+        outputs=[task_status]
     )
 
     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],
+        outputs=[task_status]
     )
 
     enhance_button.click(enhance_prompt_func, inputs=[prompt], outputs=[prompt])
 
-    
     refresh_button.click(update_server_status, outputs=[cpu_status, memory_status, disk_status, gpu_status])
     demo.load(update_server_status, outputs=[cpu_status, memory_status, disk_status, gpu_status], every=1)
 
+    refresh_tasks_button.click(update_task_status, outputs=[task_status])
+    demo.load(update_task_status, outputs=[task_status], every=5)  # 每5秒自动刷新一次
+
 if __name__ == "__main__":
-    demo.queue(max_size=10, default_concurrency_limit=1)
-    demo.launch()
+    demo.queue(max_size=10, concurrency_count=1)
+    demo.launch()