From 887a6341ed14c9904230cf55a0eabe95cd0218d3 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Fri, 21 Jun 2024 08:41:31 -0400
Subject: [PATCH] Proper ModelMergeSD3_2B node.

---
 comfy_extras/nodes_model_merging_model_specific.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/comfy_extras/nodes_model_merging_model_specific.py b/comfy_extras/nodes_model_merging_model_specific.py
index 74ab681..df111bd 100644
--- a/comfy_extras/nodes_model_merging_model_specific.py
+++ b/comfy_extras/nodes_model_merging_model_specific.py
@@ -52,7 +52,7 @@ class ModelMergeSDXL(comfy_extras.nodes_model_merging.ModelMergeBlocks):
 
         return {"required": arg_dict}
 
-class ModelMergeSD3(comfy_extras.nodes_model_merging.ModelMergeBlocks):
+class ModelMergeSD3_2B(comfy_extras.nodes_model_merging.ModelMergeBlocks):
     CATEGORY = "advanced/model_merging/model_specific"
 
     @classmethod
@@ -68,7 +68,7 @@ class ModelMergeSD3(comfy_extras.nodes_model_merging.ModelMergeBlocks):
         arg_dict["y_embedder."] = argument
         arg_dict["t_embedder."] = argument
 
-        for i in range(38):
+        for i in range(24):
             arg_dict["joint_blocks.{}.".format(i)] = argument
 
         arg_dict["final_layer."] = argument
@@ -79,5 +79,5 @@ NODE_CLASS_MAPPINGS = {
     "ModelMergeSD1": ModelMergeSD1,
     "ModelMergeSD2": ModelMergeSD1, #SD1 and SD2 have the same blocks
     "ModelMergeSDXL": ModelMergeSDXL,
-    "ModelMergeSD3": ModelMergeSD3,
+    "ModelMergeSD3_2B": ModelMergeSD3_2B,
 }