Fix Flux2 DreamBooth prior preservation prompt repeats

Author: azolotenkov

examples/dreambooth/train_dreambooth_lora_flux2.py

@@ -1740,7 +1740,9 @@ 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)
+                    # With prior preservation, prompt_embeds already contains [instance, class] embeddings,
+                    # while collate_fn doubles the prompts list. Repeat by the number of instance prompts only.
+                    num_repeat_elements = len(prompts) // 2 if args.with_prior_preservation else len(prompts)
                     prompt_embeds = prompt_embeds.repeat(num_repeat_elements, 1, 1)
                     text_ids = text_ids.repeat(num_repeat_elements, 1, 1)
 

examples/dreambooth/train_dreambooth_lora_flux2_klein.py

@@ -1680,7 +1680,9 @@ 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)
+                    # With prior preservation, prompt_embeds already contains [instance, class] embeddings,
+                    # while collate_fn doubles the prompts list. Repeat by the number of instance prompts only.
+                    num_repeat_elements = len(prompts) // 2 if args.with_prior_preservation else len(prompts)
                     prompt_embeds = prompt_embeds.repeat(num_repeat_elements, 1, 1)
                     text_ids = text_ids.repeat(num_repeat_elements, 1, 1)