Upload code/pipeline.py with huggingface_hub

code/pipeline.py

@@ -0,0 +1,691 @@
+"""
+EMOLIPS Pipeline
+================
+Emotion-Driven Lip-Sync Synthesis Pipeline
+
+Orchestrates:
+1. Audio emotion detection (automatic or manual override)
+2. Emotion intensity estimation
+3. SadTalker talking face generation
+4. Emotion-conditioned coefficient modification
+5. Output video rendering
+
+Usage:
+    pipeline = EmolipsPipeline(device="cuda")
+    pipeline.generate(
+        audio_path="speech.wav",
+        image_path="face.jpg",
+        emotion="happy",        # Optional: auto-detected if not specified
+        intensity=0.7,          # Optional: auto-estimated if not specified
+        output_path="output.mp4"
+    )
+"""
+
+import os
+import sys
+import subprocess
+import shutil
+import json
+import numpy as np
+from pathlib import Path
+from typing import Optional, Dict, List
+import warnings
+warnings.filterwarnings("ignore")
+
+from emotion_module import (
+    PracticalEmotionModifier,
+    AudioEmotionDetector,
+    EmotionIntensityEstimator,
+    EMOTION_PROFILES
+)
+
+
+class EmolipsPipeline:
+    """
+    Main EMOLIPS inference pipeline.
+
+    Wraps SadTalker backbone with emotion conditioning.
+    """
+
+    def __init__(
+        self,
+        sadtalker_dir: str = "./SadTalker",
+        device: str = "cuda",
+        checkpoint_dir: str = None
+    ):
+        self.sadtalker_dir = Path(sadtalker_dir).resolve()
+        self.device = device
+        self.checkpoint_dir = checkpoint_dir or str(self.sadtalker_dir / "checkpoints")
+
+        # Initialize emotion components
+        self.emotion_detector = AudioEmotionDetector(device=device)
+        self.intensity_estimator = EmotionIntensityEstimator()
+        self.emotion_modifier = PracticalEmotionModifier()
+
+        # Verify SadTalker installation
+        if not self.sadtalker_dir.exists():
+            print(f"⚠ SadTalker not found at {self.sadtalker_dir}")
+            print("  Run setup.sh first or specify correct path")
+
+    def detect_emotion(self, audio_path: str) -> Dict:
+        """Auto-detect emotion from audio."""
+        print("  [1/4] Detecting emotion from audio...")
+        result = self.emotion_detector.detect(audio_path)
+        print(f"        Detected: {result['detected_emotion']} "
+              f"(confidence: {result['confidence']:.2f})")
+        return result
+
+    def estimate_intensity(self, audio_path: str) -> float:
+        """Estimate emotion intensity from audio features."""
+        intensity = self.intensity_estimator.estimate(audio_path)
+        print(f"        Intensity: {intensity:.2f}")
+        return intensity
+
+    def run_sadtalker(
+        self,
+        audio_path: str,
+        image_path: str,
+        output_dir: str,
+        expression_scale: float = 1.0,
+        still_mode: bool = False,
+        preprocess: str = "crop",
+        size: int = 256,
+        pose_style: int = 0
+    ) -> Optional[str]:
+        """
+        Run SadTalker to generate base talking face video.
+
+        Returns path to generated video.
+        """
+        print("  [2/4] Running SadTalker backbone...")
+
+        # Build SadTalker command
+        inference_script = self.sadtalker_dir / "inference.py"
+
+        cmd = [
+            sys.executable, str(inference_script),
+            "--driven_audio", str(audio_path),
+            "--source_image", str(image_path),
+            "--result_dir", str(output_dir),
+            "--expression_scale", str(expression_scale),
+            "--preprocess", preprocess,
+            "--size", str(size),
+            "--pose_style", str(pose_style),
+        ]
+
+        if still_mode:
+            cmd.append("--still")
+
+        # Add checkpoint paths
+        checkpoint_dir = Path(self.checkpoint_dir)
+        if checkpoint_dir.exists():
+            cmd.extend(["--checkpoint_dir", str(checkpoint_dir)])
+
+        try:
+            env = os.environ.copy()
+            env["PYTHONPATH"] = str(self.sadtalker_dir) + ":" + env.get("PYTHONPATH", "")
+
+            result = subprocess.run(
+                cmd,
+                capture_output=True,
+                text=True,
+                cwd=str(self.sadtalker_dir),
+                env=env,
+                timeout=300  # 5 min timeout
+            )
+
+            if result.returncode != 0:
+                print(f"  ⚠ SadTalker error: {result.stderr[-500:]}")
+                return None
+
+            # Find generated video
+            output_path = Path(output_dir)
+            videos = list(output_path.rglob("*.mp4"))
+            if videos:
+                return str(sorted(videos, key=os.path.getmtime)[-1])
+
+            return None
+
+        except subprocess.TimeoutExpired:
+            print("  ⚠ SadTalker timed out (>5 min)")
+            return None
+        except Exception as e:
+            print(f"  ⚠ SadTalker failed: {e}")
+            return None
+
+    def apply_emotion_postprocess(
+        self,
+        video_path: str,
+        emotion: str,
+        intensity: float,
+        output_path: str
+    ) -> str:
+        """
+        Apply emotion-based post-processing to generated video.
+
+        This applies subtle facial modifications via:
+        1. Face landmark detection on each frame
+        2. Emotion-specific spatial warping
+        3. Color grading for emotional tone
+        """
+        print("  [3/4] Applying emotion conditioning...")
+
+        try:
+            import cv2
+            import mediapipe as mp
+
+            mp_face_mesh = mp.solutions.face_mesh
+            face_mesh = mp_face_mesh.FaceMesh(
+                static_image_mode=False,
+                max_num_faces=1,
+                min_detection_confidence=0.5
+            )
+
+            cap = cv2.VideoCapture(video_path)
+            fps = int(cap.get(cv2.CAP_PROP_FPS))
+            w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+            h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+            # Temp output (will mux audio later)
+            temp_path = output_path.replace(".mp4", "_temp.mp4")
+            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+            out = cv2.VideoWriter(temp_path, fourcc, fps, (w, h))
+
+            profile = EMOTION_PROFILES.get(emotion, EMOTION_PROFILES["neutral"])
+
+            frame_count = 0
+            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+            while cap.isOpened():
+                ret, frame = cap.read()
+                if not ret:
+                    break
+
+                # Apply emotion-specific color grading
+                frame = self._apply_color_grade(frame, emotion, intensity)
+
+                # Apply subtle face warping if emotion is strong
+                if intensity > 0.3 and emotion != "neutral":
+                    frame = self._apply_face_warp(frame, face_mesh, emotion, intensity)
+
+                out.write(frame)
+                frame_count += 1
+
+            cap.release()
+            out.release()
+            face_mesh.close()
+
+            # Mux original audio back
+            self._mux_audio(temp_path, video_path, output_path)
+
+            # Cleanup temp
+            if os.path.exists(temp_path):
+                os.remove(temp_path)
+
+            print(f"        Processed {frame_count} frames")
+            return output_path
+
+        except ImportError as e:
+            print(f"  ⚠ Post-processing skipped (missing {e}). Copying base video.")
+            shutil.copy2(video_path, output_path)
+            return output_path
+        except Exception as e:
+            print(f"  ⚠ Post-processing error: {e}. Using base video.")
+            shutil.copy2(video_path, output_path)
+            return output_path
+
+    def _apply_color_grade(
+        self, frame: np.ndarray, emotion: str, intensity: float
+    ) -> np.ndarray:
+        """Apply subtle emotion-specific color grading."""
+        import cv2
+
+        # Very subtle color shifts based on emotion
+        color_shifts = {
+            "happy": (5, 5, 15),      # Warm (slight yellow)
+            "sad": (-5, -3, -10),      # Cool (slight blue)
+            "angry": (10, -5, -5),     # Warm red
+            "fear": (-5, -5, 5),       # Cool green
+            "surprise": (5, 5, 5),     # Bright
+            "disgust": (-3, 5, -5),    # Sickly green
+            "neutral": (0, 0, 0),
+        }
+
+        shift = color_shifts.get(emotion, (0, 0, 0))
+        scale = intensity * 0.5  # Keep it very subtle
+
+        adjusted = frame.astype(np.float32)
+        adjusted[:, :, 0] += shift[0] * scale  # B
+        adjusted[:, :, 1] += shift[1] * scale  # G
+        adjusted[:, :, 2] += shift[2] * scale  # R
+
+        return np.clip(adjusted, 0, 255).astype(np.uint8)
+
+    def _apply_face_warp(
+        self,
+        frame: np.ndarray,
+        face_mesh,
+        emotion: str,
+        intensity: float
+    ) -> np.ndarray:
+        """
+        Apply subtle facial warping based on emotion.
+        Uses MediaPipe landmarks to create emotion-specific deformations.
+        """
+        import cv2
+
+        h, w = frame.shape[:2]
+        rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        results = face_mesh.process(rgb)
+
+        if not results.multi_face_landmarks:
+            return frame
+
+        landmarks = results.multi_face_landmarks[0]
+
+        # Key landmark indices for warping
+        # Brow: 70, 63, 105, 66, 107 (left), 336, 296, 334, 293, 300 (right)
+        # Mouth corners: 61, 291
+        # Jaw: 152
+
+        profile = EMOTION_PROFILES.get(emotion, {})
+        brow_shift = profile.get("brow_scale", 0) * intensity * 3  # pixels
+        mouth_shift = profile.get("mouth_scale", 0) * intensity * 2
+
+        if abs(brow_shift) < 0.5 and abs(mouth_shift) < 0.5:
+            return frame  # Not enough to notice
+
+        # Simple approach: use cv2.remap with subtle displacement
+        # This is fast and produces decent results
+        map_x = np.tile(np.arange(w, dtype=np.float32), (h, 1))
+        map_y = np.tile(np.arange(h, dtype=np.float32).reshape(-1, 1), (1, w))
+
+        # Get face center and brow/mouth regions
+        face_pts = [(int(l.x * w), int(l.y * h)) for l in landmarks.landmark]
+
+        # Brow region (top 1/3 of face)
+        brow_y = face_pts[10][1]  # Top of face
+        nose_y = face_pts[1][1]   # Nose tip
+        brow_region = (brow_y, nose_y)
+
+        # Apply brow displacement in brow region
+        for y_idx in range(max(0, brow_region[0]), min(h, brow_region[1])):
+            # Gaussian falloff from center of region
+            region_center = (brow_region[0] + brow_region[1]) // 2
+            dist = abs(y_idx - region_center) / max(1, (brow_region[1] - brow_region[0]) / 2)
+            falloff = np.exp(-dist ** 2 * 2)
+            map_y[y_idx, :] -= brow_shift * falloff
+
+        # Apply mouth displacement in lower face
+        mouth_y = face_pts[13][1]  # Upper lip
+        chin_y = face_pts[152][1]  # Chin
+        mouth_center_x = (face_pts[61][0] + face_pts[291][0]) // 2
+
+        for y_idx in range(max(0, mouth_y - 10), min(h, chin_y + 10)):
+            for x_idx in range(max(0, mouth_center_x - 40), min(w, mouth_center_x + 40)):
+                dist_y = abs(y_idx - mouth_y) / max(1, (chin_y - mouth_y))
+                dist_x = abs(x_idx - mouth_center_x) / 40.0
+                falloff = np.exp(-(dist_y ** 2 + dist_x ** 2) * 2)
+                map_x[y_idx, x_idx] += mouth_shift * falloff * (1 if x_idx > mouth_center_x else -1)
+
+        warped = cv2.remap(frame, map_x, map_y, cv2.INTER_LINEAR)
+        return warped
+
+    def _mux_audio(self, video_path: str, audio_source: str, output_path: str):
+        """Combine processed video with original audio."""
+        try:
+            subprocess.run([
+                "ffmpeg", "-y",
+                "-i", video_path,
+                "-i", audio_source,
+                "-c:v", "copy",
+                "-c:a", "aac",
+                "-map", "0:v:0",
+                "-map", "1:a:0",
+                "-shortest",
+                output_path
+            ], capture_output=True, timeout=60)
+        except Exception:
+            # If ffmpeg fails, just use the video without audio
+            shutil.copy2(video_path, output_path)
+
+    def generate(
+        self,
+        audio_path: str,
+        image_path: str,
+        emotion: Optional[str] = None,
+        intensity: Optional[float] = None,
+        output_path: str = "output.mp4",
+        expression_scale: float = 1.0,
+        still_mode: bool = False,
+        preprocess: str = "crop",
+        size: int = 256
+    ) -> Dict:
+        """
+        Full EMOLIPS generation pipeline.
+
+        Args:
+            audio_path: Path to speech audio file
+            image_path: Path to source face image
+            emotion: Target emotion (auto-detected if None)
+            intensity: Emotion intensity 0-1 (auto-estimated if None)
+            output_path: Where to save result
+            expression_scale: SadTalker expression scale
+            still_mode: Reduce head motion
+            preprocess: SadTalker preprocess mode
+            size: Output resolution
+
+        Returns:
+            Dict with generation metadata
+        """
+        print("=" * 50)
+        print("  EMOLIPS: Emotion-Driven Lip-Sync Generation")
+        print("=" * 50)
+
+        # Validate inputs
+        assert os.path.exists(audio_path), f"Audio not found: {audio_path}"
+        assert os.path.exists(image_path), f"Image not found: {image_path}"
+
+        result_meta = {
+            "audio": audio_path,
+            "image": image_path,
+            "output": output_path,
+        }
+
+        # Step 1: Emotion detection
+        if emotion is None:
+            detection = self.detect_emotion(audio_path)
+            emotion = detection["detected_emotion"]
+            result_meta["emotion_detection"] = detection
+        else:
+            print(f"  [1/4] Using specified emotion: {emotion}")
+            result_meta["emotion_detection"] = {"manual": emotion}
+
+        # Step 2: Intensity estimation
+        if intensity is None:
+            intensity = self.estimate_intensity(audio_path)
+        else:
+            print(f"        Using specified intensity: {intensity}")
+        result_meta["emotion"] = emotion
+        result_meta["intensity"] = intensity
+
+        # Adjust SadTalker expression scale based on emotion
+        emotion_expression_map = {
+            "neutral": 1.0,
+            "happy": 1.3,
+            "sad": 0.9,
+            "angry": 1.4,
+            "fear": 1.2,
+            "surprise": 1.5,
+            "disgust": 1.1
+        }
+        adjusted_scale = expression_scale * emotion_expression_map.get(emotion, 1.0) * (0.5 + 0.5 * intensity)
+
+        # Step 3: Run SadTalker
+        temp_dir = os.path.join(os.path.dirname(output_path) or ".", "temp_sadtalker")
+        os.makedirs(temp_dir, exist_ok=True)
+
+        base_video = self.run_sadtalker(
+            audio_path=audio_path,
+            image_path=image_path,
+            output_dir=temp_dir,
+            expression_scale=adjusted_scale,
+            still_mode=still_mode,
+            preprocess=preprocess,
+            size=size
+        )
+
+        if base_video is None:
+            print("  ✗ SadTalker generation failed!")
+            result_meta["success"] = False
+            return result_meta
+
+        print(f"        Base video: {base_video}")
+        result_meta["base_video"] = base_video
+
+        # Step 4: Apply emotion post-processing
+        final_video = self.apply_emotion_postprocess(
+            video_path=base_video,
+            emotion=emotion,
+            intensity=intensity,
+            output_path=output_path
+        )
+
+        result_meta["output"] = final_video
+        result_meta["success"] = True
+
+        print(f"\n  [4/4] Generation complete!")
+        print(f"        Output: {final_video}")
+        print(f"        Emotion: {emotion} (intensity: {intensity:.2f})")
+        print("=" * 50)
+
+        # Save metadata
+        meta_path = output_path.replace(".mp4", "_meta.json")
+        with open(meta_path, "w") as f:
+            json.dump(result_meta, f, indent=2, default=str)
+
+        return result_meta
+
+    def generate_all_emotions(
+        self,
+        audio_path: str,
+        image_path: str,
+        output_dir: str = "outputs",
+        intensity: float = 0.7,
+        **kwargs
+    ) -> List[Dict]:
+        """
+        Generate same audio+image across all 7 emotions.
+        This is the key demo for showing emotion conditioning works.
+        """
+        os.makedirs(output_dir, exist_ok=True)
+        results = []
+
+        emotions = ["neutral", "happy", "sad", "angry", "fear", "surprise", "disgust"]
+
+        for emotion in emotions:
+            print(f"\n{'='*50}")
+            print(f"  Generating: {emotion.upper()}")
+            print(f"{'='*50}")
+
+            out_path = os.path.join(output_dir, f"emolips_{emotion}.mp4")
+
+            result = self.generate(
+                audio_path=audio_path,
+                image_path=image_path,
+                emotion=emotion,
+                intensity=intensity,
+                output_path=out_path,
+                **kwargs
+            )
+            results.append(result)
+
+        # Create comparison grid
+        self._create_comparison_grid(output_dir, emotions)
+
+        return results
+
+    def _create_comparison_grid(self, output_dir: str, emotions: List[str]):
+        """Create side-by-side comparison video."""
+        try:
+            videos = []
+            for emotion in emotions:
+                path = os.path.join(output_dir, f"emolips_{emotion}.mp4")
+                if os.path.exists(path):
+                    videos.append(path)
+
+            if len(videos) < 2:
+                return
+
+            # Use ffmpeg to create grid
+            # 4 videos in a row, 2 rows
+            filter_parts = []
+            inputs = []
+            for i, v in enumerate(videos[:8]):  # Max 8
+                inputs.extend(["-i", v])
+                filter_parts.append(f"[{i}:v]scale=256:256[v{i}]")
+
+            n = len(videos[:8])
+            cols = min(4, n)
+            rows = (n + cols - 1) // cols
+
+            # Build xstack filter
+            layout_parts = []
+            for i in range(min(n, 8)):
+                x = (i % cols) * 256
+                y = (i // cols) * 256
+                layout_parts.append(f"{x}_{y}")
+
+            inputs_str = "".join(f"[v{i}]" for i in range(min(n, 8)))
+            filter_str = ";".join(filter_parts) + f";{inputs_str}xstack=inputs={min(n,8)}:layout={'|'.join(layout_parts)}"
+
+            grid_path = os.path.join(output_dir, "comparison_grid.mp4")
+
+            subprocess.run(
+                ["ffmpeg", "-y"] + inputs + [
+                    "-filter_complex", filter_str,
+                    "-c:v", "libx264",
+                    "-crf", "23",
+                    grid_path
+                ],
+                capture_output=True,
+                timeout=120
+            )
+
+            if os.path.exists(grid_path):
+                print(f"\n  ✓ Comparison grid: {grid_path}")
+
+        except Exception as e:
+            print(f"  ⚠ Could not create comparison grid: {e}")
+
+
+# ============================================================
+# STANDALONE MODE (without SadTalker, for testing pipeline)
+# ============================================================
+
+class EmolipsStandalone:
+    """
+    Standalone mode that works WITHOUT SadTalker.
+    Uses MediaPipe face mesh + direct warping for quick demo.
+
+    Good for:
+    - Testing the emotion module independently
+    - Quick demos without full SadTalker setup
+    - Verifying the pipeline logic
+    """
+
+    def __init__(self):
+        self.emotion_detector = AudioEmotionDetector(device="cpu")
+        self.intensity_estimator = EmotionIntensityEstimator()
+        self.emotion_modifier = PracticalEmotionModifier()
+
+    def generate_emotion_frames(
+        self,
+        image_path: str,
+        emotion: str,
+        intensity: float = 0.7,
+        num_frames: int = 30
+    ) -> List[np.ndarray]:
+        """
+        Generate emotion-modified face frames from a single image.
+        No audio needed - just shows the emotion transformation.
+        """
+        import cv2
+        import mediapipe as mp
+
+        img = cv2.imread(image_path)
+        if img is None:
+            raise ValueError(f"Could not read image: {image_path}")
+
+        mp_face_mesh = mp.solutions.face_mesh
+        face_mesh = mp_face_mesh.FaceMesh(static_image_mode=True, max_num_faces=1)
+
+        frames = []
+        for i in range(num_frames):
+            # Gradual emotion onset
+            t = min(1.0, i / (num_frames * 0.3))  # Ramp up in first 30%
+            current_intensity = intensity * t
+
+            frame = img.copy()
+
+            # Apply warping
+            if current_intensity > 0.1:
+                h, w = frame.shape[:2]
+                rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                results = face_mesh.process(rgb)
+
+                if results.multi_face_landmarks:
+                    profile = EMOTION_PROFILES.get(emotion, {})
+                    brow_shift = profile.get("brow_scale", 0) * current_intensity * 5
+                    mouth_shift = profile.get("mouth_scale", 0) * current_intensity * 4
+
+                    if abs(brow_shift) > 0.3 or abs(mouth_shift) > 0.3:
+                        map_x = np.tile(np.arange(w, dtype=np.float32), (h, 1))
+                        map_y = np.tile(np.arange(h, dtype=np.float32).reshape(-1, 1), (1, w))
+
+                        face_pts = [(int(l.x * w), int(l.y * h))
+                                   for l in results.multi_face_landmarks[0].landmark]
+
+                        brow_y = face_pts[10][1]
+                        nose_y = face_pts[1][1]
+
+                        for y_idx in range(max(0, brow_y), min(h, nose_y)):
+                            center = (brow_y + nose_y) // 2
+                            dist = abs(y_idx - center) / max(1, (nose_y - brow_y) / 2)
+                            falloff = np.exp(-dist ** 2 * 2)
+                            map_y[y_idx, :] -= brow_shift * falloff
+
+                        frame = cv2.remap(frame, map_x, map_y, cv2.INTER_LINEAR)
+
+            # Apply color grading
+            color_shifts = {
+                "happy": (5, 5, 15), "sad": (-5, -3, -10),
+                "angry": (10, -5, -5), "fear": (-5, -5, 5),
+                "surprise": (5, 5, 5), "disgust": (-3, 5, -5),
+                "neutral": (0, 0, 0)
+            }
+            shift = color_shifts.get(emotion, (0, 0, 0))
+            adjusted = frame.astype(np.float32)
+            for c in range(3):
+                adjusted[:, :, c] += shift[c] * current_intensity * 0.5
+            frame = np.clip(adjusted, 0, 255).astype(np.uint8)
+
+            frames.append(frame)
+
+        face_mesh.close()
+        return frames
+
+    def save_demo_video(
+        self,
+        image_path: str,
+        emotions: List[str] = None,
+        output_dir: str = "outputs",
+        fps: int = 30,
+        duration: float = 2.0
+    ):
+        """Save emotion demo videos from a single face image."""
+        import cv2
+
+        os.makedirs(output_dir, exist_ok=True)
+
+        if emotions is None:
+            emotions = ["neutral", "happy", "sad", "angry", "fear", "surprise", "disgust"]
+
+        num_frames = int(fps * duration)
+
+        for emotion in emotions:
+            print(f"  Generating {emotion}...")
+            frames = self.generate_emotion_frames(image_path, emotion, 0.7, num_frames)
+
+            out_path = os.path.join(output_dir, f"demo_{emotion}.mp4")
+            h, w = frames[0].shape[:2]
+            out = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
+            for f in frames:
+                out.write(f)
+            out.release()
+            print(f"    ✓ {out_path}")
+
+
+if __name__ == "__main__":
+    print("EMOLIPS Pipeline module loaded.")
+    print("Use EmolipsPipeline for full generation or EmolipsStandalone for quick demo.")