Browse Source

Merge pull request #6851 from ddPn08/master

Add `--vae-dir` argument
AUTOMATIC1111 2 years ago
parent
commit
43fd6eaab8
2 changed files with 8 additions and 0 deletions
  1. 7 0
      modules/sd_vae.py
  2. 1 0
      modules/shared.py

+ 7 - 0
modules/sd_vae.py

@@ -72,6 +72,13 @@ def refresh_vae_list():
             os.path.join(shared.cmd_opts.ckpt_dir, '**/*.vae.safetensors'),
             os.path.join(shared.cmd_opts.ckpt_dir, '**/*.vae.safetensors'),
         ]
         ]
 
 
+    if shared.cmd_opts.vae_dir is not None and os.path.isdir(shared.cmd_opts.vae_dir):
+        paths += [
+            os.path.join(shared.cmd_opts.vae_dir, '**/*.ckpt'),
+            os.path.join(shared.cmd_opts.vae_dir, '**/*.pt'),
+            os.path.join(shared.cmd_opts.vae_dir, '**/*.safetensors'),
+        ]
+
     candidates = []
     candidates = []
     for path in paths:
     for path in paths:
         candidates += glob.iglob(path, recursive=True)
         candidates += glob.iglob(path, recursive=True)

+ 1 - 0
modules/shared.py

@@ -26,6 +26,7 @@ parser = argparse.ArgumentParser()
 parser.add_argument("--config", type=str, default=os.path.join(script_path, "configs/v1-inference.yaml"), help="path to config which constructs model",)
 parser.add_argument("--config", type=str, default=os.path.join(script_path, "configs/v1-inference.yaml"), help="path to config which constructs model",)
 parser.add_argument("--ckpt", type=str, default=sd_model_file, help="path to checkpoint of stable diffusion model; if specified, this checkpoint will be added to the list of checkpoints and loaded",)
 parser.add_argument("--ckpt", type=str, default=sd_model_file, help="path to checkpoint of stable diffusion model; if specified, this checkpoint will be added to the list of checkpoints and loaded",)
 parser.add_argument("--ckpt-dir", type=str, default=None, help="Path to directory with stable diffusion checkpoints")
 parser.add_argument("--ckpt-dir", type=str, default=None, help="Path to directory with stable diffusion checkpoints")
+parser.add_argument("--vae-dir", type=str, default=None, help="Path to directory with VAE files")
 parser.add_argument("--gfpgan-dir", type=str, help="GFPGAN directory", default=('./src/gfpgan' if os.path.exists('./src/gfpgan') else './GFPGAN'))
 parser.add_argument("--gfpgan-dir", type=str, help="GFPGAN directory", default=('./src/gfpgan' if os.path.exists('./src/gfpgan') else './GFPGAN'))
 parser.add_argument("--gfpgan-model", type=str, help="GFPGAN model file name", default=None)
 parser.add_argument("--gfpgan-model", type=str, help="GFPGAN model file name", default=None)
 parser.add_argument("--no-half", action='store_true', help="do not switch the model to 16-bit floats")
 parser.add_argument("--no-half", action='store_true', help="do not switch the model to 16-bit floats")