multi-gpu/model-sharding incompatible with FluxControlNetPipeline

### Describe the bug

I have an environment with 2 A10 (each 22GB), when I follow [distributed_inference.md] to apply model sharding and run flux.1-dev, it works fine for basic scenario (ie. txt2img, img2img), even lora, **but not controlnet**.

An error got:
_Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)_

I guess the current FluxControlNetPipeline requires the FluxControlNetModel and the transformer to be on the same cuda, but this is not satisfied when I do as [distributed_inference.md] told, to distribute transformer among many cudas.

I've tried diffusers==0.31.0 and 0.32.1, both failed.

Can this hopefully be supported? Really appreciate it.

distributed_inference.md: (https://github.com/huggingface/diffusers/blob/main/docs/source/en/training/distributed_inference.md)

### Reproduction

```python
from diffusers import FluxTransformer2DModel, FluxPipeline, FluxControlNetModel, FluxControlNetPipeline
from transformers import T5EncoderModel
from diffusers.utils import load_image
import torch

transformer = FluxTransformer2DModel.from_pretrained(
"path/to/FLUX.1-dev",
subfolder="transformer",
device_map="auto", # transformer part of Flux.1-dev is huge, so distribute it among multi gpus
torch_dtype=torch.bfloat16
)
print(transformer.hf_device_map) # in my case, successfully spread transformer blocks among cuda:0 and cuda:1

text_encoder_2 = T5EncoderModel.from_pretrained(
"path/to/FLUX.1-dev",
subfolder="text_encoder_2",
device_map="auto", # text_encoder_2 of Flux.1-dev is huge, so distribute it among multi gpus
torch_dtype=torch.bfloat16
)
print(text_encoder_2.hf_device_map) # successfully spread text_encoder_2 blocks among cuda:0 and cuda:1

pipeline = FluxPipeline.from_pretrained(
"path/to/FLUX.1-dev",
transformer=transformer,
text_encoder_2=text_encoder_2,
device_map="balanced",
torch_dtype=torch.bfloat16
)
out =pipeline("a cat").images[0] # this works fine

### next do something similarly for controlnet pipeline
canny_controlnet = FluxControlNetModel.from_pretrained(
"path/to/XLabs-AI/flux-controlnet-canny-diffusers",
torch_dtype=torch.bfloat16,
use_safetensors=True,
).to("cuda")

xlabs_canny_pipe = from_pretrained(
"path/to/FLUX.1-dev",
controlnet=canny_controlnet,
transformer=transformer,
text_encoder_2=text_encoder_2,
device_map="balanced",
torch_dtype=torch.bfloat16
) # succeed at this stage, but fail later on when generating image using this contolnet pipe

control_image = load_image("path/to/whaterver/image")

image = xlabs_canny_pipe(
'A woman with scarf',
control_image=control_image,
controlnet_conditioning_scale=0.8
).images[0] # i got the error declared above
```

### Logs

```shell
RuntimeError Traceback (most recent call last)
Cell In[16], line 5
1 prompt = 'A woman with scarf'
2 control_image = load_image("//ossfs/node_50570412/workspace/test_controlnet_1223/control_image_annotated_depth.png")
----> 5 image = xlabs_canny_pipe(
6 prompt,
7 control_image=control_image,
8 controlnet_conditioning_scale=0.8,
9 num_inference_steps=25,
10 guidance_scale=3.5,
11 # generator=torch.Generator("cuda").manual_seed(128),
12 # generator=torch.manual_seed(128),
13 height=480,
14 width=480,
15 ).images[0]
17 image.save("xlabs_control_pipe_15.png")

File /opt/conda/lib/python3.10/site-packages/torch/utils/_contextlib.py:115, in context_decorator..decorate_context(*args, **kwargs)
112 @functools.wraps(func)
113 def decorate_context(*args, **kwargs):
114 with ctx_factory():
--> 115 return func(*args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_controlnet.py:908, in FluxControlNetPipeline.__call__(self, prompt, prompt_2, height, width, num_inference_steps, timesteps, guidance_scale, control_guidance_start, control_guidance_end, control_image, control_mode, controlnet_conditioning_scale, num_images_per_prompt, generator, latents, prompt_embeds, pooled_prompt_embeds, output_type, return_dict, joint_attention_kwargs, callback_on_step_end, callback_on_step_end_tensor_inputs, max_sequence_length)
905 cond_scale = controlnet_cond_scale * controlnet_keep[i]
907 # controlnet
--> 908 controlnet_block_samples, controlnet_single_block_samples = self.controlnet(
909 hidden_states=latents,
910 controlnet_cond=control_image,
911 controlnet_mode=control_mode,
912 conditioning_scale=cond_scale,
913 timestep=timestep / 1000,
914 guidance=guidance,
915 pooled_projections=pooled_prompt_embeds,
916 encoder_hidden_states=prompt_embeds,
917 txt_ids=text_ids,
918 img_ids=latent_image_ids,
919 joint_attention_kwargs=self.joint_attention_kwargs,
920 return_dict=False,
921 )
923 guidance = (
924 torch.tensor([guidance_scale], device=device) if self.transformer.config.guidance_embeds else None
925 )
926 guidance = guidance.expand(latents.shape[0]) if guidance is not None else None

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
1516 return self._compiled_call_impl(*args, **kwargs) # type: ignore[misc]
1517 else:
-> 1518 return self._call_impl(*args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
1522 # If we don't have any hooks, we want to skip the rest of the logic in
1523 # this function, and just call forward.
1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
1525 or _global_backward_pre_hooks or _global_backward_hooks
1526 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527 return forward_call(*args, **kwargs)
1529 try:
1530 result = None

File /opt/conda/lib/python3.10/site-packages/diffusers/models/controlnet_flux.py:278, in FluxControlNetModel.forward(self, hidden_states, controlnet_cond, controlnet_mode, conditioning_scale, encoder_hidden_states, pooled_projections, timestep, img_ids, txt_ids, guidance, joint_attention_kwargs, return_dict)
274 if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
275 logger.warning(
276 "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
277 )
--> 278 hidden_states = self.x_embedder(hidden_states)
280 if self.input_hint_block is not None:
281 controlnet_cond = self.input_hint_block(controlnet_cond)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
1516 return self._compiled_call_impl(*args, **kwargs) # type: ignore[misc]
1517 else:
-> 1518 return self._call_impl(*args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
1522 # If we don't have any hooks, we want to skip the rest of the logic in
1523 # this function, and just call forward.
1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
1525 or _global_backward_pre_hooks or _global_backward_hooks
1526 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527 return forward_call(*args, **kwargs)
1529 try:
1530 result = None

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/linear.py:114, in Linear.forward(self, input)
113 def forward(self, input: Tensor) -> Tensor:
--> 114 return F.linear(input, self.weight, self.bias)

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)
```

### System Info

- 🤗 Diffusers version: 0.31.0
- Platform: Linux-4.9.151-015.ali3000.alios7.x86_64-x86_64-with-glibc2.32
- Running on Google Colab?: No
- Python version: 3.10.13
- PyTorch version (GPU?): 2.1.0+cu121 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Huggingface_hub version: 0.26.2
- Transformers version: 4.44.0
- Accelerate version: 1.2.0
- PEFT version: 0.14.0
- Bitsandbytes version: 0.45.0
- Safetensors version: 0.4.5
- xFormers version: not installed
- Accelerator: NVIDIA A10, 22731 MiB
NVIDIA A10, 22731 MiB
- Using GPU in script?:
- Using distributed or parallel set-up in script?:

### Who can help?

@sayakpaul @yiyixuxu

Start with docs/source/en/training/distributed_inference.md and reproduce the failure using the FluxControlNetPipeline entry point shown in the report. Trace the call from pipeline_flux_controlnet.py into controlnet_flux.py, especially the x_embedder path, and verify that ControlNet generation works with the reported multi-GPU device maps without the cross-device error.

Scritto dal modello di indicizzazione a partire dal testo della issue.