begin to use new vae impl

modules_forge/forge_loader.py

@@ -18,6 +18,7 @@ from ldm.util import instantiate_from_config
 from modules_forge import forge_clip
 from modules_forge.unet_patcher import UnetPatcher
 from ldm_patched.modules.model_base import model_sampling, ModelType
+from backend.vae.loader import load_vae_from_state_dict
 
 import open_clip
 from transformers import CLIPTextModel, CLIPTokenizer
@@ -105,9 +106,8 @@ def load_checkpoint_guess_config(sd, output_vae=True, output_clip=True, output_c
         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 = model_config.process_vae_state_dict(vae_sd)
-        vae = VAE(sd=vae_sd)
+        vae, mapping = load_vae_from_state_dict(sd)
+        vae = VAE(model=vae, mapping=mapping)
 
     if output_clip:
         w = WeightsLoader()