mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-03-15 14:09:36 +00:00
Support loading checkpoints with the UNETLoader node.
This commit is contained in:
parent
537f35c7bc
commit
d7484ef30c
24
comfy/sd.py
24
comfy/sd.py
@ -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:
|
||||||
|
Loading…
Reference in New Issue
Block a user