From 122c9ca1cec50e78fb0fb0eb7a3d7fd015e7f037 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Thu, 14 Nov 2024 07:51:20 -0500
Subject: [PATCH] Add advanced model merging node for mochi.

---
 .../nodes_model_merging_model_specific.py     | 23 +++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/comfy_extras/nodes_model_merging_model_specific.py b/comfy_extras/nodes_model_merging_model_specific.py
index 0e847192..e3529e46 100644
--- a/comfy_extras/nodes_model_merging_model_specific.py
+++ b/comfy_extras/nodes_model_merging_model_specific.py
@@ -124,6 +124,28 @@ class ModelMergeSD35_Large(comfy_extras.nodes_model_merging.ModelMergeBlocks):
 
         return {"required": arg_dict}
 
+class ModelMergeMochiPreview(comfy_extras.nodes_model_merging.ModelMergeBlocks):
+    CATEGORY = "advanced/model_merging/model_specific"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        arg_dict = { "model1": ("MODEL",),
+                              "model2": ("MODEL",)}
+
+        argument = ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01})
+
+        arg_dict["pos_frequencies."] = argument
+        arg_dict["t_embedder."] = argument
+        arg_dict["t5_y_embedder."] = argument
+        arg_dict["t5_yproj."] = argument
+
+        for i in range(48):
+            arg_dict["blocks.{}.".format(i)] = argument
+
+        arg_dict["final_layer."] = argument
+
+        return {"required": arg_dict}
+
 NODE_CLASS_MAPPINGS = {
     "ModelMergeSD1": ModelMergeSD1,
     "ModelMergeSD2": ModelMergeSD1, #SD1 and SD2 have the same blocks
@@ -131,4 +153,5 @@ NODE_CLASS_MAPPINGS = {
     "ModelMergeSD3_2B": ModelMergeSD3_2B,
     "ModelMergeFlux1": ModelMergeFlux1,
     "ModelMergeSD35_Large": ModelMergeSD35_Large,
+    "ModelMergeMochiPreview": ModelMergeMochiPreview,
 }