Update app.py

app.py

@@ -4,6 +4,7 @@ import asyncio
 import subprocess
 import sys
 import threading
+import json
 import time
 from dataclasses import dataclass
 from typing import List, Optional
@@ -18,12 +19,15 @@ from reachy_mini.utils import create_head_pose
 
 robot_ws = None  # type: Optional[WebSocket]
 robot_loop = None  # type: Optional[asyncio.AbstractEventLoop]
-# Add after imports
+
 _black_frame = np.zeros((640, 640, 3), dtype=np.uint8)
-# Encode it to JPEG immediately
 _, _buffer = cv2.imencode('.jpg', _black_frame)
-# Set the global variable to bytes
-latest_test_frame = _buffer.tobytes()
+
+frame_lock = threading.Lock()
+latest_frame_data = {
+    "bytes": _buffer.tobytes(),
+    "timestamp": time.time()
+}
 
 @dataclass
 class Movement:
@@ -154,17 +158,28 @@ class ReachyController:
             time.sleep(0.04)
             
     def generate_mjpeg_stream(self):
-        global latest_test_frame
-        while True:
-            # Get the raw bytes (which you said are already JPEG)
-            frame_bytes = latest_test_frame
+            global latest_frame_data, frame_lock
+            last_timestamp = 0.0
             
-            if frame_bytes is not None and len(frame_bytes) > 0:
-                # Create the MJPEG boundary format
-                yield (b'--frame\r\n'
-                    b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n')
-        
-        time.sleep(0.04) # Cap at ~25 FPS
+            while True:
+                # 1. Check if frame has changed
+                with frame_lock:
+                    current_bytes = latest_frame_data["bytes"]
+                    current_timestamp = latest_frame_data["timestamp"]
+                
+                # 2. Only yield if this is a new frame
+                if current_timestamp > last_timestamp:
+                    last_timestamp = current_timestamp
+                    if current_bytes is not None:
+                        yield (b'--frame\r\n'
+                            b'Content-Type: image/jpeg\r\n\r\n' + current_bytes + b'\r\n')
+                else:
+                    # If no new frame, sleep a bit longer to save CPU
+                    time.sleep(0.02)
+                    continue
+                
+                # Cap FPS slightly to prevent saturation
+                time.sleep(0.02)
 
     def auto_start(self):
         """Auto-start daemon and robot connection"""
@@ -642,6 +657,7 @@ def send_robot_command(cmd: dict):
         return
 
     async def _send():
+        global robot_ws
         try:
             send_ws_start = time.perf_counter()
             await robot_ws.send_json(cmd)  # or send_text(json.dumps(cmd)) if you prefer
@@ -649,6 +665,7 @@ def send_robot_command(cmd: dict):
             print(f"[WebSocket /robot] Send time: {send_ws_time*1000:.2f}ms")
         except Exception as e:
             print("[Space] Error sending command to robot:", e)
+            robot_ws = None
 
     # Schedule coroutine on the main event loop from this worker thread
     fut = asyncio.run_coroutine_threadsafe(_send(), robot_loop)
@@ -660,43 +677,46 @@ def send_robot_command(cmd: dict):
         print(f"[WebSocket /robot] Total send time (including scheduling): {total_send_time*1000:.2f}ms")
     except Exception as e:
         print("[Space] Command failed:", e)
+        robot_ws = None
    
 @app.websocket("/mujoco_stream")
 async def robot_socket(ws: WebSocket):
-    global latest_test_frame
+    global latest_frame_data, frame_lock
     await ws.accept()
     print("[Space] MuJoCo stream connected")
     
-    receive_count = 0
-    receive_start_time = time.perf_counter()
-
     try:
         while True:
-            receive_start = time.perf_counter()
             msg = await ws.receive()
-            receive_time = time.perf_counter() - receive_start
             
             if msg["type"] == "websocket.receive":
-                if msg["bytes"] is not None:
-                    frame_update_start = time.perf_counter()
-                    latest_test_frame = msg["bytes"]
-                    frame_update_time = time.perf_counter() - frame_update_start
-                    receive_count += 1
-                    
-                    # Print stats every 30 frames
-                    if receive_count % 10 == 0:
-                        elapsed = time.perf_counter() - receive_start_time
-                        fps = receive_count / elapsed
-                        print(f"[WebSocket /mujoco_stream] Receive FPS: {fps:.2f} | Last receive: {receive_time*1000:.2f}ms | Frame update: {frame_update_time*1000:.2f}ms")
-                        receive_count = 0
-                        receive_start_time = time.perf_counter()
-                elif msg["text"] is not None:
-                    print(f"[WebSocket /mujoco_stream] JSON received (receive time: {receive_time*1000:.2f}ms):", msg["text"])
+                # 1. Safe retrieval using .get() to avoid KeyError
+                payload_bytes = msg.get("bytes")
+                payload_text = msg.get("text")
+
+                # 2. Handle Binary Frame
+                if payload_bytes is not None:
+                    with frame_lock:
+                        latest_frame_data["bytes"] = payload_bytes
+                        latest_frame_data["timestamp"] = time.time()
+                        
+                # 3. Handle Text Message (The Ping)
+                elif payload_text is not None:
+                    try:
+                        data = json.loads(payload_text)
+                        if data.get("type") == "ping":
+                            # Optional: Print debug only if you need to verify it's working
+                            print("[Space] Frame keep-alive received")
+                            pass
+                    except json.JSONDecodeError:
+                        print(f"[Space] Received invalid JSON: {payload_text}")
+                        
             elif msg["type"] == "websocket.disconnect":
                 print("[Space] MuJoCo stream disconnected")
                 break
+                
     except Exception as e:
-        print("[Space] MuJoCo stream disconnected", e)
+        print(f"[Space] MuJoCo stream disconnected error: {e}")
 
 @app.get("/video_feed")
 def video_feed():
@@ -713,33 +733,35 @@ async def robot_socket(ws: WebSocket):
     robot_loop = asyncio.get_running_loop()
     print("[Space] Robot connected")
     
-    receive_count = 0
-    receive_start_time = time.perf_counter()
+    # --- HEARTBEAT MECHANISM ---
+    # Define a keep-alive task that runs in the background
+    async def keep_alive():
+        try:
+            while True:
+                await asyncio.sleep(30) # Send ping every 30s
+                if robot_ws:
+                    # Send a lightweight ping
+                    await robot_ws.send_json({"type": "ping"}) 
+        except Exception as e:
+            print(f"[Space] Heartbeat stopped: {e}")
+
+    # Start the heartbeat
+    heartbeat_task = asyncio.create_task(keep_alive())
 
     try:
         while True:
-            receive_start = time.perf_counter()
             msg = await ws.receive()
-            receive_time = time.perf_counter() - receive_start
-            
             if msg["type"] == "websocket.receive":
                 if msg["text"] is not None:
-                    receive_count += 1
-                    print(f"[WebSocket /robot] JSON received (receive time: {receive_time*1000:.2f}ms):", msg["text"])
-                    
-                    # Print stats every 10 messages
-                    if receive_count % 10 == 0:
-                        elapsed = time.perf_counter() - receive_start_time
-                        msg_rate = receive_count / elapsed
-                        print(f"[WebSocket /robot] Message rate: {msg_rate:.2f} msg/s | Avg receive time: {(elapsed/receive_count)*1000:.2f}ms")
-                        receive_count = 0
-                        receive_start_time = time.perf_counter()
+                    # Handle robot messages here
+                    pass
             elif msg["type"] == "websocket.disconnect":
                 print("[Space] Robot disconnected")
                 break
     except Exception as e:
-        print("[Space] Robot disconnected", e)
+        print("[Space] Robot disconnected (Error)", e)
     finally:
+        heartbeat_task.cancel() # Stop the heartbeat when connection dies
         if robot_ws is ws:
             robot_ws = None
 
@@ -747,4 +769,10 @@ app = gr.mount_gradio_app(app, demo, path="/")
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(
+            app, 
+            host="0.0.0.0", 
+            port=7860, 
+            proxy_headers=True,      # Tell Uvicorn to trust the proxy's headers
+            forwarded_allow_ips="*"  # Trust headers from any IP (the internal HF proxy)
+        )