Update idle_generator.py

idle_generator.py

@@ -12,9 +12,6 @@ import logging
 import torch
 import cv2
 import numpy as np
-import albumentations as A
-import albumentations.pytorch.transforms as A_pytorch
-import face_alignment
 
 logger = logging.getLogger(__name__)
 
@@ -45,65 +42,42 @@ CLIP_MODES = [
 class IdleVideoGenerator:
     """Generates idle animation clips using FLOAT model with scripted pose control."""
 
-    def __init__(self, float_model, device, input_size=512):
+    def __init__(self, float_model, device, input_size=512, lipsync_instance=None):
         """
         Args:
             float_model: Loaded FLOAT model instance (already on device, eval mode)
             device: torch device
             input_size: image size (default 512)
+            lipsync_instance: FloatLipsync instance for reusing its face crop logic
         """
         self.model = float_model
         self.device = device
         self.input_size = input_size
-
-        # Face alignment for cropping
-        self.fa = face_alignment.FaceAlignment(
-            face_alignment.LandmarksType.TWO_D, flip_input=False, device=str(device)
-        )
-
-        # Image transform
-        self.transform = A.Compose([
-            A.Resize(height=input_size, width=input_size, interpolation=cv2.INTER_AREA),
-            A.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
-            A_pytorch.ToTensorV2(),
-        ])
+        self.lipsync = lipsync_instance
 
         # Store original sample function for monkey-patching
         self._original_sample = self.model.sample
 
-    def _crop_face(self, img):
-        """Detect and crop face from image."""
-        mult = 360.0 / img.shape[0]
-        interp = cv2.INTER_AREA if mult < 1.0 else cv2.INTER_CUBIC
-        resized = cv2.resize(img, (0, 0), fx=mult, fy=mult, interpolation=interp)
-
-        bboxes = self.fa.face_detector.detect_from_image(resized)
-        if not bboxes:
-            # No face found — just resize
-            return cv2.resize(img, (self.input_size, self.input_size))
-
-        x1, y1, x2, y2, score = bboxes[0]
-        mx = int((x1 + x2) / 2 / mult)
-        my = int((y1 + y2) / 2 / mult)
-        bs = int(max(y2 - y1, x2 - x1) / mult * 0.8)
-
-        img = cv2.copyMakeBorder(img, bs, bs, bs, bs, cv2.BORDER_REFLECT_101)
-        my += bs
-        mx += bs
-        crop = img[my:my + 2 * bs, mx:mx + 2 * bs]
-        return cv2.resize(crop, (self.input_size, self.input_size))
-
-    def _load_reference_image(self, image_path, no_crop=False):
-        """Load and preprocess a reference image."""
-        img = cv2.imread(image_path)
-        if img is None:
-            raise FileNotFoundError(f"Could not read image: {image_path}")
-        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-
-        if not no_crop:
-            img = self._crop_face(img)
-
-        return self.transform(image=img)['image'].unsqueeze(0).to(self.device)
+    def _load_reference_image(self, image_path):
+        """Load and preprocess a reference image using float_lipsync's proven crop logic."""
+        if self.lipsync:
+            # Use the same crop/transform as float_lipsync for consistency
+            self.lipsync._preload_reference_image(image_path)
+            return self.lipsync.preprocessed_ref_image.clone()
+        else:
+            # Fallback: basic resize without face crop
+            import albumentations as A
+            import albumentations.pytorch.transforms as A_pytorch
+            img = cv2.imread(image_path)
+            if img is None:
+                raise FileNotFoundError(f"Could not read image: {image_path}")
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            transform = A.Compose([
+                A.Resize(height=self.input_size, width=self.input_size, interpolation=cv2.INTER_AREA),
+                A.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
+                A_pytorch.ToTensorV2(),
+            ])
+            return transform(image=img)['image'].unsqueeze(0).to(self.device)
 
     def generate(self, ref_image_path, output_dir, avatar_name,
                  num_clips=NUM_CLIPS, progress_callback=None):
@@ -271,5 +245,6 @@ def get_idle_generator():
             float_model=lipsync.model,
             device=lipsync.device,
             input_size=lipsync.opt.input_size,
+            lipsync_instance=lipsync,
         )
     return _generator