Skip to content

[lora training] add caption dropout and aspect ratio buckets to krea2 and qwen-image#14089

Open
linoytsaban wants to merge 7 commits into
mainfrom
lora-caption-dropout-ar-buckets
Open

[lora training] add caption dropout and aspect ratio buckets to krea2 and qwen-image#14089
linoytsaban wants to merge 7 commits into
mainfrom
lora-caption-dropout-ar-buckets

Conversation

@linoytsaban

@linoytsaban linoytsaban commented Jun 29, 2026

Copy link
Copy Markdown
Collaborator

Adds aspect-ratio bucketing and caption dropout to the krea2 and qwen-image dreambooth LoRA scripts (neither supported these before):

  • Aspect-ratio bucketing for krea2 and qwen-image.
  • Caption dropout (--caption_dropout): randomly replace an instance caption with "" so that fraction trains unconditionally (improves CFG).

Split from the original PR; bucketing fixes to the other scripts that already support buckets are in #14158.

@github-actions github-actions Bot added examples size/L PR with diff > 200 LOC labels Jun 29, 2026
@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@linoytsaban
linoytsaban force-pushed the lora-caption-dropout-ar-buckets branch from a539d74 to 975986d Compare June 30, 2026 09:09
@linoytsaban

Copy link
Copy Markdown
Collaborator Author

krea comparison with and without buckets and dropout
image

@linoytsaban
linoytsaban marked this pull request as ready for review June 30, 2026 11:56
@linoytsaban
linoytsaban requested a review from sayakpaul June 30, 2026 11:56

@sayakpaul sayakpaul left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we do it just one and then propagate to the others? This way, it's easier to review.

@linoytsaban

Copy link
Copy Markdown
Collaborator Author

Should we do it just one and then propagate to the others? This way, it's easier to review.

I can split the krea 2 + qwen image changes (that introduce aspect ratio buckets + caption dropout since they dont currently support these at all) from the changes to the other scripts (that already support bucketing) where it's fixing bucketing issues. wdyt?

@sayakpaul

Copy link
Copy Markdown
Member

Yes sure

@linoytsaban
linoytsaban force-pushed the lora-caption-dropout-ar-buckets branch from 098cc5b to 25d9f73 Compare July 10, 2026 08:35
@linoytsaban linoytsaban changed the title [lora training] add caption dropout and aspect ratio buckets [lora training] add caption dropout and aspect ratio buckets to krea2 and qwen-image Jul 10, 2026
@linoytsaban
linoytsaban requested a review from sayakpaul July 10, 2026 10:28

@sayakpaul sayakpaul left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same comments as the other PR. WDYT?

linoytsaban and others added 3 commits July 15, 2026 17:26
Gate bucketing behind --use_aspect_ratio_buckets (a list without it errors), compute
buckets on the fly capped to each image's own resolution via the shared training_utils
helper, and drop --bucket_no_upscale and padding. Skip the opt-in bucketing/caption
dropout qwen tests, matching the flux2/kontext/z-image scripts.
@linoytsaban
linoytsaban requested a review from sayakpaul July 22, 2026 12:17
Comment on lines +530 to +549
parser.add_argument(
"--aspect_ratio_buckets",
type=str,
default=None,
help=(
"Aspect ratio buckets to use for training. Define as a string of 'h1,w1;h2,w2;...'. "
"e.g. '1024,1024;768,1360;1360,768;880,1168;1168,880;1248,832;832,1248'. "
"Requires --use_aspect_ratio_buckets. Images are resized to cover and cropped to the nearest "
"listed bucket (smaller images are upscaled). When set, --resolution is ignored."
),
)
parser.add_argument(
"--use_aspect_ratio_buckets",
action="store_true",
help=(
"Enable aspect-ratio bucketing. Without --aspect_ratio_buckets, the buckets are computed on the "
"fly from --resolution and capped to each image's own resolution, so smaller images are assigned "
"to a smaller bucket instead of being upscaled. Provide --aspect_ratio_buckets to use an explicit list."
),
)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If --aspect_ratio_buckets is provided and --use_aspect_ratio_buckets is not set, should we throw an error? We could also just keep --aspect_ratio_buckets, default it to None. If specified, then only use it for the conditionals. WDYT?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Kept the error here to stay aligned with the behavior we landed on in #14158 (discussion) - keeping --use_aspect_ratio_buckets as the single explicit switch across all the dreambooth scripts, so a bucket list without it fails fast instead of silently enabling bucketing.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

examples size/L PR with diff > 200 LOC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants