From 6f9dba7777154215e372d86e70981e865ec68b00 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Mon, 8 Jul 2024 20:50:44 +0300
Subject: [PATCH] fixes

---
 liveportrait/live_portrait_pipeline.py |  9 +++++----
 nodes.py                               | 15 +++++++++++----
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/liveportrait/live_portrait_pipeline.py b/liveportrait/live_portrait_pipeline.py
index 4375364..20d17e0 100644
--- a/liveportrait/live_portrait_pipeline.py
+++ b/liveportrait/live_portrait_pipeline.py
@@ -70,7 +70,7 @@ class LivePortraitPipeline(object):
             ]
 
     def execute(
-        self, source_np, driving_images_np, mismatch_method="repeat", reference_frame=0
+        self, source_np, driving_images_np, crop_info, mismatch_method="repeat", reference_frame=0
     ):
         inference_cfg = self.live_portrait_wrapper.cfg
         is_video = source_np.shape[0] > 1
@@ -85,7 +85,8 @@ class LivePortraitPipeline(object):
         ref_frame = self._get_source_frame(
             source_np, reference_frame, total_frames, mismatch_method
         )
-        rcrop_info = self.cropper.crop_single_image(ref_frame)
+        #rcrop_info = self.cropper.crop_single_image(ref_frame)
+        rcrop_info = crop_info
         rsource_lmk = rcrop_info["lmk_crop"]
         rimg_crop, ref_crop_256x256 = (
             rcrop_info["img_crop"],
@@ -98,7 +99,7 @@ class LivePortraitPipeline(object):
             )
             driving_frame = driving_images_np[i]
 
-            crop_info = self.cropper.crop_single_image(source_frame_rgb)
+            crop_info, _ = self.cropper.crop_single_image(source_frame_rgb)
             source_lmk = crop_info["lmk_crop"]
             img_crop, img_crop_256x256 = (
                 crop_info["img_crop"],
@@ -125,7 +126,7 @@ class LivePortraitPipeline(object):
                 c_d_lip_before_animation = [0.0]
                 combined_lip_ratio_tensor_before_animation = (
                     self.live_portrait_wrapper.calc_combined_lip_ratio(
-                        c_d_lip_before_animation, source_lmk
+                        c_d_lip_before_animation, source_lmk, inference_cfg
                     )
                 )
                 # TODO: expose lip_zero_threshold
diff --git a/nodes.py b/nodes.py
index c5e9e53..7d6ef4b 100644
--- a/nodes.py
+++ b/nodes.py
@@ -316,11 +316,13 @@ class LivePortraitProcess:
         pipeline.live_portrait_wrapper.cfg.flag_relative = relative
         pipeline.live_portrait_wrapper.cfg.flag_lip_zero = lip_zero
 
+        pipeline.cropper = crop_info['cropper']
+
         cropped_out_list = []
         full_out_list = []
 
         cropped_out_list, full_out_list = pipeline.execute(
-            source_np, driving_images_np, crop_info, mismatch_method
+            source_np, driving_images_np, crop_info['crop_info'], mismatch_method
         )
         cropped_tensors_out = (
             torch.stack([torch.from_numpy(np_array) for np_array in cropped_out_list])
@@ -385,7 +387,12 @@ class LivePortraitCropper:
         
         print(cropped_tensors.shape)
 
-        return (cropped_tensors, crop_info, keypoints_image_tensor)
+        cropper_dict = {
+            "cropper": cropper,
+            "crop_info": crop_info,
+        }
+
+        return (cropped_tensors, cropper_dict, keypoints_image_tensor)
 
 class KeypointScaler:
     @classmethod
@@ -406,7 +413,7 @@ class KeypointScaler:
 
     def process(self, crop_info, offset_x, offset_y, scale):
 
-        keypoints = crop_info['lmk_crop'].copy()
+        keypoints = crop_info['crop_info']['lmk_crop'].copy()
 
         # Create an offset array
         # Calculate the centroid of the keypoints
@@ -421,7 +428,7 @@ class KeypointScaler:
         # Translate scaled keypoints back to original position and then apply the offset
         final_keypoints = scaled_keypoints + centroid + np.array([offset_x, offset_y])
 
-        crop_info['lmk_crop'] = final_keypoints
+        crop_info['crop_info']['lmk_crop'] = final_keypoints
 
         # Draw each landmark as a circle
         width, height = 512, 512