Problem running train_dreambooth_lora_flux.py with model "black-forest-labs/FLUX.1-dev"

### Describe the bug

I am using the training script documented here https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/README_flux.md to train a LORA on my dataset.

here is the command I am running:

accelerate launch train_dreambooth_lora_flux.py
--pretrained_model_name_or_path="black-forest-labs/FLUX.1-dev"
--dataset_name="path_to_my_dataset"
--output_dir="dreambooth-flux-dev"
--mixed_precision="bf16"
--instance_prompt="a photo of sks chair"
--resolution=1024
--train_batch_size=4
--gradient_accumulation_steps=1
--gradient_checkpointing
--learning_rate=1e-5
--lr_scheduler="constant"
--lr_warmup_steps=0
--validation_prompt="A photo of a sks chair in the garden"
--validation_epochs=10
--seed="42"
--rank=16
--train_text_encoder
--max_sequence_length=512 # I am trying to use long prompts
--max_train_steps=500
--caption_column="prompt"

**if I specify the --train_text_encoder I get the following error:**

Traceback (most recent call last):
File "/workspace/diffusers/examples/dreambooth/train_dreambooth_lora_flux.py", line 1857, in
main(args)
File "/workspace/diffusers/examples/dreambooth/train_dreambooth_lora_flux.py", line 1663, in main
model_pred = transformer(
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
return forward_call(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/accelerate/utils/operations.py", line 819, in forward
return model_forward(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/accelerate/utils/operations.py", line 807, in __call__
return convert_to_fp32(self.model_forward(*args, **kwargs))
File "/usr/local/lib/python3.10/dist-packages/torch/amp/autocast_mode.py", line 16, in decorate_autocast
return func(*args, **kwargs)
File "/workspace/diffusers/src/diffusers/models/transformers/transformer_flux.py", line 385, in forward
ids = torch.cat((txt_ids, img_ids), dim=1)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument tensors in method wrapper_CUDA_cat)

This error does not appear if I leave the "--train_text_encoder" argument out.

### Reproduction

accelerate launch train_dreambooth_lora_flux.py
--pretrained_model_name_or_path="black-forest-labs/FLUX.1-dev"
--dataset_name="path_to_my_dataset"
--output_dir="dreambooth-flux-dev"
--mixed_precision="bf16"
--instance_prompt="a photo of sks object"
--resolution=1024
--train_batch_size=4
--gradient_accumulation_steps=1
--gradient_checkpointing
--learning_rate=1e-5
--lr_scheduler="constant"
--lr_warmup_steps=0
--validation_prompt="A photo of a sks object in the garden"
--validation_epochs=10
--seed="42"
--rank=16
--train_text_encoder
--max_sequence_length=512
--max_train_steps=500
--caption_column="prompt"

### Logs

_No response_

### System Info

Running on RunPod using the runpod/kohya:24.1.6 template
one GPU: **1 x A100 PCIe**
**12 vCPU 125 GB RAM**

Using the default config for accelerate accelerate config default.
No distributed training.

### Who can help?

_No response_

Start with examples/dreambooth/train_dreambooth_lora_flux.py at the training call near line 1663, then inspect src/diffusers/models/transformers/transformer_flux.py around line 385 where txt_ids and img_ids are concatenated. Reproduce with --train_text_encoder using the provided command and verify that the training step no longer raises a CPU-versus-cuda device mismatch.

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