"""MrrrMe Backend WebSocket Server - Web-Accessible Emotion AI""" import os import sys # ===== SET CACHE DIRECTORIES FIRST ===== os.environ['HF_HOME'] = '/tmp/huggingface' os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers' os.environ['HF_HUB_CACHE'] = '/tmp/huggingface/hub' os.environ['TORCH_HOME'] = '/tmp/torch' os.makedirs('/tmp/huggingface', exist_ok=True) os.makedirs('/tmp/transformers', exist_ok=True) os.makedirs('/tmp/huggingface/hub', exist_ok=True) os.makedirs('/tmp/torch', exist_ok=True) # ===== GPU FIX: Patch TensorBoard ===== class DummySummaryWriter: def __init__(self, *args, **kwargs): pass def __getattr__(self, name): return lambda *args, **kwargs: None try: import tensorboardX tensorboardX.SummaryWriter = DummySummaryWriter except: pass # ===== GPU FIX: Patch Logging to redirect /work paths ===== import logging _original_FileHandler = logging.FileHandler class RedirectingFileHandler(_original_FileHandler): def __init__(self, filename, mode='a', encoding=None, delay=False, errors=None): if isinstance(filename, str) and filename.startswith('/work'): filename = '/tmp/openface_log.txt' os.makedirs(os.path.dirname(filename) if os.path.dirname(filename) else '/tmp', exist_ok=True) super().__init__(filename, mode, encoding, delay, errors) logging.FileHandler = RedirectingFileHandler # Now import everything else import asyncio import json import base64 import numpy as np import cv2 import io import torch from fastapi import FastAPI, WebSocket, WebSocketDisconnect from fastapi.middleware.cors import CORSMiddleware import requests from PIL import Image # Check GPU if not torch.cuda.is_available(): print("[Backend] ⚠️ No GPU detected - using CPU mode") else: print(f"[Backend] ✅ GPU available: {torch.cuda.get_device_name(0)}") app = FastAPI() # CORS for browser access app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Global model variables (will be loaded after startup) face_processor = None text_analyzer = None whisper_worker = None voice_worker = None llm_generator = None fusion_engine = None models_ready = False # Avatar backend URL AVATAR_API = "http://localhost:8765" @app.on_event("startup") async def startup_event(): """Start loading models in background after server is ready""" asyncio.create_task(load_models()) async def load_models(): """Load all AI models asynchronously""" global face_processor, text_analyzer, whisper_worker, voice_worker global llm_generator, fusion_engine, models_ready print("[Backend] 🚀 Initializing MrrrMe AI models in background...") try: # Import modules from mrrrme.vision.face_processor import FaceProcessor from mrrrme.audio.voice_emotion import VoiceEmotionWorker from mrrrme.audio.whisper_transcription import WhisperTranscriptionWorker from mrrrme.nlp.text_sentiment import TextSentimentAnalyzer from mrrrme.nlp.llm_generator_groq import LLMResponseGenerator from mrrrme.config import FUSE4 # Load models print("[Backend] Loading FaceProcessor...") face_processor = FaceProcessor() print("[Backend] Loading TextSentiment...") text_analyzer = TextSentimentAnalyzer() print("[Backend] Loading Whisper...") whisper_worker = WhisperTranscriptionWorker(text_analyzer) print("[Backend] Loading VoiceEmotion...") voice_worker = VoiceEmotionWorker(whisper_worker=whisper_worker) print("[Backend] Initializing LLM...") groq_api_key = os.getenv("GROQ_API_KEY", "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") llm_generator = LLMResponseGenerator(api_key=groq_api_key) # Initialize fusion engine class FusionEngine: def __init__(self): self.alpha_face = 0.5 self.alpha_voice = 0.3 self.alpha_text = 0.2 def fuse(self, face_probs, voice_probs, text_probs): fused = ( self.alpha_face * face_probs + self.alpha_voice * voice_probs + self.alpha_text * text_probs ) fused = fused / (np.sum(fused) + 1e-8) fused_idx = int(np.argmax(fused)) fused_emotion = FUSE4[fused_idx] intensity = float(np.max(fused)) return fused_emotion, intensity fusion_engine = FusionEngine() models_ready = True print("[Backend] ✅ All models loaded!") except Exception as e: print(f"[Backend] ❌ Error loading models: {e}") import traceback traceback.print_exc() @app.get("/") async def root(): """Root endpoint""" return { "status": "running", "models_ready": models_ready, "message": "MrrrMe AI Backend" } @app.get("/health") async def health(): """Health check - responds immediately""" return { "status": "healthy", "models_ready": models_ready } @app.websocket("/ws") async def websocket_endpoint(websocket: WebSocket): await websocket.accept() print("[WebSocket] ✅ Client connected!") # Wait for models to load if needed if not models_ready: await websocket.send_json({ "type": "status", "message": "AI models are loading, please wait..." }) # Wait up to 15 minutes for models for _ in range(900): if models_ready: await websocket.send_json({ "type": "status", "message": "Models loaded! Ready to chat." }) break await asyncio.sleep(1) if not models_ready: await websocket.send_json({ "type": "error", "message": "Models failed to load. Please refresh." }) return # Session state audio_buffer = [] try: while True: data = await websocket.receive_json() msg_type = data.get("type") # ============ VIDEO FRAME ============ if msg_type == "video_frame": try: # Decode base64 image img_data = base64.b64decode(data["frame"].split(",")[1]) img = Image.open(io.BytesIO(img_data)) frame = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) # Process face emotion try: processed_frame, result = face_processor.process_frame(frame) face_emotion = face_processor.get_last_emotion() or "Neutral" face_confidence = face_processor.get_last_confidence() or 0.0 except Exception as proc_err: print(f"[FaceProcessor] Error: {proc_err}") face_emotion = "Neutral" face_confidence = 0.0 # Send face emotion to frontend await websocket.send_json({ "type": "face_emotion", "emotion": face_emotion, "confidence": face_confidence }) except Exception as e: print(f"[Video] Error: {e}") # ============ AUDIO CHUNK ============ elif msg_type == "audio_chunk": try: audio_data = base64.b64decode(data["audio"]) audio_buffer.append(audio_data) if len(audio_buffer) >= 5: voice_probs, voice_emotion = voice_worker.get_probs() await websocket.send_json({ "type": "voice_emotion", "emotion": voice_emotion }) audio_buffer = audio_buffer[-3:] except Exception as e: print(f"[Audio] Error: {e}") # ============ USER FINISHED SPEAKING ============ elif msg_type == "speech_end": transcription = data.get("text", "").strip() print(f"\n[Speech End] User said: '{transcription}'") # Filter short/meaningless transcriptions if len(transcription) < 2: continue hallucinations = {"thank you", "thanks", "okay", "ok", "you", "yeah", "yep"} if transcription.lower().strip('.,!?') in hallucinations: continue try: # Get face emotion face_emotion = face_processor.get_last_emotion() face_confidence = face_processor.get_last_confidence() # Create emotion probabilities emotion_map = {'Neutral': 0, 'Happy': 1, 'Sad': 2, 'Angry': 3} face_probs = np.array([0.25, 0.25, 0.25, 0.25], dtype=np.float32) if face_emotion in emotion_map: face_idx = emotion_map[face_emotion] face_probs[face_idx] = face_confidence face_probs = face_probs / face_probs.sum() # Get voice and text emotions voice_probs, voice_emotion = voice_worker.get_probs() text_analyzer.analyze(transcription) text_probs, _ = text_analyzer.get_probs() # Fuse emotions fused_emotion, intensity = fusion_engine.fuse( face_probs, voice_probs, text_probs ) print(f"[Fusion] Face: {face_emotion}, Voice: {voice_emotion}, Fused: {fused_emotion}") # Generate LLM response response_text = llm_generator.generate_response( fused_emotion, face_emotion, voice_emotion, transcription, force=True, intensity=intensity ) print(f"[LLM] Response: '{response_text}'") # Send to avatar for TTS try: avatar_response = requests.post( f"{AVATAR_API}/speak", data={"text": response_text}, timeout=45 ) avatar_response.raise_for_status() avatar_data = avatar_response.json() await websocket.send_json({ "type": "llm_response", "text": response_text, "emotion": fused_emotion, "intensity": intensity, "audio_url": avatar_data.get("audio_url"), "visemes": avatar_data.get("visemes") }) except Exception as avatar_err: print(f"[Avatar] Error: {avatar_err}") await websocket.send_json({ "type": "llm_response", "text": response_text, "emotion": fused_emotion, "intensity": intensity, "error": "Avatar TTS failed" }) except Exception as e: print(f"[Speech Processing] Error: {e}") import traceback traceback.print_exc() except WebSocketDisconnect: print("[WebSocket] ❌ Client disconnected") except Exception as e: print(f"[WebSocket] Error: {e}") import traceback traceback.print_exc() if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000)