add CropReferenceImage Node

README.md

@@ -23,6 +23,10 @@
 
 This repository is the official implementation of the [`HelloMeme`](https://arxiv.org/pdf/2410.22901) ComfyUI interface, featuring both image and video generation functionalities. Example workflow files can be found in the `ComfyUI_HelloMeme/workflows` directory. Test images and videos are saved in the `ComfyUI_HelloMeme/examples` directory. Below are screenshots of the interfaces for image and video generation.
 
+## 🔆 New Features/Updates
+
+- ✅ `11/6/2024` The face proportion in the reference image significantly affects the generation quality. We have encapsulated the **recommended image cropping method** used during training into a `CropReferenceImage` Node. Refer to the workflows in the `ComfyUI_HelloMeme/workflows directory`: `hellomeme_video_cropref_workflow.json` and `hellomeme_image_cropref_workflow.json`.
+
 ### Image Generation Interface
 
 <p align="center">

meme.py

@@ -32,6 +32,7 @@ if missing:
 
 from .hellomeme.utils import (face_params_to_tensor,
                              get_drive_params,
+                             crop_and_resize,
                              get_face_params,
                              load_data_list,
                              load_unet_from_safetensors)
@@ -186,6 +187,35 @@ class GetReferenceImageRT:
         ref_rot, ref_trans = face_toolkits['h3dmm'].forward_params(image_np, ref_landmark)
         return (dict(rot=ref_rot, trans=ref_trans), )
 
+class CropReferenceImage:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "face_toolkits": ("FACE_TOOLKITS",),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "crop_reference_image"
+    CATEGORY = "hellomeme"
+
+    def crop_reference_image(self, image, face_toolkits):
+        image_np = cv2.cvtColor((image[0] * 255).cpu().numpy().astype(np.uint8), cv2.COLOR_BGR2RGB)
+        image_np = cv2.resize(image_np, (512, 512))
+        print(image_np.shape)
+        face_toolkits['face_aligner'].reset_track()
+        faces = face_toolkits['face_aligner'].forward(image_np)
+        assert len(faces) > 0
+        face = sorted(faces, key=lambda x: (x['face_rect'][2] - x['face_rect'][0]) * (
+                x['face_rect'][3] - x['face_rect'][1]))[-1]
+        ref_landmark = face['pre_kpt_222']
+
+        new_image = crop_and_resize(image_np[np.newaxis, :,:,:], ref_landmark[np.newaxis, :,:], 512, crop=True)[0]
+        new_image = cv2.cvtColor(new_image[0], cv2.COLOR_RGB2BGR)
+        return (torch.from_numpy(new_image[np.newaxis, :,:,:]).float() / 255., )
+
 class GetImageDriveParams:
     @classmethod
     def INPUT_TYPES(s):
@@ -371,6 +401,7 @@ NODE_CLASS_MAPPINGS = {
     "GetVideoDriveParams": GetVideoDriveParams,
     "HMPipelineImage": HMPipelineImage,
     "HMPipelineVideo": HMPipelineVideo,
+    "CropReferenceImage": CropReferenceImage,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -382,4 +413,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "GetVideoDriveParams": "Get Drive Video Parameters",
     "HMPipelineImage": "HelloMeme Image Pipeline",
     "HMPipelineVideo": "HelloMeme Video Pipeline",
+    "CropReferenceImage": "Crop Reference Image",
 }