""" Fast Avatar Engine - MuseTalk Integration for real-time lip-sync avatar generation. Compatible with RunPod MuseTalk setup at /workspace/MuseTalk Based on Robin's working implementation with proper face detection and blending. """ import os import sys import logging import subprocess import time import copy from pathlib import Path from typing import Optional, Generator, List import numpy as np import cv2 logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Server directory SERVER_DIR = Path(__file__).parent # MuseTalk installation path (RunPod) MUSETALK_ROOT = Path(os.getenv("MUSETALK_DIR", "/workspace/MuseTalk")) sys.path.insert(0, str(MUSETALK_ROOT)) os.chdir(str(MUSETALK_ROOT)) # Default avatar video DEFAULT_AVATAR_VIDEO = SERVER_DIR / "avatar_videos" / "idle.mp4" class MuseTalkEngine: """MuseTalk-based avatar engine for real-time lip-sync video generation.""" def __init__(self, avatar_video=None, resolution=256, fps=25): self.avatar_video = avatar_video or str(DEFAULT_AVATAR_VIDEO) self.resolution = resolution # Internal processing resolution (256 for MuseTalk) self.fps = fps self._avatar_loaded = False self._models_loaded = False # Models self.audio_processor = None self.vae = None self.unet = None self.pe = None self.device = None self.timesteps = None self.face_parser = None # Avatar data self.full_frames = [] # Original full-resolution frames self.idle_frames = [] # Resized frames for processing self.idle_fps = fps self.input_latent_list = [] self.coord_list = [] # Face bounding boxes self.mask_list = [] # Pre-computed masks for blending self.mask_coords_list = [] # Pre-computed crop boxes for blending self.original_width = None self.original_height = None # MuseTalk V1.5 config (matching realtime_inference.py) self.version = "v15" self.bbox_shift = 0 # V1.5 uses 0 self.extra_margin = 10 # Extra margin for face cropping self.parsing_mode = "jaw" # Face blending mode self.left_cheek_width = 90 self.right_cheek_width = 90 self.upper_boundary_ratio = 0.5 self.expand = 1.5 logger.info("[MuseTalk] Initializing engine...") self._load_models() self._load_avatar() @property def avatar_loaded(self): return self._avatar_loaded def _load_models(self): """Load MuseTalk models including face detection and parsing.""" try: import torch from musetalk.utils.utils import load_all_model from musetalk.whisper.audio2feature import Audio2Feature from musetalk.utils.face_parsing import FaceParsing logger.info("[MuseTalk] Loading models...") start = time.time() self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") logger.info(f"[MuseTalk] Using device: {self.device}") # Load all models using MuseTalk utility (returns 3 values: vae, unet, pe) self.vae, self.unet, self.pe = load_all_model() # Load audio processor separately logger.info("[MuseTalk] Loading Audio2Feature...") self.audio_processor = Audio2Feature(model_path="tiny") # Load face parser for blending (Robin's config) logger.info("[MuseTalk] Loading FaceParsing...") self.face_parser = FaceParsing( left_cheek_width=self.left_cheek_width, right_cheek_width=self.right_cheek_width ) # Move to device and FP16 self.pe = self.pe.half().to(self.device) self.vae.vae = self.vae.vae.half().to(self.device) self.unet.model = self.unet.model.half().to(self.device) self.timesteps = torch.tensor([0], device=self.device) self._models_loaded = True logger.info(f"[MuseTalk] Models loaded in {time.time() - start:.2f}s") except Exception as e: logger.error(f"[MuseTalk] Error loading models: {e}") import traceback traceback.print_exc() self._models_loaded = False def _get_cache_path(self): """Get cache directory path based on avatar video.""" import hashlib video_hash = hashlib.md5(self.avatar_video.encode()).hexdigest()[:8] cache_dir = MUSETALK_ROOT / "results" / "v15" / "avatars" / f"cache_{video_hash}" return cache_dir def _load_avatar(self): """Load avatar video frames and detect faces (with caching).""" import pickle import torch try: if not os.path.exists(self.avatar_video): logger.error(f"[MuseTalk] Avatar video not found: {self.avatar_video}") return logger.info(f"[MuseTalk] Loading avatar from: {self.avatar_video}") # Check cache cache_dir = self._get_cache_path() cache_file = cache_dir / "avatar_cache.pkl" latents_file = cache_dir / "latents.pt" if cache_file.exists() and latents_file.exists(): logger.info(f"[MuseTalk] Loading from cache: {cache_dir}") start = time.time() with open(cache_file, 'rb') as f: cache = pickle.load(f) self.full_frames = cache['full_frames'] self.idle_frames = cache['idle_frames'] self.coord_list = cache['coord_list'] self.mask_list = cache['mask_list'] self.mask_coords_list = cache['mask_coords_list'] self.original_width = cache['original_width'] self.original_height = cache['original_height'] self.idle_fps = cache['idle_fps'] self.input_latent_list = torch.load(latents_file, weights_only=False) logger.info(f"[MuseTalk] Loaded {len(self.full_frames)} frames from cache in {time.time()-start:.1f}s") self._avatar_loaded = True return # No cache - load from video logger.info("[MuseTalk] No cache found, processing avatar (this will be cached)...") cap = cv2.VideoCapture(self.avatar_video) self.idle_fps = cap.get(cv2.CAP_PROP_FPS) or self.fps full_frames = [] while True: ret, frame = cap.read() if not ret: break full_frames.append(frame) cap.release() if not full_frames: logger.error("[MuseTalk] No frames loaded from video") return # Store original frames and dimensions self.full_frames = full_frames self.original_height, self.original_width = full_frames[0].shape[:2] logger.info(f"[MuseTalk] Loaded {len(full_frames)} frames at {self.idle_fps} fps (resolution: {self.original_width}x{self.original_height})") # Detect faces and create cropped frames for processing self._detect_faces() # Precompute latents if models loaded if self._models_loaded and self.vae and self.idle_frames: self._precompute_latents() # Save cache try: cache_dir.mkdir(parents=True, exist_ok=True) cache = { 'full_frames': self.full_frames, 'idle_frames': self.idle_frames, 'coord_list': self.coord_list, 'mask_list': self.mask_list, 'mask_coords_list': self.mask_coords_list, 'original_width': self.original_width, 'original_height': self.original_height, 'idle_fps': self.idle_fps, } with open(cache_file, 'wb') as f: pickle.dump(cache, f) torch.save(self.input_latent_list, latents_file) logger.info(f"[MuseTalk] Saved cache to: {cache_dir}") except Exception as e: logger.warning(f"[MuseTalk] Could not save cache: {e}") self._avatar_loaded = True except Exception as e: logger.error(f"[MuseTalk] Error loading avatar: {e}") import traceback traceback.print_exc() def _detect_faces(self): """Detect faces using MuseTalk's preprocessing (matching realtime_inference.py).""" try: from musetalk.utils.blending import get_image_prepare_material import numpy as np from musetalk.utils.face_detection import FaceAlignment, LandmarksType logger.info("[MuseTalk] Detecting faces using MuseTalk preprocessing...") device_str = 'cuda' if self.device is not None and self.device.type == 'cuda' else 'cpu' fa = FaceAlignment(LandmarksType._2D, flip_input=False, device=device_str) self.idle_frames = [] self.coord_list = [] self.mask_list = [] self.mask_coords_list = [] # Process in batches batch_size = 8 for batch_start in range(0, len(self.full_frames), batch_size): batch_end = min(batch_start + batch_size, len(self.full_frames)) batch_frames = self.full_frames[batch_start:batch_end] batch_array = np.stack(batch_frames, axis=0) detections = fa.get_detections_for_batch(batch_array) for i, (frame, detection) in enumerate(zip(batch_frames, detections)): frame_idx = batch_start + i h, w = frame.shape[:2] if detection is None: # Fallback: center crop size = min(h, w) // 2 center_x, center_y = w // 2, h // 2 x1 = center_x - size // 2 y1 = center_y - size // 2 x2 = x1 + size y2 = y1 + size else: x1, y1, x2, y2 = [int(v) for v in detection] # Apply bbox_shift (0 for V1.5) x1 = max(0, x1 + self.bbox_shift) y1 = max(0, y1 + self.bbox_shift) x2 = min(w, x2 + self.bbox_shift) y2 = min(h, y2 + self.bbox_shift) # For V1.5: add extra_margin to y2 (matching realtime_inference.py) if self.version == "v15": y2_extended = min(y2 + self.extra_margin, h) else: y2_extended = y2 bbox = [x1, y1, x2, y2_extended] self.coord_list.append(bbox) # Crop and resize face for latent computation (256x256) face_crop = frame[y1:y2_extended, x1:x2] if face_crop.size > 0: face_resized = cv2.resize(face_crop, (256, 256), interpolation=cv2.INTER_LANCZOS4) else: face_resized = cv2.resize(frame, (256, 256), interpolation=cv2.INTER_LANCZOS4) self.idle_frames.append(face_resized) logger.info(f"[MuseTalk] Face detection complete, {len(self.coord_list)} faces processed") # Pre-compute masks and crop_boxes for blending (matching realtime_inference.py) if self.face_parser is not None: logger.info("[MuseTalk] Pre-computing blending masks...") for i, (frame, bbox) in enumerate(zip(self.full_frames, self.coord_list)): try: mask, crop_box = get_image_prepare_material( frame, bbox, upper_boundary_ratio=self.upper_boundary_ratio, expand=self.expand, fp=self.face_parser, mode=self.parsing_mode ) self.mask_list.append(mask) self.mask_coords_list.append(crop_box) except Exception as e: logger.warning(f"[MuseTalk] Error computing mask for frame {i}: {e}") # Create empty mask as fallback self.mask_list.append(np.zeros((256, 256), dtype=np.uint8)) self.mask_coords_list.append([0, 0, 256, 256]) logger.info(f"[MuseTalk] Pre-computed {len(self.mask_list)} masks") except Exception as e: logger.error(f"[MuseTalk] Error in face detection: {e}") import traceback traceback.print_exc() # Fallback: use full frames resized self.idle_frames = [] self.coord_list = [] for frame in self.full_frames: resized = cv2.resize(frame, (256, 256)) self.idle_frames.append(resized) h, w = frame.shape[:2] self.coord_list.append([0, 0, w, h]) def _precompute_latents(self): """Precompute latents for avatar frames.""" try: import torch logger.info("[MuseTalk] Precomputing latents...") self.input_latent_list = [] for frame in self.idle_frames[:50]: # Use first 50 frames for looping frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) latents = self.vae.get_latents_for_unet(frame_rgb) self.input_latent_list.append(latents) logger.info(f"[MuseTalk] Precomputed {len(self.input_latent_list)} latents") except Exception as e: logger.error(f"[MuseTalk] Error precomputing latents: {e}") def get_idle_frames(self): """Get idle animation frames.""" return self.idle_frames if self.idle_frames else [np.zeros((self.resolution, self.resolution, 3), dtype=np.uint8)] def generate_frames_streaming(self, audio_path: str, resolution: int = None, batch_size: int = 4) -> Generator[dict, None, None]: """ Generate video frames from audio in streaming fashion. Args: audio_path: Path to audio file resolution: Output resolution (None = use original video resolution) batch_size: Batch size for inference """ # Use original resolution if not specified output_width = resolution if resolution else self.original_width output_height = resolution if resolution else self.original_height # If models not loaded, just return original full frames if not self._models_loaded or not self.audio_processor: logger.warning("[MuseTalk] Models not loaded, returning idle frames") duration = self._get_audio_duration(audio_path) num_frames = int(duration * self.fps) # Send info message first yield {"type": "info", "total_frames": num_frames, "fps": self.fps, "width": output_width, "height": output_height} for i in range(num_frames): frame_idx = i % len(self.full_frames) frame = self.full_frames[frame_idx].copy() yield {"type": "frame", "frame": frame, "index": i, "total": num_frames} return try: import torch from musetalk.utils.blending import get_image_blending logger.info(f"[MuseTalk] Processing audio: {audio_path}") # Extract audio features whisper_feature = self.audio_processor.audio2feat(audio_path) whisper_chunks = self.audio_processor.feature2chunks( feature_array=whisper_feature, fps=self.fps ) total_frames = len(whisper_chunks) logger.info(f"[MuseTalk] Generating {total_frames} frames (output: {output_width}x{output_height})") # Send info message first (required by frontend to activate canvas) yield {"type": "info", "total_frames": total_frames, "fps": self.fps, "width": output_width, "height": output_height} # Create cycled lists (like realtime_inference.py: forward + reverse) num_avatar_frames = len(self.full_frames) frame_list_cycle = self.full_frames + self.full_frames[::-1] coord_list_cycle = self.coord_list + self.coord_list[::-1] latent_list_cycle = self.input_latent_list + self.input_latent_list[::-1] mask_list_cycle = self.mask_list + self.mask_list[::-1] if self.mask_list else [] mask_coords_list_cycle = self.mask_coords_list + self.mask_coords_list[::-1] if self.mask_coords_list else [] # Generate frames for i, whisper_batch in enumerate(whisper_chunks): try: # Get corresponding data (cycling through avatar frames) cycle_idx = i % len(frame_list_cycle) latent_idx = i % len(latent_list_cycle) latent = latent_list_cycle[latent_idx] bbox = coord_list_cycle[cycle_idx] original_frame = copy.deepcopy(frame_list_cycle[cycle_idx]) # Prepare audio features audio_feat = torch.from_numpy(whisper_batch).unsqueeze(0).half().to(self.device) audio_feat = self.pe(audio_feat) # Generate face with UNet with torch.no_grad(): latent_input = latent.half().to(self.device) pred = self.unet.model( latent_input, self.timesteps, encoder_hidden_states=audio_feat ).sample # Decode to image (returns list of RGB frames) pred_faces = self.vae.decode_latents(pred) pred_face = pred_faces[0] # Get first (and only) face - RGB format # Resize predicted face to match bbox size (matching realtime_inference.py) x1, y1, x2, y2 = bbox try: pred_face_resized = cv2.resize(pred_face.astype(np.uint8), (x2 - x1, y2 - y1)) except: pred_face_resized = pred_face # Convert RGB to BGR for blending (get_image_blending expects BGR) pred_face_resized = cv2.cvtColor(pred_face_resized, cv2.COLOR_RGB2BGR) # Blend using pre-computed masks (matching realtime_inference.py) if mask_list_cycle and mask_coords_list_cycle: mask = mask_list_cycle[cycle_idx] mask_crop_box = mask_coords_list_cycle[cycle_idx] combined_frame = get_image_blending( original_frame, pred_face_resized, bbox, mask, mask_crop_box ) else: # Fallback: simple paste without blending combined_frame = original_frame.copy() combined_frame[y1:y2, x1:x2] = pred_face_resized yield {"type": "frame", "frame": combined_frame, "index": i, "total": total_frames} except Exception as e: logger.error(f"[MuseTalk] Error generating frame {i}: {e}") import traceback traceback.print_exc() # Fallback to original frame frame_idx = i % len(self.full_frames) yield {"type": "frame", "frame": self.full_frames[frame_idx].copy(), "index": i, "total": total_frames} logger.info("[MuseTalk] Frame generation complete") except Exception as e: logger.error(f"[MuseTalk] Error in streaming generation: {e}") import traceback traceback.print_exc() # Return full frames as fallback (original resolution) for i, frame in enumerate(self.full_frames[:30]): yield {"type": "frame", "frame": frame.copy(), "index": i, "total": 30} def _get_audio_duration(self, audio_path: str) -> float: """Get audio duration in seconds.""" try: result = subprocess.run( ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", audio_path], capture_output=True, text=True ) return float(result.stdout.strip()) except: return 1.0 # Default 1 second # Global engine instance _engine: Optional[MuseTalkEngine] = None def initialize_engine(avatar_path=None, resolution=256, fps=25) -> MuseTalkEngine: """Initialize and return the global MuseTalk engine.""" global _engine if _engine is None: _engine = MuseTalkEngine(avatar_video=avatar_path, resolution=resolution, fps=fps) return _engine def get_engine() -> Optional[MuseTalkEngine]: """Get the global engine instance.""" return _engine def generate_frames_streaming(audio_path: str, resolution: int = None, batch_size: int = 4) -> Generator[dict, None, None]: """Generate streaming frames from audio using the global engine.""" engine = get_engine() if engine is None: logger.error("[MuseTalk] Engine not initialized") res = resolution or 512 yield {"type": "frame", "frame": np.zeros((res, res, 3), dtype=np.uint8), "index": 0, "total": 1} return yield from engine.generate_frames_streaming(audio_path, resolution, batch_size) def get_idle_frames() -> List[np.ndarray]: """Get idle animation frames (original resolution).""" engine = get_engine() if engine is None: return [np.zeros((256, 256, 3), dtype=np.uint8)] # Return full frames if available, otherwise idle frames if engine.full_frames: return engine.full_frames return engine.get_idle_frames() if __name__ == "__main__": engine = initialize_engine() print(f"Avatar loaded: {engine.avatar_loaded}") print(f"Models loaded: {engine._models_loaded}") print(f"Idle frames: {len(engine.idle_frames)}")