OPENGIT/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2025-01-04 05:45:05 +08:00
Code Issues Packages Projects Releases Wiki Activity

use existing function for loading VAE weights from file

Browse Source
This commit is contained in:
AUTOMATIC 2023-01-09 19:58:35 +03:00
parent 99da2c5af6
commit 49c4509ce2
1 changed files with 5 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
modules/sd_vae.py
View File

@ -3,7 +3,7 @@ import safetensors.torch
import os
import collections
from collections import namedtuple
from modules import shared, devices, script_callbacks
from modules import shared, devices, script_callbacks, sd_models
from modules.paths import models_path
import glob
from copy import deepcopy
@ -172,13 +172,8 @@ def load_vae(model, vae_file=None):
assert os.path.isfile(vae_file), f"VAE file doesn't exist: {vae_file}"
print(f"Loading VAE weights from: {vae_file}")
store_base_vae(model)
_, extension = os.path.splitext(vae_file)
if extension.lower() == ".safetensors":
vae_ckpt = safetensors.torch.load_file(vae_file, device=shared.weight_load_location)
else:
vae_ckpt = torch.load(vae_file, map_location=shared.weight_load_location)
if "state_dict" in vae_ckpt:
vae_ckpt = vae_ckpt["state_dict"]
vae_ckpt = sd_models.read_state_dict(vae_file, map_location=shared.weight_load_location)
vae_dict_1 = {k: v for k, v in vae_ckpt.items() if k[0:4] != "loss" and k not in vae_ignore_keys}
_load_vae_dict(model, vae_dict_1)
@ -210,10 +205,12 @@ def _load_vae_dict(model, vae_dict_1):
model.first_stage_model.load_state_dict(vae_dict_1)
model.first_stage_model.to(devices.dtype_vae)
def clear_loaded_vae():
global loaded_vae_file
loaded_vae_file = None
def reload_vae_weights(sd_model=None, vae_file="auto"):
from modules import lowvram, devices, sd_hijack