diff --git a/README.md b/README.md index e66f6b0..89c8e24 100644 --- a/README.md +++ b/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

diff --git a/meme.py b/meme.py index 027e3ea..ed3466a 100644 --- a/meme.py +++ b/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", } \ No newline at end of file