feat: add VAE select

modules/async_worker.py

@@ -166,6 +166,7 @@ def worker():
         adaptive_cfg = args.pop()
         sampler_name = args.pop()
         scheduler_name = args.pop()
+        vae_name = args.pop()
         overwrite_step = args.pop()
         overwrite_switch = args.pop()
         overwrite_width = args.pop()
@@ -428,7 +429,7 @@ def worker():
             progressbar(async_task, 3, 'Loading models ...')
             pipeline.refresh_everything(refiner_model_name=refiner_model_name, base_model_name=base_model_name,
                                         loras=loras, base_model_additional_loras=base_model_additional_loras,
-                                        use_synthetic_refiner=use_synthetic_refiner)
+                                        use_synthetic_refiner=use_synthetic_refiner, vae_name=vae_name)
 
             progressbar(async_task, 3, 'Processing prompts ...')
             tasks = []

modules/config.py

@@ -189,6 +189,7 @@ paths_checkpoints = get_dir_or_set_default('path_checkpoints', ['../models/check
 paths_loras = get_dir_or_set_default('path_loras', ['../models/loras/'], True)
 path_embeddings = get_dir_or_set_default('path_embeddings', '../models/embeddings/')
 path_vae_approx = get_dir_or_set_default('path_vae_approx', '../models/vae_approx/')
+path_vae = get_dir_or_set_default('path_vae', '../models/vae/')
 path_upscale_models = get_dir_or_set_default('path_upscale_models', '../models/upscale_models/')
 path_inpaint = get_dir_or_set_default('path_inpaint', '../models/inpaint/')
 path_controlnet = get_dir_or_set_default('path_controlnet', '../models/controlnet/')
@@ -346,6 +347,11 @@ default_scheduler = get_config_item_or_set_default(
     default_value='karras',
     validator=lambda x: x in modules.flags.scheduler_list
 )
+default_vae = get_config_item_or_set_default(
+    key='default_vae',
+    default_value='None',
+    validator=lambda x: isinstance(x, str)
+)
 default_styles = get_config_item_or_set_default(
     key='default_styles',
     default_value=[
@@ -535,6 +541,7 @@ with open(config_example_path, "w", encoding="utf-8") as json_file:
 
 model_filenames = []
 lora_filenames = []
+vae_filenames = []
 wildcard_filenames = []
 
 sdxl_lcm_lora = 'sdxl_lcm_lora.safetensors'
@@ -546,15 +553,20 @@ def get_model_filenames(folder_paths, extensions=None, name_filter=None):
     if extensions is None:
         extensions = ['.pth', '.ckpt', '.bin', '.safetensors', '.fooocus.patch']
     files = []
+
+    if not isinstance(folder_paths, list):
+        folder_paths = [folder_paths]
     for folder in folder_paths:
         files += get_files_from_folder(folder, extensions, name_filter)
+
     return files
 
 
 def update_files():
-    global model_filenames, lora_filenames, wildcard_filenames, available_presets
+    global model_filenames, lora_filenames, vae_filenames, wildcard_filenames, available_presets
     model_filenames = get_model_filenames(paths_checkpoints)
     lora_filenames = get_model_filenames(paths_loras)
+    vae_filenames = get_model_filenames(path_vae)
     wildcard_filenames = get_files_from_folder(path_wildcards, ['.txt'])
     available_presets = get_presets()
     return

modules/core.py

@@ -35,12 +35,13 @@ opModelSamplingDiscrete = ModelSamplingDiscrete()
 
 
 class StableDiffusionModel:
-    def __init__(self, unet=None, vae=None, clip=None, clip_vision=None, filename=None):
+    def __init__(self, unet=None, vae=None, clip=None, clip_vision=None, filename=None, vae_filename=None):
         self.unet = unet
         self.vae = vae
         self.clip = clip
         self.clip_vision = clip_vision
         self.filename = filename
+        self.vae_filename = filename
         self.unet_with_lora = unet
         self.clip_with_lora = clip
         self.visited_loras = ''
@@ -142,8 +143,9 @@ def apply_controlnet(positive, negative, control_net, image, strength, start_per
 
 @torch.no_grad()
 @torch.inference_mode()
-def load_model(ckpt_filename):
-    unet, clip, vae, clip_vision = load_checkpoint_guess_config(ckpt_filename, embedding_directory=path_embeddings)
+def load_model(ckpt_filename, vae_filename=None):
+    unet, clip, vae, clip_vision = load_checkpoint_guess_config(ckpt_filename, embedding_directory=path_embeddings,
+                                                                vae_filename=vae_filename)
     return StableDiffusionModel(unet=unet, clip=clip, vae=vae, clip_vision=clip_vision, filename=ckpt_filename)
 
 

modules/default_pipeline.py

@@ -58,16 +58,20 @@ def assert_model_integrity():
 
 @torch.no_grad()
 @torch.inference_mode()
-def refresh_base_model(name):
+def refresh_base_model(name, vae_name=None):
     global model_base
 
     filename = get_file_from_folder_list(name, modules.config.paths_checkpoints)
 
-    if model_base.filename == filename:
+    vae_filename = None
+    if vae_name is not None and vae_name != 'None':
+        vae_filename = get_file_from_folder_list(vae_name, modules.config.path_vae)
+
+    if model_base.filename == filename and model_base.vae_filename == vae_filename:
         return
 
     model_base = core.StableDiffusionModel()
-    model_base = core.load_model(filename)
+    model_base = core.load_model(filename, vae_filename)
     print(f'Base model loaded: {model_base.filename}')
     return
 
@@ -216,7 +220,7 @@ def prepare_text_encoder(async_call=True):
 @torch.no_grad()
 @torch.inference_mode()
 def refresh_everything(refiner_model_name, base_model_name, loras,
-                       base_model_additional_loras=None, use_synthetic_refiner=False):
+                       base_model_additional_loras=None, use_synthetic_refiner=False, vae_name=None):
     global final_unet, final_clip, final_vae, final_refiner_unet, final_refiner_vae, final_expansion
 
     final_unet = None
@@ -227,11 +231,11 @@ def refresh_everything(refiner_model_name, base_model_name, loras,
 
     if use_synthetic_refiner and refiner_model_name == 'None':
         print('Synthetic Refiner Activated')
-        refresh_base_model(base_model_name)
+        refresh_base_model(base_model_name, vae_name)
         synthesize_refiner_model()
     else:
         refresh_refiner_model(refiner_model_name)
-        refresh_base_model(base_model_name)
+        refresh_base_model(base_model_name, vae_name)
 
     refresh_loras(loras, base_model_additional_loras=base_model_additional_loras)
     assert_model_integrity()
@@ -254,7 +258,8 @@ def refresh_everything(refiner_model_name, base_model_name, loras,
 refresh_everything(
     refiner_model_name=modules.config.default_refiner_model_name,
     base_model_name=modules.config.default_base_model_name,
-    loras=get_enabled_loras(modules.config.default_loras)
+    loras=get_enabled_loras(modules.config.default_loras),
+    vae_name=modules.config.default_vae,
 )
 
 

modules/util.py

@@ -371,6 +371,9 @@ def is_json(data: str) -> bool:
 
 
 def get_file_from_folder_list(name, folders):
+    if not isinstance(folders, list):
+        folders = [folders]
+
     for folder in folders:
         filename = os.path.abspath(os.path.realpath(os.path.join(folder, name)))
         if os.path.isfile(filename):