From 3a8786c206227addc3571539f7e2feb5ded9e8c5 Mon Sep 17 00:00:00 2001
From: ozbayb <17261091+ozbayb@users.noreply.github.com>
Date: Sat, 8 Nov 2025 11:48:31 -0700
Subject: [PATCH] DiffusionModelLoaderKJ: Allow model checkpoints to be used
 with extra state dict input

---
 nodes/model_optimization_nodes.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/nodes/model_optimization_nodes.py b/nodes/model_optimization_nodes.py
index 61941da..86cc463 100644
--- a/nodes/model_optimization_nodes.py
+++ b/nodes/model_optimization_nodes.py
@@ -391,6 +391,14 @@ class DiffusionModelLoaderKJ(BaseLoaderKJ):
     
         sd = comfy.utils.load_torch_file(unet_path)
         if extra_state_dict is not None:
+            # If the model is a checkpoint, strip additional non-diffusion model entries before adding extra state dict
+            from comfy import model_detection
+            diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd)
+            if diffusion_model_prefix == "model.diffusion_model.":
+                temp_sd = comfy.utils.state_dict_prefix_replace(sd, {diffusion_model_prefix: ""}, filter_keys=True)
+                if len(temp_sd) > 0:
+                    sd = temp_sd
+            
             extra_sd = comfy.utils.load_torch_file(extra_state_dict)
             sd.update(extra_sd)
             del extra_sd