Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add partial SDXL model #61

Merged
merged 15 commits into from
Aug 21, 2023
Merged

Add partial SDXL model #61

merged 15 commits into from
Aug 21, 2023

Conversation

jazcollins
Copy link
Contributor

This PR adds diffusion.models.models.stable_diffusion_xl, which incorporates a few ideas from SDXL. Namely:

  • Training with a larger UNet
  • Loading the pre-trained SDXL VAE

Currently we're using the madebyollin/sdxl-vae-fp16-fix VAE checkpoint rather than the official stabilityai/stable-diffusion-xl-base-1.0 VAE checkpoint. The stabilityai checkpoint leads to NaNs when training with fp16. More on that issue here.

Copy link
Collaborator

@coryMosaicML coryMosaicML left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks great! Only a minor comment.

diffusion/models/models.py Outdated Show resolved Hide resolved
diffusion/models/models.py Show resolved Hide resolved
Copy link
Contributor

@Landanjs Landanjs left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks great! The only major adjustment is the NotImplementedError for the pretrained arg. I also had a question around which VAE we should prefer if we can run it on bf16.

For the future, there should be a future PR that adds the other modeling features (in addition to the conditioning PR), so we can run a pre-trained SDXL. Sorry for going back-and-forth on whether to fully implement SDXL

diffusion/models/models.py Outdated Show resolved Hide resolved
diffusion/models/models.py Outdated Show resolved Hide resolved
diffusion/models/models.py Show resolved Hide resolved
@jazcollins jazcollins merged commit bacab36 into mosaicml:main Aug 21, 2023
7 checks passed
@jazcollins jazcollins deleted the sdxl-pr branch August 21, 2023 17:38
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants