add Z-Image

diffsynth_engine/configs/__init__.py

@@ -1,9 +1,11 @@
 from .base import PipelineConfig
 from .qwen_image import QwenImagePipelineConfig
 from .wan import WanPipelineConfig
+from .z_image import ZImagePipelineConfig
 
 __all__ = [
     "PipelineConfig",
     "QwenImagePipelineConfig",
     "WanPipelineConfig",
+    "ZImagePipelineConfig",
 ]

diffsynth_engine/configs/z_image.py

@@ -0,0 +1,8 @@
+from dataclasses import dataclass
+
+from diffsynth_engine.configs.base import PipelineConfig
+
+
+@dataclass
+class ZImagePipelineConfig(PipelineConfig):
+    pass

diffsynth_engine/models/qwen_image/transformer_qwenimage.py

@@ -697,6 +697,8 @@ class QwenImageTransformer2DModel(DiffusionModel):
             The dimensions to use for the rotary positional embeddings.
     """
 
+    _repeated_blocks = ["QwenImageTransformerBlock"]
+
     @register_to_config
     def __init__(
         self,

diffsynth_engine/models/wan/transformer_wan.py

@@ -101,7 +101,7 @@ def __init__(
             self.norm_added_k = nn.RMSNorm(dim_head * heads, eps=eps)
 
         forward_context = get_forward_context()
-        attn_cls = LocalAttention if cross_attention_dim_head is not None else USPAttention
+        attn_cls = USPAttention if cross_attention_dim_head is None else LocalAttention
         self.attn = attn_cls(
             num_heads=heads,
             head_size=dim_head,
@@ -436,6 +436,7 @@ class WanTransformer3DModel(DiffusionModel):
 
     _keep_in_fp32_modules = ["time_embedder", "scale_shift_table", "norm1", "norm2", "norm3"]
     _keys_to_ignore_on_load_unexpected = ["norm_added_q"]
+    _repeated_blocks = ["WanTransformerBlock"]
 
     @register_to_config
     def __init__(

diffsynth_engine/models/wan/transformer_wan_animate.py

@@ -545,6 +545,7 @@ class WanAnimateTransformer3DModel(DiffusionModel):
         "motion_synthesis_weight",
     ]
     _keys_to_ignore_on_load_unexpected = ["norm_added_q"]
+    _repeated_blocks = ["WanTransformerBlock"]
 
     @register_to_config
     def __init__(

diffsynth_engine/models/wan/transformer_wan_vace.py

@@ -180,6 +180,7 @@ class WanVACETransformer3DModel(DiffusionModel):
 
     _keep_in_fp32_modules = ["time_embedder", "scale_shift_table", "norm1", "norm2", "norm3"]
     _keys_to_ignore_on_load_unexpected = ["norm_added_q"]
+    _repeated_blocks = ["WanTransformerBlock", "WanVACETransformerBlock"]
 
     @register_to_config
     def __init__(

diffsynth_engine/models/z_image/__init__.py

@@ -0,0 +1,5 @@
+from .transformer_z_image import ZImageTransformer2DModel
+
+__all__ = [
+    "ZImageTransformer2DModel",
+]