Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
linoytsaban wants to merge 2 commits into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a2a6678
[lora training] fix aspect ratio bucketing and add caption dropout + …
linoytsaban Jul 10, 2026
af19b04
Merge branch 'main' into lora-bucketing-fixes
sayakpaul Jul 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 52 additions & 0 deletions examples/dreambooth/test_dreambooth_lora_flux2_klein.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -260,3 +260,55 @@ def test_dreambooth_lora_with_metadata(self):
self.assertTrue(loaded_lora_alpha == lora_alpha)
loaded_lora_rank = raw["transformer.r"]
self.assertTrue(loaded_lora_rank == rank)

def test_dreambooth_lora_flux2_aspect_ratio_buckets(self):

@sayakpaul sayakpaul Jul 10, 2026

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.

with tempfile.TemporaryDirectory() as tmpdir:
test_args = f"""
{self.script_path}
--pretrained_model_name_or_path {self.pretrained_model_name_or_path}
--instance_data_dir {self.instance_data_dir}
--instance_prompt {self.instance_prompt}
--aspect_ratio_buckets 64,64;64,128
--bucket_no_upscale
--cache_latents
--train_batch_size 1
--gradient_accumulation_steps 1
--max_train_steps 2
--learning_rate 5.0e-04
--lr_scheduler constant
--lr_warmup_steps 0
--max_sequence_length 8
--text_encoder_out_layers 1
--output_dir {tmpdir}
""".split()

run_command(self._launch_args + test_args)
self.assertTrue(os.path.isfile(os.path.join(tmpdir, "pytorch_lora_weights.safetensors")))
lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdir, "pytorch_lora_weights.safetensors"))
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):

@sayakpaul sayakpaul Jul 10, 2026

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.

with tempfile.TemporaryDirectory() as tmpdir:
test_args = f"""
{self.script_path}
--pretrained_model_name_or_path {self.pretrained_model_name_or_path}
--instance_data_dir {self.instance_data_dir}
--instance_prompt {self.instance_prompt}
--resolution 64
--caption_dropout 1.0
--train_batch_size 1
--gradient_accumulation_steps 1
--max_train_steps 2
--learning_rate 5.0e-04
--lr_scheduler constant
--lr_warmup_steps 0
--max_sequence_length 8
--text_encoder_out_layers 1
--output_dir {tmpdir}
""".split()

run_command(self._launch_args + test_args)
self.assertTrue(os.path.isfile(os.path.join(tmpdir, "pytorch_lora_weights.safetensors")))
lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdir, "pytorch_lora_weights.safetensors"))
self.assertTrue(all("lora" in k for k in lora_state_dict.keys()))
58 changes: 26 additions & 32 deletions examples/dreambooth/train_dreambooth_lora_flux2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -450,6 +450,11 @@ def parse_args(input_args=None):
"Images will be resized and cropped to fit the nearest bucket. If provided, --resolution is ignored."
),
)
parser.add_argument(
"--bucket_no_upscale",
action="store_true",
help="If set, images smaller than their aspect-ratio bucket are padded instead of upscaled.",
)
Comment on lines +453 to +457

@sayakpaul sayakpaul Jul 10, 2026

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?

parser.add_argument(
"--center_crop",
default=False,
Expand Down Expand Up @@ -890,15 +895,6 @@ def __init__(
else:
self.class_data_root = None

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]),
]
)
Comment on lines -893 to -900

@sayakpaul sayakpaul Jul 10, 2026

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?


def __len__(self):
return self._length

Expand All @@ -924,37 +920,35 @@ def __getitem__(self, index):

if not class_image.mode == "RGB":
class_image = class_image.convert("RGB")
example["class_images"] = self.image_transforms(class_image)
# Match the class image to the paired instance image's bucket so they can be stacked into one batch.
example["class_images"] = self.train_transform(
class_image, size=self.buckets[bucket_idx], center_crop=self.center_crop
)
example["class_prompt"] = self.class_prompt

return example

def train_transform(self, image, size=(224, 224), center_crop=False, random_flip=False):
# 1. Resize (deterministic)
resize = transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR)
image = resize(image)

# 2. Crop: either center or SAME random crop
def train_transform(self, image, size, center_crop=False, random_flip=False):
# Resize preserving aspect ratio so the image covers the bucket, then crop to the bucket size.
target_height, target_width = size
width, height = image.size
scale = max(target_height / height, target_width / width)
if args.bucket_no_upscale:
scale = min(scale, 1.0)
new_height, new_width = round(height * scale), round(width * scale)
image = TF.resize(image, [new_height, new_width], interpolation=transforms.InterpolationMode.BILINEAR)
# Pad to the bucket when no-upscale leaves the image smaller, so batched samples share a shape.
pad_w, pad_h = max(0, target_width - new_width), max(0, target_height - new_height)
if pad_w or pad_h:
image = TF.pad(image, [pad_w // 2, pad_h // 2, pad_w - pad_w // 2, pad_h - pad_h // 2])
if center_crop:
crop = transforms.CenterCrop(size)
image = crop(image)
image = TF.center_crop(image, size)
else:
# get_params returns (i, j, h, w)
i, j, h, w = transforms.RandomCrop.get_params(image, output_size=size)
image = TF.crop(image, i, j, h, w)

# 3. Random horizontal flip with the SAME coin flip
if random_flip:
do_flip = random.random() < 0.5
if do_flip:
image = TF.hflip(image)

# 4. ToTensor + Normalize (deterministic)
to_tensor = transforms.ToTensor()
normalize = transforms.Normalize([0.5], [0.5])
image = normalize(to_tensor(image))

return image
if random_flip and random.random() < 0.5:
image = TF.hflip(image)
return TF.normalize(TF.to_tensor(image), [0.5], [0.5])


def collate_fn(examples, with_prior_preservation=False):
Expand Down
64 changes: 30 additions & 34 deletions examples/dreambooth/train_dreambooth_lora_flux2_img2img.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -434,6 +434,11 @@ def parse_args(input_args=None):
"Images will be resized and cropped to fit the nearest bucket. If provided, --resolution is ignored."
),
)
parser.add_argument(
"--bucket_no_upscale",
action="store_true",
help="If set, images smaller than their aspect-ratio bucket are padded instead of upscaled.",
)
parser.add_argument(
"--center_crop",
default=False,
Expand Down Expand Up @@ -884,15 +889,6 @@ def __init__(
self.num_instance_images = len(self.instance_images)
self._length = self.num_instance_images

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]),
]
)

def __len__(self):
return self._length

Expand All @@ -918,40 +914,40 @@ def __getitem__(self, index):
return example

def paired_transform(self, image, dest_image=None, size=(224, 224), center_crop=False, random_flip=False):
# 1. Resize (deterministic)
resize = transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR)
image = resize(image)
# Resize preserving aspect ratio so the image covers the bucket, then crop to the bucket size.
# The same geometry is applied to the conditioning image so the pair stays aligned.
target_height, target_width = size
width, height = image.size
scale = max(target_height / height, target_width / width)
if args.bucket_no_upscale:
scale = min(scale, 1.0)
new_size = [round(height * scale), round(width * scale)]
# Pad to the bucket when no-upscale leaves the image smaller, so batched samples share a shape.
pad_w, pad_h = max(0, target_width - new_size[1]), max(0, target_height - new_size[0])
padding = [pad_w // 2, pad_h // 2, pad_w - pad_w // 2, pad_h - pad_h // 2]
image = TF.resize(image, new_size, interpolation=transforms.InterpolationMode.BILINEAR)
if pad_w or pad_h:
image = TF.pad(image, padding)
if dest_image is not None:
dest_image = resize(dest_image)

# 2. Crop: either center or SAME random crop
dest_image = TF.resize(dest_image, new_size, interpolation=transforms.InterpolationMode.BILINEAR)
if pad_w or pad_h:
dest_image = TF.pad(dest_image, padding)
if center_crop:
crop = transforms.CenterCrop(size)
image = crop(image)
image = TF.center_crop(image, size)
if dest_image is not None:
dest_image = crop(dest_image)
dest_image = TF.center_crop(dest_image, size)
else:
# get_params returns (i, j, h, w)
i, j, h, w = transforms.RandomCrop.get_params(image, output_size=size)
image = TF.crop(image, i, j, h, w)
if dest_image is not None:
dest_image = TF.crop(dest_image, i, j, h, w)

# 3. Random horizontal flip with the SAME coin flip
if random_flip:
do_flip = random.random() < 0.5
if do_flip:
image = TF.hflip(image)
if dest_image is not None:
dest_image = TF.hflip(dest_image)

# 4. ToTensor + Normalize (deterministic)
to_tensor = transforms.ToTensor()
normalize = transforms.Normalize([0.5], [0.5])
image = normalize(to_tensor(image))
if random_flip and random.random() < 0.5:
image = TF.hflip(image)
if dest_image is not None:
dest_image = TF.hflip(dest_image)
image = TF.normalize(TF.to_tensor(image), [0.5], [0.5])
if dest_image is not None:
dest_image = normalize(to_tensor(dest_image))

dest_image = TF.normalize(TF.to_tensor(dest_image), [0.5], [0.5])
return (image, dest_image) if dest_image is not None else (image, None)


Expand Down
87 changes: 55 additions & 32 deletions examples/dreambooth/train_dreambooth_lora_flux2_klein.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -450,6 +450,11 @@ def parse_args(input_args=None):
"Images will be resized and cropped to fit the nearest bucket. If provided, --resolution is ignored."
),
)
parser.add_argument(
"--bucket_no_upscale",
action="store_true",
help="If set, images smaller than their aspect-ratio bucket are padded instead of upscaled.",
)
parser.add_argument(
"--center_crop",
default=False,
Expand All @@ -464,6 +469,16 @@ def parse_args(input_args=None):
action="store_true",
help="whether to randomly flip images horizontally",
)
parser.add_argument(
"--caption_dropout",
type=float,
default=0.0,
help=(
"Probability of replacing an instance image's caption with an empty string during training, so that"
" fraction of samples is trained unconditionally. Improves classifier-free guidance. A common value is"
" 0.1. Class/prior-preservation captions are never dropped."
),
)
parser.add_argument(
"--train_batch_size", type=int, default=4, help="Batch size (per device) for the training dataloader."
)
Expand Down Expand Up @@ -885,15 +900,6 @@ def __init__(
else:
self.class_data_root = None

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]),
]
)

def __len__(self):
return self._length

Expand All @@ -919,37 +925,35 @@ def __getitem__(self, index):

if not class_image.mode == "RGB":
class_image = class_image.convert("RGB")
example["class_images"] = self.image_transforms(class_image)
# Match the class image to the paired instance image's bucket so they can be stacked into one batch.
example["class_images"] = self.train_transform(
class_image, size=self.buckets[bucket_idx], center_crop=self.center_crop
)
example["class_prompt"] = self.class_prompt

return example

def train_transform(self, image, size=(224, 224), center_crop=False, random_flip=False):
# 1. Resize (deterministic)
resize = transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR)
image = resize(image)

# 2. Crop: either center or SAME random crop
def train_transform(self, image, size, center_crop=False, random_flip=False):
# Resize preserving aspect ratio so the image covers the bucket, then crop to the bucket size.
target_height, target_width = size
width, height = image.size
scale = max(target_height / height, target_width / width)
if args.bucket_no_upscale:
scale = min(scale, 1.0)
new_height, new_width = round(height * scale), round(width * scale)
image = TF.resize(image, [new_height, new_width], interpolation=transforms.InterpolationMode.BILINEAR)
# Pad to the bucket when no-upscale leaves the image smaller, so batched samples share a shape.
pad_w, pad_h = max(0, target_width - new_width), max(0, target_height - new_height)
if pad_w or pad_h:
image = TF.pad(image, [pad_w // 2, pad_h // 2, pad_w - pad_w // 2, pad_h - pad_h // 2])
if center_crop:
crop = transforms.CenterCrop(size)
image = crop(image)
image = TF.center_crop(image, size)
else:
# get_params returns (i, j, h, w)
i, j, h, w = transforms.RandomCrop.get_params(image, output_size=size)
image = TF.crop(image, i, j, h, w)

# 3. Random horizontal flip with the SAME coin flip
if random_flip:
do_flip = random.random() < 0.5
if do_flip:
image = TF.hflip(image)

# 4. ToTensor + Normalize (deterministic)
to_tensor = transforms.ToTensor()
normalize = transforms.Normalize([0.5], [0.5])
image = normalize(to_tensor(image))

return image
if random_flip and random.random() < 0.5:
image = TF.hflip(image)
return TF.normalize(TF.to_tensor(image), [0.5], [0.5])


def collate_fn(examples, with_prior_preservation=False):
Expand Down Expand Up @@ -1505,6 +1509,13 @@ def compute_text_embeddings(prompt, text_encoding_pipeline):
class_prompt_hidden_states, class_text_ids = compute_text_embeddings(
args.class_prompt, text_encoding_pipeline
)

# When caption dropout is enabled, we precompute the empty ("") prompt embedding once and swap it in
# for randomly selected instance samples at training time (see the training loop below).
if args.caption_dropout > 0:
with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload):
empty_prompt_hidden_states, empty_text_ids = compute_text_embeddings("", text_encoding_pipeline)

validation_embeddings = {}
if args.validation_prompt is not None:
with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload):
Expand Down Expand Up @@ -1744,6 +1755,18 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
prompt_embeds = static_prompt_embeds.repeat_interleave(num_repeat_elements, dim=0)
text_ids = static_text_ids.repeat_interleave(num_repeat_elements, dim=0)

# Caption dropout: replace a sample's caption embedding with the empty-prompt embedding so it
# trains unconditionally. Both paths order embeddings as [instance..., class...], so only the
# leading instance rows are dropped, never class/prior captions.
if args.caption_dropout > 0:
n_inst = len(sample_indices)
drop_mask = torch.rand(n_inst, device=prompt_embeds.device) < args.caption_dropout
if drop_mask.any():
prompt_embeds[:n_inst][drop_mask] = empty_prompt_hidden_states.to(
device=prompt_embeds.device, dtype=prompt_embeds.dtype
)
text_ids[:n_inst][drop_mask] = empty_text_ids.to(device=text_ids.device, dtype=text_ids.dtype)

# Convert images to latent space
if args.cache_latents:
model_input = torch.cat([instance_latents_cache[idx] for idx in sample_indices], dim=0)
Expand Down
Loading
Loading