This commit is contained in:
songkey-win
2025-06-26 14:34:18 +08:00
parent bd8c9d8534
commit 19cd5d725e
10 changed files with 82 additions and 71 deletions
+3 -1
View File
@@ -42,7 +42,8 @@ class HM3ImagePipeline(HMPipeline):
self.vae_decode = copy.deepcopy(self.vae)
self.text_encoder.cpu()
self.text_encoder_ref = copy.deepcopy(self.text_encoder)
self.safety_checker.cpu()
if hasattr(self, 'safety_checker'):
del self.safety_checker
def insert_hm_modules(self, version='v3', dtype=torch.float16, modelscope=False):
self.version = version
@@ -82,6 +83,7 @@ class HM3ImagePipeline(HMPipeline):
self.vae.to(device='cpu', dtype=dtype).eval()
self.vae_decode.to(device='cpu', dtype=dtype).eval()
self.text_encoder.to(device='cpu', dtype=dtype).eval()
self.text_encoder_ref.to(device='cpu', dtype=dtype).eval()
@torch.no_grad()
def __call__(
+3 -1
View File
@@ -58,7 +58,8 @@ class HM3VideoPipeline(HMPipeline):
self.vae_decode = copy.deepcopy(self.vae)
self.text_encoder.cpu()
self.text_encoder_ref = copy.deepcopy(self.text_encoder)
self.safety_checker.cpu()
if hasattr(self, 'safety_checker'):
del self.safety_checker
def insert_hm_modules(self, version='v3', dtype=torch.float16, modelscope=False):
self.version = version
@@ -106,6 +107,7 @@ class HM3VideoPipeline(HMPipeline):
self.vae.to(device='cpu', dtype=dtype).eval()
self.vae_decode.to(device='cpu', dtype=dtype).eval()
self.text_encoder.to(device='cpu', dtype=dtype).eval()
self.text_encoder_ref.to(device='cpu', dtype=dtype).eval()
@torch.no_grad()
def gen_ref_cache(self, latent, prompt_embeds, added_cond_kwargs, device, add_axis=False):
+3 -1
View File
@@ -44,7 +44,8 @@ class HM5ImagePipeline(HMPipeline):
self.vae_decode = copy.deepcopy(self.vae)
self.text_encoder.cpu()
self.text_encoder_ref = copy.deepcopy(self.text_encoder)
self.safety_checker.cpu()
if hasattr(self, 'safety_checker'):
del self.safety_checker
def insert_hm_modules(self, version='v5', dtype=torch.float16, modelscope=False):
self.version = version
@@ -83,6 +84,7 @@ class HM5ImagePipeline(HMPipeline):
self.vae.to(device='cpu', dtype=dtype).eval()
self.vae_decode.to(device='cpu', dtype=dtype).eval()
self.text_encoder.to(device='cpu', dtype=dtype).eval()
self.text_encoder_ref.to(device='cpu', dtype=dtype).eval()
@torch.no_grad()
def __call__(
+3 -1
View File
@@ -49,7 +49,8 @@ class HM5VideoPipeline(HMPipeline):
self.vae_decode = copy.deepcopy(self.vae)
self.text_encoder.cpu()
self.text_encoder_ref = copy.deepcopy(self.text_encoder)
self.safety_checker.cpu()
if hasattr(self, 'safety_checker'):
del self.safety_checker
def insert_hm_modules(self, version='v5', dtype=torch.float16, modelscope=False):
self.version = version
@@ -89,6 +90,7 @@ class HM5VideoPipeline(HMPipeline):
self.vae.to(device='cpu', dtype=dtype).eval()
self.vae_decode.to(device='cpu', dtype=dtype).eval()
self.text_encoder.to(device='cpu', dtype=dtype).eval()
self.text_encoder_ref.to(device='cpu', dtype=dtype).eval()
@torch.no_grad()
def gen_ref_cache(self, latent, prompt_embeds, added_cond_kwargs, device, add_axis=False):
+3 -1
View File
@@ -44,7 +44,8 @@ class HMImagePipeline(HMPipeline):
self.vae_decode = copy.deepcopy(self.vae)
self.text_encoder.cpu()
self.text_encoder_ref = copy.deepcopy(self.text_encoder)
self.safety_checker.cpu()
if hasattr(self, 'safety_checker'):
del self.safety_checker
def insert_hm_modules(self, version, dtype, modelscope=False):
self.version = version
@@ -85,6 +86,7 @@ class HMImagePipeline(HMPipeline):
self.vae.to(device='cpu', dtype=dtype).eval()
self.vae_decode.to(device='cpu', dtype=dtype).eval()
self.text_encoder.to(device='cpu', dtype=dtype).eval()
self.text_encoder_ref.to(device='cpu', dtype=dtype).eval()
@torch.no_grad()
def __call__(
+4 -1
View File
@@ -95,12 +95,15 @@ class HMVideoPipeline(HMPipeline):
if version == 'v1':
self.mp_control2 = creat_model_from_cloud(HMControlNet2, hm_control2_dir, modelscope=modelscope)
else:
self.mp_control2 = creat_model_from_cloud(HMV2ControlNet2, hm2_control2_dir, modelscope=modelscope)
self.mp_control2 = creat_model_from_cloud(HMV2ControlNet2, hm_control2_dir, modelscope=modelscope)
self.mp_control2.to(device='cpu', dtype=dtype).eval()
self.vae.to(device='cpu', dtype=dtype).eval()
self.vae_decode.to(device='cpu', dtype=dtype).eval()
self.text_encoder.to(device='cpu', dtype=dtype).eval()
self.text_encoder_ref.to(device='cpu', dtype=dtype).eval()
if hasattr(self, 'safety_checker'):
del self.safety_checker
@torch.no_grad()
def __call__(
+17 -13
View File
@@ -146,19 +146,23 @@ def creat_model_from_cloud(model_cls,
cache_dir=None,
subfolder=None,
hf_token='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'):
if modelscope:
from modelscope import snapshot_download
try:
model_path = snapshot_download(model_id, cache_dir=cache_dir)
except Exception as e:
print(e)
assert False, "@@ Failed to download model from modelscope"
model = model_cls.from_pretrained(model_path, subfolder=subfolder)
if osp.isfile(model_id) and model_id.endswith('.safetensors'):
model = model_cls.from_single_file(model_id)
else:
try:
model = model_cls.from_pretrained(model_id, subfolder=subfolder, cache_dir=cache_dir, token=hf_token)
except Exception as e:
print(e)
assert False, "@@ Failed to download model from huggingface hub"
if modelscope:
from modelscope import snapshot_download
try:
model_path = snapshot_download(model_id, cache_dir=cache_dir)
except Exception as e:
print(e)
assert False, "@@ Failed to download model from modelscope"
model = model_cls.from_pretrained(model_path, subfolder=subfolder)
else:
try:
model = model_cls.from_pretrained(model_id, subfolder=subfolder, cache_dir=cache_dir, token=hf_token)
except Exception as e:
print(e)
assert False, "@@ Failed to download model from huggingface hub"
return model
+17 -32
View File
@@ -56,42 +56,29 @@ def load_face_toolkits(dtype=torch.float16, gpu_id=-1, modelscope=False):
image_encoder=image_encoder
)
def append_pipline_weights(pipeline, checkpoint_path=None,
lora_path=None, vae_path=None,
stylize='x1', lora_scale=1.0, 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:
# print("## checkpoint_path", checkpoint_path)
# print("## lora_path", lora_path)
# print("## vae_path", vae_path)
unet_stats, vae_stats = None, None
if checkpoint_path and not checkpoint_path.startswith('SD1.5'):
if osp.isfile(checkpoint_path):
print("Loading checkpoint from", checkpoint_path)
if checkpoint_path.endswith('.safetensors'):
raw_stats = load_safetensors(checkpoint_path)
unet_stats = convert_ldm_unet_checkpoint(raw_stats, pipeline.unet_ref.config)
vae_stats = convert_ldm_vae_checkpoint(raw_stats, pipeline.vae.config)
else:
raw_stats = torch.load(checkpoint_path)
unet_stats = convert_ldm_unet_checkpoint(raw_stats, pipeline.unet_ref.config)
vae_stats = convert_ldm_vae_checkpoint(raw_stats, pipeline.vae.config)
else:
tmp_pipeline = creat_model_from_cloud(HMImagePipeline, checkpoint_path, modelscope=modelscope)
unet_stats = tmp_pipeline.unet.state_dict()
vae_stats = tmp_pipeline.vae.state_dict()
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 unet_stats:
try:
if hasattr(pipeline, 'unet_pre'):
pipeline.unet_pre.load_state_dict(unet_stats, strict=False)
pipeline.unet.load_state_dict(unet_stats, strict=False)
if stylize == 'x2' and hasattr(pipeline, 'unet_ref'):
pipeline.unet_ref.load_state_dict(unet_stats, strict=False)
except:
raise ValueError("Failed to load checkpoint from", checkpoint_path)
if stylize == 'x1' and hasattr(pipeline, 'unet_ref'):
pipeline.unet_ref.load_state_dict(unet_stats, strict=True)
if stylize == 'x1' and 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)
pipeline.vae.load_state_dict(vae_stats, strict=True)
if vae_path and not vae_path.startswith('SD1.5 default vae'):
if osp.isfile(vae_path):
@@ -106,8 +93,6 @@ def append_pipline_weights(pipeline, checkpoint_path=None,
try:
if hasattr(pipeline, 'vae_decode'):
pipeline.vae_decode.load_state_dict(vae_stats, strict=True)
if stylize == 'x2':
pipeline.vae.load_state_dict(vae_stats, strict=True)
except:
raise ValueError("Failed to load vae from", vae_path)
@@ -128,7 +113,7 @@ def append_pipline_weights(pipeline, checkpoint_path=None,
if stylize == 'x2' and hasattr(pipeline, 'unet_ref'):
pipeline.load_lora_weights_sk(pipeline.unet_ref, osp.dirname(lora_path),
weight_name=osp.basename(lora_path),
adapter_name="lora", text_encoder=None,
adapter_name="lora", text_encoder=pipeline.text_encoder_ref,
lora_scale=lora_scale)
except:
+28 -19
View File
@@ -70,15 +70,7 @@ def get_models_files():
return checkpoint_files, vae_files, lora_files
def format_model_path(pipeline, config, checkpoint, vae, lora, stylize, lora_scale, deployment):
if checkpoint and not checkpoint.startswith('SD1.5'):
if checkpoint in config['sd15']['checkpoints']:
checkpoint_path = config['sd15']['checkpoints'][checkpoint]
else:
checkpoint_path = folder_paths.get_full_path_or_raise("checkpoints", checkpoint)
else:
checkpoint_path = checkpoint
def format_model_path(pipeline, config, vae, lora, stylize, lora_scale, official_id, deployment):
if vae and vae.startswith("[checkpoint] "):
vae_path = folder_paths.get_full_path_or_raise("checkpoints", vae.replace("[checkpoint] ", ""))
elif vae and vae.startswith("[vae] "):
@@ -95,9 +87,20 @@ def format_model_path(pipeline, config, checkpoint, vae, lora, stylize, lora_sca
else:
lora_path = lora
append_pipline_weights(pipeline, checkpoint_path=checkpoint_path, lora_path=lora_path, vae_path=vae_path,
stylize=stylize, lora_scale=lora_scale, modelscope=deployment == 'modelscope')
append_pipline_weights(pipeline, lora_path=lora_path, vae_path=vae_path,
stylize=stylize, lora_scale=lora_scale,
official_id=official_id,
modelscope=deployment == 'modelscope')
def get_checkpoint_path(config, checkpoint):
if checkpoint and not checkpoint.startswith('SD1.5'):
if checkpoint in config['sd15']['checkpoints']:
checkpoint_path = config['sd15']['checkpoints'][checkpoint]
else:
checkpoint_path = folder_paths.get_full_path_or_raise("checkpoints", checkpoint)
else:
checkpoint_path = checkpoint
return checkpoint_path
class HMImagePipelineLoader:
@classmethod
@@ -124,19 +127,22 @@ class HMImagePipelineLoader:
version='v2', stylize='x1', deployment='huggingface', lora_scale=1.0, dtype='fp32'):
dtype = torch.float32 if dtype == 'fp32' else torch.float16
checkpoint_path = get_checkpoint_path(MODEL_CONFIG, checkpoint)
if version == 'v3' or version == 'v4':
pipeline = creat_model_from_cloud(HM3ImagePipeline, "songkey/stable-diffusion-v1-5",
pipeline = creat_model_from_cloud(HM3ImagePipeline, checkpoint_path,
modelscope=deployment == 'modelscope')
elif version == 'v5' or version == 'v5b':
pipeline = creat_model_from_cloud(HM5ImagePipeline, "songkey/stable-diffusion-v1-5",
pipeline = creat_model_from_cloud(HM5ImagePipeline, checkpoint_path,
modelscope=deployment == 'modelscope')
else:
pipeline = creat_model_from_cloud(HMImagePipeline, "songkey/stable-diffusion-v1-5",
pipeline = creat_model_from_cloud(HMImagePipeline, checkpoint_path,
modelscope=deployment == 'modelscope')
pipeline.to(dtype=dtype)
pipeline.caryomitosis(version=version, modelscope=deployment == 'modelscope')
format_model_path(pipeline, MODEL_CONFIG, checkpoint, vae, lora, stylize, lora_scale, deployment)
format_model_path(pipeline, MODEL_CONFIG, vae, lora, stylize, lora_scale,
'songkey/stable-diffusion-v1-5', deployment)
pipeline.insert_hm_modules(version=version, dtype=dtype, modelscope=deployment == 'modelscope')
@@ -170,19 +176,22 @@ class HMVideoPipelineLoader:
version='v2', stylize='x1', deployment='huggingface', lora_scale=1.0, dtype='fp32'):
dtype = torch.float32 if dtype == 'fp32' else torch.float16
checkpoint_path = get_checkpoint_path(MODEL_CONFIG, checkpoint)
if version == 'v3' or version == 'v4':
pipeline = creat_model_from_cloud(HM3VideoPipeline, "songkey/stable-diffusion-v1-5",
pipeline = creat_model_from_cloud(HM3VideoPipeline, checkpoint_path,
modelscope=deployment == 'modelscope')
elif version == 'v5' or version == 'v5b':
pipeline = creat_model_from_cloud(HM5VideoPipeline, "songkey/stable-diffusion-v1-5",
pipeline = creat_model_from_cloud(HM5VideoPipeline, checkpoint_path,
modelscope=deployment == 'modelscope')
else:
pipeline = creat_model_from_cloud(HMVideoPipeline, "songkey/stable-diffusion-v1-5",
pipeline = creat_model_from_cloud(HMVideoPipeline, checkpoint_path,
modelscope=deployment == 'modelscope')
pipeline.to(dtype=dtype)
pipeline.caryomitosis(version=version, modelscope=deployment == 'modelscope')
format_model_path(pipeline, MODEL_CONFIG, checkpoint, vae, lora, stylize, lora_scale, deployment)
format_model_path(pipeline, MODEL_CONFIG, vae, lora, stylize, lora_scale,
'songkey/stable-diffusion-v1-5', deployment)
pipeline.insert_hm_modules(version=version, dtype=dtype, modelscope=deployment == 'modelscope')
+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.19"
version = "1.3.20"
license = {file = "LICENSE"}
[project.urls]