[Examples] Add streaming support to the ControlNet training example in JAX

[sayakpaul]

Potentially closes #2840.

Example command:

python3 train_controlnet_flax.py \
	--pretrained_model_name_or_path=$MODEL_DIR \
	--output_dir=$OUTPUT_DIR \
	--dataset_name=multimodalart/facesyntheticsspigacaptioned \
	--streaming \
	--conditioning_image_column=spiga_seg \
	--image_column=image \
	--caption_column=image_caption \
	--resolution=512 \
	--max_train_samples 50 \
	--max_train_steps 5 \
	--learning_rate=1e-5 \
	--validation_steps=2 \
	--train_batch_size=1 \
	--revision="flax" \
	--report_to="wandb"

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[patrickvonplaten]

Looks cool, @yiyixuxu can you take a look as well?

[sayakpaul]

I forgot to add a note about how to deal with num_update_steps_per_epoch when streaming is enabled.

I couldn't find any proper resource that shows how to deal with this. Most of them deal with this situation by using a pre-specified number such as done in https://github.com/borisdayma/dalle-mini/blob/main/src/dalle_mini/data.py. I think this is fine, no?

[yiyixuxu]

@sayakpaul

Testing now here, got this error
https://wandb.ai/yiyixu/train_controlnet_flax/runs/izb6xux1/logs?workspace=user-yiyixu

[yiyixuxu]

@sayakpaul

I think it's en easy fix:)
I will get it working and probably going to push one commit to your PR

[yiyixuxu]

@sayakpaul

ok, tested here https://wandb.ai/yiyixu/train_controlnet_flax/runs/np6ktkj4/logs?workspace=user-yiyixu
Thanks a lot for the PR! super cool!

feel free to merge whenever:)

[sayakpaul]

Thanks so much, @yiyixuxu, for taking care of the issue.

Merging.