Intergrate Native AutoEncoderKL

backend/vae.py

@@ -1,38 +1,14 @@
-from diffusers import AutoencoderKL
-from backend.state_dict import split_state_dict_with_prefix, compile_state_dict
+from backend.state_dict import filter_state_dict_with_prefix
 from backend.operations import using_forge_operations
-from backend.attention import AttentionProcessorForge
-from diffusers.loaders.single_file_model import convert_ldm_vae_checkpoint
+from backend.nn.autoencoder_kl import IntegratedAutoencoderKL
 
 
-class BaseAutoencoderKL(AutoencoderKL):
-    def __init__(self, *args, **kwargs):
+def load_vae(state_dict, config_path):
+    config = IntegratedAutoencoderKL.load_config(config_path)
 
-        super().__init__(*args, **kwargs)
-        self.state_dict_mapping = {}
-
-    def encode(self, x, regulation=None, mode=False):
-        latent_dist = super().encode(x).latent_dist
-        if mode:
-            return latent_dist.mode()
-        elif regulation is not None:
-            return regulation(latent_dist)
-        else:
-            return latent_dist.sample()
-
-    def decode(self, x):
-        return super().decode(x).sample
-
-
-def load_vae(state_dict, config):
     with using_forge_operations():
-        model = BaseAutoencoderKL(**config)
+        model = IntegratedAutoencoderKL.from_config(config)
 
-    vae_state_dict = split_state_dict_with_prefix(state_dict, "first_stage_model.")
-    vae_state_dict = convert_ldm_vae_checkpoint(vae_state_dict, config)
-    vae_state_dict, mapping = compile_state_dict(vae_state_dict)
+    vae_state_dict = filter_state_dict_with_prefix(state_dict, "first_stage_model.")
     model.load_state_dict(vae_state_dict, strict=True)
-    model.set_attn_processor(AttentionProcessorForge())
-    model.state_dict_mapping = mapping
-
     return model