Put VAE key name in model config.

comfyanonymous · Jan 30, 2024 · da7a8df · da7a8df
1 parent ed2fa10
commit da7a8df
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/comfy/sd.py b/comfy/sd.py
@@ -462,7 +462,7 @@ class WeightsLoader(torch.nn.Module):
         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
@@ -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