flux controlnet mode to take into account batch size

Author: christopher-beckham

src/diffusers/pipelines/flux/pipeline_flux_controlnet.py

@@ -748,9 +748,11 @@ def __call__(
             )
 
             # set control mode
+            orig_mode_type = type(control_mode)
             if control_mode is not None:
-                control_mode = torch.tensor(control_mode).to(device, dtype=torch.long)
-                control_mode = control_mode.reshape([-1, 1])
+                control_mode = torch.tensor(control_mode).to(device, dtype=torch.long).view(-1,1)
+                if orig_mode_type == int:
+                    control_mode = control_mode.repeat(control_image.shape[0], 1)
 
         elif isinstance(self.controlnet, FluxMultiControlNetModel):
             control_images = []
@@ -793,8 +795,10 @@ def __call__(
                         control_mode_.append(-1)
                     else:
                         control_mode_.append(cmode)
-            control_mode = torch.tensor(control_mode_).to(device, dtype=torch.long)
-            control_mode = control_mode.reshape([-1, 1])
+                control_mode = torch.tensor(control_mode_).to(device, dtype=torch.long)
+                control_mode = control_mode.view(-1, 1)
+            else:
+                raise ValueError("For multi-controlnet, control_mode should be a list")
 
         # 4. Prepare latent variables
         num_channels_latents = self.transformer.config.in_channels // 4