Added legacy paths

toolkit/models/v2/vae/autoencoder_kl_flux2.py

@@ -0,0 +1,10 @@
+from diffusers import AutoencoderKLFlux2
+
+from .._mixin import OstrisModelMixin
+
+
+class Flux2KLVAE(AutoencoderKLFlux2, OstrisModelMixin):
+    """The diffusers Flux2 KL VAE (ernie_image; distinct from the hand-rolled
+    BFL AutoEncoder in vae/flux2_kl.py that flux2/ideogram4 use)."""
+
+    aitk_subfolder = "vae"

toolkit/models/v2/vae/flux2_kl.py

@@ -20,6 +20,8 @@ import torch.utils.checkpoint as ckpt
 from einops import rearrange
 from torch import Tensor, nn
 
+from .._mixin import OstrisModelMixin
+
 
 @dataclass
 class AutoEncoderParams:
@@ -354,7 +356,14 @@ class Decoder(nn.Module):
         return h
 
 
-class AutoEncoder(nn.Module):
+class AutoEncoder(nn.Module, OstrisModelMixin):
+    @classmethod
+    def aitk_config_from_state_dict(cls, state_dict):
+        # small decoder builds report 96ch at the first decoder block
+        if state_dict["decoder.up.0.block.0.conv1.bias"].shape[0] == 96:
+            return AutoEncoderSmallDecoderParams()
+        return AutoEncoderParams()
+
     def __init__(self, params: AutoEncoderParams):
         super().__init__()
         self.params = params

toolkit/models/v2/vae/ltx2.py

@@ -0,0 +1,16 @@
+from diffusers.models.autoencoders import (
+    AutoencoderKLLTX2Audio as DiffusersAutoencoderKLLTX2Audio,
+)
+from diffusers.models.autoencoders import (
+    AutoencoderKLLTX2Video as DiffusersAutoencoderKLLTX2Video,
+)
+
+from .._mixin import OstrisModelMixin
+
+
+class LTX2VideoVAE(DiffusersAutoencoderKLLTX2Video, OstrisModelMixin):
+    aitk_subfolder = "vae"
+
+
+class LTX2AudioVAE(DiffusersAutoencoderKLLTX2Audio, OstrisModelMixin):
+    aitk_subfolder = "audio_vae"

toolkit/models/v2/vae/wan.py

@@ -0,0 +1,9 @@
+from diffusers import AutoencoderKLWan
+
+from .._mixin import OstrisModelMixin
+
+
+class WanVAE(AutoencoderKLWan, OstrisModelMixin):
+    """The wan 2.1/2.2 causal video VAE."""
+
+    aitk_subfolder = "vae"