Add VAE selector to the Parameter Generator node #15

receyuki · Nov 21, 2023 · fc11266 · fc11266
1 parent d02265d
commit fc11266
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/nodes.py b/nodes.py
@@ -483,6 +483,10 @@ def INPUT_TYPES(s):
                 "ckpt_name": (folder_paths.get_filename_list("checkpoints"),),
             },
             "optional": {
+                "vae_name": (
+                    ["baked VAE"] + folder_paths.get_filename_list("vae"),
+                    {"default": "baked VAE"},
+                ),
                 "model_version": (
                     list(SDParameterGenerator.MODEL_SCALING_FACTOR.keys()),
                     {"default": "SDv1 512px"},
@@ -586,8 +590,9 @@ def INPUT_TYPES(s):
 
     def generate_parameter(
         self,
-        model_version,
         ckpt_name,
+        vae_name,
+        model_version,
         config_name,
         seed,
         steps,
@@ -622,6 +627,12 @@ def generate_parameter(
                 embedding_directory=folder_paths.get_folder_paths("embeddings"),
             )[:3]
 
+        if vae_name != "baked VAE":
+            vae_path = folder_paths.get_full_path("vae", vae_name)
+            sd = comfy.utils.load_torch_file(vae_path)
+            vae = comfy.sd.VAE(sd=sd)
+            checkpoint = (*checkpoint[:2], vae)
+
         if aspect_ratio != "custom":
             aspect_ratio_value = aspect_ratio.split(" - ")[0]
             width = int(