diff --git a/comfy/sd.py b/comfy/sd.py index f331feefb..848b727a4 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -1743,17 +1743,16 @@ def load_diffusion_model_state_dict(sd, model_options={}, metadata=None, disable """ dtype = model_options.get("dtype", None) - custom_operations = model_options.get("custom_operations", None) - if custom_operations is None: - sd, metadata = comfy.utils.convert_old_quants(sd, "", metadata=metadata) - #Allow loading unets from checkpoint files diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd) + + custom_operations = model_options.get("custom_operations", None) + if custom_operations is None: + sd, metadata = comfy.utils.convert_old_quants(sd, diffusion_model_prefix, metadata=metadata) + temp_sd = comfy.utils.state_dict_prefix_replace(sd, {diffusion_model_prefix: ""}, filter_keys=True) if len(temp_sd) > 0: sd = temp_sd - if custom_operations is None: - sd, metadata = comfy.utils.convert_old_quants(sd, "", metadata=metadata) parameters = comfy.utils.calculate_parameters(sd) weight_dtype = comfy.utils.weight_dtype(sd)