diff --git a/README.md b/README.md index 0de4a6bb5..6366280e7 100644 --- a/README.md +++ b/README.md @@ -273,6 +273,8 @@ You can install ComfyUI in Apple Mac silicon (M1 or M2) with any recent macOS ve #### DirectML (AMD Cards on Windows) +This is very badly supported and is not recommended. There are some unofficial builds of pytorch ROCm on windows that exist that will give you a much better experience than this. This readme will be updated once official pytorch ROCm builds for windows come out. + ```pip install torch-directml``` Then you can launch ComfyUI with: ```python main.py --directml``` #### Ascend NPUs diff --git a/comfy_extras/nodes_model_merging_model_specific.py b/comfy_extras/nodes_model_merging_model_specific.py index dc3411947..2c93cd84f 100644 --- a/comfy_extras/nodes_model_merging_model_specific.py +++ b/comfy_extras/nodes_model_merging_model_specific.py @@ -268,6 +268,52 @@ class ModelMergeWAN2_1(comfy_extras.nodes_model_merging.ModelMergeBlocks): return {"required": arg_dict} +class ModelMergeCosmosPredict2_2B(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_embedder."] = argument + arg_dict["x_embedder."] = argument + arg_dict["t_embedder."] = argument + arg_dict["t_embedding_norm."] = argument + + + for i in range(28): + arg_dict["blocks.{}.".format(i)] = argument + + arg_dict["final_layer."] = argument + + return {"required": arg_dict} + +class ModelMergeCosmosPredict2_14B(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_embedder."] = argument + arg_dict["x_embedder."] = argument + arg_dict["t_embedder."] = argument + arg_dict["t_embedding_norm."] = argument + + + for i in range(36): + 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 @@ -281,4 +327,6 @@ NODE_CLASS_MAPPINGS = { "ModelMergeCosmos7B": ModelMergeCosmos7B, "ModelMergeCosmos14B": ModelMergeCosmos14B, "ModelMergeWAN2_1": ModelMergeWAN2_1, + "ModelMergeCosmosPredict2_2B": ModelMergeCosmosPredict2_2B, + "ModelMergeCosmosPredict2_14B": ModelMergeCosmosPredict2_14B, }