Fix Flux2 DreamBooth prior preservation prompt repeats

Author: azolotenkov

examples/dreambooth/train_dreambooth_lora_flux2.py

@@ -1740,9 +1740,12 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
                     prompt_embeds = prompt_embeds_cache[step]
                     text_ids = text_ids_cache[step]
                 else:
-                    num_repeat_elements = len(prompts)
-                    prompt_embeds = prompt_embeds.repeat(num_repeat_elements, 1, 1)
-                    text_ids = text_ids.repeat(num_repeat_elements, 1, 1)
+                    # With prior preservation, prompt_embeds/text_ids already contain [instance, class] entries,
+                    # while collate_fn orders batches as [inst1..instB, class1..classB]. Repeat each entry along
+                    # dim 0 to preserve that grouping instead of interleaving [inst, class, inst, class, ...].
+                    num_repeat_elements = len(prompts) // 2 if args.with_prior_preservation else len(prompts)
+                    prompt_embeds = prompt_embeds.repeat_interleave(num_repeat_elements, dim=0)
+                    text_ids = text_ids.repeat_interleave(num_repeat_elements, dim=0)
 
                 # Convert images to latent space
                 if args.cache_latents:

examples/dreambooth/train_dreambooth_lora_flux2_klein.py

@@ -1680,9 +1680,12 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
                     prompt_embeds = prompt_embeds_cache[step]
                     text_ids = text_ids_cache[step]
                 else:
-                    num_repeat_elements = len(prompts)
-                    prompt_embeds = prompt_embeds.repeat(num_repeat_elements, 1, 1)
-                    text_ids = text_ids.repeat(num_repeat_elements, 1, 1)
+                    # With prior preservation, prompt_embeds/text_ids already contain [instance, class] entries,
+                    # while collate_fn orders batches as [inst1..instB, class1..classB]. Repeat each entry along
+                    # dim 0 to preserve that grouping instead of interleaving [inst, class, inst, class, ...].
+                    num_repeat_elements = len(prompts) // 2 if args.with_prior_preservation else len(prompts)
+                    prompt_embeds = prompt_embeds.repeat_interleave(num_repeat_elements, dim=0)
+                    text_ids = text_ids.repeat_interleave(num_repeat_elements, dim=0)
 
                 # Convert images to latent space
                 if args.cache_latents: