diff --git a/examples/dreambooth/test_dreambooth_lora_flux2_klein.py b/examples/dreambooth/test_dreambooth_lora_flux2_klein.py index 1bc8da075374..9492f14b2d8d 100644 --- a/examples/dreambooth/test_dreambooth_lora_flux2_klein.py +++ b/examples/dreambooth/test_dreambooth_lora_flux2_klein.py @@ -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): + 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): + 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())) diff --git a/examples/dreambooth/train_dreambooth_lora_flux2.py b/examples/dreambooth/train_dreambooth_lora_flux2.py index 5643fb95a910..c47a9e8eed4f 100644 --- a/examples/dreambooth/train_dreambooth_lora_flux2.py +++ b/examples/dreambooth/train_dreambooth_lora_flux2.py @@ -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, @@ -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]), - ] - ) - def __len__(self): return self._length @@ -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): diff --git a/examples/dreambooth/train_dreambooth_lora_flux2_img2img.py b/examples/dreambooth/train_dreambooth_lora_flux2_img2img.py index 514832dd31db..fed31bf200bc 100644 --- a/examples/dreambooth/train_dreambooth_lora_flux2_img2img.py +++ b/examples/dreambooth/train_dreambooth_lora_flux2_img2img.py @@ -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, @@ -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 @@ -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) diff --git a/examples/dreambooth/train_dreambooth_lora_flux2_klein.py b/examples/dreambooth/train_dreambooth_lora_flux2_klein.py index 7b7e1634f8a7..3bf617b257d1 100644 --- a/examples/dreambooth/train_dreambooth_lora_flux2_klein.py +++ b/examples/dreambooth/train_dreambooth_lora_flux2_klein.py @@ -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, @@ -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." ) @@ -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 @@ -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): @@ -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): @@ -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) diff --git a/examples/dreambooth/train_dreambooth_lora_flux2_klein_img2img.py b/examples/dreambooth/train_dreambooth_lora_flux2_klein_img2img.py index ca29cf2dbaff..1c1373d63255 100644 --- a/examples/dreambooth/train_dreambooth_lora_flux2_klein_img2img.py +++ b/examples/dreambooth/train_dreambooth_lora_flux2_klein_img2img.py @@ -435,6 +435,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, @@ -880,15 +885,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 @@ -914,40 +910,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) diff --git a/examples/dreambooth/train_dreambooth_lora_flux_kontext.py b/examples/dreambooth/train_dreambooth_lora_flux_kontext.py index 846ce6336e54..a5019831041e 100644 --- a/examples/dreambooth/train_dreambooth_lora_flux_kontext.py +++ b/examples/dreambooth/train_dreambooth_lora_flux_kontext.py @@ -452,6 +452,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, @@ -914,15 +919,6 @@ def __init__( else: self.class_data_root = None - self.image_transforms = transforms.Compose( - [ - transforms.Resize(self.size, interpolation=transforms.InterpolationMode.BILINEAR), - transforms.CenterCrop(self.size) if center_crop else transforms.RandomCrop(self.size), - transforms.ToTensor(), - transforms.Normalize([0.5], [0.5]), - ] - ) - def __len__(self): return self._length @@ -951,46 +947,49 @@ 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.paired_transform( + class_image, size=self.buckets[bucket_idx], center_crop=self.center_crop + )[0] example["class_prompt"] = self.class_prompt 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) diff --git a/examples/dreambooth/train_dreambooth_lora_z_image.py b/examples/dreambooth/train_dreambooth_lora_z_image.py index de3e01e4f2de..e956b28057f4 100644 --- a/examples/dreambooth/train_dreambooth_lora_z_image.py +++ b/examples/dreambooth/train_dreambooth_lora_z_image.py @@ -444,6 +444,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, @@ -879,15 +884,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 @@ -912,37 +908,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):