aboutsummaryrefslogtreecommitdiffstats
path: root/modules/sd_vae.py
diff options
context:
space:
mode:
authorAUTOMATIC1111 <16777216c@gmail.com>2023-05-18 07:12:17 +0000
committerGitHub <noreply@github.com>2023-05-18 07:12:17 +0000
commit182330ae40c93b761fced8d03bbdd8523f1739d0 (patch)
treec95f2e531cc34475eb39266add6d12a2a5e1e2b3 /modules/sd_vae.py
parent0d31f20cbd556ea4ba3d8ad9254bcce71c32088c (diff)
parent983f2c494ad8fed2f08193681ba0bf5dda01555a (diff)
downloadstable-diffusion-webui-gfx803-182330ae40c93b761fced8d03bbdd8523f1739d0.tar.gz
stable-diffusion-webui-gfx803-182330ae40c93b761fced8d03bbdd8523f1739d0.tar.bz2
stable-diffusion-webui-gfx803-182330ae40c93b761fced8d03bbdd8523f1739d0.zip
Merge branch 'dev' into ngrok-py
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