FluxKlein Training Scripts - CFG issue #13762

Activity

1. chinmayirm commented on Sep 22, 2026

   

   chinmayirm

   on Sep 22, 2026

   More actions

   Hi @cryptexis, the split you describe is real, and it is easy to miss because the two Klein variants share one pipeline.

   Root cause

   Flux2KleinPipeline.do_classifier_free_guidance is _guidance_scale > 1 and not self.config.is_distilled. On the base checkpoint that is true, so inference runs a conditional and an unconditional branch. That only matches training if some steps were trained on an empty caption.

   train_dreambooth_lora_flux2_klein.py can do that, through --caption_dropout, and the default is 0.0. train_dreambooth_lora_flux2_klein_img2img.py has no caption dropout at all. Its only dropout is lora_dropout. guidance_embeds is a different switch: Klein does not use the FLUX.1 dev guidance embedding, so the training loop passes guidance=None and the guidance the pipeline applies is entirely the unconditional branch.

   A LoRA trained with the image-to-image script, or with the text-to-image script left at the default, never sees an empty caption. guidance_scale > 1 on the base model then extrapolates from a model that was only trained conditionally. The distilled checkpoint is unaffected, because the pipeline ignores guidance_scale there.

   Proposed approach

   Bring the text-to-image caption-dropout block into the image-to-image script: cache the empty-prompt embedding, and replace only instance rows, never prior-preservation captions. Keep the default at 0.0 so existing runs do not change. In both scripts' help text, state that the non-distilled base model needs a non-zero value (0.1 is the usual one) before guidance_scale > 1 at inference corresponds to anything the LoRA was trained on. Leave the distilled path alone.

   Tests

   A tiny run of the dropout mask: with the flag set, dropped rows equal the empty embedding and class rows are unchanged. With the flag at 0, embeddings are unchanged. The image-to-image script accepts the new argument.

   I would be glad to open a PR for this once the approach looks right to you.