diff --git a/README.md b/README.md index 836d7be..21ab920 100644 --- a/README.md +++ b/README.md @@ -18,10 +18,12 @@ + ## πŸ”† New Features/Updates +- βœ… `12/17/2024` Support modelscope ([Modelscope Demo](https://www.modelscope.cn/studios/songkey/HelloMeme)). - βœ… `12/08/2024` Added **HelloMemeV2** (select "v2" in the version option of the LoadHelloMemeImage/Video Node). Its features include: a. Improved expression consistency between the generated video and the driving video. b. Better compatibility with third-party checkpoints (*we will continuously collect compatible free third-party checkpoints and share them on this page. If you'd like to recommend one, please post it on the issue page or email us at songkey@pku.edu.cn*). diff --git a/hellomeme/pipelines/pipline_hm_image.py b/hellomeme/pipelines/pipline_hm_image.py index edfef92..f52b59a 100644 --- a/hellomeme/pipelines/pipline_hm_image.py +++ b/hellomeme/pipelines/pipline_hm_image.py @@ -41,12 +41,28 @@ class HMImagePipeline(StableDiffusionImg2ImgPipeline): self.vae_decode = copy.deepcopy(self.vae) - def insert_hm_modules(self, version, dtype): + def insert_hm_modules(self, version, dtype, modelscope=False): + if modelscope: + from modelscope import snapshot_download + hm_reference_dir = snapshot_download('songkey/hm_reference') + hm2_reference_dir = snapshot_download('songkey/hm2_reference') + hm_control_dir = snapshot_download('songkey/hm_control') + hm_control2_dir = snapshot_download('songkey/hm_control2') + hm2_control_dir = snapshot_download('songkey/hm2_control') + hm2_control2_dir = snapshot_download('songkey/hm2_control2') + else: + hm_reference_dir = 'songkey/hm_reference' + hm2_reference_dir = 'songkey/hm2_reference' + hm_control_dir = 'songkey/hm_control' + hm_control2_dir = 'songkey/hm_control2' + hm2_control_dir = 'songkey/hm2_control' + hm2_control2_dir = 'songkey/hm2_control2' + if isinstance(self.unet, HMDenoising3D): if version == 'v1': - hm_adapter = HMReferenceAdapter.from_pretrained('songkey/hm_reference') + hm_adapter = HMReferenceAdapter.from_pretrained(hm_reference_dir) else: - hm_adapter = HMReferenceAdapter.from_pretrained('songkey/hm2_reference') + hm_adapter = HMReferenceAdapter.from_pretrained(hm2_reference_dir) self.unet.insert_reference_adapter(hm_adapter) self.unet.to(device='cpu', dtype=dtype).eval() @@ -56,17 +72,17 @@ class HMImagePipeline(StableDiffusionImg2ImgPipeline): if hasattr(self, "mp_control"): del self.mp_control if version == 'v1': - self.mp_control = HMControlNet.from_pretrained('songkey/hm_control') + self.mp_control = HMControlNet.from_pretrained(hm_control_dir) else: - self.mp_control = HMV2ControlNet.from_pretrained('songkey/hm2_control') + self.mp_control = HMV2ControlNet.from_pretrained(hm2_control_dir) self.mp_control.to(device='cpu', dtype=dtype).eval() if hasattr(self, "mp_control2"): del self.mp_control2 if version == 'v1': - self.mp_control2 = HMControlNet2.from_pretrained('songkey/hm_control2') + self.mp_control2 = HMControlNet2.from_pretrained(hm_control2_dir) else: - self.mp_control2 = HMV2ControlNet2.from_pretrained('songkey/hm2_control2') + self.mp_control2 = HMV2ControlNet2.from_pretrained(hm2_control2_dir) self.mp_control2.to(device='cpu', dtype=dtype).eval() self.vae.to(device='cpu', dtype=dtype).eval() diff --git a/hellomeme/pipelines/pipline_hm_video.py b/hellomeme/pipelines/pipline_hm_video.py index c1680af..d814be7 100644 --- a/hellomeme/pipelines/pipline_hm_video.py +++ b/hellomeme/pipelines/pipline_hm_video.py @@ -26,7 +26,7 @@ from ..models import HMDenoising3D, HMDenoisingMotion, HMControlNet, HMControlNe from ..models import HMReferenceAdapter class HMVideoPipeline(StableDiffusionImg2ImgPipeline): - def caryomitosis(self, version, **kwargs): + def caryomitosis(self, version, modelscope=False, **kwargs): if hasattr(self, "unet_ref"): del self.unet_ref self.unet_ref = HMDenoising3D.from_unet2d(self.unet) @@ -37,11 +37,19 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline): self.unet_pre = HMDenoising3D.from_unet2d(self.unet) self.unet_pre.cpu() + if modelscope: + from modelscope import snapshot_download + hm_animatediff_dir = snapshot_download('songkey/hm_animatediff_frame12') + hm2_animatediff_dir = snapshot_download('songkey/hm2_animatediff_frame12') + else: + hm_animatediff_dir = 'songkey/hm_animatediff_frame12' + hm2_animatediff_dir = 'songkey/hm2_animatediff_frame12' + self.num_frames = 12 if version == 'v1': - adapter = MotionAdapter.from_pretrained("songkey/hm_animatediff_frame12", torch_dtype=torch.float16) + adapter = MotionAdapter.from_pretrained(hm_animatediff_dir, torch_dtype=torch.float16) else: - adapter = MotionAdapter.from_pretrained("songkey/hm2_animatediff_frame12", torch_dtype=torch.float16) + adapter = MotionAdapter.from_pretrained(hm2_animatediff_dir, torch_dtype=torch.float16) unet = HMDenoisingMotion.from_unet2d(unet=self.unet, motion_adapter=adapter, load_weights=True) # todo: δΈε€ŸδΌ˜ι›… del self.unet @@ -49,11 +57,27 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline): self.vae_decode = copy.deepcopy(self.vae) - def insert_hm_modules(self, version, dtype): - if version == 'v1': - hm_adapter = HMReferenceAdapter.from_pretrained('songkey/hm_reference') + def insert_hm_modules(self, version, dtype, modelscope=False): + if modelscope: + from modelscope import snapshot_download + hm_reference_dir = snapshot_download('songkey/hm_reference') + hm2_reference_dir = snapshot_download('songkey/hm2_reference') + hm_control_dir = snapshot_download('songkey/hm_control') + hm_control2_dir = snapshot_download('songkey/hm_control2') + hm2_control_dir = snapshot_download('songkey/hm2_control') + hm2_control2_dir = snapshot_download('songkey/hm2_control2') else: - hm_adapter = HMReferenceAdapter.from_pretrained('songkey/hm2_reference') + hm_reference_dir = 'songkey/hm_reference' + hm2_reference_dir = 'songkey/hm2_reference' + hm_control_dir = 'songkey/hm_control' + hm_control2_dir = 'songkey/hm_control2' + hm2_control_dir = 'songkey/hm2_control' + hm2_control2_dir = 'songkey/hm2_control2' + + if version == 'v1': + hm_adapter = HMReferenceAdapter.from_pretrained(hm_reference_dir) + else: + hm_adapter = HMReferenceAdapter.from_pretrained(hm2_reference_dir) if isinstance(self.unet, HMDenoisingMotion): self.unet.insert_reference_adapter(hm_adapter) self.unet.to(device='cpu', dtype=dtype).eval() @@ -68,18 +92,18 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline): if hasattr(self, "mp_control"): del self.mp_control if version == 'v1': - self.mp_control = HMControlNet.from_pretrained('songkey/hm_control') + self.mp_control = HMControlNet.from_pretrained(hm_control_dir) else: - self.mp_control = HMV2ControlNet.from_pretrained('songkey/hm2_control') + self.mp_control = HMV2ControlNet.from_pretrained(hm2_control_dir) self.mp_control.to(device='cpu', dtype=dtype).eval() if hasattr(self, "mp_control2"): del self.mp_control2 if version == 'v1': - self.mp_control2 = HMControlNet2.from_pretrained('songkey/hm_control2') + self.mp_control2 = HMControlNet2.from_pretrained(hm_control2_dir) else: - self.mp_control2 = HMV2ControlNet2.from_pretrained('songkey/hm2_control2') + self.mp_control2 = HMV2ControlNet2.from_pretrained(hm2_control2_dir) self.mp_control2.to(device='cpu', dtype=dtype).eval() self.vae.to(device='cpu', dtype=dtype).eval() diff --git a/hellomeme/tools/hello_3dmm.py b/hellomeme/tools/hello_3dmm.py index 7f8fa2e..4584855 100644 --- a/hellomeme/tools/hello_3dmm.py +++ b/hellomeme/tools/hello_3dmm.py @@ -11,8 +11,8 @@ import numpy as np import cv2 +import os.path as osp -from huggingface_hub import hf_hub_download from .utils import get_warp_mat_bbox_by_gt_pts_float, create_onnx_session def crop_transl_to_full_transl(crop_trans, crop_center, scale, full_center, focal_length): @@ -36,8 +36,14 @@ def crop_transl_to_full_transl(crop_trans, crop_center, scale, full_center, foca return full_trans class Hello3DMMPred(object): - def __init__(self, gpu_id=None): - self.deep3d_pred_net = create_onnx_session(hf_hub_download('songkey/hello_group_facemodel', filename='hello_3dmm.onnx'), gpu_id=gpu_id) + def __init__(self, gpu_id=None, modelscope=False): + if modelscope: + from modelscope import snapshot_download + model_path = osp.join(snapshot_download('songkey/hello_group_facemodel'), 'hello_3dmm.onnx') + else: + from huggingface_hub import hf_hub_download + model_path = hf_hub_download('songkey/hello_group_facemodel', filename='hello_3dmm.onnx') + self.deep3d_pred_net = create_onnx_session(model_path, gpu_id=gpu_id) self.deep3d_pred_net_input_name = self.deep3d_pred_net.get_inputs()[0].name self.deep3d_pred_net_output_name = [output.name for output in self.deep3d_pred_net.get_outputs()] diff --git a/hellomeme/tools/hello_arkit.py b/hellomeme/tools/hello_arkit.py index 2c8b2f8..b92cc81 100644 --- a/hellomeme/tools/hello_arkit.py +++ b/hellomeme/tools/hello_arkit.py @@ -8,12 +8,19 @@ import numpy as np import cv2 -from huggingface_hub import hf_hub_download +import os.path as osp from .utils import create_onnx_session, get_warp_mat_bbox_by_gt_pts_float class HelloARKitBSPred(object): - def __init__(self, gpu_id=0): - self.face_rig_net = create_onnx_session(hf_hub_download('songkey/hello_group_facemodel', filename='hello_arkit_blendshape.onnx'), gpu_id=gpu_id) + def __init__(self, gpu_id=0, modelscope=False): + if modelscope: + from modelscope import snapshot_download + model_path = osp.join(snapshot_download('songkey/hello_group_facemodel'), 'hello_arkit_blendshape.onnx') + else: + from huggingface_hub import hf_hub_download + model_path = hf_hub_download('songkey/hello_group_facemodel', filename='hello_arkit_blendshape.onnx') + + self.face_rig_net = create_onnx_session(model_path, gpu_id=gpu_id) self.onnx_input_name = self.face_rig_net.get_inputs()[0].name self.onnx_output_name = [output.name for output in self.face_rig_net.get_outputs()] self.image_size = 224 diff --git a/hellomeme/tools/hello_face_alignment.py b/hellomeme/tools/hello_face_alignment.py index bd9b527..607622b 100644 --- a/hellomeme/tools/hello_face_alignment.py +++ b/hellomeme/tools/hello_face_alignment.py @@ -7,21 +7,31 @@ """ import cv2 +import os.path as osp import numpy as np from .hello_face_det import HelloFaceDet from .utils import get_warp_mat_bbox, get_warp_mat_bbox_by_gt_pts_float, transform_points from .utils import create_onnx_session -from huggingface_hub import hf_hub_download class HelloFaceAlignment(object): - def __init__(self, gpu_id=None): + def __init__(self, gpu_id=None, modelscope=False): expand_ratio = 0.15 + + if modelscope: + from modelscope import snapshot_download + alignment_model_path = osp.join(snapshot_download('songkey/hello_group_facemodel'), 'hello_face_landmark.onnx') + det_model_path = osp.join(snapshot_download('songkey/hello_group_facemodel'), 'hello_face_det.onnx') + else: + from huggingface_hub import hf_hub_download + alignment_model_path = hf_hub_download('songkey/hello_group_facemodel', filename='hello_face_landmark.onnx') + det_model_path = hf_hub_download('songkey/hello_group_facemodel', filename='hello_face_det.onnx') self.face_alignment_net_222 = ( - create_onnx_session(hf_hub_download('songkey/hello_group_facemodel', filename='hello_face_landmark.onnx'), gpu_id=gpu_id)) + create_onnx_session(alignment_model_path, gpu_id=gpu_id)) self.onnx_input_name_222 = self.face_alignment_net_222.get_inputs()[0].name self.onnx_output_name_222 = [output.name for output in self.face_alignment_net_222.get_outputs()] self.face_image_size = 128 - self.face_detector = HelloFaceDet(hf_hub_download('songkey/hello_group_facemodel', filename='hello_face_det.onnx'), gpu_id=gpu_id) + + self.face_detector = HelloFaceDet(det_model_path, gpu_id=gpu_id) self.expand_ratio = expand_ratio def onnx_infer(self, input_uint8): diff --git a/hellomeme/tools/pdf.py b/hellomeme/tools/pdf.py index c1900ed..92e5770 100644 --- a/hellomeme/tools/pdf.py +++ b/hellomeme/tools/pdf.py @@ -106,7 +106,7 @@ class HourGlass(nn.Module): low3 = self._modules['b3_' + str(level)](low3) up1size = up1.size() rescale_size = (up1size[2], up1size[3]) - up2 = F.upsample(low3, size=rescale_size, mode='bilinear') + up2 = F.interpolate(low3, size=rescale_size, mode='bilinear') return up1 + up2 @@ -173,7 +173,7 @@ class FAN_use(nn.Module): class FanEncoder(ModelMixin, ConfigMixin): @register_to_config - def __init__(self,pose_dim=6, eye_dim=6): + def __init__(self, pose_dim=6, eye_dim=6): super().__init__() self.model = FAN_use() @@ -205,4 +205,4 @@ class FanEncoder(ModelMixin, ConfigMixin): emo_embed = self.emo_embed(emo_feat) return torch.cat([eye_embed, emo_embed, mouth_feat], dim=1) - # return headpose_emb, eye_embed, emo_embed, mouth_feat \ No newline at end of file + # return headpose_emb, eye_embed, emo_embed, mouth_feat diff --git a/hellomeme/utils.py b/hellomeme/utils.py index 91d2e6f..84da9ba 100644 --- a/hellomeme/utils.py +++ b/hellomeme/utils.py @@ -12,6 +12,7 @@ import os import cv2 import os.path as osp import numpy as np +import random from tqdm import tqdm import torch from PIL import Image @@ -26,24 +27,38 @@ from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (convert_ldm_ from .tools import Hello3DMMPred, HelloARKitBSPred, HelloFaceAlignment, HelloCameraDemo, FanEncoder from transformers import CLIPVisionModelWithProjection +def generate_random_string(length=8): + characters = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz' + result = '' + for i in range(0, length): + result += random.choice(characters) + return result + def get_torch_device(gpu_id=-1): if gpu_id < 0: return torch.device("cpu") gpu_id = min(max(-1, gpu_id), torch.cuda.device_count()) return torch.device(f"cuda:{gpu_id}") -def load_face_toolkits(dtype=torch.float16, gpu_id=-1): - image_encoder = CLIPVisionModelWithProjection.from_pretrained( - 'h94/IP-Adapter', subfolder='models/image_encoder') +def load_face_toolkits(dtype=torch.float16, gpu_id=-1, modelscope=False): + if modelscope: + from modelscope import snapshot_download + pd_fpg_motion = FanEncoder.from_pretrained(snapshot_download('songkey/pd_fgc_motion')) + image_encoder = CLIPVisionModelWithProjection.from_pretrained( + snapshot_download('songkey/IP-Adapter'), subfolder='models/image_encoder') + else: + pd_fpg_motion = FanEncoder.from_pretrained("songkey/pd_fgc_motion") + image_encoder = CLIPVisionModelWithProjection.from_pretrained( + 'h94/IP-Adapter', subfolder='models/image_encoder') + image_encoder.to(dtype=dtype).cpu() - pd_fpg_motion = FanEncoder.from_pretrained("songkey/pd_fgc_motion") pd_fpg_motion.to(dtype=dtype).cpu() return dict( device=get_torch_device(gpu_id), dtype=dtype, pd_fpg_motion=pd_fpg_motion, - face_aligner=HelloCameraDemo(face_alignment_module=HelloFaceAlignment(gpu_id=gpu_id), reset=False), - harkit_bs=HelloARKitBSPred(gpu_id=gpu_id), - h3dmm=Hello3DMMPred(gpu_id=gpu_id), + face_aligner=HelloCameraDemo(face_alignment_module=HelloFaceAlignment(gpu_id=gpu_id, modelscope=modelscope), reset=False), + harkit_bs=HelloARKitBSPred(gpu_id=gpu_id, modelscope=modelscope), + h3dmm=Hello3DMMPred(gpu_id=gpu_id, modelscope=modelscope), image_encoder=image_encoder ) diff --git a/meme.py b/meme.py index 0de1e61..6201882 100644 --- a/meme.py +++ b/meme.py @@ -70,17 +70,18 @@ class HMImagePipelineLoader: "vae": (vae_files, ), "version": (['v1', 'v2'], ), "stylize": (['x1', 'x2'], ), + "deployment": (['diffusers', 'modelscope'], ), } } RETURN_TYPES = ("HMIMAGEPIPELINE", ) RETURN_NAMES = ("hm_image_pipeline", ) FUNCTION = "load_pipeline" CATEGORY = "hellomeme" - def load_pipeline(self, checkpoint=None, lora=None, vae=None, version='v2', stylize='x1'): + def load_pipeline(self, checkpoint=None, lora=None, vae=None, version='v2', stylize='x1', deployment='diffusers'): dtype = torch.float16 pipeline = HMImagePipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") pipeline.to(dtype=dtype) - pipeline.caryomitosis(version=version) + pipeline.caryomitosis(version=version, modelscope=deployment=='modelscope') if checkpoint and not checkpoint.startswith('SD1.5'): checkpoint_path = folder_paths.get_full_path_or_raise("checkpoints", checkpoint) @@ -101,7 +102,7 @@ class HMImagePipelineLoader: append_pipline_weights(pipeline, checkpoint_path=checkpoint_path, lora_path=lora_path, vae_path=vae_path, stylize=stylize) - pipeline.insert_hm_modules(version=version, dtype=dtype) + pipeline.insert_hm_modules(version=version, dtype=dtype, modelscope=deployment=='modelscope') return (pipeline, ) @@ -118,6 +119,7 @@ class HMVideoPipelineLoader: "vae": (vae_files, ), "version": (['v1', 'v2'], ), "stylize": (['x1', 'x2'], ), + "deployment": (['diffusers', 'modelscope'], ), } } @@ -126,11 +128,11 @@ class HMVideoPipelineLoader: FUNCTION = "load_pipeline" CATEGORY = "hellomeme" - def load_pipeline(self, checkpoint=None, lora=None, vae=None, version='v2', stylize='x1'): + def load_pipeline(self, checkpoint=None, lora=None, vae=None, version='v2', stylize='x1', deployment='diffusers'): dtype = torch.float16 pipeline = HMVideoPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") pipeline.to(dtype=dtype) - pipeline.caryomitosis(version=version) + pipeline.caryomitosis(version=version, modelscope=deployment=='modelscope') if checkpoint and not checkpoint.startswith('SD1.5'): checkpoint_path = folder_paths.get_full_path_or_raise("checkpoints", checkpoint) @@ -150,7 +152,7 @@ class HMVideoPipelineLoader: append_pipline_weights(pipeline, checkpoint_path=checkpoint_path, lora_path=lora_path, vae_path=vae_path, stylize=stylize) - pipeline.insert_hm_modules(version=version, dtype=dtype) + pipeline.insert_hm_modules(version=version, dtype=dtype, modelscope=deployment=='modelscope') return (pipeline,) @@ -161,6 +163,7 @@ class HMFaceToolkitsLoader: return { "required": { "gpu_id": ("INT", {"default": 0, "min": -1, "max": 16}, ), + "deployment": (['diffusers', 'modelscope'], ), } } @@ -169,9 +172,9 @@ class HMFaceToolkitsLoader: FUNCTION = "load_face_toolkits" CATEGORY = "hellomeme" - def load_face_toolkits(self, gpu_id): + def load_face_toolkits(self, gpu_id, deployment='diffusers'): dtype = torch.float16 - face_toolkits = load_face_toolkits(dtype=dtype, gpu_id=gpu_id) + face_toolkits = load_face_toolkits(dtype=dtype, gpu_id=gpu_id, modelscope=deployment=='modelscope') return (face_toolkits, )