Enable regional compilation on z-image transformer model (#12736)

up
2025-12-06 12:34:13 +08:00 · 2025-11-27 22:48:00 +05:30
parent 6bf668c4d2
commit 01e355516b
1 changed files with 1 additions and 0 deletions
--- a/src/diffusers/models/transformers/transformer_z_image.py
+++ b/src/diffusers/models/transformers/transformer_z_image.py
@@ -328,6 +328,7 @@ class RopeEmbedder:
 class ZImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
    _supports_gradient_checkpointing = True
    _no_split_modules = ["ZImageTransformerBlock"]
+    _repeated_blocks = ["ZImageTransformerBlock"]
    _skip_layerwise_casting_patterns = ["t_embedder", "cap_embedder"]  # precision sensitive layers

    @register_to_config