Update app.py

app.py

@@ -1,4 +1,7 @@
 import os
+
+os.environ["GRADIO_TEMP_DIR"] = os.path.join(os.getcwd(), ".tmp_outputs")
+
 import torch
 from openai import OpenAI
 from time import time
@@ -10,10 +13,6 @@ from videosys import CogVideoConfig, VideoSysEngine
 from videosys.models.cogvideo.pipeline import CogVideoPABConfig
 import psutil
 import GPUtil
-import queue
-import threading
-
-os.environ["GRADIO_TEMP_DIR"] = os.path.join(os.getcwd(), ".tmp_outputs")
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -32,15 +31,6 @@ 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
@@ -111,68 +101,30 @@ 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('/')
-    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"
+    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"
+        })
     
     return {
         'cpu': f"{cpu_percent}%",
         'memory': f"{memory.percent}%",
         'disk': f"{disk.percent}%",
-        'gpu_memory': gpu_memory
+        'gpu': gpu_info
     }
 
-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 [{'id': task['id'], 'name': task['name'], 'status': task['status']} for task in tasks]
 
 css = """
 body {
@@ -301,27 +253,56 @@ 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)):
-        task_id = add_task(generate, (load_model(), prompt, num_inference_steps, guidance_scale), f"Generate: {prompt[:20]}...")
-        return get_task_status()
+        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
 
     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)
-        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()
+        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
 
     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 (
@@ -331,44 +312,25 @@ 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=[task_status],
-        concurrency_limit=1
+        outputs=[video_output, download_video_button, elapsed_time],
     )
 
     generate_button_vs.click(
         generate_vs,
         inputs=[prompt, num_inference_steps, guidance_scale, pab_threshold, pab_gap],
-        outputs=[task_status],
-        concurrency_limit=1
+        outputs=[video_output_vs, download_video_button_vs, elapsed_time_vs],
     )
 
-    enhance_button.click(
-        enhance_prompt_func,
-        inputs=[prompt],
-        outputs=[prompt],
-        concurrency_limit=1
-    )
+    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],
-        concurrency_limit=1
-    )
+    
+    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],
-        concurrency_limit=1
-    )
-    demo.load(update_task_status, outputs=[task_status], every=5)  # 每5秒自动刷新一次
-
 if __name__ == "__main__":
-    demo.queue(max_size=10)
-    demo.launch(max_threads=4)  # 设置最大线程数为4，您可以根据需要调整这个值
+    demo.queue(max_size=10, default_concurrency_limit=1)
+    demo.launch()