Fix bugs for controlnet after merging the main branch new feature.

Author: JerryWu-code

src/diffusers/models/transformers/transformer_z_image.py

@@ -916,16 +916,20 @@ def forward(
 
         if omni_mode:
             return self._forward_omni(
-                x, t, cap_feats, cond_latents, siglip_feats, 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(x, t, cap_feats, patch_size, f_patch_size, return_dict)
+            return self._forward_basic(
+                x, t, cap_feats, controlnet_block_samples, patch_size, f_patch_size, return_dict
+            )
 
     def _forward_basic(
         self,
         x: List[torch.Tensor],
         t,
         cap_feats: List[torch.Tensor],
+        controlnet_block_samples: Optional[Dict[int, torch.Tensor]],
         patch_size: int,
         f_patch_size: int,
         return_dict: bool,
@@ -1053,6 +1057,7 @@ def _forward_omni(
         cap_feats: List[List[torch.Tensor]],
         cond_latents: List[List[torch.Tensor]],
         siglip_feats: List[List[torch.Tensor]],
+        controlnet_block_samples: Optional[Dict[int, torch.Tensor]],
         patch_size: int,
         f_patch_size: int,
         return_dict: bool,
@@ -1241,17 +1246,23 @@ def _forward_omni(
         unified_noise_mask_tensor = unified_noise_mask_tensor[:, : unified.shape[1]]
 
         if torch.is_grad_enabled() and self.gradient_checkpointing:
-            for layer in self.layers:
+            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
                 )
+                if controlnet_block_samples is not None:
+                    if layer_idx in controlnet_block_samples:
+                        unified = unified + controlnet_block_samples[layer_idx]
         else:
-            for layer in self.layers:
+            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
                 )
+                if controlnet_block_samples is not None:
+                    if layer_idx in controlnet_block_samples:
+                        unified = unified + controlnet_block_samples[layer_idx]
 
         unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"](
             unified, noise_mask=unified_noise_mask_tensor, c_noisy=t_noisy_x, c_clean=t_clean_x

src/diffusers/pipelines/z_image/pipeline_z_image_omni.py

@@ -659,12 +659,12 @@ def __call__(
                 latent_model_input_list = list(latent_model_input.unbind(dim=0))
 
                 model_out_list = self.transformer(
-                    latent_model_input_list, 
-                    timestep_model_input, 
-                    prompt_embeds_model_input, 
-                    condition_latents_model_input, 
-                    condition_siglip_embeds_model_input,
-                    return_dict=False
+                    x=latent_model_input_list,
+                    t=timestep_model_input,
+                    cap_feats=prompt_embeds_model_input,
+                    cond_latents=condition_latents_model_input,
+                    siglip_feats=condition_siglip_embeds_model_input,
+                    return_dict=False,
                 )[0]
 
                 if apply_cfg: