Phase 1

toolkit/models/v2/diffusion_models/flux.py

@@ -0,0 +1,14 @@
+from diffusers import FluxTransformer2DModel as DiffusersFluxTransformer2DModel
+
+from .._mixin import OstrisModelMixin
+
+
+class FluxTransformer2DModel(DiffusersFluxTransformer2DModel, OstrisModelMixin):
+    """Flux1-family DiT (flux, flux_kontext, chroma-adjacent finetunes in
+    diffusers layout)."""
+
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["transformer_blocks", "single_transformer_blocks"]

toolkit/models/v2/diffusion_models/hidream.py

@@ -0,0 +1,18 @@
+from diffusers.models import (
+    HiDreamImageTransformer2DModel as DiffusersHiDreamImageTransformer2DModel,
+)
+
+from .._mixin import OstrisModelMixin
+
+
+class HiDreamImageTransformer2DModel(
+    DiffusersHiDreamImageTransformer2DModel, OstrisModelMixin
+):
+    """The diffusers HiDream DiT (hidream_e1; the base hidream arch uses the
+    vendored copy in the hidream extension)."""
+
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["double_stream_blocks", "single_stream_blocks"]

toolkit/models/v2/diffusion_models/nucleus_image.py

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

toolkit/models/v2/diffusion_models/qwen_image.py

@@ -0,0 +1,29 @@
+from diffusers import (
+    QwenImageTransformer2DModel as DiffusersQwenImageTransformer2DModel,
+)
+
+from .._mixin import OstrisModelMixin
+
+
+class QwenImageTransformer2DModel(
+    DiffusersQwenImageTransformer2DModel, OstrisModelMixin
+):
+    aitk_subfolder = "transformer"
+    aitk_config_repo = "Qwen/Qwen-Image"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["transformer_blocks"]
+
+    @classmethod
+    def _load_single_file(cls, file_path, dtype, config_path=None, subfolder=None):
+        # single-file checkpoints in the wild carry diffusers or original key
+        # layouts; diffusers' single-file machinery owns that conversion
+        model = cls.from_single_file(
+            file_path,
+            config=config_path if config_path is not None else cls.aitk_config_repo,
+            subfolder="transformer",
+            torch_dtype=dtype,
+        )
+        model.to(dtype)
+        return model

toolkit/models/v2/diffusion_models/wan.py

@@ -0,0 +1,13 @@
+from diffusers import WanTransformer3DModel as DiffusersWanTransformer3DModel
+
+from .._mixin import OstrisModelMixin
+
+
+class WanTransformer3DModel(DiffusersWanTransformer3DModel, OstrisModelMixin):
+    """Wan 2.1/2.2 video DiT (wan22 loads two of these into its dual wrapper)."""
+
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["blocks"]