""" MuseTalk Engine - Keeps models loaded in memory for fast inference """ import os import sys import torch import numpy as np from pathlib import Path # Add parent directory to path sys.path.insert(0, str(Path(__file__).parent.parent)) # Global model instances _engine = None class MuseTalkEngine: """Singleton engine that keeps models loaded in memory""" def __init__(self): print("Initializing MuseTalk Engine...") self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.loaded = False # Paths self.base_dir = Path(__file__).parent.parent self.avatar_video = self.base_dir / "data" / "video" / "avatar.mp4" # Models will be loaded on first use self.vae = None self.unet = None self.pe = None self.timesteps = None self.audio_processor = None self.whisper = None # Cached avatar data self.avatar_latents = None self.avatar_coords = None self.avatar_frames = None self.avatar_fps = 25 def load_models(self): """Load all models into memory""" if self.loaded: return print("Loading MuseTalk models...") from musetalk.utils.utils import load_all_model from musetalk.whisper.audio2feature import Audio2Feature # Load models self.vae, self.unet, self.pe = load_all_model() self.audio_processor = Audio2Feature(model_path=str(self.base_dir / "models" / "whisper" / "tiny.pt")) self.timesteps = torch.tensor([0], device=self.device) print("Models loaded successfully!") self.loaded = True def preprocess_avatar(self): """Pre-process avatar video (run once)""" if self.avatar_latents is not None: print("Avatar already preprocessed") return print("Preprocessing avatar video...") import glob import pickle import cv2 from musetalk.utils.preprocessing import get_landmark_and_bbox, coord_placeholder from musetalk.utils.utils import get_video_fps, datagen # Create temp dir for frames temp_dir = self.base_dir / "results" / "server" / "avatar_frames" temp_dir.mkdir(parents=True, exist_ok=True) coord_path = temp_dir / "coords.pkl" # Check if already cached if coord_path.exists() and list(temp_dir.glob("*.png")): print("Loading cached avatar data...") with open(coord_path, "rb") as f: self.avatar_coords = pickle.load(f) self.avatar_frames = sorted(glob.glob(str(temp_dir / "*.png"))) self.avatar_fps = get_video_fps(str(self.avatar_video)) self._compute_latents() return # Extract frames import imageio reader = imageio.get_reader(str(self.avatar_video)) self.avatar_fps = get_video_fps(str(self.avatar_video)) frame_paths = [] for i, frame in enumerate(reader): frame_path = temp_dir / f"{i:08d}.png" imageio.imwrite(str(frame_path), frame) frame_paths.append(str(frame_path)) self.avatar_frames = frame_paths # Get landmarks and bounding boxes print("Computing face landmarks...") self.avatar_coords = get_landmark_and_bbox(frame_paths, 0) # Save coords with open(coord_path, "wb") as f: pickle.dump(self.avatar_coords, f) self._compute_latents() print("Avatar preprocessing complete!") def _compute_latents(self): """Compute VAE latents for avatar frames""" print("Computing VAE latents...") from musetalk.utils.preprocessing import read_imgs, coord_placeholder import copy input_latent_list = [] coord_placeholder_list = coord_placeholder(self.avatar_coords, copy.deepcopy(self.avatar_frames), self.avatar_coords[0]["bbox"]) for bbox, frame, _, _ in coord_placeholder_list: if bbox is None: continue x1, y1, x2, y2 = bbox crop_frame = frame[y1:y2, x1:x2] crop_frame = cv2.resize(crop_frame, (256, 256), interpolation=cv2.INTER_LANCZOS4) latents = self.vae.get_latents_for_unet(crop_frame) input_latent_list.append(latents) self.avatar_latents = input_latent_list print(f"Computed {len(input_latent_list)} latents") def generate_video(self, audio_path: str, output_path: str) -> str: """Generate lip-sync video from audio""" import cv2 import copy from tqdm import tqdm self.load_models() self.preprocess_avatar() print(f"Generating video for: {audio_path}") # Get audio features whisper_feature = self.audio_processor.audio2feat(audio_path) whisper_chunks = self.audio_processor.feature2chunks( feature_array=whisper_feature, fps=self.avatar_fps ) print(f"Audio chunks: {len(whisper_chunks)}") # Cycle latents to match audio length from itertools import cycle latent_cycle = cycle(self.avatar_latents) coord_cycle = cycle(self.avatar_coords) frame_cycle = cycle(self.avatar_frames) # Generate frames batch_size = 8 gen_frames = [] for i in range(0, len(whisper_chunks), batch_size): whisper_batch = whisper_chunks[i:i+batch_size] latent_batch = [next(latent_cycle) for _ in range(len(whisper_batch))] # Convert to tensors audio_tensors = [torch.FloatTensor(w).to(self.device) for w in whisper_batch] audio_tensor = torch.stack(audio_tensors) latent_tensor = torch.cat(latent_batch, dim=0) # Generate with UNet with torch.no_grad(): pred = self.unet(latent_tensor, self.timesteps, encoder_hidden_states=audio_tensor).sample recon = self.vae.decode_latents(pred) for frame in recon: gen_frames.append(frame) # Compose final video print(f"Composing {len(gen_frames)} frames...") from musetalk.utils.blending import get_image output_frames = [] coord_cycle = cycle(self.avatar_coords) frame_cycle = cycle(self.avatar_frames) for i, gen_frame in enumerate(gen_frames[:len(whisper_chunks)]): coord = next(coord_cycle) orig_frame_path = next(frame_cycle) orig_frame = cv2.imread(orig_frame_path) if coord.get("bbox"): x1, y1, x2, y2 = coord["bbox"] gen_resized = cv2.resize(gen_frame, (x2-x1, y2-y1)) # Blend result = get_image(orig_frame, gen_resized, coord) output_frames.append(result) else: output_frames.append(orig_frame) # Write video print(f"Writing video to: {output_path}") h, w = output_frames[0].shape[:2] fourcc = cv2.VideoWriter_fourcc(*'mp4v') out = cv2.VideoWriter(output_path, fourcc, self.avatar_fps, (w, h)) for frame in output_frames: out.write(frame) out.release() # Add audio temp_video = output_path.replace(".mp4", "_temp.mp4") os.rename(output_path, temp_video) import subprocess cmd = [ "ffmpeg", "-y", "-i", temp_video, "-i", audio_path, "-c:v", "libx264", "-c:a", "aac", "-shortest", output_path ] subprocess.run(cmd, capture_output=True) os.remove(temp_video) print(f"Video generated: {output_path}") return output_path def get_engine() -> MuseTalkEngine: """Get or create the singleton engine""" global _engine if _engine is None: _engine = MuseTalkEngine() return _engine # Pre-load on import if running directly if __name__ == "__main__": engine = get_engine() engine.load_models() engine.preprocess_avatar() print("Engine ready!")