Skip to content

[lora training] fix aspect ratio bucketing in dreambooth scripts (+ caption dropout, --bucket_no_upscale)#14158

Open
linoytsaban wants to merge 2 commits into
mainfrom
lora-bucketing-fixes
Open

[lora training] fix aspect ratio bucketing in dreambooth scripts (+ caption dropout, --bucket_no_upscale)#14158
linoytsaban wants to merge 2 commits into
mainfrom
lora-bucketing-fixes

Conversation

@linoytsaban

Copy link
Copy Markdown
Collaborator

Split from #14089. Bucketing fixes for the scripts that already support aspect-ratio buckets (flux2, flux2_img2img, flux2_klein(+img2img), flux_kontext, z_image):

  • resize-to-cover then crop instead of resizing to exact bucket dims (which distorted images and made the crop a no-op)
  • bucket class/prior-preservation images too (were forced to a fixed square)
  • --bucket_no_upscale: pad images smaller than their bucket instead of upscaling
  • caption dropout added to flux2_klein

@github-actions github-actions Bot added size/L PR with diff > 200 LOC examples and removed size/L PR with diff > 200 LOC labels Jul 10, 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.

@github-actions github-actions Bot added the size/L PR with diff > 200 LOC label Jul 10, 2026

@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.

Left some comments. LMK.

loaded_lora_rank = raw["transformer.r"]
self.assertTrue(loaded_lora_rank == rank)

def test_dreambooth_lora_flux2_aspect_ratio_buckets(self):

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.

I think we can remove this test for now. Let's see if this feature becomes more and more used and then add it back. So, we could just skip it.

self.assertTrue(all("lora" in k for k in lora_state_dict.keys()))
self.assertTrue(all(key.startswith("transformer") for key in lora_state_dict.keys()))

def test_dreambooth_lora_flux2_caption_dropout(self):

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 for this.

Comment on lines +453 to +457
parser.add_argument(
"--bucket_no_upscale",
action="store_true",
help="If set, images smaller than their aspect-ratio bucket are padded instead of upscaled.",
)

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.

Does this help at all? If so, I would maybe keep it simple and just upscale. Also, should we enable bucketing through an arg?

Comment on lines -893 to -900
self.image_transforms = transforms.Compose(
[
transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR),
transforms.CenterCrop(size) if center_crop else transforms.RandomCrop(size),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]),
]
)

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 we're going with the above suggestion i.e. make aspect-ratio bucketing conditional on an arg, then I guess this should stay?

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