v5c
This commit is contained in:
@@ -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()
|
||||
|
||||
@@ -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
@@ -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 = {}
|
||||
|
||||
@@ -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
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user