jinx/ComfyUI
1
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-08-02 23:14:49 +08:00
Code Issues Activity

Cleaner support for loading different diffusion model types.

Browse Source
This commit is contained in:
comfyanonymous
2024-07-11 11:37:31 -04:00
parent ffe0bb0a33
commit 5e1fced639
2 changed files with 7 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
comfy/model_detection.py
View File

@@ -105,6 +105,9 @@ def detect_unet_config(state_dict, key_prefix):
unet_config["audio_model"] = "dit1.0"
return unet_config
if '{}input_blocks.0.0.weight'.format(key_prefix) not in state_dict_keys:
return None
unet_config = {
"use_checkpoint": False,
"image_size": 32,
@@ -239,6 +242,8 @@ def model_config_from_unet_config(unet_config, state_dict=None):
def model_config_from_unet(state_dict, unet_key_prefix, use_base_if_no_match=False):
unet_config = detect_unet_config(state_dict, unet_key_prefix)
if unet_config is None:
return None
model_config = model_config_from_unet_config(unet_config, state_dict)
if model_config is None and use_base_if_no_match:
return comfy.supported_models_base.BASE(unet_config)