Add controlnet_blocks_repeat to Flux forward

Author: Anghellia

src/diffusers/models/transformers/transformer_flux.py

@@ -402,6 +402,7 @@ def forward(
         controlnet_block_samples=None,
         controlnet_single_block_samples=None,
         return_dict: bool = True,
+        controlnet_blocks_repeat: bool = False,
     ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
         """
         The [`FluxTransformer2DModel`] forward method.
@@ -509,8 +510,8 @@ def custom_forward(*inputs):
                 interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
                 interval_control = int(np.ceil(interval_control))
                 # For Xlabs ControlNet.
-                if len(controlnet_block_samples) == 2:
-                    hidden_states = hidden_states + controlnet_block_samples[index_block % 2]
+                if controlnet_blocks_repeat:
+                    hidden_states = hidden_states + controlnet_block_samples[index_block % len(controlnet_block_samples)]
                 else:
                     hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
 

src/diffusers/pipelines/flux/pipeline_flux_controlnet.py

@@ -739,6 +739,7 @@ def __call__(
         )
 
         # 3. Prepare control image
+        controlnet_blocks_repeat = False
         num_channels_latents = self.transformer.config.in_channels // 4
         if isinstance(self.controlnet, FluxControlNetModel):
             control_image = self.prepare_image(
@@ -766,6 +767,8 @@ def __call__(
                     height_control_image,
                     width_control_image,
                 )
+            else:
+                controlnet_blocks_repeat = True
 
             # Here we ensure that `control_mode` has the same length as the control_image.
             if control_mode is not None:
@@ -926,6 +929,7 @@ def __call__(
                     img_ids=latent_image_ids,
                     joint_attention_kwargs=self.joint_attention_kwargs,
                     return_dict=False,
+                    controlnet_blocks_repeat=controlnet_blocks_repeat,
                 )[0]
 
                 # compute the previous noisy sample x_t -> x_t-1