Modularize Dreambooth LoRA SDXL inferencing during and after training

[rootonchair]

What does this PR do?

Part of #6545

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[yiyixuxu]

cc @sayakpaul

[sayakpaul]

Very nice changes! As with the other PR, could you provide some training commands with which we could test it?

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[sayakpaul]

@rootonchair let's fix the code quality issues. I just ran your changes and I am running into the following problem:

Traceback (most recent call last):
  File "/home/sayak/collabs/diffusers/examples/dreambooth/train_dreambooth_lora_sdxl.py", line 1791, in <module>
    main(args)
  File "/home/sayak/collabs/diffusers/examples/dreambooth/train_dreambooth_lora_sdxl.py", line 1700, in main
    images = log_validation(
  File "/home/sayak/collabs/diffusers/examples/dreambooth/train_dreambooth_lora_sdxl.py", line 186, in log_validation
    phase_name: [
  File "/home/sayak/collabs/diffusers/examples/dreambooth/train_dreambooth_lora_sdxl.py", line 187, in <listcomp>
    wandb.Image(image, caption=f"{i}: {args.validation_prompt}") for i, image in enumerate(images)
NameError: name 'wandb' is not defined

I ran it with:

accelerate launch train_dreambooth_lora_sdxl.py   --pretrained_model_name_or_path=$MODEL_NAME    --instance_data_dir=$INSTANCE_DIR   --pretrained_vae_model_name_or_path=$VAE_PATH   --output_dir=$OUTPUT_DIR   --mixed_precision="fp16"   --instance_prompt="a photo of sks dog"   --resolution=1024   --train_batch_size=1   --gradient_accumulation_steps=4   --learning_rate=1e-4   --report_to="wandb"   --lr_scheduler="constant"   --lr_warmup_steps=0   --max_train_steps=500   --validation_prompt="A photo of sks dog in a bucket"   --validation_epochs=25   --seed="0"   --push_to_hub

Am I missing out on something?

[rootonchair]

@sayakpaul seem like python doesn't allow a module imported inside a function to be used in another function. You could reproduce with this small script

>>> def foo():
...     print(os.getcwd())
... 
>>> def main():
...     import os
...     foo()
... 
>>> main()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 3, in main
  File "<stdin>", line 2, in foo
NameError: name 'os' is not defined

[sayakpaul]

But this doesn't happen in the other PR you raised for SD v1.5. Can we unify? Without the error being fixed, we cannot merge this change, I am afraid.

[rootonchair]

But this doesn't happen in the other PR you raised for SD v1.5. Can we unify? Without the error being fixed, we cannot merge this change, I am afraid.

Ahh, my bad, I forgot to revert back to the old wandb import in the other PR. Should we keep the new way of handling wadb?

[sayakpaul]

Yeah that I think that is indeed the best bet. Sorry for the trouble here. I will update my suggestions in the other PR.

[rootonchair]

No problem. I wasn't sure it was the best way until you raise that error. Let's me redo the changes and fix code quality problem

[rootonchair]

@sayakpaul could you help to check if there are any problem. Thank you

[sayakpaul]

It's working as expected now, thanks!