Skip to content

Commit

Permalink
Add advanced model merge node for Flux model.
Browse files Browse the repository at this point in the history
  • Loading branch information
comfyanonymous committed Aug 3, 2024
1 parent ea03c9d commit 7cd0cdf
Showing 1 changed file with 27 additions and 0 deletions.
27 changes: 27 additions & 0 deletions comfy_extras/nodes_model_merging_model_specific.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,9 +75,36 @@ def INPUT_TYPES(s):

return {"required": arg_dict}

class ModelMergeFlux1(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["img_in."] = argument
arg_dict["time_in."] = argument
arg_dict["guidance_in"] = argument
arg_dict["vector_in."] = argument
arg_dict["txt_in."] = argument

for i in range(19):
arg_dict["double_blocks.{}.".format(i)] = argument

for i in range(38):
arg_dict["single_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
"ModelMergeSDXL": ModelMergeSDXL,
"ModelMergeSD3_2B": ModelMergeSD3_2B,
"ModelMergeFlux1": ModelMergeFlux1,
}

0 comments on commit 7cd0cdf

Please sign in to comment.