From 862b61830c3bfc48d92c9df3c0270e676753579e Mon Sep 17 00:00:00 2001 From: lllyasviel Date: Tue, 30 Jan 2024 21:58:16 -0800 Subject: [PATCH] i --- ldm_patched/modules/sd.py | 2 +- ldm_patched/modules/supported_models_base.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/ldm_patched/modules/sd.py b/ldm_patched/modules/sd.py index e197c39c..a8413307 100644 --- a/ldm_patched/modules/sd.py +++ b/ldm_patched/modules/sd.py @@ -462,7 +462,7 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o model.load_model_weights(sd, "model.diffusion_model.") if output_vae: - vae_sd = ldm_patched.modules.utils.state_dict_prefix_replace(sd, {"first_stage_model.": ""}, filter_keys=True) + vae_sd = ldm_patched.modules.utils.state_dict_prefix_replace(sd, {k: "" for k in model_config.vae_key_prefix}, filter_keys=True) vae_sd = model_config.process_vae_state_dict(vae_sd) vae = VAE(sd=vae_sd) diff --git a/ldm_patched/modules/supported_models_base.py b/ldm_patched/modules/supported_models_base.py index 5baf4bca..58535a9f 100644 --- a/ldm_patched/modules/supported_models_base.py +++ b/ldm_patched/modules/supported_models_base.py @@ -21,6 +21,7 @@ class BASE: noise_aug_config = None sampling_settings = {} latent_format = latent_formats.LatentFormat + vae_key_prefix = ["first_stage_model."] manual_cast_dtype = None