Support saving some more modelspec types.

comfyanonymous · May 14, 2024 · 2de3b69 · 2de3b69
1 parent cf6e1ef
commit 2de3b69
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/comfy_extras/nodes_model_merging.py b/comfy_extras/nodes_model_merging.py
@@ -175,9 +175,14 @@ def save_checkpoint(model, clip=None, vae=None, clip_vision=None, filename_prefi
 
     enable_modelspec = True
     if isinstance(model.model, comfy.model_base.SDXL):
-        metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-base"
+        if isinstance(model.model, comfy.model_base.SDXL_instructpix2pix):
+            metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-edit"
+        else:
+            metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-base"
     elif isinstance(model.model, comfy.model_base.SDXLRefiner):
         metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-refiner"
+    elif isinstance(model.model, comfy.model_base.SVD_img2vid):
+        metadata["modelspec.architecture"] = "stable-video-diffusion-img2vid-v1"
     else:
         enable_modelspec = False