diff --git a/comfy_extras/nodes_model_merging_model_specific.py b/comfy_extras/nodes_model_merging_model_specific.py index e3529e46..73091626 100644 --- a/comfy_extras/nodes_model_merging_model_specific.py +++ b/comfy_extras/nodes_model_merging_model_specific.py @@ -75,6 +75,34 @@ class ModelMergeSD3_2B(comfy_extras.nodes_model_merging.ModelMergeBlocks): return {"required": arg_dict} + +class ModelMergeAuraflow(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["init_x_linear."] = argument + arg_dict["positional_encoding"] = argument + arg_dict["cond_seq_linear."] = argument + arg_dict["register_tokens"] = argument + arg_dict["t_embedder."] = argument + + for i in range(4): + arg_dict["double_layers.{}.".format(i)] = argument + + for i in range(32): + arg_dict["single_layers.{}.".format(i)] = argument + + arg_dict["modF."] = argument + arg_dict["final_linear."] = argument + + return {"required": arg_dict} + class ModelMergeFlux1(comfy_extras.nodes_model_merging.ModelMergeBlocks): CATEGORY = "advanced/model_merging/model_specific" @@ -151,6 +179,7 @@ NODE_CLASS_MAPPINGS = { "ModelMergeSD2": ModelMergeSD1, #SD1 and SD2 have the same blocks "ModelMergeSDXL": ModelMergeSDXL, "ModelMergeSD3_2B": ModelMergeSD3_2B, + "ModelMergeAuraflow": ModelMergeAuraflow, "ModelMergeFlux1": ModelMergeFlux1, "ModelMergeSD35_Large": ModelMergeSD35_Large, "ModelMergeMochiPreview": ModelMergeMochiPreview,