Text-to-Image Dataset and Flux Transform

[calvinpelletier]

Context

This adds a text-to-image dataset class, as well as a Flux transform that converts generic text-to-image data to Flux-specific data.

The text-to-image dataset class also has the optional ability to yield a row ID along with the text and image data. This will be used to preprocess the dataset before Flux finetuning.

The Flux transform also exposes a public function just for tokenizing text. This will be used by the Flux finetuning recipe to tokenize the prompts used to sample images from the model during training.

Changelog

• Text-to-image dataset class and builder
• Flux transform
• Unit test for the dataset class
• Unit test for the Flux transform

Test plan

Manual testing

• created a local dataset of labeled images and finetuned Flux on the dataset

Checklist

• run pre-commit hooks and linters (make sure you've first installed via pre-commit install)
• add unit tests for any new functionality
• update docstrings for any new or updated methods or classes
• run unit tests via pytest tests

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2311

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 3 New Failures, 6 Cancelled Jobs

As of commit 44a770e with merge base 90fd2d3 ():

NEW FAILURES - The following jobs have failed:

• GPU tests / gpu_test (3.9, stable) (gh)
  Process completed with exit code 2.
• Recipe Tests / recipe_test (3.9) (gh)
  Process completed with exit code 2.
• Unit Test / unit_tests (3.9) (gh)
  Process completed with exit code 2.

CANCELLED JOBS - The following jobs were cancelled. Please retry:

• GPU tests / gpu_test (3.10, stable) (gh)
  ##[error]The operation was canceled.
• GPU tests / gpu_test (3.11, stable) (gh)
• Recipe Tests / recipe_test (3.10) (gh)
• Recipe Tests / recipe_test (3.11) (gh)
  ##[error]The operation was canceled.
• Unit Test / unit_tests (3.10) (gh)
  ##[error]The operation was canceled.
• Unit Test / unit_tests (3.11) (gh)
  ##[error]The operation was canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.