implement ram checking

app.py

@@ -1,5 +1,6 @@
 import random
 import threading
+import psutil  # Add this import for RAM monitoring
 
 import fastapi
 import gradio as gr
@@ -9,6 +10,20 @@ from viser_proxy_manager import ViserProxyManager
 from vis_st4rtrack import visualize_st4rtrack
 
 
+def check_ram_usage(threshold_percent=90):
+    """Check if RAM usage is above the threshold.
+    
+    Args:
+        threshold_percent: Maximum RAM usage percentage allowed
+        
+    Returns:
+        bool: True if RAM usage is below threshold, False otherwise
+    """
+    ram_percent = psutil.virtual_memory().percent
+    print(f"Current RAM usage: {ram_percent}%")
+    return ram_percent < threshold_percent
+
+
 def main() -> None:
     app = fastapi.FastAPI()
     viser_manager = ViserProxyManager(app)
@@ -21,10 +36,21 @@ def main() -> None:
         # Add the iframe with a border
         # add_sphere_btn = gr.Button("Add Random Sphere")
         iframe_html = gr.HTML("")
+        status_text = gr.Markdown("")  # Add status text component
 
-        @demo.load(outputs=[iframe_html])
+        @demo.load(outputs=[iframe_html, status_text])
         def start_server(request: gr.Request):
             assert request.session_hash is not None
+            
+            # Check RAM usage before starting visualization
+            if not check_ram_usage(threshold_percent=90):
+                return """
+                <div style="text-align: center; padding: 20px; background-color: #ffeeee; border-radius: 5px;">
+                    <h2>⚠️ Server is currently under high load</h2>
+                    <p>Please try again later when resources are available.</p>
+                </div>
+                """, "**System Status:** High memory usage detected. Visualization not loaded to prevent server overload."
+            
             viser_manager.start_server(request.session_hash)
 
             # Use the request's base URL if available
@@ -46,8 +72,16 @@ def main() -> None:
             ).start()
 
             return f"""
-            <iframe src="{protocol}://{host}/viser/{request.session_hash}/" width="100%" height="500px" frameborder="0" style="display: block;"></iframe>
-            """
+            <iframe 
+                src="{protocol}://{host}/viser/{request.session_hash}/" 
+                width="100%" 
+                height="500px" 
+                frameborder="0" 
+                style="display: block;"
+                allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture"
+                loading="lazy"
+            ></iframe>
+            """, "**System Status:** Visualization loaded successfully."
         
         # @demo.load(outputs=[iframe_html])
         # def start_server(request: gr.Request):

requirements.txt

@@ -12,3 +12,4 @@ opencv-python>=4.5.0
 imageio>=2.25.0
 matplotlib>=3.5.0
 pyliblzfse>=0.1.0
+psutil>=5.9.0