298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-03-15 14:09:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

Support loading checkpoints with the UNETLoader node.

Browse Source
This commit is contained in:
comfyanonymous 2024-07-03 11:34:32 -04:00
parent 537f35c7bc
commit d7484ef30c
1 changed files with 16 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
comfy/sd.py
View File

@ -563,24 +563,32 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
return (model_patcher, clip, vae, clipvision) return (model_patcher, clip, vae, clipvision)
def load_unet_state_dict(sd): #load unet in diffusers format def load_unet_state_dict(sd): #load unet in diffusers or regular format
#Allow loading unets from checkpoint files
checkpoint = False
diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd)
temp_sd = comfy.utils.state_dict_prefix_replace(sd, {diffusion_model_prefix: ""}, filter_keys=True)
if len(temp_sd) > 0:
sd = temp_sd
checkpoint = True
parameters = comfy.utils.calculate_parameters(sd) parameters = comfy.utils.calculate_parameters(sd)
unet_dtype = model_management.unet_dtype(model_params=parameters) unet_dtype = model_management.unet_dtype(model_params=parameters)
load_device = model_management.get_torch_device() load_device = model_management.get_torch_device()
if 'transformer_blocks.0.attn.add_q_proj.weight' in sd: #MMDIT SD3 if checkpoint or "input_blocks.0.0.weight" in sd or 'clf.1.weight' in sd: #ldm or stable cascade
model_config = model_detection.model_config_from_unet(sd, "")
if model_config is None:
return None
new_sd = sd
elif 'transformer_blocks.0.attn.add_q_proj.weight' in sd: #MMDIT SD3
new_sd = model_detection.convert_diffusers_mmdit(sd, "") new_sd = model_detection.convert_diffusers_mmdit(sd, "")
if new_sd is None: if new_sd is None:
return None return None
model_config = model_detection.model_config_from_unet(new_sd, "") model_config = model_detection.model_config_from_unet(new_sd, "")
if model_config is None: if model_config is None:
return None return None
elif "input_blocks.0.0.weight" in sd or 'clf.1.weight' in sd: #ldm or stable cascade
model_config = model_detection.model_config_from_unet(sd, "")
if model_config is None:
return None
new_sd = sd
else: #diffusers else: #diffusers
model_config = model_detection.model_config_from_diffusers_unet(sd) model_config = model_detection.model_config_from_diffusers_unet(sd)
if model_config is None: if model_config is None: