From da7a8df0d2582c8dc91e5afafe51300899c91392 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 30 Jan 2024 02:24:38 -0500 Subject: [PATCH] Put VAE key name in model config. --- comfy/sd.py | 2 +- comfy/supported_models_base.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/comfy/sd.py b/comfy/sd.py index 9ca9d1d1..c15d73fe 100644 --- a/comfy/sd.py +++ b/comfy/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 = comfy.utils.state_dict_prefix_replace(sd, {"first_stage_model.": ""}, filter_keys=True) + vae_sd = comfy.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/comfy/supported_models_base.py b/comfy/supported_models_base.py index 5baf4bca..58535a9f 100644 --- a/comfy/supported_models_base.py +++ b/comfy/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