[Examples] Discussions on using streaming datasets for ControlNet training in JAX

[sayakpaul]

As we discussed how well a large dataset would play in conjunction with the ControlNet training script in JAX, I considered consolidating what I have been doing in this regard.

To mimic a mid-scale dataset for training, I created this dummy dataset: https://huggingface.co/datasets/sayakpaul/dummy-controlnet-100000-samples. It has 100000 samples.

Then I followed different training examples that use JAX to come up with a dummy script that would load the dataset in the streaming model and perform a dummy forward pass (no models needed for this):

• https://github.com/huggingface/transformers/blob/main/examples/research_projects/jax-projects/dataset-streaming/run_mlm_flax_stream.py
• https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_flax.py

This is the dummy script I used for testing: https://gist.github.com/sayakpaul/b7ba991bf15e302ad0c9bbfe6e1e7f83#file-train_dummy-py. Things seem to work fine unless I am missing out on something. In that case, please let me know.

Additionally, one can find this gist that has all the instructions and scripts I used to perform this testing.

Ccing @patrickvonplaten @yiyixuxu

[sayakpaul]

Thanks to @patil-suraj for pointing me to https://github.com/borisdayma/dalle-mini/blob/main/src/dalle_mini/data.py.

I adapted it and did something like this: https://gist.github.com/sayakpaul/f24f287b9f7443ef44c65705d3003b34.

Now, to test the data loader, I did:

from transformers import AutoTokenizer
from data import Dataset 


def main():
    dataset_id = "multimodalart/facesyntheticsspigacaptioned"
    model_id = "runwayml/stable-diffusion-v1-5"

    tokenizer = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer")
    dataset_repo_or_path = dataset_id
    resolution = 256
    batch_size = 16
    conditioning_image_column = "spiga_seg"
    caption_column = "image_caption"

    dataset = Dataset(
        dataset_repo_or_path=dataset_repo_or_path,
        resolution=resolution,
        batch_size=batch_size,
        conditioning_image_column=conditioning_image_column,
        caption_column=caption_column,
        streaming=True
    )

    dataset.preprocess(tokenizer=tokenizer)

    train_loader = dataset.dataloader("train", 16, 1)

    print(list(next(iter(train_loader))))


if __name__ == "__main__":
    main()

Seems to work as expected. Btw, we can actually enable streaming in @yiyixuxu's current PR: #2818. Just enabling streaming=True in load_dataset should work out of the box along with this change instead of with_transform().

However, as we can see here as well, we cannot set steps_per_epoch then. In this case, it's either set as None or is computed from max_train_samples.

[patrickvonplaten]

That's super helpful! Thanks a lot for compiling everything here :-)