From aac5509484371d32c13f44a289e1ef9be47e71c1 Mon Sep 17 00:00:00 2001 From: songkey-win Date: Thu, 26 Jun 2025 20:22:12 +0800 Subject: [PATCH] v5c --- hellomeme/pipelines/pipline_hm5_image.py | 43 ++++++++++++++++++++---- hellomeme/tools/utils.py | 4 +-- hellomeme/utils.py | 43 +++++++++++------------- meme.py | 27 +++++++++------ pyproject.toml | 2 +- 5 files changed, 75 insertions(+), 44 deletions(-) diff --git a/hellomeme/pipelines/pipline_hm5_image.py b/hellomeme/pipelines/pipline_hm5_image.py index 33659d0..1e7ede9 100644 --- a/hellomeme/pipelines/pipline_hm5_image.py +++ b/hellomeme/pipelines/pipline_hm5_image.py @@ -50,18 +50,21 @@ class HM5ImagePipeline(HMPipeline): def insert_hm_modules(self, version='v5', dtype=torch.float16, modelscope=False): self.version = version - if version == 'v5b': - hm_reference_dir = 'songkey/hm5b_reference' - else: - hm_reference_dir = 'songkey/hm5_reference' hm_control_dir = 'songkey/hm5_control_base' hm_control_proj_dir = 'songkey/hm5_control_proj' + if version == 'v5': + hm_reference_dir = 'songkey/hm5_reference' + elif version == 'v5b': + hm_reference_dir = 'songkey/hm5b_reference' + else: + hm_reference_dir = 'songkey/hm5c_reference' + hm_control_proj_dir = 'songkey/hm5c_control_proj' if isinstance(self.unet, HM3Denoising3D): - if version == 'v5b': - hm_adapter = creat_model_from_cloud(HM5bReferenceAdapter, hm_reference_dir, modelscope=modelscope) - else: + if version == 'v5': hm_adapter = creat_model_from_cloud(HM5ReferenceAdapter, hm_reference_dir, modelscope=modelscope) + else: + hm_adapter = creat_model_from_cloud(HM5bReferenceAdapter, hm_reference_dir, modelscope=modelscope) self.unet.insert_reference_adapter(hm_adapter) self.unet.to(device='cpu', dtype=dtype).eval() @@ -91,6 +94,7 @@ class HM5ImagePipeline(HMPipeline): self, prompt: Union[str, List[str]] = None, image: PipelineImageInput = None, + ref_params: Dict[str, Any] = None, drive_params: Dict[str, Any] = None, strength: float = 0.8, num_inference_steps: Optional[int] = 50, @@ -242,28 +246,52 @@ class HM5ImagePipeline(HMPipeline): c, h, w = ref_latents.shape[1:] condition = drive_params['condition'].clone().to(device=device) + if self.version == 'v5c': + ref_condition = ref_params['condition'].clone().to(device=device) if self.do_classifier_free_guidance: condition = torch.cat([torch.ones_like(condition) * -1, condition], dim=0) + if self.version == 'v5c': + ref_condition = torch.cat([torch.ones_like(ref_condition) * -1, ref_condition], dim=0) + if self.version == 'v5c': + ref_control_latents = {} control_latents = {} self.mp_control.to(device=device) self.mp_control_proj.to(device=device) if 'drive_coeff' in drive_params: drive_coeff = drive_params['drive_coeff'].clone().to(device=device) face_parts = drive_params['face_parts'].clone().to(device=device) + if self.version == 'v5c': + ref_drive_coeff = ref_params['drive_coeff'].clone().to(device=device) + ref_face_parts = ref_params['face_parts'].clone().to(device=device) if self.do_classifier_free_guidance: drive_coeff = torch.cat([torch.zeros_like(drive_coeff), drive_coeff], dim=0) face_parts = torch.cat([torch.zeros_like(face_parts), face_parts], dim=0) + if self.version == 'v5c': + ref_drive_coeff = torch.cat([torch.zeros_like(ref_drive_coeff), ref_drive_coeff], dim=0) + ref_face_parts = torch.cat([torch.zeros_like(ref_face_parts), ref_face_parts], dim=0) control_latents1 = self.mp_control(condition=condition, drive_coeff=drive_coeff, face_parts=face_parts) control_latents1 = self.mp_control_proj(control_latents1) control_latents.update(control_latents1) + if self.version == 'v5c': + ref_control_latents1 = self.mp_control(condition=ref_condition, drive_coeff=ref_drive_coeff, face_parts=ref_face_parts) + ref_control_latents1 = self.mp_control_proj(ref_control_latents1) + ref_control_latents.update(ref_control_latents1) elif 'pd_fpg' in drive_params: pd_fpg = drive_params['pd_fpg'].clone().to(device=device) + if self.version == 'v5c': + ref_pd_fpg = ref_params['pd_fpg'].clone().to(device=device) if self.do_classifier_free_guidance: pd_fpg = torch.cat([torch.zeros_like(pd_fpg), pd_fpg], dim=0) + if self.version == 'v5c': + ref_pd_fpg = torch.cat([torch.zeros_like(ref_pd_fpg), ref_pd_fpg], dim=0) control_latents2 = self.mp_control(condition=condition, emo_embedding=pd_fpg) control_latents2 = self.mp_control_proj(control_latents2) control_latents.update(control_latents2) + if self.version == 'v5c': + ref_control_latents2 = self.mp_control(condition=ref_condition, emo_embedding=ref_pd_fpg) + ref_control_latents2 = self.mp_control_proj(ref_control_latents2) + ref_control_latents.update(ref_control_latents2) self.mp_control.cpu() self.mp_control_proj.cpu() @@ -286,6 +314,7 @@ class HM5ImagePipeline(HMPipeline): latent_model_input.unsqueeze(2), 0, encoder_hidden_states=prompt_embeds_ref, + control_hidden_states=ref_control_latents if self.version == 'v5c' else None, return_dict=False, )[1] self.unet_ref.cpu() diff --git a/hellomeme/tools/utils.py b/hellomeme/tools/utils.py index de30af3..595c59a 100644 --- a/hellomeme/tools/utils.py +++ b/hellomeme/tools/utils.py @@ -123,7 +123,7 @@ def download_file_from_cloud(model_id, file_name, modelscope=False, cache_dir=None, - hf_token='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',): + hf_token='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',): if modelscope: from modelscope import snapshot_download try: @@ -149,7 +149,7 @@ def creat_model_from_cloud(model_cls, modelscope=False, cache_dir=None, subfolder=None, - hf_token='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'): + hf_token='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'): if osp.isfile(model_id) and model_id.endswith('.safetensors'): model = model_cls.from_single_file(model_id) diff --git a/hellomeme/utils.py b/hellomeme/utils.py index 6ee3d08..dc79371 100644 --- a/hellomeme/utils.py +++ b/hellomeme/utils.py @@ -25,8 +25,8 @@ from safetensors import safe_open from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (convert_ldm_unet_checkpoint, convert_ldm_vae_checkpoint) from .tools import Hello3DMMPred, HelloARKitBSPred, HelloFaceAlignment, HelloCameraDemo, FanEncoder -from .pipelines import HMImagePipeline -from transformers import CLIPVisionModelWithProjection +from transformers import CLIPVisionModelWithProjection, CLIPTextModel +from diffusers import AutoencoderKL, UNet2DConditionModel def generate_random_string(length=8): characters = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz' @@ -58,29 +58,24 @@ def load_face_toolkits(dtype=torch.float16, gpu_id=-1, modelscope=False): def append_pipline_weights(pipeline, lora_path=None, vae_path=None, stylize='x1', lora_scale=1.0, official_id=None, modelscope=False): - ### load customized checkpoint or lora here: + text_encoder_stats = creat_model_from_cloud(CLIPTextModel, official_id, subfolder='text_encoder', modelscope=modelscope).state_dict() + if hasattr(pipeline, 'text_encoder_ref'): + pipeline.text_encoder_ref.load_state_dict(text_encoder_stats, strict=True) - # print("## checkpoint_path", checkpoint_path) - # print("## lora_path", lora_path) - # print("## vae_path", vae_path) - - unet_stats, vae_stats, text_encoder_stats = None, None, None - if not official_id is None: - tmp_pipeline = creat_model_from_cloud(HMImagePipeline, official_id, modelscope=modelscope) - unet_stats = tmp_pipeline.unet.state_dict() - vae_stats = tmp_pipeline.vae.state_dict() - text_encoder_stats = tmp_pipeline.text_encoder.state_dict() - - if stylize == 'x1' and hasattr(pipeline, 'unet_ref'): + if stylize == 'x1': + unet_stats = creat_model_from_cloud(UNet2DConditionModel, official_id, subfolder='unet', modelscope=modelscope).state_dict() + if hasattr(pipeline, 'unet_ref'): pipeline.unet_ref.load_state_dict(unet_stats, strict=True) - if stylize == 'x1' and hasattr(pipeline, 'unet_pre'): + if hasattr(pipeline, 'unet_pre'): pipeline.unet_pre.load_state_dict(unet_stats, strict=True) - if stylize == 'x1' and hasattr(pipeline, 'text_encoder_ref'): - pipeline.text_encoder_ref.load_state_dict(text_encoder_stats, strict=True) + pipeline.text_encoder.load_state_dict(text_encoder_stats, strict=True) + + vae_stats = creat_model_from_cloud(AutoencoderKL, official_id, subfolder='vae', modelscope=modelscope).state_dict() pipeline.vae.load_state_dict(vae_stats, strict=True) - if vae_path and not vae_path.startswith('SD1.5 default vae'): + if vae_path: + vae_stats = None if osp.isfile(vae_path): print("Loading vae from", vae_path) if vae_path.endswith('.safetensors'): @@ -88,13 +83,13 @@ def append_pipline_weights(pipeline, lora_path=None, vae_path=None, else: raw_vae_stats = torch.load(vae_path) vae_stats = convert_ldm_vae_checkpoint(raw_vae_stats, pipeline.vae.config) - + elif vae_path.startswith('SD1.5 default vae') and hasattr(pipeline, 'vae_decode') and stylize != 'x1' and official_id: + vae_stats = creat_model_from_cloud(AutoencoderKL, official_id, subfolder='vae', modelscope=modelscope) if vae_stats: try: - if hasattr(pipeline, 'vae_decode'): - pipeline.vae_decode.load_state_dict(vae_stats, strict=True) + pipeline.vae_decode.load_state_dict(vae_stats, strict=True) except: - raise ValueError("Failed to load vae from", vae_path) + raise ValueError("@@ Failed to load vae from", vae_path) ### lora if lora_path and not lora_path.startswith('None'): @@ -117,7 +112,7 @@ def append_pipline_weights(pipeline, lora_path=None, vae_path=None, lora_scale=lora_scale) except: - raise ValueError("Failed to load lora from", lora_path) + raise ValueError("@@ Failed to load lora from", lora_path) def load_safetensors(model_path): tensors = {} diff --git a/meme.py b/meme.py index 56d10ec..00ff698 100644 --- a/meme.py +++ b/meme.py @@ -112,7 +112,7 @@ class HMImagePipelineLoader: "checkpoint": (checkpoint_files, ), "lora": (lora_files, ), "vae": (vae_files, ), - "version": (['v5b', 'v5', 'v4', 'v3', 'v2', 'v1'], ), + "version": (['v5c', 'v5b', 'v5', 'v4', 'v3', 'v2', 'v1'], ), "stylize": (['x1', 'x2'], ), "deployment": (['huggingface', 'modelscope'], ), "lora_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1}), @@ -129,10 +129,10 @@ class HMImagePipelineLoader: checkpoint_path = get_checkpoint_path(MODEL_CONFIG, checkpoint) - if version == 'v3' or version == 'v4': + if version in ['v3', 'v4']: pipeline = creat_model_from_cloud(HM3ImagePipeline, checkpoint_path, modelscope=deployment == 'modelscope') - elif version == 'v5' or version == 'v5b': + elif version in ['v5', 'v5b', 'v5c']: pipeline = creat_model_from_cloud(HM5ImagePipeline, checkpoint_path, modelscope=deployment == 'modelscope') else: @@ -178,10 +178,10 @@ class HMVideoPipelineLoader: checkpoint_path = get_checkpoint_path(MODEL_CONFIG, checkpoint) - if version == 'v3' or version == 'v4': + if version in ['v3', 'v4']: pipeline = creat_model_from_cloud(HM3VideoPipeline, checkpoint_path, modelscope=deployment == 'modelscope') - elif version == 'v5' or version == 'v5b': + elif version in ['v5', 'v5b', 'v5c']: pipeline = creat_model_from_cloud(HM5VideoPipeline, checkpoint_path, modelscope=deployment == 'modelscope') else: @@ -352,20 +352,26 @@ class HMPipelineImage: device = get_torch_device(gpu_id) dtype = hm_image_pipeline.dtype - PROMPT = DEFAULT_PROMPT_NEW if hm_image_pipeline.version == 'v5b' else DEFAULT_PROMPT + PROMPT = DEFAULT_PROMPT_NEW if hm_image_pipeline.version in ['v5b', 'v5c'] else DEFAULT_PROMPT prompt = PROMPT if prompt == '' else prompt + ", " + PROMPT image_np = cv2.cvtColor(ref_head_pose['frame'][0], cv2.COLOR_BGR2RGB) image_pil = Image.fromarray(image_np) - ref_trans = ref_head_pose['trans'][0] - + ref_rot, ref_trans = ref_head_pose['rot'], ref_head_pose['trans'] drive_rot, drive_trans = drive_head_pose['rot'], drive_head_pose['trans'] - condition = gen_control_heatmaps(drive_rot, drive_trans, ref_trans, 512, trans_ratio) + + condition = gen_control_heatmaps(drive_rot, drive_trans, ref_trans[0], 512, trans_ratio) + ref_condition = gen_control_heatmaps(ref_rot, ref_trans, ref_trans[0], 512, 0.0) drive_params = dict(condition=condition.unsqueeze(0)) + ref_params = dict(condition=ref_condition.unsqueeze(0)) drive_params.update(drive_expression) + ref_params.update(ref_expression) + for k, v in drive_params.items(): - drive_params[k] = v.to(dtype=dtype, device='cpu') + drive_params[k] = v.to(dtype=dtype) + for k, v in ref_params.items(): + ref_params[k] = v.to(dtype=dtype, device='cpu') generator = torch.Generator().manual_seed(seed) @@ -374,6 +380,7 @@ class HMPipelineImage: strength=1.0, image=image_pil, drive_params=drive_params, + ref_params=ref_params, num_inference_steps=steps, negative_prompt=[negative_prompt], guidance_scale=guidance_scale, diff --git a/pyproject.toml b/pyproject.toml index f3dc1fb..79257f1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui_hellomeme" description = "This repository is the official implementation of the [a/HelloMeme](https://arxiv.org/pdf/2410.22901) ComfyUI interface" -version = "1.3.21" +version = "1.3.22" license = {file = "LICENSE"} [project.urls]