This commit is contained in:
songkey-win
2025-06-26 20:22:12 +08:00
parent f659a25cfc
commit aac5509484
5 changed files with 75 additions and 44 deletions
+36 -7
View File
@@ -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()
+2 -2
View File
@@ -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)
+19 -24
View File
@@ -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 = {}
+17 -10
View File
@@ -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,
+1 -1
View File
@@ -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]