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