Update app.py

app.py

@@ -99,6 +99,31 @@ 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"
+        })
+    
+    return {
+        'cpu': f"{cpu_percent}%",
+        'memory': f"{memory.percent}%",
+        'disk': f"{disk.percent}%",
+        'gpu': gpu_info
+    }
+
+
+
 css = """
 body {
     font-family: Arial, sans-serif;
@@ -132,19 +157,43 @@ body {
     margin: 0 auto;
 }
 
+.server-status {
+    position: fixed;
+    bottom: 10px;
+    left: 10px;
+    background-color: rgba(240, 240, 240, 0.9);
+    padding: 5px;
+    border-radius: 5px;
+    box-shadow: 0 1px 3px rgba(0,0,0,0.1);
+    font-size: 0.8em;
+    max-width: 340px;
+}
+.server-status h4 {
+    margin: 0 0 5px 0;
+    font-size: 0.9em;
+}
+.server-status-details {
+    display: none;
+}
+.server-status-summary {
+    cursor: pointer;
+}
 @media (max-width: 768px) {
     .row {
         flex-direction: column;
     }
-    
     .column {
         width: 100%;
     }
-    
     .video-output {
         width: 100%;
         height: auto;
     }
+    .server-status {
+        position: static;
+        margin-top: 20px;
+        max-width: 100%;
+    }
 }
 """
 
@@ -200,6 +249,18 @@ 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.Column(elem_classes="server-status"):
+        gr.Markdown("#### Server Status")
+        with gr.Row():
+            cpu_status = gr.Textbox(label="CPU", scale=1, container=False)
+            memory_status = gr.Textbox(label="Memory", scale=1, container=False)
+            disk_status = gr.Textbox(label="Disk", scale=1, container=False)
+        with gr.Accordion("GPU Details", open=False):
+            gpu_status = gr.JSON(container=False)
+        refresh_button = gr.Button("Refresh", scale=1, size="sm")
+         
+
     def generate_vanilla(prompt, num_inference_steps, guidance_scale, progress=gr.Progress(track_tqdm=True)):
         engine = load_model()
         t = time()
@@ -225,6 +286,16 @@ with gr.Blocks(css=css) as demo:
     def enhance_prompt_func(prompt):
         return convert_prompt(prompt, retry_times=1)
 
+    def update_server_status():
+        status = get_server_status()
+        return (
+            f"{status['cpu']}",
+            f"{status['memory']}",
+            f"{status['disk']}",
+            status['gpu']
+        )
+
+
     generate_button.click(
         generate_vanilla,
         inputs=[prompt, num_inference_steps, guidance_scale],
@@ -239,6 +310,10 @@ with gr.Blocks(css=css) as demo:
 
     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=60)
+
 if __name__ == "__main__":
     demo.queue(max_size=10, default_concurrency_limit=1)
     demo.launch()