ZImageControlTransformer2DModel

Author: hlky

src/diffusers/__init__.py

@@ -277,6 +277,7 @@
             "WanTransformer3DModel",
             "WanVACETransformer3DModel",
             "ZImageTransformer2DModel",
+            "ZImageControlTransformer2DModel",
             "ZImageControlNetModel",
             "attention_backend",
         ]
@@ -1006,6 +1007,7 @@
             WanTransformer3DModel,
             WanVACETransformer3DModel,
             ZImageControlNetModel,
+            ZImageControlTransformer2DModel,
             ZImageTransformer2DModel,
             attention_backend,
         )

src/diffusers/loaders/peft.py

@@ -65,6 +65,7 @@
     "QwenImageTransformer2DModel": lambda model_cls, weights: weights,
     "Flux2Transformer2DModel": lambda model_cls, weights: weights,
     "ZImageTransformer2DModel": lambda model_cls, weights: weights,
+    "ZImageControlTransformer2DModel": lambda model_cls, weights: weights,
 }
 
 

src/diffusers/models/__init__.py

@@ -117,6 +117,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_control"] = ["ZImageControlTransformer2DModel"]
     _import_structure["unets.unet_1d"] = ["UNet1DModel"]
     _import_structure["unets.unet_2d"] = ["UNet2DModel"]
     _import_structure["unets.unet_2d_condition"] = ["UNet2DConditionModel"]
@@ -231,6 +232,7 @@
             WanAnimateTransformer3DModel,
             WanTransformer3DModel,
             WanVACETransformer3DModel,
+            ZImageControlTransformer2DModel,
             ZImageTransformer2DModel,
         )
         from .unets import (