diff --git a/hellomeme/pipelines/pipline_hm3_image.py b/hellomeme/pipelines/pipline_hm3_image.py index 1a266ea..5408cbe 100644 --- a/hellomeme/pipelines/pipline_hm3_image.py +++ b/hellomeme/pipelines/pipline_hm3_image.py @@ -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__( diff --git a/hellomeme/pipelines/pipline_hm3_video.py b/hellomeme/pipelines/pipline_hm3_video.py index 0c42bde..88bc95d 100644 --- a/hellomeme/pipelines/pipline_hm3_video.py +++ b/hellomeme/pipelines/pipline_hm3_video.py @@ -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): diff --git a/hellomeme/pipelines/pipline_hm5_image.py b/hellomeme/pipelines/pipline_hm5_image.py index 2c105dc..33659d0 100644 --- a/hellomeme/pipelines/pipline_hm5_image.py +++ b/hellomeme/pipelines/pipline_hm5_image.py @@ -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__( diff --git a/hellomeme/pipelines/pipline_hm5_video.py b/hellomeme/pipelines/pipline_hm5_video.py index 12d4b9e..6321297 100644 --- a/hellomeme/pipelines/pipline_hm5_video.py +++ b/hellomeme/pipelines/pipline_hm5_video.py @@ -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): diff --git a/hellomeme/pipelines/pipline_hm_image.py b/hellomeme/pipelines/pipline_hm_image.py index 899eb39..275c35b 100644 --- a/hellomeme/pipelines/pipline_hm_image.py +++ b/hellomeme/pipelines/pipline_hm_image.py @@ -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__( diff --git a/hellomeme/pipelines/pipline_hm_video.py b/hellomeme/pipelines/pipline_hm_video.py index f106acb..8a52291 100644 --- a/hellomeme/pipelines/pipline_hm_video.py +++ b/hellomeme/pipelines/pipline_hm_video.py @@ -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__( diff --git a/hellomeme/tools/utils.py b/hellomeme/tools/utils.py index 5786218..22ea7a4 100644 --- a/hellomeme/tools/utils.py +++ b/hellomeme/tools/utils.py @@ -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 \ No newline at end of file diff --git a/hellomeme/utils.py b/hellomeme/utils.py index e6eced5..6ee3d08 100644 --- a/hellomeme/utils.py +++ b/hellomeme/utils.py @@ -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: diff --git a/meme.py b/meme.py index bca6e48..56d10ec 100644 --- a/meme.py +++ b/meme.py @@ -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') diff --git a/pyproject.toml b/pyproject.toml index 0ab8d3c..71fce1d 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.19" +version = "1.3.20" license = {file = "LICENSE"} [project.urls]