HuggingFaceM4/reachy_mini_remote_control
11
1"""2Reachy Mini Controller3A centralized server that listens for Robot connections and hosts a Gradio control interface.4"""5 6import asyncio7import threading8import time9import queue10from dataclasses import dataclass11from typing import Optional, Tuple, Dict12 13import cv214import gradio as gr15import numpy as np16from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request, HTTPException17from fastapi.responses import StreamingResponse18from fastapi.staticfiles import StaticFiles19import os20import uvicorn21from fastrtc import StreamHandler22from huggingface_hub import get_token23import httpx24 25from reachy_mini.utils import create_head_pose26 27# token = get_token()28 29# if not token:30# raise ValueError("No token found. Please set the HF_TOKEN environment variable or login to Hugging Face.")31 32# -------------------------------------------------------------------33# 1. Configuration34# -------------------------------------------------------------------35 36AUDIO_SAMPLE_RATE = 16000 # respeaker samplerate37 38# Audio queue configuration39MAX_AUDIO_QUEUE_SIZE = 240 41# Movement step sizes42NUDGE_ANGLE = 5.0 # degrees for head roll / yaw43NUDGE_BODY = 0.3 # degrees for body_yaw44NUDGE_PITCH = 5.0 # degrees for pitch45 46# Video loop timing47FRAME_SLEEP_S = 0.04 # 25 fps48 49# TURN config50TURN_TTL_SERVER_MS = 360_00051 52turn_credentials = None53server_turn_credentials = None54# while turn_credentials is None or server_turn_credentials is None:55# try:56# if turn_credentials is None:57# turn_credentials = get_cloudflare_turn_credentials(hf_token=token)58# if server_turn_credentials is None:59# server_turn_credentials = get_cloudflare_turn_credentials(ttl=TURN_TTL_SERVER_MS, hf_token=token)60# except Exception as e:61# print(f"[Video] Error getting turn credentials: {e!r}")62# time.sleep(1)63 64 65# -------------------------------------------------------------------66# 2. Data Models67# -------------------------------------------------------------------68 69@dataclass70class Movement:71 name: str72 x: float = 073 y: float = 074 z: float = 075 roll: float = 076 pitch: float = 077 yaw: float = 078 body_yaw: float = 079 left_antenna: Optional[float] = None80 right_antenna: Optional[float] = None81 duration: float = 1.082 83 84# -------------------------------------------------------------------85# 2b. Multi-User Authentication86# -------------------------------------------------------------------87 88# Token cache to prevent rate limiting: {token -> (username, expiry_time)}89_token_cache: Dict[str, Tuple[Optional[str], float]] = {}90_token_cache_lock = threading.Lock()91TOKEN_CACHE_TTL = 300 # 5 minutes92 93async def validate_hf_token(token: str) -> Optional[str]:94 """95 Validate HuggingFace token and return username.96 Uses caching to prevent rate limiting from HuggingFace API.97 98 Args:99 token: HuggingFace API token100 101 Returns:102 Username if token is valid, None otherwise103 """104 # Check cache first105 current_time = time.time()106 with _token_cache_lock:107 if token in _token_cache:108 username, expiry = _token_cache[token]109 if current_time < expiry:110 print(f"[Auth] Using cached token validation for user: {username}")111 return username112 else:113 # Cache expired, remove it114 del _token_cache[token]115 116 # Validate with HuggingFace API117 try:118 async with httpx.AsyncClient() as client:119 response = await client.get(120 "https://huggingface.co/api/whoami-v2",121 headers={"Authorization": f"Bearer {token}"},122 timeout=5.0123 )124 if response.status_code == 200:125 data = response.json()126 username = data.get("name")127 128 # Cache the result129 with _token_cache_lock:130 _token_cache[token] = (username, current_time + TOKEN_CACHE_TTL)131 132 print(f"[Auth] Token validated for user: {username}")133 return username134 elif response.status_code == 429:135 print(f"[Auth] Rate limited by HuggingFace API! Status: {response.status_code}")136 return None137 except Exception as e:138 print(f"[Auth] Token validation error: {e}")139 return None140 141 142async def get_user_from_websocket(websocket: WebSocket) -> Optional[str]:143 """144 Extract and validate HF token from WebSocket headers or query parameters.145 146 Args:147 websocket: WebSocket connection148 149 Returns:150 Username if authenticated, None otherwise151 """152 # Check for Authorization header (for robot connections)153 auth_header = websocket.headers.get("authorization", "")154 if auth_header.startswith("Bearer "):155 token = auth_header[7:] # Remove "Bearer " prefix156 return await validate_hf_token(token)157 158 # Check for token in query parameters (for browser connections)159 token_param = websocket.query_params.get("token")160 if token_param:161 return await validate_hf_token(token_param)162 163 # Check for token in cookies (HuggingFace OAuth)164 cookie_header = websocket.headers.get("cookie", "")165 if cookie_header:166 cookies = {}167 for cookie in cookie_header.split(";"):168 if "=" in cookie:169 name, value = cookie.strip().split("=", 1)170 cookies[name] = value171 172 token = cookies.get("token") or cookies.get("hf_token")173 if token:174 return await validate_hf_token(token)175 176 return None177 178 179async def get_user_from_request(request: Request, token_param: Optional[str] = None) -> Optional[str]:180 """181 Extract and validate HF token from HTTP request.182 183 Args:184 request: FastAPI Request object185 token_param: Optional token from query parameter (from Gradio OAuth)186 187 Returns:188 Username if authenticated, None otherwise189 """190 # First check query parameter (passed from Gradio OAuth)191 if token_param:192 print(f"[DEBUG] Using token from query parameter")193 return await validate_hf_token(token_param)194 195 # Check Authorization header (for robot connections)196 auth_header = request.headers.get("authorization", "")197 if auth_header.startswith("Bearer "):198 token = auth_header[7:]199 print(f"[DEBUG] Found token in Authorization header, validating...")200 return await validate_hf_token(token)201 202 # Check for token in cookies203 token = request.cookies.get("token") or request.cookies.get("hf_token")204 if token:205 print(f"[DEBUG] Found token in cookies, validating...")206 return await validate_hf_token(token)207 208 print("[DEBUG] No authentication found")209 return None210 211 212class UserSession:213 """214 Per-user session state for multi-user support.215 Stores all connection handles and state for a single user robot.216 """217 def __init__(self, user_id: str):218 self.user_id = user_id219 220 # Connection handles221 self.robot_ws: Optional[WebSocket] = None222 self.robot_loop: Optional[asyncio.AbstractEventLoop] = None223 224 # Video Stream Data225 self.frame_lock = threading.Lock()226 self.black_frame = np.zeros((640, 640, 3), dtype=np.uint8)227 _, buffer = cv2.imencode(".jpg", self.black_frame)228 self.latest_frame_bytes = buffer.tobytes()229 self.latest_frame_ts = time.time()230 231 # Latency tracking for video232 self.video_latencies = []233 self.video_latency_window = 100234 235 # Audio from robot -> browser236 self.audio_queue: "queue.Queue[Tuple[int, bytes, float]]" = queue.Queue()237 238 # Audio from operator -> robot239 self.audio_to_robot_queue: "queue.Queue[bytes]" = queue.Queue()240 241 # Latency tracking for audio242 self.audio_latencies = []243 self.audio_latency_window = 100244 245 # Live pose state246 self.pose_lock = threading.Lock()247 self.current_pose = Movement(248 name="Current",249 x=0, y=0, z=0,250 roll=0, pitch=0, yaw=0,251 body_yaw=0,252 left_antenna=0, right_antenna=0,253 duration=0.2,254 )255 256 # Robot state (joint positions)257 self.robot_state_lock = threading.Lock()258 self.latest_robot_state: Optional[dict] = None259 self.latest_robot_state_ts: float = 0.0260 261 # --- Connection management ---262 263 def set_robot_connection(self, ws: WebSocket, loop: asyncio.AbstractEventLoop) -> None:264 self.robot_ws = ws265 self.robot_loop = loop266 267 def clear_robot_connection(self) -> None:268 self.robot_ws = None269 self.robot_loop = None270 271 # --- Video ---272 273 def update_frame(self, frame_bytes: bytes, robot_timestamp: Optional[float] = None) -> None:274 """Update the latest video frame."""275 receive_time = time.time()276 with self.frame_lock:277 self.latest_frame_bytes = frame_bytes278 self.latest_frame_ts = receive_time279 280 if robot_timestamp is not None:281 latency_ms = (receive_time - robot_timestamp) * 1000282 self.video_latencies.append(latency_ms)283 284 if len(self.video_latencies) > self.video_latency_window:285 self.video_latencies.pop(0)286 287 def get_video_latency_stats(self) -> dict:288 """Get video latency statistics in milliseconds."""289 if not self.video_latencies:290 return {"min": 0, "max": 0, "avg": 0, "latest": 0, "count": 0}291 292 return {293 "min": min(self.video_latencies),294 "max": max(self.video_latencies),295 "avg": sum(self.video_latencies) / len(self.video_latencies),296 "latest": self.video_latencies[-1],297 "count": len(self.video_latencies)298 }299 300 # --- Audio queues ---301 302 @staticmethod303 def _push_bounded(q: queue.Queue, item, max_size: int, description: str) -> None:304 while q.qsize() >= max_size:305 try:306 dropped = q.get_nowait()307 del dropped308 except queue.Empty:309 break310 q.put(item)311 312 def push_audio_from_robot(self, audio_bytes: bytes, robot_timestamp: Optional[float] = None) -> None:313 """Push audio data from robot to the queue for browser playback."""314 self._push_bounded(315 self.audio_queue,316 (AUDIO_SAMPLE_RATE, audio_bytes, robot_timestamp if robot_timestamp is not None else time.time()),317 MAX_AUDIO_QUEUE_SIZE,318 "FROM robot",319 )320 321 def push_audio_to_robot(self, audio_bytes: bytes) -> None:322 self._push_bounded(323 self.audio_to_robot_queue,324 audio_bytes,325 MAX_AUDIO_QUEUE_SIZE,326 "TO robot",327 )328 329 def get_audio_to_robot_blocking(self) -> bytes:330 try:331 return self.audio_to_robot_queue.get(timeout=0.2)332 except queue.Empty:333 return None334 335 def track_audio_latency(self, robot_timestamp: float) -> None:336 """Track audio latency when audio is about to be played."""337 playback_time = time.time()338 latency_ms = (playback_time - robot_timestamp) * 1000339 self.audio_latencies.append(latency_ms)340 341 if len(self.audio_latencies) > self.audio_latency_window:342 self.audio_latencies.pop(0)343 344 def get_audio_latency_stats(self) -> dict:345 """Get audio latency statistics in milliseconds."""346 if not self.audio_latencies:347 return {"min": 0, "max": 0, "avg": 0, "latest": 0, "count": 0}348 349 return {350 "min": min(self.audio_latencies),351 "max": max(self.audio_latencies),352 "avg": sum(self.audio_latencies) / len(self.audio_latencies),353 "latest": self.audio_latencies[-1],354 "count": len(self.audio_latencies)355 }356 357 # --- Status ---358 359 def get_connection_status(self) -> str:360 return "✅ Robot Connected" if self.robot_ws else "🔴 Waiting for Robot..."361 362 # --- Robot state ---363 364 def update_robot_state(self, data: dict) -> None:365 with self.robot_state_lock:366 self.latest_robot_state = data367 self.latest_robot_state_ts = time.time()368 369 def get_latency_display(self) -> str:370 """Get formatted latency statistics for display."""371 video_stats = self.get_video_latency_stats()372 audio_stats = self.get_audio_latency_stats()373 374 lines = []375 376 if video_stats["count"] > 0:377 lines.append(378 f"📹 Video: {video_stats['latest']:.0f}ms "379 f"(avg: {video_stats['avg']:.0f}ms, max: {video_stats['max']:.0f}ms)"380 )381 382 if audio_stats["count"] > 0:383 lines.append(384 f"🎵 Audio: {audio_stats['latest']:.0f}ms "385 f"(avg: {audio_stats['avg']:.0f}ms, max: {audio_stats['max']:.0f}ms)"386 )387 388 if not lines:389 return "⏱️ Latency: Waiting for data..."390 391 return "\n".join(lines)392 393 # --- Pose management ---394 395 def update_pose(396 self,397 dx: float = 0,398 dy: float = 0,399 dz: float = 0,400 droll: float = 0,401 dpitch: float = 0,402 dyaw: float = 0,403 dbody_yaw: float = 0,404 ) -> Movement:405 with self.pose_lock:406 p = self.current_pose407 408 new = Movement(409 name="Current",410 x=p.x + dx,411 y=p.y + dy,412 z=p.z + dz,413 roll=p.roll + droll,414 pitch=p.pitch + dpitch,415 yaw=p.yaw + dyaw,416 body_yaw=p.body_yaw + dbody_yaw,417 left_antenna=p.left_antenna,418 right_antenna=p.right_antenna,419 duration=0.1,420 )421 422 # Clamp posed values423 new.pitch = float(np.clip(new.pitch, -30, 30))424 new.yaw = float(np.clip(new.yaw, -180, 180))425 new.roll = float(np.clip(new.roll, -40, 40))426 new.body_yaw = float(np.clip(new.body_yaw, -3, 3))427 new.z = float(np.clip(new.z, -20, 50))428 new.x = float(np.clip(new.x, -50, 50))429 new.y = float(np.clip(new.y, -50, 50))430 431 self.current_pose = new432 return new433 434 def reset_pose(self) -> Movement:435 with self.pose_lock:436 self.current_pose = Movement(437 name="Current",438 x=0, y=0, z=0,439 roll=0, pitch=0, yaw=0,440 body_yaw=0,441 left_antenna=0, right_antenna=0,442 duration=0.3,443 )444 return self.current_pose445 446 def get_pose_text(self) -> str:447 with self.pose_lock:448 p = self.current_pose449 return (450 "Head position:\n"451 f" x={p.x:.1f}, y={p.y:.1f}, z={p.z:.1f}\n"452 f" roll={p.roll:.1f}, pitch={p.pitch:.1f}, yaw={p.yaw:.1f}\n"453 "Body:\n"454 f" body_yaw={p.body_yaw:.1f}"455 )456 457 458# -------------------------------------------------------------------459# 3. Global State460# -------------------------------------------------------------------461 462class GlobalState:463 """464 Multi-user state manager.465 Manages per-user sessions, each with its own robot connection and state.466 """467 def __init__(self):468 # Multi-user sessions: user_id -> UserSession469 self.sessions: Dict[str, UserSession] = {}470 self.sessions_lock = threading.Lock()471 472 def get_or_create_session(self, user_id: str) -> UserSession:473 """Get existing session or create a new one for the user."""474 with self.sessions_lock:475 if user_id not in self.sessions:476 print(f"[MultiUser] Creating new session for user: {user_id}")477 self.sessions[user_id] = UserSession(user_id)478 return self.sessions[user_id]479 480 def get_session(self, user_id: str) -> Optional[UserSession]:481 """Get session for user if it exists."""482 with self.sessions_lock:483 return self.sessions.get(user_id)484 485 def remove_session(self, user_id: str) -> None:486 """Remove a user session."""487 with self.sessions_lock:488 if user_id in self.sessions:489 print(f"[MultiUser] Removing session for user: {user_id}")490 del self.sessions[user_id]491 492 # Multi-user helper methods493 def get_all_sessions(self) -> list:494 """Get all active sessions."""495 with self.sessions_lock:496 return list(self.sessions.values())497 498 def get_session_count(self) -> int:499 """Get count of active sessions."""500 with self.sessions_lock:501 return len(self.sessions)502 503 504state = GlobalState()505 506 507# -------------------------------------------------------------------508# 4. Robot commands509# -------------------------------------------------------------------510 511def send_pose_to_robot(session: UserSession, mov: Movement, msg: str = "Move sent"):512 """Send pose command to robot for a specific user session."""513 print(f"[DEBUG] send_pose_to_robot called for user: {session.user_id}")514 print(f"[DEBUG] robot_ws: {session.robot_ws}, robot_loop: {session.robot_loop}")515 516 if not (session.robot_ws and session.robot_loop):517 print(f"[DEBUG] Robot not connected for user: {session.user_id}")518 return get_pose_string_for_session(session), "⚠️ Robot not connected"519 520 pose = create_head_pose(521 x=mov.x,522 y=mov.y,523 z=mov.z,524 roll=mov.roll,525 pitch=mov.pitch,526 yaw=mov.yaw,527 degrees=True,528 mm=True,529 )530 531 payload = {532 "type": "movement",533 "movement": {534 "head": pose.tolist(),535 "body_yaw": mov.body_yaw,536 "duration": mov.duration,537 },538 }539 540 if mov.left_antenna is not None and mov.right_antenna is not None:541 payload["movement"]["antennas"] = [542 np.deg2rad(mov.right_antenna),543 np.deg2rad(mov.left_antenna),544 ]545 546 print(f"[DEBUG] Sending payload to robot: {payload['type']}, body_yaw={payload['movement']['body_yaw']}")547 548 # Send to robot asynchronously to avoid blocking UI callbacks.549 try:550 fut = asyncio.run_coroutine_threadsafe(551 session.robot_ws.send_json(payload),552 session.robot_loop,553 )554 555 def _log_send_error(done_fut):556 try:557 done_fut.result()558 print(f"[DEBUG] Successfully sent movement to robot for {session.user_id}")559 except Exception as e:560 print(f"[Move] Failed to send movement to robot for {session.user_id}: {e!r}")561 562 fut.add_done_callback(_log_send_error)563 except Exception as e:564 print(f"[Move] Failed to queue movement for {session.user_id}: {e!r}")565 return get_pose_string_for_session(session), "❌ Failed to send movement"566 567 return get_pose_string_for_session(session), f"✅ {msg}"568 569 570# -------------------------------------------------------------------571# 6. FastAPI endpoints572# -------------------------------------------------------------------573 574app = FastAPI()575 576viz_dir = os.path.join(os.path.dirname(__file__), "dist")577if not os.path.exists(viz_dir) and os.path.exists(os.path.join(os.path.dirname(__file__), "dist.zip")):578 # unzip the dist.zip file579 import zipfile580 zip_path = os.path.join(os.path.dirname(__file__), "dist.zip")581 with zipfile.ZipFile(zip_path, 'r') as zip_ref:582 zip_ref.extractall(os.path.dirname(__file__))583 viz_dir = os.path.join(os.path.dirname(__file__), "dist") 584 585if os.path.exists(viz_dir):586 app.mount(587 "/viz",588 StaticFiles(directory=viz_dir, html=True),589 name="viz",590 )591 592 593@app.websocket("/robot")594async def robot_endpoint(ws: WebSocket):595 """Endpoint for the Robot to connect to (control channel)."""596 await ws.accept()597 598 # Authenticate robot599 user_id = await get_user_from_websocket(ws)600 if not user_id:601 print("[Auth] Robot connection rejected - no valid token")602 await ws.close(code=1008, reason="Authentication required")603 return604 605 # Get or create user session606 session = state.get_or_create_session(user_id)607 session.set_robot_connection(ws, asyncio.get_running_loop())608 print(f"[System] Robot Connected for user: {user_id}")609 610 try:611 while True:612 msg = await ws.receive()613 if msg.get("type") == "websocket.disconnect":614 break615 except (WebSocketDisconnect, Exception):616 print(f"[System] Robot Disconnected for user: {user_id}")617 finally:618 session.clear_robot_connection()619 620 621@app.websocket("/robot_state")622async def robot_state_endpoint(ws: WebSocket):623 """Endpoint for the Robot to publish joint state."""624 await ws.accept()625 626 # Authenticate robot627 user_id = await get_user_from_websocket(ws)628 if not user_id:629 print("[Auth] Robot state rejected - no valid token")630 await ws.close(code=1008, reason="Authentication required")631 return632 633 # Get user session634 session = state.get_or_create_session(user_id)635 print(f"[System] Robot State Connected for user: {user_id}")636 637 try:638 while True:639 data = await ws.receive_json()640 if isinstance(data, dict) and data.get("type") == "robot_state":641 session.update_robot_state(data)642 except (WebSocketDisconnect, Exception):643 print(f"[System] Robot State Disconnected for user: {user_id}")644 645 646@app.websocket("/joint_states")647async def joint_states_endpoint(ws: WebSocket):648 """Endpoint for the Browser to receive joint state."""649 await ws.accept()650 651 # Authenticate browser user652 user_id = await get_user_from_websocket(ws)653 if not user_id:654 print("[Auth] Joint states rejected - no valid token")655 await ws.close(code=1008, reason="Authentication required")656 return657 658 # Get user session659 session = state.get_or_create_session(user_id)660 print(f"[System] Browser Joint State Connected for user: {user_id}")661 662 try:663 while True:664 with session.robot_state_lock:665 data = session.latest_robot_state666 667 if data and (head_vals := data.get("head")):668 payload = {669 "body_yaw": head_vals[0],670 "angles": head_vals[1:7],671 "antennas": data.get("antennas"),672 "timestamp": data.get("timestamp"),673 }674 await ws.send_json(payload)675 676 await asyncio.sleep(0.03)677 except (WebSocketDisconnect, Exception):678 print(f"[System] Browser Joint State Disconnected for user: {user_id}")679 680 681@app.get("/video_feed")682async def video_feed(request: Request, token: Optional[str] = None):683 """Video feed endpoint - requires authentication via token parameter."""684 # Authenticate user685 user_id = await get_user_from_request(request, token)686 print(f"[video_feed] Authenticated user_id: {user_id}")687 688 if not user_id:689 raise HTTPException(status_code=401, detail="Authentication required. Please log in with Hugging Face.")690 691 # Get user's session692 session = state.get_session(user_id)693 if not session:694 print(f"[video_feed] No robot connected for user: {user_id}")695 raise HTTPException(status_code=404, detail="No robot connected for your account")696 697 # Create user-specific video generator698 def generate_user_video():699 last_timestamp = 0.0700 frame_count = 0701 start_time = time.time()702 while True:703 with session.frame_lock:704 current_bytes = session.latest_frame_bytes705 current_timestamp = session.latest_frame_ts706 707 if current_timestamp > last_timestamp and current_bytes is not None:708 last_timestamp = current_timestamp709 frame_count += 1710 elapsed = time.time() - start_time711 if elapsed > 1.0:712 fps = frame_count / elapsed713 print(f"[video_feed] User {user_id} FPS: {fps:.2f}")714 frame_count = 0715 start_time = time.time()716 yield (717 b"--frame\r\n"718 b"Content-Type: image/jpeg\r\n\r\n" + current_bytes + b"\r\n"719 )720 else:721 time.sleep(FRAME_SLEEP_S)722 continue723 724 time.sleep(FRAME_SLEEP_S)725 726 return StreamingResponse(727 generate_user_video(),728 media_type="multipart/x-mixed-replace; boundary=frame",729 )730 731 732@app.get("/audio_feed")733async def audio_feed(request: Request, token: Optional[str] = None):734 """Audio feed endpoint - requires authentication via token parameter."""735 # Authenticate user736 user_id = await get_user_from_request(request, token)737 if not user_id:738 raise HTTPException(status_code=401, detail="Authentication required. Please log in with Hugging Face.")739 740 # Get user's session741 session = state.get_session(user_id)742 if not session:743 raise HTTPException(status_code=404, detail="No robot connected for your account")744 745 # Create user-specific audio generator746 def generate_user_audio():747 # Clear old data to start fresh748 with session.audio_queue.mutex:749 session.audio_queue.queue.clear()750 751 TARGET_SAMPLES = 512752 byte_buffer = bytearray()753 754 while True:755 try:756 sample_rate, chunk_bytes, robot_timestamp = session.audio_queue.get(timeout=1.0)757 758 if robot_timestamp is not None:759 session.track_audio_latency(robot_timestamp)760 761 if chunk_bytes:762 byte_buffer.extend(chunk_bytes)763 except queue.Empty:764 continue765 766 chunk_size = TARGET_SAMPLES * 2767 while len(byte_buffer) >= chunk_size:768 out_bytes = byte_buffer[:chunk_size]769 byte_buffer = byte_buffer[chunk_size:]770 yield bytes(out_bytes)771 772 return StreamingResponse(773 generate_user_audio(),774 media_type="application/octet-stream",775 headers={776 "Cache-Control": "no-cache",777 "X-Content-Type-Options": "nosniff",778 }779 )780 781@app.websocket("/video_stream")782async def stream_endpoint(ws: WebSocket):783 """784 Endpoint for Robot/Sim to send video frames.785 786 Expected message formats:787 1. Binary only (legacy): Just the JPEG frame bytes788 2. JSON with timestamp: {"timestamp": <float>, "frame": <base64 encoded frame>}789 """790 await ws.accept()791 792 # Authenticate robot793 user_id = await get_user_from_websocket(ws)794 if not user_id:795 print("[Auth] Video stream rejected - no valid token")796 await ws.close(code=1008, reason="Authentication required")797 return798 799 # Get user session800 session = state.get_or_create_session(user_id)801 print(f"[Video] Stream connected for user: {user_id}")802 803 frame_count = 0804 start_time = time.time()805 latency_report_interval = 5.0 # Report latency stats every 5 seconds806 last_latency_report = time.time()807 808 try:809 while True:810 msg = await ws.receive()811 812 # Handle binary-only messages (legacy mode, no timestamp)813 data = msg.get("bytes")814 if data:815 session.update_frame(data, robot_timestamp=None)816 frame_count += 1817 818 # Handle text/JSON messages with timestamp819 text_data = msg.get("text")820 if text_data:821 import base64822 import json823 try:824 json_data = json.loads(text_data)825 timestamp = json_data.get("timestamp")826 frame_b64 = json_data.get("frame")827 if frame_b64:828 frame_bytes = base64.b64decode(frame_b64)829 session.update_frame(frame_bytes, robot_timestamp=timestamp)830 frame_count += 1831 except (json.JSONDecodeError, KeyError) as e:832 print(f"[Video] Error parsing JSON: {e}")833 834 # FPS reporting835 elapsed = time.time() - start_time836 if elapsed > 1.0:837 fps = frame_count / elapsed838 print(f"[Video] Receiving FPS: {fps:.2f}")839 frame_count = 0840 start_time = time.time()841 842 # Latency reporting843 if time.time() - last_latency_report > latency_report_interval:844 stats = session.get_video_latency_stats()845 if stats["count"] > 0:846 print(847 f"[Video Latency] User: {user_id}, "848 f"Latest: {stats['latest']:.1f}ms, "849 f"Avg: {stats['avg']:.1f}ms, "850 f"Min: {stats['min']:.1f}ms, "851 f"Max: {stats['max']:.1f}ms "852 f"(over {stats['count']} frames)"853 )854 last_latency_report = time.time()855 856 except WebSocketDisconnect as e:857 print(f"[Video] WebSocketDisconnect: code={e.code}, reason={e.reason}")858 except asyncio.CancelledError:859 print("[Video] stream_endpoint cancelled")860 except Exception as e:861 print(f"[Video] stream_endpoint closed with error: {e!r}")862 finally:863 print("[Video] stream_endpoint closed (finally)")864 865@app.websocket("/audio_stream")866async def audio_endpoint(ws: WebSocket):867 """868 Full duplex audio channel between Robot/Sim and server.869 870 Expected message formats from robot:871 1. Binary only (legacy): Just the audio bytes872 2. JSON with timestamp: {"timestamp": <float>, "audio": <base64 encoded audio>}873 """874 await ws.accept()875 876 # Authenticate robot877 user_id = await get_user_from_websocket(ws)878 if not user_id:879 print("[Auth] Audio stream rejected - no valid token")880 await ws.close(code=1008, reason="Authentication required")881 return882 883 # Get user session884 session = state.get_or_create_session(user_id)885 print(f"[Audio] Stream Connected for user: {user_id}")886 887 latency_report_interval = 5.0 # Report latency stats every 5 seconds888 last_latency_report = time.time()889 890 async def robot_to_server():891 nonlocal last_latency_report892 try:893 while True:894 data = await ws.receive()895 t = data.get("type")896 if t == "websocket.disconnect":897 print(f"[Audio] Disconnected (recv) for user: {user_id}")898 break899 900 if t == "websocket.receive":901 # Handle binary-only messages (legacy mode, no timestamp)902 if data.get("bytes"):903 session.push_audio_from_robot(data["bytes"], robot_timestamp=None)904 905 # Handle JSON messages with timestamp906 elif data.get("text"):907 text_data = data.get("text")908 if text_data == "ping":909 print("[Audio] Received ping")910 else:911 import json912 import base64913 try:914 json_data = json.loads(text_data)915 timestamp = json_data.get("timestamp")916 audio_b64 = json_data.get("audio")917 if audio_b64:918 audio_bytes = base64.b64decode(audio_b64)919 session.push_audio_from_robot(audio_bytes, robot_timestamp=timestamp)920 except (json.JSONDecodeError, KeyError):921 pass922 923 # Latency reporting924 if time.time() - last_latency_report > latency_report_interval:925 stats = session.get_audio_latency_stats()926 if stats["count"] > 0:927 print(928 f"[Audio Latency] "929 f"Latest: {stats['latest']:.1f}ms, "930 f"Avg: {stats['avg']:.1f}ms, "931 f"Min: {stats['min']:.1f}ms, "932 f"Max: {stats['max']:.1f}ms "933 f"(over {stats['count']} chunks)"934 )935 last_latency_report = time.time()936 937 except asyncio.CancelledError:938 print("[Audio] robot_to_server cancelled")939 except Exception as e:940 print(f"[Audio] robot_to_server error: {e}")941 942 async def server_to_robot():943 loop = asyncio.get_running_loop()944 try:945 while True:946 chunk: bytes = await loop.run_in_executor(947 None, session.get_audio_to_robot_blocking948 )949 if chunk is not None:950 await ws.send_bytes(chunk)951 except asyncio.CancelledError:952 print("[Audio] server_to_robot cancelled")953 except Exception as e:954 print(f"[Audio] server_to_robot error: {e}")955 956 try:957 await asyncio.gather(robot_to_server(), server_to_robot())958 except asyncio.CancelledError:959 print("[Audio] audio_endpoint cancelled")960 finally:961 print("[Audio] Stream Closed")962 963 964@app.websocket("/browser_stream")965async def browser_stream_endpoint(ws: WebSocket):966 """967 Bi-directional connection for the Browser.968 - Sends Microphone data (Browser -> Robot)969 - Receives Speaker data (Robot -> Browser)970 """971 await ws.accept()972 973 # Authenticate browser user974 user_id = await get_user_from_websocket(ws)975 if not user_id:976 print("[Auth] Browser stream rejected - no valid token")977 await ws.close(code=1008, reason="Authentication required")978 return979 980 # Get user session981 session = state.get_or_create_session(user_id)982 print(f"[Browser] WebSocket Connected for user: {user_id}")983 984 # Task: Send Audio FROM Robot TO Browser985 async def send_to_browser():986 while True:987 # Get audio from the robot queue (non-blocking check)988 if not session.audio_queue.empty():989 _, chunk_bytes, robot_timestamp = session.audio_queue.get(timeout=0.5)990 991 # Track latency if we have a timestamp992 if robot_timestamp is not None:993 session.track_audio_latency(robot_timestamp)994 try:995 # Send as binary message996 await ws.send_bytes(chunk_bytes)997 except Exception:998 break999 else:1000 await asyncio.sleep(0.005) # Tiny sleep to prevent CPU burn1001 1002 # Task: Receive Audio FROM Browser TO Robot1003 async def receive_from_browser():1004 try:1005 while True:1006 data = await ws.receive_bytes()1007 # Push directly to the robot's input queue1008 session.push_audio_to_robot(data)1009 except Exception as e:1010 print(f"[Browser] Input stream ended: {e}")1011 1012 try:1013 # Run both tasks concurrently1014 await asyncio.gather(send_to_browser(), receive_from_browser())1015 except Exception as e:1016 print(f"[Browser] WebSocket Closed: {e}")1017 finally:1018 print("[Browser] Disconnected")1019 1020 1021@app.websocket("/control")1022async def control_endpoint(ws: WebSocket):1023 """Keyboard control endpoint for browser -> robot movement commands."""1024 await ws.accept()1025 1026 user_id = await get_user_from_websocket(ws)1027 if not user_id:1028 print("[Auth] Control stream rejected - no valid token")1029 await ws.close(code=1008, reason="Authentication required")1030 return1031 1032 session = state.get_or_create_session(user_id)1033 print(f"[Control] Keyboard stream connected for user: {user_id}")1034 1035 try:1036 while True:1037 msg = await ws.receive()1038 1039 if msg.get("type") == "websocket.disconnect":1040 break1041 1042 payload = msg.get("text")1043 if not payload:1044 continue1045 1046 action = ""1047 try:1048 import json1049 data = json.loads(payload)1050 action = str(data.get("action", "")).lower().strip()1051 except Exception:1052 action = payload.lower().strip()1053 1054 if not action:1055 continue1056 1057 pose = apply_keyboard_action(session, action)1058 await ws.send_json({"pose": pose})1059 except (WebSocketDisconnect, Exception):1060 print(f"[Control] Keyboard stream disconnected for user: {user_id}")1061 1062# -------------------------------------------------------------------1063# 8. Movement UI helpers1064# -------------------------------------------------------------------1065 1066def get_pose_string():1067 """Returns pose in format JS can parse: pitch:X,yaw:Y,roll:Z,body:B1068 NOTE: This is a legacy function that returns empty if no sessions exist.1069 Should be replaced with per-user version."""1070 sessions = state.get_all_sessions()1071 if sessions:1072 session = sessions[0] # Use first session for now (temporary)1073 return get_pose_string_for_session(session)1074 return "pitch:0.0,yaw:0.0,roll:0.0,body:0.0"1075 1076 1077def get_connection_status_for_user(profile: gr.OAuthProfile | None) -> str:1078 """Get connection status for authenticated user."""1079 if profile is None:1080 return "🔴 Not authenticated"1081 1082 user_id = profile.username1083 session = state.get_session(user_id)1084 return session.get_connection_status() if session else "🔴 Waiting for Robot..."1085 1086 1087def get_latency_display_for_user(profile: gr.OAuthProfile | None) -> str:1088 """Get latency display for authenticated user."""1089 if profile is None:1090 return "Not authenticated"1091 1092 user_id = profile.username1093 session = state.get_session(user_id)1094 return session.get_latency_display() if session else "No latency data available"1095 1096 1097def nudge_pose(profile: gr.OAuthProfile | None, dpitch=0, dyaw=0, droll=0, dbody_yaw=0, label="Move"):1098 """Modified to return pose string instead of tuple. Requires OAuth profile."""1099 if profile is None:1100 return "Not authenticated"1101 1102 user_id = profile.username1103 session = state.get_session(user_id)1104 if not session:1105 return "No robot connected"1106 1107 mov = session.update_pose(1108 dpitch=dpitch,1109 dyaw=dyaw,1110 droll=droll,1111 dbody_yaw=dbody_yaw,1112 )1113 send_pose_to_robot(session, mov, label)1114 return get_pose_string_for_session(session)1115 1116 1117def center_pose(profile: gr.OAuthProfile | None):1118 """Modified to return pose string. Requires OAuth profile."""1119 if profile is None:1120 return "Not authenticated"1121 1122 user_id = profile.username1123 session = state.get_session(user_id)1124 if not session:1125 return "No robot connected"1126 1127 mov = session.reset_pose()1128 send_pose_to_robot(session, mov, "Reset pose")1129 return get_pose_string_for_session(session)1130 1131 1132def get_pose_string_for_session(session: UserSession) -> str:1133 """Get pose string for a specific session."""1134 p = session.current_pose1135 return f"pitch:{p.pitch:.1f},yaw:{p.yaw:.1f},roll:{p.roll:.1f},body:{p.body_yaw:.1f}"1136 1137 1138def apply_keyboard_action(session: UserSession, action: str) -> str:1139 """Apply one keyboard action to a user session and forward to robot."""1140 print(f"[Control] Action '{action}' for user: {session.user_id}")1141 print(f"[DEBUG] Session has robot_ws: {session.robot_ws is not None}, robot_loop: {session.robot_loop is not None}")1142 action_map = {1143 "w": (dict(dpitch=-NUDGE_PITCH), "W"),1144 "s": (dict(dpitch=NUDGE_PITCH), "S"),1145 "a": (dict(dyaw=NUDGE_ANGLE * 2), "A"),1146 "d": (dict(dyaw=-NUDGE_ANGLE * 2), "D"),1147 "q": (dict(droll=-NUDGE_ANGLE), "Q"),1148 "e": (dict(droll=NUDGE_ANGLE), "E"),1149 "j": (dict(dbody_yaw=NUDGE_BODY), "J"),1150 "l": (dict(dbody_yaw=-NUDGE_BODY), "L"),1151 }1152 1153 if action == "h":1154 mov = session.reset_pose()1155 send_pose_to_robot(session, mov, "Reset pose")1156 return get_pose_string_for_session(session)1157 1158 config = action_map.get(action)1159 if config is None:1160 return get_pose_string_for_session(session)1161 1162 kwargs, label = config1163 mov = session.update_pose(**kwargs)1164 send_pose_to_robot(session, mov, label)1165 return get_pose_string_for_session(session)1166 1167 1168# Quick action button wrapper functions1169def qa_center(profile: gr.OAuthProfile | None):1170 return center_pose(profile)1171 1172def qa_look_up(profile: gr.OAuthProfile | None):1173 return nudge_pose(profile, dpitch=-15, label="Look Up")1174 1175def qa_curious(profile: gr.OAuthProfile | None):1176 return nudge_pose(profile, dpitch=-10, droll=15, label="Curious")1177 1178def qa_excited(profile: gr.OAuthProfile | None):1179 return nudge_pose(profile, dpitch=-5, droll=-10, label="Excited")1180 1181 1182# -------------------------------------------------------------------1183# 9. Gradio UI1184# -------------------------------------------------------------------1185 1186CUSTOM_CSS = """1187/* Login view styling */1188.login-container {1189 position: fixed !important;1190 top: 0 !important;1191 left: 0 !important;1192 width: 100vw !important;1193 height: 100vh !important;1194 display: flex !important;1195 align-items: center !important;1196 justify-content: center !important;1197 padding: 20px !important;1198 background: linear-gradient(135deg, #0f0f1a 0%, #1a1a2e 100%) !important;1199 z-index: 9999 !important;1200 overflow: auto !important;