38 lines
1.3 KiB
Python
38 lines
1.3 KiB
Python
import folder_paths
|
|
import comfy.utils
|
|
from comfy.sd import CLIP
|
|
|
|
|
|
class MUReplaceModelWeights:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"model": ("MODEL",), "ckpt_name": (folder_paths.get_filename_list("checkpoints"),)}}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP")
|
|
FUNCTION = "do"
|
|
|
|
CATEGORY = "utils"
|
|
|
|
def do(self, model, ckpt_name):
|
|
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
|
sd = comfy.utils.load_torch_file(ckpt_path, safe_load=True)
|
|
to_load = {}
|
|
model = model.clone()
|
|
# Reload clip
|
|
clip_target = model.model.model_config.clip_target()
|
|
clip_sd = model.model.model_config.process_clip_state_dict(sd)
|
|
clip = CLIP(clip_target, embedding_directory=folder_paths.get_folder_paths("embeddings"))
|
|
clip.load_sd(clip_sd, full_model=True)
|
|
del clip_sd
|
|
|
|
pfx = "model.diffusion_model."
|
|
for k in list(sd.keys()):
|
|
if k.startswith(pfx):
|
|
to_load[k[len(pfx) :]] = sd.pop(k)
|
|
print("Replacing weights with", ckpt_path)
|
|
to_load = model.model.model_config.process_unet_state_dict(to_load)
|
|
model.model.diffusion_model.load_state_dict(to_load, strict=False)
|
|
del to_load
|
|
|
|
return (model, clip)
|