Adding VQGAN Training script

[isamu-isozaki]

What does this PR do?

This is a vqgan training script ported from taming-transformers and from lucidrian's muse-maskgit repo here and open-muse. I'm planning to test this on the cifar10 dataset to confirm it works

Some steps missing/need confirmation are

• Confirm einops and timm can be external dependencies. If not convert these ops to native pytorch
• Test on cifar10
• Add in test to test_models_vq and test_models_vae for the slight modification

Fixes #4702

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.

[isamu-isozaki]

Once confirmed it works with cifar10 will remove the draft part

[isamu-isozaki]

I was able to start training this script. And I removed the einops dependencies. The only additional dependency so far is timm. I plan to run this overnight on cifar with 128 image resolution and then remove the draft from this pr. Also let me know if anyone knows a good VQModel config that's easy to train/fast

[isamu-isozaki]

Ok! Training seems to work. Here's a wandb run on cifar 10. In 6gb vram, command to run this is

accelerate launch train_vqgan.py --dataset_name=cifar10 --image_column=img --validation_images images/bird.jpg images/car.jpg images/dog.jpg images/frog.jpg images/horse.jpg images/ship.jpg --resolution=128 --train_batch_size=2 --gradient_accumulation_steps=8 --report_to=wandb

For the validation images, they will be shown like so for each validation image provided. The left is the input image and the right is the generated image

The remaining parts that I can think of are

• Make log_validation support trackers other than wandb
• Make tqdm updates similar to other examples

I did find a bug where global step doesn't seem to go above 3000 but once that is fixed I'll open for review

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[isamu-isozaki]

The main logic is done so I think it's ready for review. For the 3000 step bug I'm currently running training to see if it happens again after the fixes.

[isamu-isozaki]

Ok! Seems like it was a hardware issue(I think). Got steps 3100. Script should be ready for review.

[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.

[yqy2001]

Hi there, what is the current status of this PR? It seems that everything works well. Will this be merged?

[isamu-isozaki]

@sayakpaul I tried fixing by following the @require_torch format by making a @require_timm. Let me know what you think

[yiyixuxu]

maybe we don't need this special return_loss flag
I don't think it would break, no?
it is a tuple, we would usually use out[0], if it is a DecoderOutput, usually we do out.sample; I think just adding the loss to the output should be fine cc @DN6 to confirm if it's non-breaking

[isamu-isozaki]

Hmm, wouldn't it be possible for some people to do out[-1] for tuple? I think that's the only time it'll break

[sayakpaul]

Even if they do that, I think the error message would be fairly easy to digest but I don't think it will be breaking. WDYT? I like the idea of not introducing return_loss.

[isamu-isozaki]

Good point. I'll remove return_loss

[sayakpaul]

@isamu-isozaki I pushed a few things and I hope you don't mind.

• Moved is_timm_available and require_timm to proper modules.
• Added timm as a dependency in our workflows.
• Decorated the trainer test class with require_timm instead of doing it per test methods.

[isamu-isozaki]

@sayakpaul No worries thanks a bunch for doing that. I did forget the proper way to add modules for the tests 😅

[sayakpaul]

Ah all tests passing. Sight to the sore eyes, eh!

[isamu-isozaki]

@sayakpaul awesome! I just removed return_loss(and hopefully tests still pass). I did do the tests on my end

[sayakpaul]

@isamu-isozaki could you resolve the conflicts so that it's ready for merging? We would like to include in our upcoming release. Sorry for the delay on my end.

@yiyixuxu could give the changes introduced to the library components a bit?

[isamu-isozaki]

@sayakpaul tnx I think I resolved the conflicts but let me do tests to make sure

[sayakpaul]

Okay the code quality issues should be easy to fix I think. But LMK if you find difficulties. What I would do:

• Create a fresh Python env.
• From diffusers root, run pip install -e .[quality].
• Run make style && make quality.
• Push the changes.

[isamu-isozaki]

@sayakpaul tnx a bunch. I think I fixed the ruff format error but one question I have is when I try locally the doc-builder always fails even in a fresh environment with the above steps. But when I just fix all the tests before that, the checks in the ci usually passes. It might be a bug on my part but is that a common issue?
The error is

ruff check examples scripts src tests utils benchmarks setup.py
ruff format --check examples scripts src tests utils benchmarks setup.py
918 files left unchanged
doc-builder style src/diffusers docs/source --max_len 119 --check_only
Traceback (most recent call last):
  File "/home/isamu/miniconda3/envs/diffusers/bin/doc-builder", line 8, in <module>
    sys.exit(main())
  File "/home/isamu/miniconda3/envs/diffusers/lib/python3.10/site-packages/doc_builder/commands/doc_builder_cli.py", line 47, in main
    args.func(args)
  File "/home/isamu/miniconda3/envs/diffusers/lib/python3.10/site-packages/doc_builder/commands/style.py", line 28, in style_command
    raise ValueError(f"{len(changed)} files should be restyled!")
ValueError: 284 files should be restyled!
Makefile:43: recipe for target 'quality' failed
make: *** [quality] Error 1

locally but I think it'll pass here

[sayakpaul]

That's weird. Could be a setup related problem :/

[sayakpaul]

Alright merging this now!

[sayakpaul]

Thanks a lot for shipping this super cool script, @isamu-isozaki. Really appreciate your hard work and patience!

[isamu-isozaki]

@sayakpaul np! No worries at all and thanks for the support!