For qwen-image training file, Maybe "shuffle" of dataloader should be "False" when custom_instance_prompts is not None and cache_latents is False?

[yinguoweiOvO]

Describe the bug

I think "shuffle" of dataloader should be "False" when custom_instance_prompts is not None and cache_latents is False. Otherwise, it will lead to errors in the correspondence between prompt embedding and image during training, and prompt will not be followed when performing the task of T2I.

Reproduction

None

Logs

System Info

None

Who can help?

No response

[Aishwarya0811]

Hi @yinguoweiOvO!

Thanks for reporting this! Could you provide a bit more detail to help us understand the issue better?

Questions:

1. Which specific training script are you using? (e.g., train_dreambooth_lora.py)
2. Can you share your training command/config that reproduces this?
3. What specific symptoms do you see? (poor image quality, prompts ignored, etc.)

To investigate properly, it would be helpful to see:

• The exact error or unexpected behavior
• Your training arguments/config
• Maybe a simple reproduction case

This will help us identify the root cause and implement the right fix. Thanks! 🙂

[yinguoweiOvO]

1. It's in the qwen-image training script, but I'm not sure if there are the same problems in other scripts
2. I have modified the training script based on my own needs, so I am unable to provide you with the training configuration. However, you can try to set the training set so that the images and annotations correspond one-to-one instead of using the same annotation for all images to reproduce this issue. Note: To reproduce this issue, you need to set cache_latents to False.
3. In the following code, "prompts" is pre-encoded and stored but "pixel_values" is not when train_dataset have custom_instance_prompts and "cache_latents" is False

precompute_latents = args.cache_latents or train_dataset.custom_instance_prompts
if precompute_latents:
    prompt_embeds_cache = []
    prompt_embeds_mask_cache = []
    latents_cache = []
    for batch in tqdm(train_dataloader, desc="Caching latents"):
        with torch.no_grad():
            if args.cache_latents:
                with offload_models(vae, device=accelerator.device, offload=args.offload):
                    batch["pixel_values"] = batch["pixel_values"].to(
                        accelerator.device, non_blocking=True, dtype=vae.dtype
                    )
                latents_cache.append(vae.encode(batch["pixel_values"]).latent_dist)
            if train_dataset.custom_instance_prompts:
                with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload):
                    prompt_embeds, prompt_embeds_mask = compute_text_embeddings(
                        batch["prompts"], text_encoding_pipeline
                    )
                prompt_embeds_cache.append(prompt_embeds)
                prompt_embeds_mask_cache.append(prompt_embeds_mask)

In the following code, if "shuffle" of dataloader is True, the order of the data sampled from dataloader in each epoch is different but the order of prompt_embeds sampled from prompt_embeds_cache is same. This will cause an error in the correspondence between "prompts" and "pixel_values", thereby resulting in the generation of images not following the prompt.

for step, batch in enumerate(train_dataloader):
    models_to_accumulate = [transformer]
    prompts = batch["prompts"]

    with accelerator.accumulate(models_to_accumulate):
        # encode batch prompts when custom prompts are provided for each image -
        if train_dataset.custom_instance_prompts:
            prompt_embeds = prompt_embeds_cache[step]
            prompt_embeds_mask = prompt_embeds_mask_cache[step]
        else:
            num_repeat_elements = len(prompts)
            prompt_embeds = prompt_embeds.repeat(num_repeat_elements, 1, 1)
            prompt_embeds_mask = prompt_embeds_mask.repeat(num_repeat_elements, 1)
        # Convert images to latent space
        if args.cache_latents:
            model_input = latents_cache[step].sample()
        else:
            with offload_models(vae, device=accelerator.device, offload=args.offload):
                pixel_values = batch["pixel_values"].to(dtype=vae.dtype)
            model_input = vae.encode(pixel_values).latent_dist.sample()