""" MuseTalk Fast Speech-to-Speech Server With pre-loaded models, avatar caching, and streaming chunks """ import os import sys import asyncio import tempfile import uuid import time import json from pathlib import Path from typing import Optional import subprocess import threading sys.path.insert(0, str(Path(__file__).parent.parent)) from fastapi import FastAPI, UploadFile, File, Form, HTTPException, BackgroundTasks, WebSocket, WebSocketDisconnect from fastapi.staticfiles import StaticFiles from fastapi.responses import FileResponse, StreamingResponse from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel import uvicorn import httpx import base64 import cv2 import numpy as np import queue # H.264 encoder for efficient streaming try: from server.h264_encoder import H264StreamEncoder H264_AVAILABLE = True except ImportError: H264_AVAILABLE = False print("[WARNING] H.264 encoder not available, using JPEG only") # API Keys GROQ_API_KEY = "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" ELEVENLABS_API_KEY = "sk_857e9e6f2412ddf3ff5334b736e4b571641d26225c0d8d62" ELEVENLABS_VOICE_ID = "21m00Tcm4TlvDq8ikWAM" # Paths BASE_DIR = Path(__file__).parent.parent RESULTS_DIR = BASE_DIR / "results" / "server" RESULTS_DIR.mkdir(parents=True, exist_ok=True) app = FastAPI(title="MuseTalk Fast API") app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) app.mount("/static", StaticFiles(directory=Path(__file__).parent / "static"), name="static") app.mount("/avatar_videos", StaticFiles(directory=Path(__file__).parent / "avatar_videos"), name="avatar_videos") # React app static files (built with Vite) REACT_BUILD_DIR = Path(__file__).parent / "web" / "dist" if REACT_BUILD_DIR.exists(): app.mount("/app/assets", StaticFiles(directory=REACT_BUILD_DIR / "assets"), name="react_assets") @app.middleware("http") async def add_cache_headers(request, call_next): response = await call_next(request) if request.url.path.endswith(".mp4") or request.url.path.endswith(".jpg"): response.headers["Cache-Control"] = "public, max-age=3600, immutable" return response # Session Metrics class SessionMetrics: def __init__(self, request_id: str): self.request_id = request_id self.start_time = time.time() self.timings = {} self.frame_timings = [] def mark(self, event: str): elapsed = time.time() - self.start_time self.timings[event] = round(elapsed * 1000) # ms print(f"[METRICS] {event}: {elapsed:.3f}s") def frame_sent(self, index: int): self.frame_timings.append({"i": index, "t": round((time.time() - self.start_time) * 1000)}) def to_dict(self): return { "request_id": self.request_id, "total_ms": round((time.time() - self.start_time) * 1000), "timings": self.timings, "frames": len(self.frame_timings) } # Global state engine = None conversation_history = [] video_jobs = {} # job_id -> {"status": str, "progress": int, "video_path": str} class TextRequest(BaseModel): text: str def convert_webm_to_wav(input_path: str, output_path: str) -> str: cmd = ["ffmpeg", "-y", "-v", "quiet", "-i", input_path, "-ar", "16000", "-ac", "1", "-c:a", "pcm_s16le", output_path] subprocess.run(cmd, capture_output=True) return output_path DISTIL_WHISPER_URL = "http://localhost:8766" # Distil-Whisper PT-BR (~50ms) async def transcribe_audio_local(audio_path: str) -> str: """Transcribe using Distil-Whisper PT-BR (local, ~50ms).""" async with httpx.AsyncClient(timeout=30.0) as client: with open(audio_path, "rb") as f: response = await client.post( f"{DISTIL_WHISPER_URL}/inference", files={"file": (os.path.basename(audio_path), f, "audio/wav")}, data={"response_format": "json", "language": "pt"} ) if response.status_code != 200: raise HTTPException(status_code=500, detail=f"Distil-Whisper error: {response.text}") result = response.json() return result.get("text", "").strip() async def transcribe_audio_groq(audio_path: str) -> str: """Fallback to Groq API if local Whisper fails.""" async with httpx.AsyncClient(timeout=30.0) as client: with open(audio_path, "rb") as f: response = await client.post( "https://api.groq.com/openai/v1/audio/transcriptions", headers={"Authorization": f"Bearer {GROQ_API_KEY}"}, files={"file": ("audio.wav", f, "audio/wav")}, data={"model": "whisper-large-v3", "response_format": "text", "language": "pt"} ) if response.status_code != 200: raise HTTPException(status_code=500, detail=f"Groq STT error: {response.text}") return response.text.strip() async def transcribe_audio(audio_path: str) -> str: """Transcribe audio - tries local Whisper first, falls back to Groq.""" try: return await transcribe_audio_local(audio_path) except Exception as e: print(f"[STT] Local Whisper failed ({e}), falling back to Groq") return await transcribe_audio_groq(audio_path) VLLM_URL = "http://localhost:8000" VLLM_MODEL = "Qwen/Qwen2.5-0.5B-Instruct" OLLAMA_URL = "http://localhost:11434" OLLAMA_MODEL = "gemma3:1b-it-q4_K_M" # System prompt for short, direct responses LLM_SYSTEM_PROMPT = """Responda em portugues de forma muito curta e direta. Maximo 1 frase.""" async def chat_vllm(messages: list) -> str: """Chat using vLLM (fastest, ~65ms).""" import time start = time.time() # Build messages with system prompt llm_messages = [ {"role": "system", "content": LLM_SYSTEM_PROMPT} ] for msg in messages: if msg.get("role") == "user": llm_messages.append({"role": "user", "content": msg.get("content", "")}) async with httpx.AsyncClient(timeout=30.0) as client: response = await client.post( f"{VLLM_URL}/v1/chat/completions", json={ "model": VLLM_MODEL, "messages": llm_messages, "max_tokens": 50, "temperature": 0.7 } ) if response.status_code != 200: raise Exception(f"vLLM error: {response.text}") result = response.json() elapsed = (time.time() - start) * 1000 content = result.get("choices", [{}])[0].get("message", {}).get("content", "").strip() print(f"[LLM] vLLM local: {elapsed:.0f}ms") return content async def chat_ollama(messages: list) -> str: """Chat using local Ollama/Gemma (fast, ~340ms).""" import time start = time.time() # Get user message user_msg = "" for msg in messages: if msg.get("role") == "user": user_msg = msg.get("content", "") # Simple prompt for Gemma prompt = f"Responda em portugues, muito curto (1 frase): {user_msg}" async with httpx.AsyncClient(timeout=30.0) as client: response = await client.post( f"{OLLAMA_URL}/api/generate", json={ "model": OLLAMA_MODEL, "prompt": prompt, "stream": False, "options": {"num_predict": 50, "temperature": 0.7, "stop": ["\n", "."]} } ) if response.status_code != 200: raise Exception(f"Ollama error: {response.text}") result = response.json() elapsed = (time.time() - start) * 1000 print(f"[LLM] Ollama local: {elapsed:.0f}ms") return result.get("response", "").strip() async def chat_groq(messages: list) -> str: """Fallback to Groq API.""" async with httpx.AsyncClient(timeout=60.0) as client: response = await client.post( "https://api.groq.com/openai/v1/chat/completions", headers={"Authorization": f"Bearer {GROQ_API_KEY}", "Content-Type": "application/json"}, json={"model": "llama-3.1-8b-instant", "messages": messages, "max_tokens": 500, "temperature": 0.7} ) if response.status_code != 200: raise HTTPException(status_code=500, detail=f"Groq LLM error: {response.text}") return response.json()["choices"][0]["message"]["content"] async def chat_llm(messages: list) -> str: """Chat - tries vLLM first, then Ollama, then Groq.""" # Try vLLM first (fastest, ~65ms) try: return await chat_vllm(messages) except Exception as e: print(f"[LLM] vLLM failed ({e}), trying Ollama") # Try Ollama (fast, ~340ms) try: return await chat_ollama(messages) except Exception as e: print(f"[LLM] Ollama failed ({e}), falling back to Groq") # Fallback to Groq API return await chat_groq(messages) def text_to_speech_espeak(text: str, output_path: str, target_duration_ms: int = None) -> str: """ Ultra-fast TTS with espeak-ng (~15ms). Used to generate video while ElevenLabs loads in parallel. If target_duration_ms is provided, adjusts speed to match. """ base_speed = 175 if target_duration_ms: # First generate with base speed to calculate ratio temp_path = output_path + ".temp.wav" subprocess.run([ "espeak-ng", "-v", "pt-br", "-s", str(base_speed), "-w", temp_path, text ], capture_output=True) # Get duration and calculate required speed result = subprocess.run([ 'ffprobe', '-v', 'error', '-show_entries', 'format=duration', '-of', 'default=noprint_wrappers=1:nokey=1', temp_path ], capture_output=True, text=True) base_dur_ms = int(float(result.stdout.strip()) * 1000) if base_dur_ms > 0: required_speed = int(base_speed * (base_dur_ms / target_duration_ms)) required_speed = max(80, min(400, required_speed)) else: required_speed = base_speed os.remove(temp_path) else: required_speed = base_speed # Generate final audio subprocess.run([ "espeak-ng", "-v", "pt-br", "-s", str(required_speed), "-w", output_path, text ], capture_output=True) return output_path async def text_to_speech_elevenlabs(text: str, output_path: str) -> str: """ Text-to-speech with ElevenLabs. Uses eleven_flash_v2_5 model (fastest available). """ async with httpx.AsyncClient(timeout=60.0) as client: response = await client.post( f"https://api.elevenlabs.io/v1/text-to-speech/{ELEVENLABS_VOICE_ID}/stream", headers={ "xi-api-key": ELEVENLABS_API_KEY, "Content-Type": "application/json" }, json={ "text": text, "model_id": "eleven_flash_v2_5", "voice_settings": {"stability": 0.5, "similarity_boost": 0.75} } ) if response.status_code != 200: raise HTTPException(status_code=500, detail=f"ElevenLabs error: {response.text}") with open(output_path, "wb") as f: f.write(response.content) return output_path def generate_video_background(job_id: str, audio_path: str, resolution: int = 256, batch_size: int = 8): """Background task to generate video""" global engine, video_jobs def progress_callback(progress, message): video_jobs[job_id]["progress"] = progress video_jobs[job_id]["message"] = message try: video_jobs[job_id] = {"status": "processing", "progress": 0, "message": "Iniciando..."} output_path = str(RESULTS_DIR / f"video_{job_id}.mp4") engine.generate_video_fast(audio_path, output_path, resolution=resolution, batch_size=batch_size, callback=progress_callback) video_jobs[job_id]["status"] = "completed" video_jobs[job_id]["video_path"] = output_path video_jobs[job_id]["progress"] = 100 except Exception as e: video_jobs[job_id]["status"] = "error" video_jobs[job_id]["error"] = str(e) @app.on_event("startup") async def startup_event(): """Initialize engine on startup""" global engine print("\n" + "=" * 60) print("STARTING FAST MUSETALK SERVER") print("=" * 60) from server.fast_engine import initialize_engine engine = initialize_engine() print("\n✓ Server ready!") print("=" * 60 + "\n") @app.get("/") async def root(): return FileResponse(Path(__file__).parent / "static" / "index_rtc.html") @app.get("/websocket") async def websocket_page(): """Old WebSocket version (backup)""" return FileResponse(Path(__file__).parent / "static" / "index_video.html") @app.get("/video") async def video_page(): return FileResponse(Path(__file__).parent / "static" / "index_video.html") @app.get("/api/health") async def health(): return {"status": "ok", "engine_ready": engine is not None and engine.avatar_loaded} @app.post("/api/chat") async def chat(request: TextRequest): global conversation_history conversation_history.append({"role": "user", "content": request.text}) messages = [{"role": "system", "content": "Voce e uma assistente virtual amigavel. Responda de forma concisa (1-3 frases)."}] + conversation_history[-10:] response_text = await chat_llm(messages) conversation_history.append({"role": "assistant", "content": response_text}) return {"text": response_text} @app.post("/api/tts") async def tts(request: TextRequest): output_path = str(RESULTS_DIR / f"tts_{uuid.uuid4().hex[:8]}.mp3") await text_to_speech_elevenlabs(request.text, output_path) return FileResponse(output_path, media_type="audio/mpeg") @app.post("/api/transcribe") async def transcribe(audio: UploadFile = File(...)): temp_audio = tempfile.NamedTemporaryFile(suffix=".wav", delete=False) content = await audio.read() temp_audio.write(content) temp_audio.close() try: text = await transcribe_audio(temp_audio.name) return {"text": text} finally: os.unlink(temp_audio.name) @app.post("/api/conversation") async def conversation_fast( audio: UploadFile = File(...), resolution: int = Form(256), batch_size: int = Form(8), background_tasks: BackgroundTasks = None ): """ Fast conversation: generates video with audio synced """ start_time = time.time() job_id = uuid.uuid4().hex[:8] print(f"\n{'='*50}") print(f"NEW CONVERSATION - Job: {job_id}") print(f"Settings: resolution={resolution}, batch_size={batch_size}") print(f"{'='*50}") # Save and convert audio temp_webm = tempfile.NamedTemporaryFile(suffix=".webm", delete=False) content = await audio.read() temp_webm.write(content) temp_webm.close() temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False) temp_wav.close() convert_webm_to_wav(temp_webm.name, temp_wav.name) try: # STT stt_start = time.time() user_text = await transcribe_audio(temp_wav.name) stt_time = time.time() - stt_start print(f"STT ({stt_time:.2f}s): {user_text}") # LLM llm_start = time.time() global conversation_history conversation_history.append({"role": "user", "content": user_text}) messages = [{"role": "system", "content": "Voce e uma assistente virtual amigavel. Responda de forma concisa (1-3 frases)."}] + conversation_history[-10:] assistant_text = await chat_llm(messages) conversation_history.append({"role": "assistant", "content": assistant_text}) llm_time = time.time() - llm_start print(f"LLM ({llm_time:.2f}s): {assistant_text}") # TTS tts_start = time.time() tts_path = str(RESULTS_DIR / f"tts_{job_id}.mp3") await text_to_speech_elevenlabs(assistant_text, tts_path) tts_time = time.time() - tts_start print(f"TTS ({tts_time:.2f}s)") total_time = time.time() - start_time print(f"Fast response: {total_time:.2f}s (resolution={resolution}, batch={batch_size})") # Start video generation in background tts_wav = str(RESULTS_DIR / f"tts_{job_id}.wav") subprocess.run(["ffmpeg", "-y", "-v", "quiet", "-i", tts_path, "-ar", "16000", "-ac", "1", tts_wav], capture_output=True) video_jobs[job_id] = {"status": "queued", "progress": 0} thread = threading.Thread(target=generate_video_background, args=(job_id, tts_wav, resolution, batch_size)) thread.start() return { "job_id": job_id, "user_text": user_text, "assistant_text": assistant_text, "audio_url": f"/api/audio/{job_id}", "video_status": "generating", "timing": {"stt": round(stt_time, 2), "llm": round(llm_time, 2), "tts": round(tts_time, 2), "total": round(total_time, 2)} } finally: try: os.unlink(temp_webm.name) os.unlink(temp_wav.name) except: pass @app.get("/api/audio/{job_id}") async def get_audio(job_id: str): audio_path = RESULTS_DIR / f"tts_{job_id}.mp3" if not audio_path.exists(): raise HTTPException(status_code=404, detail="Audio not found") return FileResponse(audio_path, media_type="audio/mpeg") @app.get("/api/video-status/{job_id}") async def video_status(job_id: str): """Check video generation status""" if job_id not in video_jobs: raise HTTPException(status_code=404, detail="Job not found") return video_jobs[job_id] @app.get("/api/video/{job_id}") async def get_video(job_id: str): """Get generated video""" if job_id not in video_jobs: raise HTTPException(status_code=404, detail="Job not found") job = video_jobs[job_id] if job["status"] != "completed": raise HTTPException(status_code=202, detail=f"Video still generating: {job.get('progress', 0)}%") return FileResponse(job["video_path"], media_type="video/mp4") @app.get("/api/video-stream/{job_id}") async def video_stream(job_id: str): """SSE stream for video generation progress""" async def event_generator(): while True: if job_id not in video_jobs: yield f"data: {json.dumps({'error': 'Job not found'})}\n\n" break job = video_jobs[job_id] yield f"data: {json.dumps(job)}\n\n" if job["status"] in ["completed", "error"]: break await asyncio.sleep(0.5) return StreamingResponse(event_generator(), media_type="text/event-stream") @app.post("/api/conversation-with-video") async def conversation_with_video(audio: UploadFile = File(...)): """ Full conversation with video - waits for video to complete """ start_time = time.time() job_id = uuid.uuid4().hex[:8] # Save and convert audio temp_webm = tempfile.NamedTemporaryFile(suffix=".webm", delete=False) content = await audio.read() temp_webm.write(content) temp_webm.close() temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False) temp_wav.close() convert_webm_to_wav(temp_webm.name, temp_wav.name) try: # STT stt_start = time.time() user_text = await transcribe_audio(temp_wav.name) stt_time = time.time() - stt_start # LLM llm_start = time.time() global conversation_history conversation_history.append({"role": "user", "content": user_text}) messages = [{"role": "system", "content": "Voce e uma assistente virtual amigavel. Responda de forma concisa (1-3 frases)."}] + conversation_history[-10:] assistant_text = await chat_llm(messages) conversation_history.append({"role": "assistant", "content": assistant_text}) llm_time = time.time() - llm_start # TTS tts_start = time.time() tts_path = str(RESULTS_DIR / f"tts_{job_id}.mp3") await text_to_speech_elevenlabs(assistant_text, tts_path) tts_time = time.time() - tts_start # Convert to WAV for MuseTalk tts_wav = str(RESULTS_DIR / f"tts_{job_id}.wav") subprocess.run(["ffmpeg", "-y", "-v", "quiet", "-i", tts_path, "-ar", "16000", "-ac", "1", tts_wav], capture_output=True) # Generate video (blocking) video_start = time.time() video_path = str(RESULTS_DIR / f"video_{job_id}.mp4") engine.generate_video_fast(tts_wav, video_path) video_time = time.time() - video_start total_time = time.time() - start_time return { "job_id": job_id, "user_text": user_text, "assistant_text": assistant_text, "audio_url": f"/api/audio/{job_id}", "video_url": f"/api/video-file/{job_id}", "timing": { "stt": round(stt_time, 2), "llm": round(llm_time, 2), "tts": round(tts_time, 2), "video": round(video_time, 2), "total": round(total_time, 2) } } finally: try: os.unlink(temp_webm.name) os.unlink(temp_wav.name) except: pass @app.get("/api/video-file/{job_id}") async def get_video_file(job_id: str): video_path = RESULTS_DIR / f"video_{job_id}.mp4" if not video_path.exists(): raise HTTPException(status_code=404, detail="Video not found") return FileResponse(video_path, media_type="video/mp4") @app.post("/api/clear-history") async def clear_history(): global conversation_history conversation_history = [] return {"status": "ok"} @app.get("/streaming") async def streaming_page(): """Streaming video page""" return FileResponse(Path(__file__).parent / "static" / "index_streaming.html") # WebSocket streaming state streaming_sessions = {} @app.websocket("/ws/stream") async def websocket_stream(websocket: WebSocket): """ WebSocket endpoint for real-time frame streaming. """ await websocket.accept() session_id = uuid.uuid4().hex[:8] print(f"\n[WebSocket] Client connected: {session_id}") try: while True: data = await websocket.receive_json() if data.get("type") == "ping": await websocket.send_json({"type": "pong"}) continue if data.get("type") == "start": # Unique request ID for this specific request request_id = uuid.uuid4().hex[:8] metrics = SessionMetrics(request_id) audio_base64 = data.get("audio_base64") resolution = data.get("resolution", 256) batch_size = data.get("batch_size", 8) codec = data.get("codec", "jpeg").lower() # "jpeg" or "h264" # Validate codec if codec == "h264" and not H264_AVAILABLE: codec = "jpeg" await websocket.send_json({ "type": "warning", "message": "H.264 not available, using JPEG" }) print(f"\n{'='*50}") print(f"[WebSocket] New request: {request_id}") print(f"[WebSocket] Settings: resolution={resolution}, batch={batch_size}, codec={codec}") print(f"{'='*50}") # Decode and save audio audio_bytes = base64.b64decode(audio_base64) temp_audio = tempfile.NamedTemporaryFile(suffix=".webm", delete=False) temp_audio.write(audio_bytes) temp_audio.close() temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False) temp_wav.close() convert_webm_to_wav(temp_audio.name, temp_wav.name) try: # STT await websocket.send_json({"type": "status", "message": "Transcrevendo..."}) user_text = await transcribe_audio(temp_wav.name) await websocket.send_json({"type": "transcription", "text": user_text}) print(f"[STT] {user_text}") metrics.mark("stt_done") # LLM await websocket.send_json({"type": "status", "message": "Pensando..."}) global conversation_history conversation_history.append({"role": "user", "content": user_text}) messages = [{"role": "system", "content": "Voce e uma assistente virtual amigavel. Responda de forma concisa (1-3 frases)."}] + conversation_history[-10:] assistant_text = await chat_llm(messages) conversation_history.append({"role": "assistant", "content": assistant_text}) await websocket.send_json({"type": "response", "text": assistant_text}) print(f"[LLM] {assistant_text}") metrics.mark("llm_done") # TTS - PARALLEL PIPELINE # 1. Generate espeak immediately (fast ~15ms) for video timing # 2. Start ElevenLabs in parallel for high quality audio await websocket.send_json({"type": "status", "message": "Gerando voz..."}) espeak_wav = str(RESULTS_DIR / f"stream_espeak_{request_id}.wav") tts_path = str(RESULTS_DIR / f"stream_tts_{request_id}.mp3") # Fast espeak for video generation espeak_start = time.time() text_to_speech_espeak(assistant_text, espeak_wav) espeak_time = (time.time() - espeak_start) * 1000 metrics.mark("espeak_done") print(f"[TTS] espeak: {espeak_time:.0f}ms") # Start ElevenLabs in parallel (don't await yet) elevenlabs_start = time.time() elevenlabs_task = asyncio.create_task( text_to_speech_elevenlabs(assistant_text, tts_path) ) # Convert espeak for MuseTalk (16kHz) tts_wav = str(RESULTS_DIR / f"stream_tts_{request_id}.wav") subprocess.run(["ffmpeg", "-y", "-v", "quiet", "-i", espeak_wav, "-ar", "16000", "-ac", "1", tts_wav], capture_output=True) metrics.mark("tts_done") # Send audio URL (will serve ElevenLabs when ready, espeak as fallback) await websocket.send_json({"type": "audio", "url": f"/api/stream-audio/{request_id}"}) # Stream frames (using espeak timing, video starts immediately!) await websocket.send_json({"type": "status", "message": "Gerando video..."}) print(f"[VIDEO] Starting frame generation (codec={codec})...") frame_count = 0 total_frames = 0 gen_start = time.time() h264_encoder = None total_bytes_sent = 0 for frame_data in engine.generate_frames_streaming(tts_wav, resolution=resolution, batch_size=batch_size): if frame_data["type"] == "info": total_frames = frame_data["total_frames"] fps = frame_data["fps"] # Initialize H.264 encoder if needed if codec == "h264": h264_encoder = H264StreamEncoder( width=resolution, height=resolution, fps=fps ) # Send codec config (SPS/PPS) for decoder init codec_config = h264_encoder.start() if codec_config: await websocket.send_bytes(b'\x00' + codec_config) # 0x00 = config print(f"[H264] Sent codec config: {len(codec_config)} bytes") await websocket.send_json({ "type": "info", "total_frames": total_frames, "fps": fps, "codec": codec }) print(f"[VIDEO] Total frames: {total_frames}, FPS: {fps}") elif frame_data["type"] == "frame": frame = frame_data["frame"] if codec == "h264" and h264_encoder: # H.264 encoding (delta frames) h264_data, meta = h264_encoder.encode(frame) if h264_data: # Send as binary: 0x01=keyframe, 0x02=delta frame_type = b'\x01' if meta['keyframe'] else b'\x02' await websocket.send_bytes(frame_type + h264_data) total_bytes_sent += len(h264_data) else: # JPEG encoding (full frames) _, buffer = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 80]) frame_base64 = base64.b64encode(buffer).decode('utf-8') total_bytes_sent += len(buffer) await websocket.send_json({ "type": "frame", "frame": frame_base64, "index": frame_data["index"], "total": total_frames }) frame_count += 1 # Yield to event loop every frame to ensure immediate send await asyncio.sleep(0) # Log first 5 frames to show timing if frame_count == 1: metrics.mark("first_frame") if frame_count <= 5: elapsed = time.time() - gen_start if codec == "h264": frame_type_str = "I" if meta.get('keyframe') else "P" else: frame_type_str = "JPEG" print(f"[VIDEO] Frame {frame_count} ({frame_type_str}) sent at {elapsed:.2f}s") # Then log every batch elif frame_count % batch_size == 0: elapsed = time.time() - gen_start print(f"[VIDEO] Batch complete: {frame_count}/{total_frames} frames at {elapsed:.2f}s") # Cleanup H.264 encoder if h264_encoder: h264_encoder.stop() gen_time = time.time() - gen_start avg_frame_size = total_bytes_sent / max(1, frame_count) print(f"[VIDEO] Complete: {frame_count} frames in {gen_time:.2f}s ({frame_count/gen_time:.1f} fps)") print(f"[VIDEO] Bandwidth: {total_bytes_sent/1024:.1f}KB total, {avg_frame_size:.0f}B/frame avg ({codec})") # Ensure ElevenLabs audio is ready for playback try: await elevenlabs_task metrics.mark("elevenlabs_done") elevenlabs_time = (time.time() - elevenlabs_start) * 1000 print(f"[TTS] ElevenLabs ready: {elevenlabs_time:.0f}ms") except Exception as e: print(f"[TTS] ElevenLabs error (using espeak fallback): {e}") metrics.mark("all_done") await websocket.send_json({ "type": "done", "total_frames": frame_count, "gen_time": round(gen_time, 2), "codec": codec, "total_bytes": total_bytes_sent, "avg_frame_bytes": round(avg_frame_size), "metrics": metrics.to_dict() }) finally: try: os.unlink(temp_audio.name) os.unlink(temp_wav.name) except: pass elif data.get("type") == "ping": await websocket.send_json({"type": "pong"}) except WebSocketDisconnect: print(f"[WebSocket] Client disconnected: {session_id}") except Exception as e: print(f"[WebSocket] Error: {e}") import traceback traceback.print_exc() try: await websocket.send_json({"type": "error", "message": str(e)}) except: pass @app.get("/api/stream-audio/{session_id}") async def get_stream_audio(session_id: str): """Get TTS audio for streaming session (ElevenLabs preferred, espeak fallback)""" # Try ElevenLabs first (high quality) elevenlabs_path = RESULTS_DIR / f"stream_tts_{session_id}.mp3" if elevenlabs_path.exists(): return FileResponse(elevenlabs_path, media_type="audio/mpeg") # Fallback to espeak (if ElevenLabs not ready yet) espeak_path = RESULTS_DIR / f"stream_espeak_{session_id}.wav" if espeak_path.exists(): return FileResponse(espeak_path, media_type="audio/wav") raise HTTPException(status_code=404, detail="Audio not found") @app.get("/api/idle-frames") async def get_idle_frames(): """Get idle animation frames from idle.mp4 (for avatar listening state)""" global engine if engine is None: raise HTTPException(status_code=503, detail="Engine not initialized") idle_frames, idle_fps = engine.get_idle_frames() if not idle_frames: raise HTTPException(status_code=404, detail="No idle frames available") # Encode frames as base64 JPEG (RGB to BGR for cv2) frames_b64 = [] for frame in idle_frames: # Convert RGB to BGR for cv2.imencode frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) _, buffer = cv2.imencode('.jpg', frame_bgr, [cv2.IMWRITE_JPEG_QUALITY, 85]) frames_b64.append(base64.b64encode(buffer).decode('utf-8')) return { "frames": frames_b64, "fps": idle_fps, "count": len(frames_b64) } # ============================================ # WebRTC Streaming # ============================================ from server.webrtc_stream import webrtc_manager class RTCOfferRequest(BaseModel): sdp: str type: str session_id: str @app.post("/api/rtc/offer") async def rtc_offer(request: RTCOfferRequest): """Handle WebRTC offer and return answer.""" try: answer = await webrtc_manager.handle_offer( session_id=request.session_id, sdp=request.sdp, type=request.type, fps=30 ) return answer except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @app.post("/api/rtc/close/{session_id}") async def rtc_close(session_id: str): """Close WebRTC connection.""" await webrtc_manager.close_connection(session_id) return {"status": "closed"} @app.get("/rtc") async def rtc_page(): """Serve WebRTC test page.""" return FileResponse(Path(__file__).parent / "static" / "index_rtc.html") @app.get("/ws") async def ws_page(): """Serve WebSocket-only streaming page (proxy-friendly, no WebRTC).""" return FileResponse(Path(__file__).parent / "static" / "index_ws.html") @app.get("/h264") async def h264_page(): """Serve H.264 WebCodecs streaming page (efficient delta frames).""" return FileResponse(Path(__file__).parent / "static" / "index_h264.html") @app.get("/auto") async def auto_page(): """Auto-detect best streaming mode (WebRTC if UDP available, else H.264/JPEG).""" return FileResponse(Path(__file__).parent / "static" / "index_auto.html") @app.get("/app/{full_path:path}") async def react_app(full_path: str = ""): """Serve React SPA - all routes return index.html for client-side routing.""" react_index = Path(__file__).parent / "web" / "dist" / "index.html" if react_index.exists(): return FileResponse(react_index) return FileResponse(Path(__file__).parent / "static" / "index_auto.html") @app.get("/app") async def react_app_root(): """React app root.""" return await react_app("") @app.websocket("/ws/rtc") async def websocket_rtc(websocket: WebSocket): """ WebSocket endpoint for WebRTC signaling + audio processing. Video is sent via WebRTC, signaling via WebSocket. """ print(f"\n[WebRTC-WS] Accepting connection...") try: await websocket.accept() except Exception as e: print(f"[WebRTC-WS] Accept error: {e}") raise session_id = uuid.uuid4().hex[:8] print(f"[WebRTC-WS] Client connected: {session_id}") try: while True: data = await websocket.receive_json() if data.get("type") == "ping": await websocket.send_json({"type": "pong"}) continue if data.get("type") == "offer": # Handle WebRTC offer answer = await webrtc_manager.handle_offer( session_id=session_id, sdp=data["sdp"], type="offer", fps=30 ) await websocket.send_json({ "type": "answer", "sdp": answer["sdp"] }) print(f"[WebRTC-WS] Sent answer to {session_id}") elif data.get("type") == "start": # Process audio and stream video via WebRTC request_id = uuid.uuid4().hex[:8] metrics = SessionMetrics(request_id) audio_base64 = data.get("audio_base64") resolution = data.get("resolution", 256) batch_size = data.get("batch_size", 16) print(f"\n{'='*50}") print(f"[WebRTC] New request: {request_id}") print(f"[WebRTC] Settings: resolution={resolution}, batch={batch_size}") print(f"{'='*50}") # Decode audio audio_bytes = base64.b64decode(audio_base64) temp_audio = tempfile.NamedTemporaryFile(suffix=".webm", delete=False) temp_audio.write(audio_bytes) temp_audio.close() temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False) temp_wav.close() convert_webm_to_wav(temp_audio.name, temp_wav.name) try: # STT await websocket.send_json({"type": "status", "message": "Transcrevendo..."}) user_text = await transcribe_audio(temp_wav.name) metrics.mark("stt_done") await websocket.send_json({"type": "transcription", "text": user_text}) print(f"[STT] {user_text}") # LLM await websocket.send_json({"type": "status", "message": "Pensando..."}) global conversation_history conversation_history.append({"role": "user", "content": user_text}) messages = [{"role": "system", "content": "Voce e uma assistente virtual amigavel. Responda de forma concisa (1-3 frases)."}] + conversation_history[-10:] assistant_text = await chat_llm(messages) conversation_history.append({"role": "assistant", "content": assistant_text}) metrics.mark("llm_done") await websocket.send_json({"type": "response", "text": assistant_text}) print(f"[LLM] {assistant_text}") # TTS - PARALLEL PIPELINE # 1. Generate espeak immediately (fast ~15ms) for video timing # 2. Start ElevenLabs in parallel for high quality audio await websocket.send_json({"type": "status", "message": "Gerando voz..."}) espeak_wav = str(RESULTS_DIR / f"stream_espeak_{request_id}.wav") tts_path = str(RESULTS_DIR / f"stream_tts_{request_id}.mp3") # Fast espeak for video generation espeak_start = time.time() text_to_speech_espeak(assistant_text, espeak_wav) espeak_time = (time.time() - espeak_start) * 1000 metrics.mark("espeak_done") print(f"[TTS] espeak: {espeak_time:.0f}ms") # Start ElevenLabs in parallel (don't await yet) elevenlabs_start = time.time() elevenlabs_task = asyncio.create_task( text_to_speech_elevenlabs(assistant_text, tts_path) ) # Convert espeak for MuseTalk (16kHz) tts_wav = str(RESULTS_DIR / f"stream_tts_{request_id}.wav") subprocess.run(["ffmpeg", "-y", "-v", "quiet", "-i", espeak_wav, "-ar", "16000", "-ac", "1", tts_wav], capture_output=True) metrics.mark("tts_done") # Send audio URL (will serve ElevenLabs when ready) await websocket.send_json({"type": "audio", "url": f"/api/stream-audio/{request_id}"}) # Check WebRTC connection or use WebSocket fallback use_websocket_video = data.get("ws_video", False) # Client can request WS video rtc_connected = webrtc_manager.is_connected(session_id) if not use_websocket_video and not rtc_connected: await websocket.send_json({"type": "status", "message": "Conectando WebRTC..."}) print(f"[WebRTC] Waiting for connection to be established...") rtc_connected = await webrtc_manager.wait_for_connection(session_id, timeout=5.0) if not rtc_connected: use_websocket_video = True await websocket.send_json({"type": "ws_video_mode", "enabled": True}) print(f"[WS-VIDEO] WebRTC not available, using WebSocket video fallback") # Stream frames await websocket.send_json({"type": "status", "message": "Gerando video..."}) mode = "WS-VIDEO" if use_websocket_video else "WebRTC" print(f"[{mode}] Starting frame generation...") frame_count = 0 total_frames = 0 gen_start = time.time() for frame_data in engine.generate_frames_streaming(tts_wav, resolution=resolution, batch_size=batch_size): if frame_data["type"] == "info": total_frames = frame_data["total_frames"] fps = frame_data["fps"] await websocket.send_json({ "type": "info", "total_frames": total_frames, "fps": fps }) print(f"[{mode}] Total frames: {total_frames}, FPS: {fps}") elif frame_data["type"] == "frame": frame = frame_data["frame"] if use_websocket_video: # Send frame as JPEG base64 via WebSocket _, jpeg_data = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 80]) frame_b64 = base64.b64encode(jpeg_data).decode('utf-8') await websocket.send_json({ "type": "video_frame", "frame": frame_b64, "index": frame_count }) else: # Send frame via WebRTC await webrtc_manager.send_frame(session_id, frame) frame_count += 1 if frame_count == 1: metrics.mark("first_frame") # Small delay to match video FPS await asyncio.sleep(1.0 / 30.0) if frame_count <= 5: elapsed = time.time() - gen_start print(f"[{mode}] Frame {frame_count} sent at {elapsed:.2f}s") elif frame_count % batch_size == 0: elapsed = time.time() - gen_start print(f"[{mode}] Batch complete: {frame_count}/{total_frames} frames at {elapsed:.2f}s") gen_time = time.time() - gen_start print(f"[{mode}] Complete: {frame_count} frames in {gen_time:.2f}s ({frame_count/gen_time:.1f} fps)") # Ensure ElevenLabs audio is ready for playback try: await elevenlabs_task metrics.mark("elevenlabs_done") elevenlabs_time = (time.time() - elevenlabs_start) * 1000 print(f"[TTS] ElevenLabs ready: {elevenlabs_time:.0f}ms") except Exception as e: print(f"[TTS] ElevenLabs error (using espeak fallback): {e}") metrics.mark("all_done") await websocket.send_json({ "type": "done", "total_frames": frame_count, "gen_time": round(gen_time, 2), "metrics": metrics.to_dict() }) finally: try: os.unlink(temp_audio.name) os.unlink(temp_wav.name) except: pass except WebSocketDisconnect: print(f"[WebRTC-WS] Client disconnected: {session_id}") await webrtc_manager.close_connection(session_id) except Exception as e: print(f"[WebRTC-WS] Error: {e}") await webrtc_manager.close_connection(session_id) # ============================================ # Main entry point # ============================================ if __name__ == "__main__": import sys port = int(sys.argv[1]) if len(sys.argv) > 1 else 8000 uvicorn.run(app, host="0.0.0.0", port=port)