Fix for auto_pipeline, Add Styling.

Author: JerryWu-code

src/diffusers/__init__.py

@@ -675,8 +675,8 @@
             "ZImageControlNetInpaintPipeline",
             "ZImageControlNetPipeline",
             "ZImageImg2ImgPipeline",
-            "ZImagePipeline",
             "ZImageOmniPipeline",
+            "ZImagePipeline",
         ]
     )
 
@@ -1387,8 +1387,8 @@
             ZImageControlNetInpaintPipeline,
             ZImageControlNetPipeline,
             ZImageImg2ImgPipeline,
-            ZImagePipeline,
             ZImageOmniPipeline,
+            ZImagePipeline,
         )
 
     try:

src/diffusers/models/transformers/transformer_z_image.py

@@ -222,7 +222,7 @@ def forward(
         if self.modulation:
             if noise_mask is not None and adaln_noisy is not None and adaln_clean is not None:
                 # Per-token modulation based on noise_mask, (batch, seq_len), 1 for noisy tokens, 0 for clean tokens
-                batch_size, seq_len = x.shape[0], x.shape[1]
+                _, seq_len = x.shape[0], x.shape[1]
 
                 mod_noisy = self.adaLN_modulation(adaln_noisy)
                 mod_clean = self.adaLN_modulation(adaln_clean)
@@ -260,7 +260,9 @@ def forward(
             else:
                 # Original global modulation
                 assert adaln_input is not None
-                scale_msa, gate_msa, scale_mlp, gate_mlp = self.adaLN_modulation(adaln_input).unsqueeze(1).chunk(4, dim=2)
+                scale_msa, gate_msa, scale_mlp, gate_mlp = (
+                    self.adaLN_modulation(adaln_input).unsqueeze(1).chunk(4, dim=2)
+                )
                 gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh()
                 scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp
 
@@ -297,7 +299,7 @@ def __init__(self, hidden_size, out_channels):
     def forward(self, x, c=None, noise_mask=None, c_noisy=None, c_clean=None):
         if noise_mask is not None and c_noisy is not None and c_clean is not None:
             # Per-token modulation based on noise_mask
-            batch_size, seq_len = x.shape[0], x.shape[1]
+            _, seq_len = x.shape[0], x.shape[1]
             scale_noisy = 1.0 + self.adaLN_modulation(c_noisy)
             scale_clean = 1.0 + self.adaLN_modulation(c_clean)
 
@@ -916,8 +918,15 @@ def forward(
 
         if omni_mode:
             return self._forward_omni(
-                x, t, cap_feats, cond_latents, siglip_feats,
-                controlnet_block_samples, patch_size, f_patch_size, return_dict
+                x,
+                t,
+                cap_feats,
+                cond_latents,
+                siglip_feats,
+                controlnet_block_samples,
+                patch_size,
+                f_patch_size,
+                return_dict,
             )
         else:
             return self._forward_basic(
@@ -1130,14 +1139,23 @@ def _forward_omni(
         if torch.is_grad_enabled() and self.gradient_checkpointing:
             for layer in self.noise_refiner:
                 x = self._gradient_checkpointing_func(
-                    layer, x, x_attn_mask, x_freqs_cis,
-                    noise_mask=x_noise_mask_tensor, adaln_noisy=t_noisy_x, adaln_clean=t_clean_x
+                    layer,
+                    x,
+                    x_attn_mask,
+                    x_freqs_cis,
+                    noise_mask=x_noise_mask_tensor,
+                    adaln_noisy=t_noisy_x,
+                    adaln_clean=t_clean_x,
                 )
         else:
             for layer in self.noise_refiner:
                 x = layer(
-                    x, x_attn_mask, x_freqs_cis,
-                    noise_mask=x_noise_mask_tensor, adaln_noisy=t_noisy_x, adaln_clean=t_clean_x
+                    x,
+                    x_attn_mask,
+                    x_freqs_cis,
+                    noise_mask=x_noise_mask_tensor,
+                    adaln_noisy=t_noisy_x,
+                    adaln_clean=t_clean_x,
                 )
 
         # cap embed & refine (no modulation)
@@ -1208,9 +1226,7 @@ def _forward_omni(
                 x_len = x_item_seqlens[i]
                 cap_len = cap_item_seqlens[i]
                 siglip_len = siglip_item_seqlens[i]
-                unified.append(
-                    torch.cat([cap_feats[i][:cap_len], x[i][:x_len], siglip_feats[i][:siglip_len]])
-                )
+                unified.append(torch.cat([cap_feats[i][:cap_len], x[i][:x_len], siglip_feats[i][:siglip_len]]))
                 unified_freqs_cis.append(
                     torch.cat([cap_freqs_cis[i][:cap_len], x_freqs_cis[i][:x_len], siglip_freqs_cis[i][:siglip_len]])
                 )
@@ -1221,7 +1237,9 @@ def _forward_omni(
                         device=device,
                     )
                 )
-            unified_item_seqlens = [a + b + c for a, b, c in zip(cap_item_seqlens, x_item_seqlens, siglip_item_seqlens)]
+            unified_item_seqlens = [
+                a + b + c for a, b, c in zip(cap_item_seqlens, x_item_seqlens, siglip_item_seqlens)
+            ]
         else:
             for i in range(bsz):
                 x_len = x_item_seqlens[i]
@@ -1248,17 +1266,26 @@ def _forward_omni(
         if torch.is_grad_enabled() and self.gradient_checkpointing:
             for layer_idx, layer in enumerate(self.layers):
                 unified = self._gradient_checkpointing_func(
-                    layer, unified, unified_attn_mask, unified_freqs_cis,
-                    noise_mask=unified_noise_mask_tensor, adaln_noisy=t_noisy_x, adaln_clean=t_clean_x
+                    layer,
+                    unified,
+                    unified_attn_mask,
+                    unified_freqs_cis,
+                    noise_mask=unified_noise_mask_tensor,
+                    adaln_noisy=t_noisy_x,
+                    adaln_clean=t_clean_x,
                 )
                 if controlnet_block_samples is not None:
                     if layer_idx in controlnet_block_samples:
                         unified = unified + controlnet_block_samples[layer_idx]
         else:
             for layer_idx, layer in enumerate(self.layers):
                 unified = layer(
-                    unified, unified_attn_mask, unified_freqs_cis,
-                    noise_mask=unified_noise_mask_tensor, adaln_noisy=t_noisy_x, adaln_clean=t_clean_x
+                    unified,
+                    unified_attn_mask,
+                    unified_freqs_cis,
+                    noise_mask=unified_noise_mask_tensor,
+                    adaln_noisy=t_noisy_x,
+                    adaln_clean=t_clean_x,
                 )
                 if controlnet_block_samples is not None:
                     if layer_idx in controlnet_block_samples:

src/diffusers/pipelines/auto_pipeline.py

@@ -119,7 +119,13 @@
 )
 from .wan import WanImageToVideoPipeline, WanPipeline, WanVideoToVideoPipeline
 from .wuerstchen import WuerstchenCombinedPipeline, WuerstchenDecoderPipeline
-from .z_image import ZImageImg2ImgPipeline, ZImagePipeline
+from .z_image import (
+    ZImageControlNetInpaintPipeline,
+    ZImageControlNetPipeline,
+    ZImageImg2ImgPipeline,
+    ZImageOmniPipeline,
+    ZImagePipeline,
+)
 
 
 AUTO_TEXT2IMAGE_PIPELINES_MAPPING = OrderedDict(
@@ -164,6 +170,9 @@
         ("qwenimage", QwenImagePipeline),
         ("qwenimage-controlnet", QwenImageControlNetPipeline),
         ("z-image", ZImagePipeline),
+        ("z-image-controlnet", ZImageControlNetPipeline),
+        ("z-image-controlnet-inpaint", ZImageControlNetInpaintPipeline),
+        ("z-image-omni", ZImageOmniPipeline),
     ]
 )
 

src/diffusers/pipelines/z_image/pipeline_z_image_omni.py

@@ -16,22 +16,19 @@
 from typing import Any, Callable, Dict, List, Optional, Union
 
 import PIL
-
 import torch
 from transformers import AutoTokenizer, PreTrainedModel, Siglip2ImageProcessorFast, Siglip2VisionModel
 
-from ...image_processor import VaeImageProcessor
 from ...loaders import FromSingleFileMixin, ZImageLoraLoaderMixin
 from ...models.autoencoders import AutoencoderKL
+from ...models.transformers import ZImageTransformer2DModel
 from ...pipelines.pipeline_utils import DiffusionPipeline
 from ...schedulers import FlowMatchEulerDiscreteScheduler
 from ...utils import logging, replace_example_docstring
 from ...utils.torch_utils import randn_tensor
-from .pipeline_output import ZImagePipelineOutput
-
 from ..flux2.image_processor import Flux2ImageProcessor
+from .pipeline_output import ZImagePipelineOutput
 
-from ...models.transformers import ZImageTransformer2DModel
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
@@ -224,7 +221,6 @@ def _encode_prompt(
             prompt = [prompt]
 
         for i, prompt_item in enumerate(prompt):
-
             if num_condition_images == 0:
                 prompt[i] = ["<|im_start|>user\n" + prompt_item + "<|im_end|>\n<|im_start|>assistant\n"]
             elif num_condition_images > 0:
@@ -236,7 +232,7 @@ def _encode_prompt(
 
         flattened_prompt = []
         prompt_list_lengths = []
-        
+
         for i in range(len(prompt)):
             prompt_list_lengths.append(len(prompt[i]))
             flattened_prompt.extend(prompt[i])
@@ -304,14 +300,15 @@ def prepare_image_latents(
         image_latents = []
         for image in images:
             image = image.to(device=device, dtype=dtype)
-            image_latent = (self.vae.encode(image.bfloat16()).latent_dist.mode()[0] - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+            image_latent = (
+                self.vae.encode(image.bfloat16()).latent_dist.mode()[0] - self.vae.config.shift_factor
+            ) * self.vae.config.scaling_factor
             image_latent = image_latent.unsqueeze(1).to(dtype)
             image_latents.append(image_latent)  # (16, 128, 128)
 
         # image_latents = [image_latents] * batch_size
         image_latents = [image_latents.copy() for _ in range(batch_size)]
 
-
         return image_latents
 
     def prepare_siglip_embeds(
@@ -327,7 +324,7 @@ def prepare_siglip_embeds(
             shape = siglip_inputs.spatial_shapes[0]
             hidden_state = self.siglip(**siglip_inputs).last_hidden_state
             B, N, C = hidden_state.shape
-            hidden_state = hidden_state[:, :shape[0] * shape[1]]
+            hidden_state = hidden_state[:, : shape[0] * shape[1]]
             hidden_state = hidden_state.view(shape[0], shape[1], C)
             siglip_embeds.append(hidden_state.to(dtype))
 
@@ -529,7 +526,7 @@ def __call__(
                 image_height = (image_height // multiple_of) * multiple_of
                 img = self.image_processor.preprocess(img, height=image_height, width=image_width, resize_mode="crop")
                 condition_images.append(img)
-            
+
             if len(condition_images) > 0:
                 height = height or image_height
                 width = width or image_width
@@ -591,7 +588,9 @@ def __call__(
                 negative_prompt_embeds = [npe for npe in negative_prompt_embeds for _ in range(num_images_per_prompt)]
 
         condition_siglip_embeds = [None if sels == [] else sels + [None] for sels in condition_siglip_embeds]
-        negative_condition_siglip_embeds = [None if sels == [] else sels + [None] for sels in negative_condition_siglip_embeds]
+        negative_condition_siglip_embeds = [
+            None if sels == [] else sels + [None] for sels in negative_condition_siglip_embeds
+        ]
 
         actual_batch_size = batch_size * num_images_per_prompt
         image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] // 2)