Added legacy paths

toolkit/models/v2/diffusion_models/auraflow.py

@@ -0,0 +1,13 @@
+from diffusers import AuraFlowTransformer2DModel as DiffusersAuraFlowTransformer2DModel
+
+from .._mixin import OstrisModelMixin
+
+
+class AuraFlowTransformer2DModel(
+    DiffusersAuraFlowTransformer2DModel, OstrisModelMixin
+):
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["joint_transformer_blocks", "single_transformer_blocks"]

toolkit/models/v2/diffusion_models/cosmos.py

@@ -0,0 +1,15 @@
+from diffusers.models import (
+    CosmosTransformer3DModel as DiffusersCosmosTransformer3DModel,
+)
+
+from .._mixin import OstrisModelMixin
+
+
+class CosmosTransformer3DModel(DiffusersCosmosTransformer3DModel, OstrisModelMixin):
+    """Cosmos video/image DiT (anima)."""
+
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["transformer_blocks"]

toolkit/models/v2/diffusion_models/ltx2.py

@@ -0,0 +1,34 @@
+from diffusers.models.transformers import (
+    LTX2VideoTransformer3DModel as DiffusersLTX2VideoTransformer3DModel,
+)
+from diffusers.pipelines.ltx2 import (
+    LTX2TextConnectors as DiffusersLTX2TextConnectors,
+)
+from diffusers.pipelines.ltx2 import LTX2Vocoder as DiffusersLTX2Vocoder
+from diffusers.pipelines.ltx2 import (
+    LTX2VocoderWithBWE as DiffusersLTX2VocoderWithBWE,
+)
+
+from .._mixin import OstrisModelMixin
+
+
+class LTX2VideoTransformer3DModel(
+    DiffusersLTX2VideoTransformer3DModel, OstrisModelMixin
+):
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["transformer_blocks"]
+
+
+class LTX2TextConnectors(DiffusersLTX2TextConnectors, OstrisModelMixin):
+    aitk_subfolder = "connectors"
+
+
+class LTX2Vocoder(DiffusersLTX2Vocoder, OstrisModelMixin):
+    aitk_subfolder = "vocoder"
+
+
+class LTX2VocoderWithBWE(DiffusersLTX2VocoderWithBWE, OstrisModelMixin):
+    aitk_subfolder = "vocoder"

toolkit/models/v2/diffusion_models/lumina2.py

@@ -0,0 +1,13 @@
+from diffusers import Lumina2Transformer2DModel as DiffusersLumina2Transformer2DModel
+
+from .._mixin import OstrisModelMixin
+
+
+class Lumina2Transformer2DModel(
+    DiffusersLumina2Transformer2DModel, OstrisModelMixin
+):
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["layers"]

toolkit/models/v2/diffusion_models/pixart.py

@@ -0,0 +1,24 @@
+from diffusers import PixArtTransformer2DModel as DiffusersPixArtTransformer2DModel
+from diffusers import Transformer2DModel as DiffusersTransformer2DModel
+
+from .._mixin import OstrisModelMixin
+
+
+class PixArtTransformer2DModel(
+    DiffusersPixArtTransformer2DModel, OstrisModelMixin
+):
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["transformer_blocks"]
+
+
+class Transformer2DModel(DiffusersTransformer2DModel, OstrisModelMixin):
+    """The generic diffusers DiT the pixart sigma path loads."""
+
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["transformer_blocks"]

toolkit/models/v2/diffusion_models/sd3.py

@@ -0,0 +1,11 @@
+from diffusers import SD3Transformer2DModel as DiffusersSD3Transformer2DModel
+
+from .._mixin import OstrisModelMixin
+
+
+class SD3Transformer2DModel(DiffusersSD3Transformer2DModel, OstrisModelMixin):
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["transformer_blocks"]

toolkit/models/v2/diffusion_models/unet.py

@@ -0,0 +1,13 @@
+from diffusers import UNet2DConditionModel as DiffusersUNet2DConditionModel
+
+from .._mixin import OstrisModelMixin
+
+
+class UNet2DConditionModel(DiffusersUNet2DConditionModel, OstrisModelMixin):
+    """The SD1/SD2/SDXL-family UNet."""
+
+    aitk_subfolder = "unet"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["down_blocks", "up_blocks"]