server/musetalk_engine.py

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