Phase 1

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"]