[lora training] add caption dropout and aspect ratio buckets to krea2 and qwen-image#14089
[lora training] add caption dropout and aspect ratio buckets to krea2 and qwen-image#14089linoytsaban wants to merge 7 commits into
Conversation
|
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. |
a539d74 to
975986d
Compare
sayakpaul
left a comment
There was a problem hiding this comment.
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? |
|
Yes sure |
098cc5b to
25d9f73
Compare
sayakpaul
left a comment
There was a problem hiding this comment.
Same comments as the other PR. WDYT?
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.
| 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." | ||
| ), | ||
| ) |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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.

Adds aspect-ratio bucketing and caption dropout to the krea2 and qwen-image dreambooth LoRA scripts (neither supported these before):
--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.