Skip to content

Conversation

Light-yzc
Copy link

PR Title: feat: Add aspect ratio bucketing to SDXL LoRA training scripts

What does this PR do?

This PR adds support for aspect ratio bucketing to the train_text_to_image_lora_sdxl.py and train_dreambooth_lora_sdxl.py scripts.

This allows training on images with different resolutions and aspect ratios. It avoids forcing all images to be cropped into a square, which helps preserve the original composition.

Changes:

  • Added a BucketBatchSampler that creates batches of images with similar resolutions.
  • Modified the data preprocessing logic to assign each image to a bucket and resize/crop it accordingly.
  • Updated compute_time_ids to handle the dynamic target sizes from the buckets.
  • For the DreamBooth script, the bucketing logic now also applies to class images to ensure instance/prior image pairs have matching dimensions.

New command-line arguments like --enable_bucketing and --bucket_size are added to control this feature.

Fixes # (issue)

Before submitting

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant