diff --git a/advanced_model_merging.py b/advanced_model_merging.py new file mode 100644 index 0000000..2a69acf --- /dev/null +++ b/advanced_model_merging.py @@ -0,0 +1,30 @@ +import comfy_extras.nodes_model_merging + +class ModelMergeBlockNumber(comfy_extras.nodes_model_merging.ModelMergeBlocks): + @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["time_embed."] = argument + arg_dict["label_emb."] = argument + + for i in range(12): + arg_dict["input_blocks.{}".format(i)] = argument + + for i in range(3): + arg_dict["middle_block.{}".format(i)] = argument + + for i in range(12): + arg_dict["output_blocks.{}".format(i)] = argument + + arg_dict["out."] = argument + + return {"required": arg_dict} + + +NODE_CLASS_MAPPINGS = { + "ModelMergeBlockNumber": ModelMergeBlockNumber, +}