From 5fe4119d53a5bf672b45838fdf47c17efd89cfab Mon Sep 17 00:00:00 2001 From: Aseem Saxena Date: Thu, 24 Apr 2025 12:54:32 -0700 Subject: [PATCH] Update utils.py --- comfy/utils.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/comfy/utils.py b/comfy/utils.py index ed40e3f26..062920173 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -118,16 +118,16 @@ def state_dict_key_replace(state_dict, keys_to_replace): return state_dict def state_dict_prefix_replace(state_dict, replace_prefix, filter_keys=False): - out = {} if filter_keys else state_dict - - for old_prefix, new_prefix in replace_prefix.items(): - keys_to_replace = [key for key in state_dict if key.startswith(old_prefix)] - - for key in keys_to_replace: - new_key = new_prefix + key[len(old_prefix):] - out[new_key] = state_dict.pop(key) - - return out + if not filter_keys: + return {} + else: + out = state_dict + for old_prefix, new_prefix in replace_prefix.items(): + keys_to_replace = [key for key in state_dict if key.startswith(old_prefix)] + for key in keys_to_replace: + new_key = new_prefix + key[len(old_prefix):] + out[new_key] = state_dict.pop(key) + return out def transformers_convert(sd, prefix_from, prefix_to, number):