Add z-image-omni-base implementation

src/diffusers/__init__.py

@@ -279,6 +279,7 @@
             "WanTransformer3DModel",
             "WanVACETransformer3DModel",
             "ZImageTransformer2DModel",
+            "ZImageOmniTransformer2DModel",
             "attention_backend",
         ]
     )
@@ -668,6 +669,7 @@
             "WuerstchenPriorPipeline",
             "ZImageImg2ImgPipeline",
             "ZImagePipeline",
+            "ZImageOmniPipeline",
         ]
     )
 
@@ -1013,6 +1015,7 @@
             WanTransformer3DModel,
             WanVACETransformer3DModel,
             ZImageTransformer2DModel,
+            ZImageOmniTransformer2DModel,
             attention_backend,
         )
         from .modular_pipelines import ComponentsManager, ComponentSpec, ModularPipeline, ModularPipelineBlocks
@@ -1371,6 +1374,7 @@
             WuerstchenPriorPipeline,
             ZImageImg2ImgPipeline,
             ZImagePipeline,
+            ZImageOmniPipeline,
         )
 
     try:

src/diffusers/models/__init__.py

@@ -116,6 +116,7 @@
     _import_structure["transformers.transformer_wan_animate"] = ["WanAnimateTransformer3DModel"]
     _import_structure["transformers.transformer_wan_vace"] = ["WanVACETransformer3DModel"]
     _import_structure["transformers.transformer_z_image"] = ["ZImageTransformer2DModel"]
+    _import_structure["transformers.transformer_z_image_omni"] = ["ZImageOmniTransformer2DModel"]
     _import_structure["unets.unet_1d"] = ["UNet1DModel"]
     _import_structure["unets.unet_2d"] = ["UNet2DModel"]
     _import_structure["unets.unet_2d_condition"] = ["UNet2DConditionModel"]

src/diffusers/models/transformers/__init__.py

@@ -48,3 +48,4 @@
     from .transformer_wan_animate import WanAnimateTransformer3DModel
     from .transformer_wan_vace import WanVACETransformer3DModel
     from .transformer_z_image import ZImageTransformer2DModel
+    from .transformer_z_image_omni import ZImageOmniTransformer2DModel