server/fast_engine.py

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