aboutsummaryrefslogtreecommitdiffstats
path: root/modules/sd_vae.py
diff options
context:
space:
mode:
authorAUTOMATIC1111 <16777216c@gmail.com>2023-07-08 13:50:23 +0000
committerGitHub <noreply@github.com>2023-07-08 13:50:23 +0000
commit993dd9a8927407de8d19142cacb07e6f76686a67 (patch)
tree89df31c33ecf054c9cb70aaf38fc3382f306f450 /modules/sd_vae.py
parentff6acd35d0807a4e0c3ee86cdb1520a4a3a11cdd (diff)
parentd7d6e8cfc8b85a99a48f82975ee213d487783c28 (diff)
downloadstable-diffusion-webui-gfx803-993dd9a8927407de8d19142cacb07e6f76686a67.tar.gz
stable-diffusion-webui-gfx803-993dd9a8927407de8d19142cacb07e6f76686a67.tar.bz2
stable-diffusion-webui-gfx803-993dd9a8927407de8d19142cacb07e6f76686a67.zip
Merge branch 'dev' into patch-1
Diffstat (limited to 'modules/sd_vae.py')
-rw-r--r--modules/sd_vae.py11
1 files changed, 4 insertions, 7 deletions
diff --git a/modules/sd_vae.py b/modules/sd_vae.py
index 521e485a..e4ff2994 100644
--- a/modules/sd_vae.py
+++ b/modules/sd_vae.py
@@ -1,8 +1,5 @@
-import torch
-import safetensors.torch
import os
import collections
-from collections import namedtuple
from modules import paths, shared, devices, script_callbacks, sd_models
import glob
from copy import deepcopy
@@ -88,10 +85,10 @@ def refresh_vae_list():
def find_vae_near_checkpoint(checkpoint_file):
- checkpoint_path = os.path.splitext(checkpoint_file)[0]
- for vae_location in [f"{checkpoint_path}.vae.pt", f"{checkpoint_path}.vae.ckpt", f"{checkpoint_path}.vae.safetensors"]:
- if os.path.isfile(vae_location):
- return vae_location
+ checkpoint_path = os.path.basename(checkpoint_file).rsplit('.', 1)[0]
+ for vae_file in vae_dict.values():
+ if os.path.basename(vae_file).startswith(checkpoint_path):
+ return vae_file
return None