Move controlnet sample initialization from transformer to pipeline

Author: dimitribarbot

src/diffusers/models/transformers/transformer_qwenimage.py

@@ -639,9 +639,7 @@ def forward(
             if controlnet_block_samples is not None:
                 interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
                 interval_control = int(np.ceil(interval_control))
-                sample = controlnet_block_samples[index_block // interval_control]
-                sample_size = min(sample.size(1), hidden_states.size(1))
-                hidden_states[:, :sample_size] = hidden_states[:, :sample_size] + sample[:, :sample_size]
+                hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
 
         # Use only the image part (hidden_states) from the dual-stream blocks
         hidden_states = self.norm_out(hidden_states, temb)

src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_controlnet.py

@@ -18,14 +18,15 @@
 
 import numpy as np
 import torch
+import torch.nn.functional as F
 from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor
 
 from ...image_processor import PipelineImageInput, VaeImageProcessor
 from ...loaders import QwenImageLoraLoaderMixin
 from ...models import AutoencoderKLQwenImage, QwenImageTransformer2DModel
 from ...models.controlnets.controlnet_qwenimage import QwenImageControlNetModel, QwenImageMultiControlNetModel
 from ...schedulers import FlowMatchEulerDiscreteScheduler
-from ...utils import deprecate, is_torch_xla_available, logging, replace_example_docstring
+from ...utils import is_torch_xla_available, logging, replace_example_docstring
 from ...utils.torch_utils import randn_tensor
 from ..pipeline_utils import DiffusionPipeline
 from .pipeline_output import QwenImagePipelineOutput
@@ -970,6 +971,14 @@ def __call__(
                     return_dict=False,
                 )
 
+                if image_latents is not None:
+                    padding_size = image_latents.shape[1]
+                    for i, sample in enumerate(controlnet_block_samples):
+                        # Pad right with padding_size zeros at dimension 1 of each sample
+                        pad_tuple = [0] * (2 * sample.dim())
+                        pad_tuple[-3] = padding_size
+                        controlnet_block_samples[i] = F.pad(sample, pad_tuple, mode="constant", value=0)
+
                 with self.transformer.cache_context("cond"):
                     noise_pred = self.transformer(
                         hidden_states=latent_model_input,