Repository: huggingface/diffusers
File: examples/dreambooth/train_dreambooth_lora_flux_kontext.py (Python)

Code before:
            image = crop(image)
            if dest_image is not None:
                dest_image = crop(dest_image)
        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 dest_image is not None:
            dest_image = normalize(to_tensor(dest_image))

        return (image, dest_image) if dest_image is not None else (image, None)


def collate_fn(examples, with_prior_preservation=False):
    pixel_values = [example["instance_images"] for example in examples]
    prompts = [example["instance_prompt"] for example in examples]

    # Concat class and instance examples for prior preservation.
    # We do this to avoid doing two forward passes.
    if with_prior_preservation:
        pixel_values += [example["class_images"] for example in examples]
        prompts += [example["class_prompt"] for example in examples]

    pixel_values = torch.stack(pixel_values)
    pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float()

    batch = {"pixel_values": pixel_values, "prompts": prompts}
    if any("cond_images" in example for example in examples):
        cond_pixel_values = [example["cond_images"] for example in examples]
        cond_pixel_values = torch.stack(cond_pixel_values)
        cond_pixel_values = cond_pixel_values.to(memory_format=torch.contiguous_format).float()
        batch.update({"cond_pixel_values": cond_pixel_values})
    return batch


class BucketBatchSampler(BatchSampler):

Review comment (suggestion):
Should initialize one-time only. All deterministic transformations should be initialized only once. Future PR.

Code after:
            dest_image = resize(dest_image)

        # 2. Crop: either center or SAME random crop
        if center_crop:
            crop = transforms.CenterCrop(size)
            image = crop(image)
            if dest_image is not None:
                dest_image = crop(dest_image)
        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 dest_image is not None:
            dest_image = normalize(to_tensor(dest_image))

        return (image, dest_image) if dest_image is not None else (image, None)


def collate_fn(examples, with_prior_preservation=False):
    pixel_values = [example["instance_images"] for example in examples]
    prompts = [example["instance_prompt"] for example in examples]

    # Concat class and instance examples for prior preservation.
    # We do this to avoid doing two forward passes.
    if with_prior_preservation:
        pixel_values += [example["class_images"] for example in examples]
        prompts += [example["class_prompt"] for example in examples]

    pixel_values = torch.stack(pixel_values)
    pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float()

    batch = {"pixel_values": pixel_values, "prompts": prompts}
    if any("cond_images" in example for example in examples):
        cond_pixel_values = [example["cond_images"] for example in examples]
        cond_pixel_values = torch.stack(cond_pixel_values)
        cond_pixel_values = cond_pixel_values.to(memory_format=torch.contiguous_format).float()