mrrrme/backend_server.py

"""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)