diff --git a/__init__.py b/__init__.py index 19c95c7..0e4d63f 100644 --- a/__init__.py +++ b/__init__.py @@ -1,15 +1,19 @@ from .merge.dare import DareModelMerger +from .merge.dareMBW import DareModelMergerMBW from .merge.mag import MagnitudePruningModelMerger from .merge.block import BlockModelMergerAdv + NODE_CLASS_MAPPINGS = { "DareModelMerger": DareModelMerger, + "DareModelMergerMBW": DareModelMergerMBW, "MagnitudeModelMerger": MagnitudePruningModelMerger, "BlockModelMergerAdv": BlockModelMergerAdv, } NODE_DISPLAY_NAME_MAPPINGS = { "DareModelMerger": "ModelMergeByDARE", + "DareModelMergerMBW": "ModelMergeByDAREMBW", "MagnitudeModelMerger": "ModelMergeByMagnitudePruning", "BlockModelMergerAdv": "ModelMergeByBlock (Advanced)", }