From 2e4e530b7849b7be286fc2ef1e69022e52db646a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafa=C5=82=20T=C5=82uchowski?= <rafal@tluchowski.pl>
Date: Sun, 23 Feb 2025 19:34:27 +0100
Subject: [PATCH] fp_pipe normalization across different workflows

---
 nodes/face_fit_and_restore.py | 356 +++++++++++++++++++---------------
 nodes/face_wrapper.py         | 103 ++++++++--
 2 files changed, 288 insertions(+), 171 deletions(-)

diff --git a/nodes/face_fit_and_restore.py b/nodes/face_fit_and_restore.py
index c289b6b..f5c681c 100644
--- a/nodes/face_fit_and_restore.py
+++ b/nodes/face_fit_and_restore.py
@@ -8,6 +8,8 @@ from ..core.image_processor import ImageProcessor
 
 
 class FaceFitAndRestore:
+    """ComfyUI node for processing face images in Fit or Restore mode."""
+
     def __init__(self):
         self.face_detector = FaceDetector()
         self.image_processor = ImageProcessor()
@@ -50,245 +52,293 @@ class FaceFitAndRestore:
 
     def process_image(self, mode, workflow, padding_percent=0.0, bbox_size="1024",
                       image=None, fp_pipe=None, image_sequence=None):
+        """Process images in either Fit or Restore mode with unified fp_pipe structure."""
 
+        # Validate inputs
+        if image is None and (workflow == "single_image" or mode == "Restore"):
+            print("Error: Image is required for single_image workflow and Restore mode")
+            return None, fp_pipe, None, int(bbox_size)
 
-        # Check if we should use workflow from fp_pipe in Restore mode
-        if mode == "Restore":
-            if fp_pipe is None:
-                print("Error: fp_pipe is required in Restore mode")
-                return None, {}, None, int(bbox_size)
+        if workflow == "image_sequence" and mode == "Fit" and (not image_sequence or "frames" not in image_sequence):
+            print("Error: Valid image sequence data is required for Fit mode")
+            return None, fp_pipe, None, int(bbox_size)
 
-            if "workflow" in fp_pipe:
-                workflow = fp_pipe["workflow"]
-                print(f"Using workflow from pipe: {workflow}")
+        if mode == "Restore" and (not fp_pipe or "frames" not in fp_pipe):
+            print("Error: Valid fp_pipe with frames is required for Restore mode")
+            return None, fp_pipe, None, int(bbox_size)
 
-            # Additional check for sequence restore
-            if workflow == "image_sequence" and "frames" not in fp_pipe:
-                print("Error: fp_pipe does not contain frames data required for image sequence restore")
-                return None, {}, None, int(bbox_size)
+        # Initialize fp_pipe if None
+        if fp_pipe is None:
+            fp_pipe = {
+                "workflow": workflow,
+                "current_frame": 0,
+                "padding_percent": padding_percent,
+                "target_lm": {},  # Will be populated by face_wrapper
+                "frames": {}
+            }
+        else:
+            # Update common parameters
+            fp_pipe["workflow"] = workflow
+            fp_pipe["padding_percent"] = padding_percent
 
-        if workflow == "single_image":
-            if image is None:
-                print("Error: Image is required for single_image workflow")
-                return None, {}, None, int(bbox_size)
+        try:
+            if workflow == "single_image":
+                return self._process_frame(
+                    mode=mode,
+                    image=image,
+                    frame_number=0,
+                    padding_percent=padding_percent,
+                    bbox_size=bbox_size,
+                    fp_pipe=fp_pipe
+                )
+            else:  # image_sequence
+                # W trybie Restore używamy current_frame z fp_pipe
+                current_frame = (fp_pipe.get("current_frame", 0) if mode == "Restore"
+                                 else image_sequence.get("current_frame", 0))
+                fp_pipe["current_frame"] = current_frame
 
-            if mode == "Fit":
-                result = self._fit(image, padding_percent, bbox_size)
-            else:  # Restore
-                result = self._restore(image, fp_pipe)
+                # W trybie Restore używamy liczby klatek z fp_pipe
+                if mode == "Restore":
+                    total_frames = len(fp_pipe["frames"])
+                    frames_to_process = {
+                        idx: None for idx in range(total_frames)
+                    }
+                else:  # Fit mode
+                    total_frames = len(image_sequence["frames"])
+                    frames_to_process = image_sequence["frames"]
 
-            # Add workflow info to settings
-            result_settings = result[1]
-            result_settings["workflow"] = "single_image"
-            return result
-
-        elif workflow == "image_sequence":
-            # In Fit mode, we need image_sequence input
-            if mode == "Fit":
-                if image_sequence is None or not image_sequence.get("frames"):
-                    print("Error: Valid image_sequence is required for image_sequence workflow in Fit mode")
-                    return None, {}, None, int(bbox_size)
-                sequence_settings = {
-                    "workflow": "image_sequence",
-                    "frames": {},
-                    "current_frame": image_sequence.get("current_frame", 0)
-                }
-
-                selected_result = None
-                total_frames = len(image_sequence["frames"])
                 print(f"Processing sequence of {total_frames} frames...")
-                print(f"Current frame index: {sequence_settings['current_frame']}")
-                for frame_idx, frame_path in image_sequence["frames"].items():
+                current_frame_result = None
+
+                # Process each frame
+                for frame_idx in range(total_frames):
                     print(f"Processing frame {frame_idx + 1}/{total_frames}")
-                    try:
-                        frame_tensor = self._load_image_from_path(frame_path)
 
-                        if frame_tensor is None:
-                            print(f"Warning: Could not load frame {frame_idx} from {frame_path}")
+                    # Przygotowanie obrazu wejściowego
+                    if mode == "Fit":
+                        frame_path = frames_to_process[frame_idx]
+                        frame_image = self._load_image_from_path(frame_path)
+                        if frame_image is None:
+                            print(f"Warning: Could not load frame {frame_idx}")
                             continue
+                    else:  # Restore
+                        frame_image = image  # Ten sam obraz dla każdej klatki w trybie Restore
+                        frame_path = None
 
-                        result = self._fit(frame_tensor, padding_percent, bbox_size)
+                    # Process frame
+                    result = self._process_frame(
+                        mode=mode,
+                        image=frame_image,
+                        frame_number=frame_idx,
+                        padding_percent=padding_percent,
+                        bbox_size=bbox_size,
+                        fp_pipe=fp_pipe,
+                        original_path=frame_path if mode == "Fit" else None
+                    )
 
-                        if result:  # Check if result was created successfully
-                            result[1]["original_image_path"] = frame_path
-                            sequence_settings["frames"][f"frame_{frame_idx}"] = result[1]
-                            if frame_idx == sequence_settings["current_frame"]:
-                                selected_result = result
+                    # Store results for current frame
+                    if frame_idx == current_frame:
+                        current_frame_result = result
 
-                    except Exception as e:
-                        print(f"Error processing frame {frame_idx}: {str(e)}")
-                        continue
+                # After processing all frames, return results for current frame
+                if current_frame_result is None:
+                    print(f"Warning: Frame {current_frame} not found in sequence")
+                    return None, fp_pipe, None, int(bbox_size)
 
+                return current_frame_result
+
+        except Exception as e:
+            print(f"Error in process_image: {str(e)}")
+            import traceback
+            traceback.print_exc()
+            return None, fp_pipe, None, int(bbox_size)
+
+    def _process_frame(self, mode, image, frame_number, padding_percent, bbox_size, fp_pipe, original_path=None):
+        """Process a single frame in either Fit or Restore mode."""
+        try:
+            frame_key = f"frame_{frame_number}"
+
+            if mode == "Fit":
+                # Process frame in Fit mode
+                result_image, frame_settings, result_mask = self._fit(
+                    image=image,
+                    padding_percent=padding_percent,
+                    bbox_size=bbox_size
+                )
+
+                if frame_settings:
+                    # Create standardized frame data structure
+                    frame_data = {
+                        "bbox_size": int(bbox_size),
+                        "crop_bbox": frame_settings["crop_bbox"],
+                        "original_image_shape": frame_settings["original_image_shape"],
+                        "rotation_angle": frame_settings["rotation_angle"],
+                        "detected_lm": {}  # Will be populated by face_wrapper
+                    }
+
+                    # Add original path for sequence processing
+                    if original_path:
+                        frame_data["original_image_path"] = original_path
+
+                    # Store frame data
+                    fp_pipe["frames"][frame_key] = frame_data
 
             else:  # Restore mode
-                if image is None:
-                    print("Error: Input image is required for Restore mode")
-                    return None, {}, None, int(bbox_size)
-
-                current_frame = fp_pipe.get("current_frame", 0)
-                frame_settings = fp_pipe.get("frames", {}).get(f"frame_{current_frame}")
-
+                frame_settings = fp_pipe["frames"].get(frame_key)
                 if frame_settings is None:
-                    print(f"Error: No processor settings for frame {current_frame}")
-                    return None, {}, None, int(bbox_size)
-                print(f"Processing restore for frame {current_frame} with settings: {frame_settings}")
-                result = self._restore(image, frame_settings)
-
-                if result is None:
-                    print(f"Error: Failed to restore frame {current_frame}")
+                    print(f"Error: No settings found for frame {frame_number}")
                     return None, fp_pipe, None, int(bbox_size)
-                return result[0], fp_pipe, result[2], int(bbox_size)
 
-            if selected_result is None:
-                print(f"Warning: Could not process current frame, using default empty result")
-                return None, sequence_settings, None, int(bbox_size)
+                result = self._restore(image, frame_settings)
+                if result is None:
+                    return None, fp_pipe, None, int(bbox_size)
 
-            return selected_result[0], sequence_settings, selected_result[2], int(bbox_size)
+                result_image, result_mask = result[0], result[2]
+
+            return result_image, fp_pipe, result_mask, int(bbox_size)
+
+        except Exception as e:
+            print(f"Error processing frame {frame_number}: {str(e)}")
+            import traceback
+            traceback.print_exc()
+            return None, fp_pipe, None, int(bbox_size)
 
     def _fit(self, image, padding_percent, bbox_size):
         """Fit mode: Crop and process the face."""
+        # Convert image to numpy array
         image_np = self.image_processor.convert_to_numpy(image)
         if image_np is None:
-            return image, {}, self._create_empty_mask(image), int(bbox_size)
+            return image, {}, self._create_empty_mask(image)
 
-        # Detect facial landmarks using mediapipe
+        # Detect facial landmarks
         landmarks_df = self.face_detector.detect_landmarks_mp(image_np)
         if landmarks_df is None:
             print("No face detected, returning original image")
-            return image, {}, self._create_empty_mask(image), int(bbox_size)
+            return image, {}, self._create_empty_mask(image)
 
-        # Calculate rotation angle and rotate the image
+        # Calculate rotation and rotate image
         rotation_angle = self.image_processor.calculate_rotation_angle(landmarks_df)
         rotated_image, updated_landmarks = self.image_processor.rotate_image(image_np, landmarks_df)
         if rotated_image is None:
             print("Failed to rotate image, returning original")
-            return image, {}, self._create_empty_mask(image), int(bbox_size)
+            return image, {}, self._create_empty_mask(image)
 
-        # Crop face region to a square (1:1)
-        cropped_face, crop_bbox = self.image_processor.crop_face_to_square(rotated_image, updated_landmarks,
-                                                                           padding_percent)
+        # Crop face region
+        cropped_face, crop_bbox = self.image_processor.crop_face_to_square(
+            rotated_image, updated_landmarks, padding_percent
+        )
         if cropped_face is None:
             print("Failed to crop face, returning original image")
-            return image, {}, self._create_empty_mask(image), int(bbox_size)
+            return image, {}, self._create_empty_mask(image)
 
         # Resize to target size
         target_size = int(bbox_size)
         final_image = self.image_processor.resize_image(cropped_face, target_size)
         if final_image is None:
             print("Failed to resize image, returning original")
-            return image, {}, self._create_empty_mask(image), int(bbox_size)
+            return image, {}, self._create_empty_mask(image)
 
-        # Convert to tensor format
+        # Convert to tensor
         final_image = self.image_processor.numpy_to_tensor(final_image)
 
-        # Save processor settings for restoration
-        fp_pipe = {
+        # Create frame settings
+        frame_settings = {
             "original_image_shape": image_np.shape,
             "rotation_angle": rotation_angle,
-            "crop_bbox": crop_bbox,  # (x, y, w, h)
-            "padding_percent": padding_percent,
-            "bbox_size": target_size,
-            "original_image_path": None
+            "crop_bbox": crop_bbox,
+            "bbox_size": target_size
         }
 
-        # Create a mask for the face area
-        mask = np.ones((target_size, target_size), dtype=np.float32)
+        # Create mask
+        mask = self._create_mask(
+            image_shape=frame_settings["original_image_shape"],
+            crop_bbox=frame_settings["crop_bbox"],
+            rotation_angle=frame_settings["rotation_angle"]
+        )
 
-        if rotation_angle != 0:
-            original_mask = np.ones(image_np.shape[:2], dtype=np.float32)
-            height, width = image_np.shape[:2]
-            center = (width // 2, height // 2)
-            rotation_matrix = cv2.getRotationMatrix2D(center, rotation_angle, 1.0)
-            rotated_mask = cv2.warpAffine(original_mask, rotation_matrix, (width, height), flags=cv2.INTER_LINEAR)
-
-            x1, y1, w, h = crop_bbox
-            cropped_mask = rotated_mask[y1:y1 + h, x1:x1 + w]
-            resized_mask = cv2.resize(cropped_mask, (target_size, target_size), interpolation=cv2.INTER_LINEAR)
-            mask = resized_mask
-
-        mask = self.image_processor.convert_mask_to_tensor(mask)
-        return final_image, fp_pipe, mask, int(bbox_size)
-
-    def _restore(self, image, fp_pipe):
-        """Restore mode: Restore the face to the original image."""
-
-        if fp_pipe is None:
-            print("Error in restore: fp_pipe is None")
-            return image, {}, self._create_empty_mask(image)
-
-        processed_face_np = self.image_processor.convert_to_numpy(image)
-        if processed_face_np is None:
-            print("Error in restore: Failed to convert image to numpy array")
-            return image, {}, self._create_empty_mask(image)
-
-        original_image_shape = fp_pipe.get("original_image_shape")
-        rotation_angle = fp_pipe.get("rotation_angle")
-        crop_bbox = fp_pipe.get("crop_bbox")
-        padding_percent = fp_pipe.get("padding_percent")
-        bbox_size = fp_pipe.get("bbox_size")
-
-        # Debug info
-        print(f"Restore settings: shape={original_image_shape}, rotation={rotation_angle}, bbox={crop_bbox}")
-
-        if not all([original_image_shape, crop_bbox]):
-            missing = []
-            if not original_image_shape:
-                missing.append("original_image_shape")
-            if not crop_bbox:
-                missing.append("crop_bbox")
-            print(f"Error in restore: Missing required settings: {', '.join(missing)}")
-            return image, {}, self._create_empty_mask(image)
+        return final_image, frame_settings, mask
 
+    def _restore(self, image, frame_settings):
+        """Restore mode: Restore the face back to original position."""
         try:
-            restored_image = np.zeros(original_image_shape, dtype=np.uint8)
+            # Validate frame settings
+            required_keys = ["original_image_shape", "rotation_angle", "crop_bbox"]
+            if not all(key in frame_settings for key in required_keys):
+                print("Error: Missing required frame settings")
+                return None
 
-            x1, y1, w, h = crop_bbox
+            # Convert input image to numpy
+            processed_face_np = self.image_processor.convert_to_numpy(image)
+            if processed_face_np is None:
+                return None
+
+            # Get settings
+            original_shape = frame_settings["original_image_shape"]
+            rotation_angle = frame_settings["rotation_angle"]
+            x1, y1, w, h = frame_settings["crop_bbox"]
+
+            # Create output image
+            restored_image = np.zeros(original_shape, dtype=np.uint8)
+
+            # Resize face to original crop size
             resized_face = cv2.resize(processed_face_np, (w, h), interpolation=cv2.INTER_LANCZOS4)
 
+            # Place face in position
             restored_image[y1:y1 + h, x1:x1 + w] = resized_face
 
+            # Rotate back if needed
             if rotation_angle != 0:
-                height, width = original_image_shape[:2]
+                height, width = original_shape[:2]
                 center = (width // 2, height // 2)
                 rotation_matrix = cv2.getRotationMatrix2D(center, -rotation_angle, 1.0)
-                restored_image = cv2.warpAffine(restored_image, rotation_matrix, (width, height),
-                                                flags=cv2.INTER_LANCZOS4)
+                restored_image = cv2.warpAffine(
+                    restored_image,
+                    rotation_matrix,
+                    (width, height),
+                    flags=cv2.INTER_LANCZOS4
+                )
 
-            # Convert to tensor format
+            # Convert to tensor
             restored_image = self.image_processor.numpy_to_tensor(restored_image)
-            mask = self._create_mask(original_image_shape, crop_bbox, rotation_angle)
 
-            return restored_image, fp_pipe, mask
+            # Create mask
+            mask = self._create_mask(original_shape, (x1, y1, w, h), rotation_angle)
+
+            return restored_image, frame_settings, mask
 
         except Exception as e:
             print(f"Error in restore: {str(e)}")
-            return image, fp_pipe, self._create_empty_mask(image)
+            return None
 
     def _create_mask(self, image_shape, crop_bbox, rotation_angle):
         """Create a mask for the face area."""
         mask = np.zeros(image_shape[:2], dtype=np.float32)
         x1, y1, w, h = crop_bbox
 
-        # Draw a white rectangle for the face area
+        # Create face area mask
         mask[y1:y1 + h, x1:x1 + w] = 1.0
 
-        # Rotate the mask back to the original orientation
+        # Rotate mask if needed
         if rotation_angle != 0:
             height, width = image_shape[:2]
             center = (width // 2, height // 2)
             rotation_matrix = cv2.getRotationMatrix2D(center, -rotation_angle, 1.0)
-            mask = cv2.warpAffine(mask, rotation_matrix, (width, height), flags=cv2.INTER_LINEAR)
+            mask = cv2.warpAffine(
+                mask,
+                rotation_matrix,
+                (width, height),
+                flags=cv2.INTER_LINEAR
+            )
 
         return self.image_processor.convert_mask_to_tensor(mask)
 
     def _create_empty_mask(self, image):
-        """Create an empty mask with the same shape as the input image."""
+        """Create an empty mask matching the image dimensions."""
         if torch.is_tensor(image):
-            image_shape = image.shape[1:3]  # (H, W)
+            shape = image.shape[1:3]
         else:
-            image_shape = image.shape[:2]  # (H, W)
-        mask = torch.zeros((1, *image_shape), dtype=torch.float32)
-        return mask
+            shape = image.shape[:2]
+        return torch.zeros((1, *shape), dtype=torch.float32)
 
     def _load_image_from_path(self, image_path):
         """Load and convert image from file path to tensor format."""
diff --git a/nodes/face_wrapper.py b/nodes/face_wrapper.py
index 38eaf00..c49a52e 100644
--- a/nodes/face_wrapper.py
+++ b/nodes/face_wrapper.py
@@ -90,6 +90,7 @@ class FaceWrapper:
             if det_overlay is not None:
                 overlays.append(det_overlay)
 
+        # Get base landmarks
         base_landmarks = MediapipeBaseLandmarks.get_base_landmarks(
             (width, height), x_scale=x_scale, y_translation=y_transform
         )
@@ -117,12 +118,16 @@ class FaceWrapper:
 
         output_image = self.image_processor.numpy_to_tensor(result_image)
         output_mask = self.image_processor.convert_mask_to_tensor(mask_np)
-        landmarks_data = self._prepare_landmarks_data(landmarks_df, base_landmarks)
 
-        return output_image, self._update_pipe(fp_pipe, landmarks_data), output_mask
+        # Prepare and update landmarks data
+        landmarks_data = self._prepare_landmarks_data(landmarks_df, base_landmarks)
+        updated_pipe = self._update_pipe(fp_pipe, landmarks_data)
+
+        return output_image, updated_pipe, output_mask
 
     def _unwrap_mode(self, image_np, landmarks_df, width, height, device,
                      x_scale, y_transform, fp_pipe, mask_np):
+        # Get base landmarks
         base_landmarks = MediapipeBaseLandmarks.get_base_landmarks(
             (width, height), x_scale=x_scale, y_translation=y_transform
         )
@@ -141,8 +146,11 @@ class FaceWrapper:
 
         output_image = self.image_processor.pil_to_tensor(warped_image)
 
+        # Prepare and update landmarks data
         landmarks_data = self._prepare_landmarks_data(landmarks_df, base_landmarks)
-        return output_image, self._update_pipe(fp_pipe, landmarks_data), warped_mask
+        updated_pipe = self._update_pipe(fp_pipe, landmarks_data)
+
+        return output_image, updated_pipe, warped_mask
 
     def _wrap_mode(self, image_np, landmarks_df, width, height, device,
                    x_scale, y_transform, fp_pipe, mask_np):
@@ -152,12 +160,25 @@ class FaceWrapper:
             output_mask = self.image_processor.convert_mask_to_tensor(mask_np)
             return output_image, {}, output_mask
 
+        current_frame = fp_pipe.get("current_frame", 0)
+        frame_key = f"frame_{current_frame}"
+
+        # Check if we have required data
+        if frame_key not in fp_pipe["frames"] or "detected_lm" not in fp_pipe["frames"][frame_key]:
+            print(f"No detected landmarks found for frame {current_frame}")
+            output_image = self.image_processor.numpy_to_tensor(image_np)
+            output_mask = self.image_processor.convert_mask_to_tensor(mask_np)
+            return output_image, fp_pipe, output_mask
+
+        # Get source landmarks from target_lm at root level
         source_x = fp_pipe['target_lm']['x']
         source_y = fp_pipe['target_lm']['y']
         source_landmarks = np.column_stack((source_x, source_y))[:468]
 
-        detected_x = fp_pipe['detected_lm']['x']
-        detected_y = fp_pipe['detected_lm']['y']
+        # Get target landmarks from detected_lm in current frame
+        frame_data = fp_pipe["frames"][frame_key]
+        detected_x = frame_data['detected_lm']['x']
+        detected_y = frame_data['detected_lm']['y']
         target_landmarks = np.column_stack((detected_x, detected_y))[:468]
 
         # Process image
@@ -183,18 +204,64 @@ class FaceWrapper:
             return CPUDeformer.warp_face(image, source_landmarks, target_landmarks)
 
     def _prepare_landmarks_data(self, detected_df, target_lm):
-        return {
-            'detected_lm': {
-                'x': detected_df['x'].tolist(),
-                'y': detected_df['y'].tolist(),
-                'indices': detected_df['index'].tolist()
-            },
-            'target_lm': {
-                'x': target_lm[:, 0].tolist(),
-                'y': target_lm[:, 1].tolist(),
-                'indices': list(range(len(target_lm)))
-            }
+        """
+        Prepare landmarks data in the correct format for fp_pipe
+
+        Args:
+            detected_df: DataFrame with detected landmarks
+            target_lm: Target landmarks array
+
+        Returns:
+            dict: Landmarks data in the format compatible with fp_pipe
+        """
+        # Prepare detected landmarks data
+        detected_data = {
+            'x': detected_df['x'].tolist(),
+            'y': detected_df['y'].tolist(),
+            'indices': detected_df['index'].tolist()
         }
 
-    def _update_pipe(self, pipe, new_data):
-        return {**(pipe or {}), **new_data}
+        # Prepare target landmarks data
+        target_data = {
+            'x': target_lm[:, 0].tolist(),
+            'y': target_lm[:, 1].tolist(),
+            'indices': list(range(len(target_lm)))
+        }
+
+        return detected_data, target_data
+
+    def _update_pipe(self, pipe, landmarks_data):
+        """
+        Update fp_pipe with new landmarks data
+
+        Args:
+            pipe: Existing fp_pipe dictionary or None
+            landmarks_data: Tuple of (detected_landmarks, target_landmarks)
+
+        Returns:
+            dict: Updated fp_pipe structure
+        """
+        if pipe is None:
+            pipe = {
+                "workflow": "single_image",
+                "current_frame": 0,
+                "frames": {}
+            }
+
+        detected_lm, target_lm = landmarks_data
+
+        # Update target landmarks at the root level
+        pipe["target_lm"] = target_lm
+
+        # Update detected landmarks for the current frame
+        current_frame = pipe.get("current_frame", 0)
+        frame_key = f"frame_{current_frame}"
+
+        # Ensure the frame exists in the structure
+        if frame_key not in pipe["frames"]:
+            pipe["frames"][frame_key] = {}
+
+        # Update detected landmarks for the current frame
+        pipe["frames"][frame_key]["detected_lm"] = detected_lm
+
+        return pipe
\ No newline at end of file