feat: add MuseTalk fast_engine integration and model download script

- Created server/fast_engine.py with MuseTalk integration for lip-sync avatar
- Added scripts/download_models.sh to download all required models from HuggingFace
- Added assets/avatar.png placeholder image

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

scripts/download_models.sh

@@ -0,0 +1,171 @@
+#!/bin/bash
+#
+# Download MuseTalk models from HuggingFace
+#
+
+set -e
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
+MODELS_DIR="${PROJECT_DIR}/models"
+
+GREEN='\033[0;32m'
+CYAN='\033[0;36m'
+YELLOW='\033[1;33m'
+NC='\033[0m'
+
+success() { echo -e "${GREEN}✓${NC} $1"; }
+header() { echo -e "\n${CYAN}$1${NC}"; }
+warn() { echo -e "${YELLOW}!${NC} $1"; }
+
+echo ""
+echo "╔══════════════════════════════════════════════════════════════╗"
+echo "║          MUSETALK - MODEL DOWNLOAD                           ║"
+echo "╚══════════════════════════════════════════════════════════════╝"
+echo ""
+
+# Check for huggingface-cli
+if ! command -v huggingface-cli &> /dev/null; then
+    echo "Installing huggingface_hub..."
+    pip install huggingface_hub -q
+fi
+
+# Create directories
+mkdir -p "${MODELS_DIR}/avatar/musetalk"
+mkdir -p "${MODELS_DIR}/avatar/vae"
+mkdir -p "${MODELS_DIR}/whisper"
+mkdir -p "${MODELS_DIR}/dwpose"
+mkdir -p "${MODELS_DIR}/face_parse"
+
+# ============================================================
+# 1. MUSETALK MODELS
+# ============================================================
+header "[1/5] Downloading MuseTalk UNet..."
+
+if [ -f "${MODELS_DIR}/avatar/musetalk/pytorch_model.bin" ]; then
+    success "MuseTalk UNet already exists"
+else
+    huggingface-cli download TMElyralab/MuseTalk \
+        --include "models/musetalk/*" \
+        --local-dir "${MODELS_DIR}/temp_musetalk" \
+        --local-dir-use-symlinks False
+
+    # Move files to correct location
+    if [ -d "${MODELS_DIR}/temp_musetalk/models/musetalk" ]; then
+        cp -r "${MODELS_DIR}/temp_musetalk/models/musetalk/"* "${MODELS_DIR}/avatar/musetalk/"
+        rm -rf "${MODELS_DIR}/temp_musetalk"
+        success "MuseTalk UNet downloaded"
+    else
+        warn "MuseTalk download may have failed"
+    fi
+fi
+
+# ============================================================
+# 2. VAE (SD-VAE-FT-MSE)
+# ============================================================
+header "[2/5] Downloading VAE..."
+
+if [ -f "${MODELS_DIR}/avatar/vae/config.json" ]; then
+    success "VAE already exists"
+else
+    huggingface-cli download stabilityai/sd-vae-ft-mse \
+        --local-dir "${MODELS_DIR}/avatar/vae" \
+        --local-dir-use-symlinks False
+
+    success "VAE downloaded"
+fi
+
+# ============================================================
+# 3. WHISPER
+# ============================================================
+header "[3/5] Downloading Whisper..."
+
+if [ -f "${MODELS_DIR}/whisper/tiny.pt" ]; then
+    success "Whisper already exists"
+else
+    # Download from OpenAI
+    wget -q -O "${MODELS_DIR}/whisper/tiny.pt" \
+        "https://openaipublic.azureedge.net/main/whisper/models/65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/tiny.pt" \
+        || curl -s -o "${MODELS_DIR}/whisper/tiny.pt" \
+        "https://openaipublic.azureedge.net/main/whisper/models/65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/tiny.pt"
+
+    if [ -f "${MODELS_DIR}/whisper/tiny.pt" ]; then
+        success "Whisper downloaded"
+    else
+        warn "Whisper download failed"
+    fi
+fi
+
+# ============================================================
+# 4. DWPOSE
+# ============================================================
+header "[4/5] Downloading DWPose..."
+
+if [ -f "${MODELS_DIR}/dwpose/dw-ll_ucoco_384.pth" ]; then
+    success "DWPose already exists"
+else
+    huggingface-cli download yzd-v/DWPose \
+        --include "dw-ll_ucoco_384.pth" \
+        --local-dir "${MODELS_DIR}/dwpose" \
+        --local-dir-use-symlinks False
+
+    success "DWPose downloaded"
+fi
+
+# ============================================================
+# 5. FACE PARSING
+# ============================================================
+header "[5/5] Downloading Face Parsing models..."
+
+if [ -f "${MODELS_DIR}/face_parse/79999_iter.pth" ]; then
+    success "Face parsing already exists"
+else
+    # Download from MuseTalk repo
+    huggingface-cli download TMElyralab/MuseTalk \
+        --include "models/face-parse-bisenet/*" \
+        --local-dir "${MODELS_DIR}/temp_fp" \
+        --local-dir-use-symlinks False
+
+    if [ -d "${MODELS_DIR}/temp_fp/models/face-parse-bisenet" ]; then
+        cp -r "${MODELS_DIR}/temp_fp/models/face-parse-bisenet/"* "${MODELS_DIR}/face_parse/"
+        rm -rf "${MODELS_DIR}/temp_fp"
+        success "Face parsing downloaded"
+    fi
+
+    # Also need resnet18
+    if [ ! -f "${MODELS_DIR}/face_parse/resnet18-5c106cde.pth" ]; then
+        wget -q -O "${MODELS_DIR}/face_parse/resnet18-5c106cde.pth" \
+            "https://download.pytorch.org/models/resnet18-5c106cde.pth" \
+            || curl -s -o "${MODELS_DIR}/face_parse/resnet18-5c106cde.pth" \
+            "https://download.pytorch.org/models/resnet18-5c106cde.pth"
+    fi
+fi
+
+# ============================================================
+# SUMMARY
+# ============================================================
+echo ""
+echo "╔══════════════════════════════════════════════════════════════╗"
+echo "║          DOWNLOAD COMPLETE                                   ║"
+echo "╚══════════════════════════════════════════════════════════════╝"
+echo ""
+
+echo "Models directory: ${MODELS_DIR}"
+echo ""
+
+# Check all models
+check_model() {
+    if [ -f "$1" ] || [ -d "$1" ]; then
+        echo -e "  ${GREEN}✓${NC} $2"
+    else
+        echo -e "  ✗ $2 (missing)"
+    fi
+}
+
+echo "Installed models:"
+check_model "${MODELS_DIR}/avatar/musetalk/pytorch_model.bin" "MuseTalk UNet"
+check_model "${MODELS_DIR}/avatar/vae/config.json" "VAE"
+check_model "${MODELS_DIR}/whisper/tiny.pt" "Whisper Tiny"
+check_model "${MODELS_DIR}/dwpose/dw-ll_ucoco_384.pth" "DWPose"
+check_model "${MODELS_DIR}/face_parse/79999_iter.pth" "Face Parsing"
+echo ""

server/fast_engine.py

@@ -0,0 +1,487 @@
+"""
+Fast Avatar Engine - MuseTalk Integration for real-time lip-sync avatar generation.
+Provides streaming video frame generation from audio input.
+"""
+
+import os
+import sys
+import logging
+import asyncio
+import tempfile
+from pathlib import Path
+from typing import Optional, Generator, List, Callable
+import numpy as np
+import cv2
+import torch
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Add server directory to path for imports
+SERVER_DIR = Path(__file__).parent
+sys.path.insert(0, str(SERVER_DIR))
+
+# Model paths
+MODELS_DIR = Path(os.getenv("MODELS_DIR", SERVER_DIR.parent / "models"))
+AVATAR_DIR = MODELS_DIR / "avatar"
+MUSETALK_DIR = AVATAR_DIR / "musetalk"
+VAE_DIR = AVATAR_DIR / "vae"
+WHISPER_DIR = MODELS_DIR / "whisper"
+DWPOSE_DIR = MODELS_DIR / "dwpose"
+FACE_PARSE_DIR = MODELS_DIR / "face_parse"
+
+# Default avatar image
+DEFAULT_AVATAR = os.getenv("AVATAR_IMAGE", str(SERVER_DIR.parent / "assets" / "avatar.png"))
+
+
+class MuseTalkEngine:
+    """
+    MuseTalk-based avatar engine for real-time lip-sync video generation.
+    """
+
+    def __init__(self,
+                 avatar_image: Optional[str] = None,
+                 resolution: int = 256,
+                 fps: int = 25,
+                 use_float16: bool = True):
+        """
+        Initialize the MuseTalk engine.
+
+        Args:
+            avatar_image: Path to the avatar image
+            resolution: Output resolution (default 256x256)
+            fps: Frames per second (default 25)
+            use_float16: Use FP16 for faster inference
+        """
+        self.avatar_image = avatar_image or DEFAULT_AVATAR
+        self.resolution = resolution
+        self.fps = fps
+        self.use_float16 = use_float16
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        self._avatar_loaded = False
+        self._models_loaded = False
+
+        # Models
+        self.vae = None
+        self.unet = None
+        self.audio2feature = None
+        self.face_parser = None
+
+        # Cached avatar data
+        self.avatar_frame = None
+        self.avatar_latents = None
+        self.face_bbox = None
+        self.idle_frames = []
+
+        logger.info(f"[MuseTalk] Initialized with device: {self.device}")
+
+    @property
+    def avatar_loaded(self) -> bool:
+        """Check if avatar is loaded and ready."""
+        return self._avatar_loaded and self._models_loaded
+
+    def _check_models(self) -> bool:
+        """Check if all required models exist."""
+        required_paths = [
+            (VAE_DIR, "VAE model"),
+            (MUSETALK_DIR / "musetalk.json", "MuseTalk config"),
+            (MUSETALK_DIR / "pytorch_model.bin", "MuseTalk weights"),
+        ]
+
+        for path, name in required_paths:
+            if not path.exists():
+                logger.warning(f"[MuseTalk] Missing {name}: {path}")
+                return False
+
+        return True
+
+    def load_models(self) -> bool:
+        """Load all required models."""
+        if self._models_loaded:
+            return True
+
+        if not self._check_models():
+            logger.error("[MuseTalk] Required models not found")
+            return False
+
+        try:
+            logger.info("[MuseTalk] Loading models...")
+
+            # Load VAE
+            from musetalk.models.vae import VAE
+            self.vae = VAE(
+                model_path=str(VAE_DIR),
+                resized_img=self.resolution,
+                use_float16=self.use_float16
+            )
+            logger.info("[MuseTalk] VAE loaded")
+
+            # Load UNet
+            from musetalk.models.unet import UNet
+            self.unet = UNet(
+                unet_config=str(MUSETALK_DIR / "musetalk.json"),
+                model_path=str(MUSETALK_DIR / "pytorch_model.bin"),
+                use_float16=self.use_float16,
+                device=self.device
+            )
+            logger.info("[MuseTalk] UNet loaded")
+
+            # Load Audio2Feature (Whisper)
+            whisper_path = WHISPER_DIR / "tiny.pt"
+            if whisper_path.exists():
+                from musetalk.whisper.audio2feature import Audio2Feature
+                self.audio2feature = Audio2Feature(
+                    whisper_model_type="tiny",
+                    model_path=str(whisper_path)
+                )
+                logger.info("[MuseTalk] Audio2Feature loaded")
+            else:
+                logger.warning(f"[MuseTalk] Whisper model not found: {whisper_path}")
+
+            self._models_loaded = True
+            logger.info("[MuseTalk] All models loaded successfully")
+            return True
+
+        except Exception as e:
+            logger.error(f"[MuseTalk] Error loading models: {e}")
+            import traceback
+            traceback.print_exc()
+            return False
+
+    def load_avatar(self, avatar_path: Optional[str] = None) -> bool:
+        """
+        Load and preprocess the avatar image.
+
+        Args:
+            avatar_path: Path to avatar image (uses default if None)
+        """
+        if avatar_path:
+            self.avatar_image = avatar_path
+
+        if not os.path.exists(self.avatar_image):
+            logger.error(f"[MuseTalk] Avatar image not found: {self.avatar_image}")
+            return False
+
+        try:
+            logger.info(f"[MuseTalk] Loading avatar: {self.avatar_image}")
+
+            # Load and resize avatar
+            avatar = cv2.imread(self.avatar_image)
+            if avatar is None:
+                logger.error(f"[MuseTalk] Failed to read avatar image")
+                return False
+
+            self.avatar_frame = cv2.resize(
+                avatar,
+                (self.resolution, self.resolution),
+                interpolation=cv2.INTER_LANCZOS4
+            )
+
+            # Precompute avatar latents if models are loaded
+            if self._models_loaded and self.vae:
+                self.avatar_latents = self.vae.get_latents_for_unet(self.avatar_frame)
+                logger.info("[MuseTalk] Avatar latents computed")
+
+            # Generate idle frames (just the static avatar for now)
+            self.idle_frames = [self.avatar_frame.copy() for _ in range(10)]
+
+            self._avatar_loaded = True
+            logger.info("[MuseTalk] Avatar loaded successfully")
+            return True
+
+        except Exception as e:
+            logger.error(f"[MuseTalk] Error loading avatar: {e}")
+            import traceback
+            traceback.print_exc()
+            return False
+
+    def get_idle_frames(self) -> List[np.ndarray]:
+        """Get idle animation frames."""
+        if not self.idle_frames and self.avatar_frame is not None:
+            return [self.avatar_frame.copy()]
+        return self.idle_frames
+
+    def _process_audio(self, audio_path: str) -> Optional[np.ndarray]:
+        """Extract audio features using Whisper."""
+        if not self.audio2feature:
+            logger.warning("[MuseTalk] Audio2Feature not loaded")
+            return None
+
+        try:
+            feature_array = self.audio2feature.audio2feat(audio_path)
+            return feature_array
+        except Exception as e:
+            logger.error(f"[MuseTalk] Error processing audio: {e}")
+            return None
+
+    def generate_frames_streaming(
+        self,
+        audio_path: str,
+        resolution: int = 256,
+        batch_size: int = 4
+    ) -> Generator[np.ndarray, None, None]:
+        """
+        Generate video frames from audio in a streaming fashion.
+
+        Args:
+            audio_path: Path to audio file
+            resolution: Output resolution
+            batch_size: Batch size for inference
+
+        Yields:
+            Video frames as numpy arrays (BGR format)
+        """
+        if not self.avatar_loaded:
+            logger.error("[MuseTalk] Engine not ready - avatar not loaded")
+            # Yield a black frame as fallback
+            yield np.zeros((resolution, resolution, 3), dtype=np.uint8)
+            return
+
+        # If models aren't fully loaded, return static avatar
+        if not self._models_loaded or not self.audio2feature:
+            logger.warning("[MuseTalk] Models not fully loaded, returning static avatar")
+            # Estimate frame count from audio duration
+            import subprocess
+            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
+                )
+                duration = float(result.stdout.strip())
+                num_frames = int(duration * self.fps)
+            except:
+                num_frames = 25  # Default 1 second
+
+            for _ in range(num_frames):
+                yield self.avatar_frame.copy()
+            return
+
+        try:
+            # Extract audio features
+            logger.info(f"[MuseTalk] Processing audio: {audio_path}")
+            feature_array = self._process_audio(audio_path)
+
+            if feature_array is None:
+                logger.error("[MuseTalk] Failed to extract audio features")
+                yield self.avatar_frame.copy()
+                return
+
+            # Get whisper chunks for each frame
+            whisper_chunks = self.audio2feature.feature2chunks(
+                feature_array,
+                fps=self.fps,
+                audio_feat_length=[2, 2]
+            )
+
+            logger.info(f"[MuseTalk] Generating {len(whisper_chunks)} frames")
+
+            # Generate frames in batches
+            for i in range(0, len(whisper_chunks), batch_size):
+                batch_chunks = whisper_chunks[i:i + batch_size]
+
+                for chunk in batch_chunks:
+                    try:
+                        # Convert chunk to tensor
+                        audio_feat = torch.from_numpy(chunk).unsqueeze(0)
+                        audio_feat = audio_feat.to(self.device)
+
+                        if self.use_float16:
+                            audio_feat = audio_feat.half()
+
+                        # Add positional encoding
+                        audio_feat = self.unet.pe(audio_feat)
+
+                        # Generate frame using UNet
+                        with torch.no_grad():
+                            latents = self.unet.model(
+                                self.avatar_latents.to(self.device),
+                                0,  # timestep
+                                encoder_hidden_states=audio_feat
+                            ).sample
+
+                        # Decode latents to image
+                        frame = self.vae.decode_latents(latents)[0]
+
+                        # Resize if needed
+                        if frame.shape[0] != resolution or frame.shape[1] != resolution:
+                            frame = cv2.resize(frame, (resolution, resolution))
+
+                        yield frame
+
+                    except Exception as e:
+                        logger.error(f"[MuseTalk] Error generating frame: {e}")
+                        yield self.avatar_frame.copy()
+
+            logger.info("[MuseTalk] Frame generation complete")
+
+        except Exception as e:
+            logger.error(f"[MuseTalk] Error in streaming generation: {e}")
+            import traceback
+            traceback.print_exc()
+            yield self.avatar_frame.copy()
+
+    async def generate_video_fast(
+        self,
+        audio_path: str,
+        output_path: str,
+        resolution: int = 256,
+        batch_size: int = 8,
+        callback: Optional[Callable[[int, int], None]] = None
+    ) -> bool:
+        """
+        Generate a complete video from audio.
+
+        Args:
+            audio_path: Path to input audio
+            output_path: Path for output video
+            resolution: Output resolution
+            batch_size: Batch size for inference
+            callback: Progress callback (current_frame, total_frames)
+
+        Returns:
+            True if successful
+        """
+        try:
+            frames = []
+            frame_count = 0
+
+            for frame in self.generate_frames_streaming(audio_path, resolution, batch_size):
+                frames.append(frame)
+                frame_count += 1
+                if callback:
+                    callback(frame_count, -1)  # Unknown total
+
+            if not frames:
+                logger.error("[MuseTalk] No frames generated")
+                return False
+
+            # Write video
+            temp_video = output_path.replace('.mp4', '_noaudio.mp4')
+            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+            out = cv2.VideoWriter(temp_video, fourcc, self.fps, (resolution, resolution))
+
+            for frame in frames:
+                out.write(frame)
+            out.release()
+
+            # Add audio
+            import subprocess
+            subprocess.run([
+                'ffmpeg', '-y',
+                '-i', temp_video,
+                '-i', audio_path,
+                '-c:v', 'libx264',
+                '-c:a', 'aac',
+                '-shortest',
+                output_path
+            ], check=True, capture_output=True)
+
+            # Cleanup
+            if os.path.exists(temp_video):
+                os.remove(temp_video)
+
+            logger.info(f"[MuseTalk] Video saved: {output_path}")
+            return True
+
+        except Exception as e:
+            logger.error(f"[MuseTalk] Error generating video: {e}")
+            import traceback
+            traceback.print_exc()
+            return False
+
+
+# Global engine instance
+_engine: Optional[MuseTalkEngine] = None
+
+
+def initialize_engine(
+    avatar_path: Optional[str] = None,
+    resolution: int = 256,
+    fps: int = 25
+) -> MuseTalkEngine:
+    """
+    Initialize and return the global MuseTalk engine.
+
+    Args:
+        avatar_path: Path to avatar image
+        resolution: Output resolution
+        fps: Frames per second
+    """
+    global _engine
+
+    if _engine is not None:
+        return _engine
+
+    _engine = MuseTalkEngine(
+        avatar_image=avatar_path,
+        resolution=resolution,
+        fps=fps
+    )
+
+    # Try to load models (may fail if not installed)
+    _engine.load_models()
+
+    # Try to load avatar (may fail if image not found)
+    _engine.load_avatar()
+
+    return _engine
+
+
+def get_engine() -> Optional[MuseTalkEngine]:
+    """Get the global engine instance."""
+    return _engine
+
+
+# Convenience exports
+def generate_video_fast(
+    audio_path: str,
+    output_path: str,
+    resolution: int = 256,
+    batch_size: int = 8,
+    callback: Optional[Callable[[int, int], None]] = None
+) -> bool:
+    """Generate video from audio using the global engine."""
+    engine = get_engine()
+    if engine is None:
+        logger.error("[MuseTalk] Engine not initialized")
+        return False
+
+    import asyncio
+    return asyncio.run(engine.generate_video_fast(
+        audio_path, output_path, resolution, batch_size, callback
+    ))
+
+
+def generate_frames_streaming(
+    audio_path: str,
+    resolution: int = 256,
+    batch_size: int = 4
+) -> Generator[np.ndarray, None, None]:
+    """Generate streaming frames from audio using the global engine."""
+    engine = get_engine()
+    if engine is None:
+        logger.error("[MuseTalk] Engine not initialized")
+        yield np.zeros((resolution, resolution, 3), dtype=np.uint8)
+        return
+
+    yield from engine.generate_frames_streaming(audio_path, resolution, batch_size)
+
+
+def get_idle_frames() -> List[np.ndarray]:
+    """Get idle animation frames."""
+    engine = get_engine()
+    if engine is None:
+        return [np.zeros((256, 256, 3), dtype=np.uint8)]
+    return engine.get_idle_frames()
+
+
+if __name__ == "__main__":
+    # Test the engine
+    engine = initialize_engine()
+    print(f"Avatar loaded: {engine.avatar_loaded}")
+    print(f"Models loaded: {engine._models_loaded}")
+
+    idle = get_idle_frames()
+    print(f"Idle frames: {len(idle)}")