improve vae key mapping

backend/state_dict.py

@@ -24,16 +24,20 @@ def split_state_dict_with_prefix(sd, prefix):
     return vae_sd
 
 
-def shrink_last_key(t):
-    ts = t.split('.')
-    del ts[-1]
-    return '.'.join(ts)
-
-
 def compile_state_dict(state_dict):
     sd = {}
     mapping = {}
     for k, v in state_dict.items():
         sd[k] = v.value
-        mapping[shrink_last_key(v.key)] = shrink_last_key(k)
+        mapping[v.key] = (k, v.advanced_indexing)
     return sd, mapping
+
+
+def map_state_dict(sd, mapping):
+    new_sd = {}
+    for k, v in sd.items():
+        k, indexing = mapping.get(k, (k, None))
+        if indexing is not None:
+            v = v[indexing]
+        new_sd[k] = v
+    return new_sd

backend/vae/loader.py

@@ -7,6 +7,11 @@ from diffusers.loaders.single_file_model import convert_ldm_vae_checkpoint
 
 
 class BaseVAE(AutoencoderKL):
+    def __init__(self, *args, **kwargs):
+
+        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:
@@ -31,5 +36,6 @@ def load_vae_from_state_dict(state_dict):
     vae_state_dict, mapping = compile_state_dict(vae_state_dict)
     model.load_state_dict(vae_state_dict, strict=True)
     model.set_attn_processor(AttentionProcessorForge())
+    model.state_dict_mapping = mapping
 
-    return model, mapping
+    return model