[WIP]Flax training script for controlnet

[yiyixuxu]

wandb test run here: https://wandb.ai/yiyixu/train_controlnet_flax/runs/jye8xcs9?workspace=user-yiyixu

wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_1.png

export MODEL_DIR="runwayml/stable-diffusion-v1-5"
export OUTPUT_DIR="control_out"

python3 train_controlnet_flax.py \
 --pretrained_model_name_or_path=$MODEL_DIR \
 --output_dir=$OUTPUT_DIR \
 --dataset_name=fusing/fill50k \
 --resolution=512 \
 --learning_rate=1e-5 \
 --validation_image "./conditioning_image_1.png"  \
 --validation_prompt "red circle with blue background" \
 --validation_steps=1000 \
 --train_batch_size=1 \
 --revision="flax"\
 --report_to="wandb"

[HuggingFaceDocBuilderDev]

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

[jingyangcarl]

Hi, awesome works!! I run the code and find the following @yiyixuxu, any chance you can share some ideas on this? #2822

[patrickvonplaten]

Super nice! Feel free to merge :-)

[patrickvonplaten]

BTW, could you also add this:

export MODEL_DIR="runwayml/stable-diffusion-v1-5"
export OUTPUT_DIR="control_out"

python3 train_controlnet_flax.py \
 --pretrained_model_name_or_path=$MODEL_DIR \
 --output_dir=$OUTPUT_DIR \
 --dataset_name=fusing/fill50k \
 --resolution=512 \
 --learning_rate=1e-5 \
 --validation_image "./conditioning_image_1.png"  \
 --validation_prompt "red circle with blue background" \
 --validation_steps=12500 \
 --train_batch_size=1 \
 --revision="flax"\
 --report_to="wandb"

to the README here: https://github.com/huggingface/diffusers/blob/d9cf213e9b8cc59be953e1da7df916346612aef3/examples/controlnet/README.md

[yiyixuxu]

I've been running into this issue this afternoon #2851
but for the last 2 times, I was able to save the model to the hub (despite the error message) - I will leave that issue open for now in case we run into this issue in the future

see an example repo here with an automatically generated model card here (checkpoint with 4k steps)
https://huggingface.co/YiYiXu/fill-circle-controlnet

I will merge this in now so that more people can start to test it. Will open a separate PR to add intermediate checkpointing and make other improvements based on feedback.

[sayakpaul]

Thanks, @yiyixuxu. Do you plan on integrating the streaming support as we were discussing in #2840?

[yiyixuxu]

Thanks @sayakpaul
I'm taking a look right now, but you're much more of an expert in this than me:) I think it's up to you and @patrickvonplaten

[sayakpaul]

I'm taking a look right now, but you're much more of an expert in this than me:)

Actually, I don't believe the changes for enabling streaming will be long. I already mentioned a couple of pointers in #2840 (comment). Maybe you could try following the changes I made in here in your next PR? Happy to guide you if need be.

But happy to open a PR and have you test the changes. Would that make sense?

[yiyixuxu]

@sayakpaul yeah a PR please! happy to test it out :)

[sayakpaul]

Perf! Will have it opened as soon as possible.