From 2c4e92a98b8338f754855a0db7dce164945e366e Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 2 Jan 2024 14:41:33 -0500 Subject: [PATCH] Fix regression. --- comfy/model_detection.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/comfy/model_detection.py b/comfy/model_detection.py index ad16c0fb..ea824c44 100644 --- a/comfy/model_detection.py +++ b/comfy/model_detection.py @@ -48,7 +48,12 @@ def detect_unet_config(state_dict, key_prefix, dtype): unet_config["dtype"] = dtype model_channels = state_dict['{}input_blocks.0.0.weight'.format(key_prefix)].shape[0] in_channels = state_dict['{}input_blocks.0.0.weight'.format(key_prefix)].shape[1] - out_channels = state_dict['{}out.2.weight'.format(key_prefix)].shape[0] + + out_key = '{}out.2.weight'.format(key_prefix) + if out_key in state_dict: + out_channels = state_dict[out_key].shape[0] + else: + out_channels = 4 num_res_blocks = [] channel_mult = []