From 310ad09258cd6ba80c44d31e183063455726dcb8 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Sat, 17 Aug 2024 21:31:15 -0400
Subject: [PATCH] Add a ModelSave node.

---
 comfy_extras/nodes_model_merging.py | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/comfy_extras/nodes_model_merging.py b/comfy_extras/nodes_model_merging.py
index 136f9a98..8fb8bf79 100644
--- a/comfy_extras/nodes_model_merging.py
+++ b/comfy_extras/nodes_model_merging.py
@@ -333,6 +333,25 @@ def save(self, vae, filename_prefix, prompt=None, extra_pnginfo=None):
         comfy.utils.save_torch_file(vae.get_sd(), output_checkpoint, metadata=metadata)
         return {}
 
+class ModelSave:
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": { "model": ("MODEL",),
+                              "filename_prefix": ("STRING", {"default": "diffusion_models/ComfyUI"}),},
+                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},}
+    RETURN_TYPES = ()
+    FUNCTION = "save"
+    OUTPUT_NODE = True
+
+    CATEGORY = "advanced/model_merging"
+
+    def save(self, model, filename_prefix, prompt=None, extra_pnginfo=None):
+        save_checkpoint(model, filename_prefix=filename_prefix, output_dir=self.output_dir, prompt=prompt, extra_pnginfo=extra_pnginfo)
+        return {}
+
 NODE_CLASS_MAPPINGS = {
     "ModelMergeSimple": ModelMergeSimple,
     "ModelMergeBlocks": ModelMergeBlocks,
@@ -344,4 +363,5 @@ def save(self, vae, filename_prefix, prompt=None, extra_pnginfo=None):
     "CLIPMergeAdd": CLIPAdd,
     "CLIPSave": CLIPSave,
     "VAESave": VAESave,
+    "ModelSave": ModelSave,
 }