Skip to content

Allow low precision vae sd xl#4083

Merged
patrickvonplaten merged 5 commits intomainfrom
allow_low_precision_vae_sd_xl
Jul 14, 2023
Merged

Allow low precision vae sd xl#4083
patrickvonplaten merged 5 commits intomainfrom
allow_low_precision_vae_sd_xl

Conversation

@patrickvonplaten
Copy link
Copy Markdown
Contributor

@patrickvonplaten patrickvonplaten commented Jul 13, 2023 •
edited
Loading

What does this PR do?

Ahllow low precision VAEs such as https://huggingface.co/madebyollin/sdxl-vae-fp16-fix

The amazing @madebyollin has fine-tuned a SD-XL VAE checkpoint that works in full fp16! You can try it out as follows:

from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline, AutoencoderKL

vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, force_upcast=False)
pipe = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-0.9", vae=vae, torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
pipe.to("cuda")

refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-0.9", vae=vae, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
refiner.to("cuda")

This saves 3GB of GPU VRAM and improves speed by 5% for 30 steps

@HuggingFaceDocBuilderDev
Copy link
Copy Markdown

HuggingFaceDocBuilderDev commented Jul 13, 2023 •
edited
Loading

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

@patrickvonplaten
Copy link
Copy Markdown
Contributor Author

patrickvonplaten commented Jul 13, 2023

For this PR to work we need to merge: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/discussions/3

@patrickvonplaten patrickvonplaten mentioned this pull request Jul 14, 2023
Merged
3 tasks
@sayakpaul
Copy link
Copy Markdown
Member

sayakpaul commented Jul 14, 2023

Contingent on: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/discussions/3#64b0725108f7eb472ef01fbf

sayakpaul
sayakpaul reviewed Jul 14, 2023
View reviewed changes
src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py
self.vae.to(dtype)
init_latents = init_latents.to(dtype)
if self.vae.config.force_upcast:
self.vae.to(dtype)
Copy link
Copy Markdown
Member

@sayakpaul sayakpaul Jul 14, 2023

Choose a reason for hiding this comment

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

If force_upcast is true, then shouldn't we use upcast_vae()?

Copy link
Copy Markdown
Contributor Author

@patrickvonplaten patrickvonplaten Jul 14, 2023

Choose a reason for hiding this comment

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

We have to possibly move it back here since for the decoder not all layers are upcasted (so we should move it back to fp16 here)

@sayakpaul
Copy link
Copy Markdown
Member

sayakpaul commented Jul 14, 2023 •
edited
Loading

Looks good to me.

Will all the compatible VAEs have force_upcast in their configs? Maybe a better check would be to first verify if that flag exists in the config no?

if force_upcast in self.vae.config:
...

Am I missing out on something? Should we also add a couple of tests here?

patrickvonplaten
patrickvonplaten commented Jul 14, 2023
View reviewed changes
src/diffusers/models/autoencoder_kl.py
norm_num_groups: int = 32,
sample_size: int = 32,
scaling_factor: float = 0.18215,
force_upcast: float = True,
Copy link
Copy Markdown
Contributor Author

@patrickvonplaten patrickvonplaten Jul 14, 2023

Choose a reason for hiding this comment

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

cc @sayakpaul every config has force_upcast

@patrickvonplaten
Copy link
Copy Markdown
Contributor Author

patrickvonplaten commented Jul 14, 2023

self.vae.to(dtype)
Member

It's there by default because it's in the constructor arguments: https://github.com/huggingface/diffusers/pull/4083/files#r1263476807

@patrickvonplaten
Copy link
Copy Markdown
Contributor Author

patrickvonplaten commented Jul 14, 2023

Merging since it's also needed for inpaint

@patrickvonplaten patrickvonplaten mentioned this pull request Jul 14, 2023
Closed
@patrickvonplaten
Copy link
Copy Markdown
Contributor Author

patrickvonplaten commented Jul 14, 2023

Looks good to me.

Will all the compatible VAEs have force_upcast in their configs? Maybe a better check would be to first verify if that flag exists in the config no?

if force_upcast in self.vae.config:
...

Am I missing out on something? Should we also add a couple of tests here?

Tests are a bit tricky here since we don't yet can do slow tests as checkpoint is behind a gate and dummy weights don't have the fp16 issue => I'll set a reminder issue: #4100

@patrickvonplaten patrickvonplaten merged commit ad8f985 into main Jul 14, 2023
@patrickvonplaten patrickvonplaten deleted the allow_low_precision_vae_sd_xl branch July 14, 2023 12:51
@AmericanPresidentJimmyCarter
Copy link
Copy Markdown
Contributor

AmericanPresidentJimmyCarter commented Jul 14, 2023

Related: #4102

Me and @Birch-san have been just using bfloat16 in our pipelines for the VAE and it works perfectly well without all the extra legwork that the upcast requires, and is faster.

@sayakpaul sayakpaul mentioned this pull request Jul 14, 2023
Merged
sayakpaul added a commit that referenced this pull request Jul 18, 2023
@sayakpaul
incorporate fix from #4083
695c3ac
sayakpaul added a commit that referenced this pull request Jul 18, 2023
@sayakpaul @patrickvonplaten
[Core] add: controlnet support for SDXL (#4038)
3eb498e 
* add: controlnet sdxl.

* modifications to controlnet.

* run styling.

* add: __init__.pys

* incorporate #4019 changes.

* run make fix-copies.

* resize the conditioning images.

* remove autocast.

* run styling.

* disable autocast.

* debugging

* device placement.

* back to autocast.

* remove comment.

* save some memory by reusing the vae and unet in the pipeline.

* apply styling.

* Allow low precision sd xl

* finish

* finish

* changes to accommodate the improved VAE.

* modifications to how we handle vae encoding in the training.

* make style

* make existing controlnet fast tests pass.

* change vae checkpoint cli arg.

* fix: vae pretrained paths.

* fix: steps in get_scheduler().

* debugging.

* debugging./

* fix: weight conversion.

* add: docs.

* add: limited tests./

* add: datasets to the requirements.

* update docstrings and incorporate the usage of watermarking.

* incorporate fix from #4083

* fix watermarking dependency handling.

* run make-fix-copies.

* Empty-Commit

* Update requirements_sdxl.txt

* remove vae upcasting part.

* Apply suggestions from code review

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

* run make style

* run make fix-copies.

* disable suppot for multicontrolnet.

* Apply suggestions from code review

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

* run make fix-copies.

* dtyle/.

* fix-copies.

---------

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>
@wilfrediscoming
Copy link
Copy Markdown

wilfrediscoming commented Jul 25, 2023

Does this work for other pipeline? How do I know which VAE will work?

pipe = StableDiffusionInpaintPipeline.from_pretrained(
    runwayml/stable-diffusion-inpainting,
    revision="fp16",
    torch_dtype=torch.float16,
)

@sayakpaul
Copy link
Copy Markdown
Member

sayakpaul commented Jul 25, 2023 •
edited
Loading

AFAIK The VAE is specifically suited for SDXL.

orpatashnik pushed a commit to orpatashnik/diffusers that referenced this pull request Aug 1, 2023
@patrickvonplaten @orpatashnik
Allow low precision vae sd xl (huggingface#4083)
562b7e1 
* Allow low precision sd xl

* finish

* finish

* make style
orpatashnik pushed a commit to orpatashnik/diffusers that referenced this pull request Aug 1, 2023
@sayakpaul @patrickvonplaten @orpatashnik
[Core] add: controlnet support for SDXL (huggingface#4038)
2206088 
* add: controlnet sdxl.

* modifications to controlnet.

* run styling.

* add: __init__.pys

* incorporate huggingface#4019 changes.

* run make fix-copies.

* resize the conditioning images.

* remove autocast.

* run styling.

* disable autocast.

* debugging

* device placement.

* back to autocast.

* remove comment.

* save some memory by reusing the vae and unet in the pipeline.

* apply styling.

* Allow low precision sd xl

* finish

* finish

* changes to accommodate the improved VAE.

* modifications to how we handle vae encoding in the training.

* make style

* make existing controlnet fast tests pass.

* change vae checkpoint cli arg.

* fix: vae pretrained paths.

* fix: steps in get_scheduler().

* debugging.

* debugging./

* fix: weight conversion.

* add: docs.

* add: limited tests./

* add: datasets to the requirements.

* update docstrings and incorporate the usage of watermarking.

* incorporate fix from huggingface#4083

* fix watermarking dependency handling.

* run make-fix-copies.

* Empty-Commit

* Update requirements_sdxl.txt

* remove vae upcasting part.

* Apply suggestions from code review

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

* run make style

* run make fix-copies.

* disable suppot for multicontrolnet.

* Apply suggestions from code review

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

* run make fix-copies.

* dtyle/.

* fix-copies.

---------

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>
orpatashnik pushed a commit to orpatashnik/diffusers that referenced this pull request Aug 1, 2023
@patrickvonplaten @orpatashnik
Allow low precision vae sd xl (huggingface#4083)
6683847 
* Allow low precision sd xl

* finish

* finish

* make style
orpatashnik pushed a commit to orpatashnik/diffusers that referenced this pull request Aug 1, 2023
@sayakpaul @patrickvonplaten @orpatashnik
[Core] add: controlnet support for SDXL (huggingface#4038)
0681072 
* add: controlnet sdxl.

* modifications to controlnet.

* run styling.

* add: __init__.pys

* incorporate huggingface#4019 changes.

* run make fix-copies.

* resize the conditioning images.

* remove autocast.

* run styling.

* disable autocast.

* debugging

* device placement.

* back to autocast.

* remove comment.

* save some memory by reusing the vae and unet in the pipeline.

* apply styling.

* Allow low precision sd xl

* finish

* finish

* changes to accommodate the improved VAE.

* modifications to how we handle vae encoding in the training.

* make style

* make existing controlnet fast tests pass.

* change vae checkpoint cli arg.

* fix: vae pretrained paths.

* fix: steps in get_scheduler().

* debugging.

* debugging./

* fix: weight conversion.

* add: docs.

* add: limited tests./

* add: datasets to the requirements.

* update docstrings and incorporate the usage of watermarking.

* incorporate fix from huggingface#4083

* fix watermarking dependency handling.

* run make-fix-copies.

* Empty-Commit

* Update requirements_sdxl.txt

* remove vae upcasting part.

* Apply suggestions from code review

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

* run make style

* run make fix-copies.

* disable suppot for multicontrolnet.

* Apply suggestions from code review

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

* run make fix-copies.

* dtyle/.

* fix-copies.

---------

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>
orpatashnik pushed a commit to orpatashnik/diffusers that referenced this pull request Aug 1, 2023
@patrickvonplaten @orpatashnik
Allow low precision vae sd xl (huggingface#4083)
4fd54e9 
* Allow low precision sd xl

* finish

* finish

* make style
orpatashnik pushed a commit to orpatashnik/diffusers that referenced this pull request Aug 1, 2023
@sayakpaul @patrickvonplaten @orpatashnik
[Core] add: controlnet support for SDXL (huggingface#4038)
7dc3f6b 
* add: controlnet sdxl.

* modifications to controlnet.

* run styling.

* add: __init__.pys

* incorporate huggingface#4019 changes.

* run make fix-copies.

* resize the conditioning images.

* remove autocast.

* run styling.

* disable autocast.

* debugging

* device placement.

* back to autocast.

* remove comment.

* save some memory by reusing the vae and unet in the pipeline.

* apply styling.

* Allow low precision sd xl

* finish

* finish

* changes to accommodate the improved VAE.

* modifications to how we handle vae encoding in the training.

* make style

* make existing controlnet fast tests pass.

* change vae checkpoint cli arg.

* fix: vae pretrained paths.

* fix: steps in get_scheduler().

* debugging.

* debugging./

* fix: weight conversion.

* add: docs.

* add: limited tests./

* add: datasets to the requirements.

* update docstrings and incorporate the usage of watermarking.

* incorporate fix from huggingface#4083

* fix watermarking dependency handling.

* run make-fix-copies.

* Empty-Commit

* Update requirements_sdxl.txt

* remove vae upcasting part.

* Apply suggestions from code review

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

* run make style

* run make fix-copies.

* disable suppot for multicontrolnet.

* Apply suggestions from code review

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

* run make fix-copies.

* dtyle/.

* fix-copies.

---------

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>
yoonseokjin pushed a commit to yoonseokjin/diffusers that referenced this pull request Dec 25, 2023
@patrickvonplaten
Allow low precision vae sd xl (huggingface#4083)
beb330b 
* Allow low precision sd xl

* finish

* finish

* make style
yoonseokjin pushed a commit to yoonseokjin/diffusers that referenced this pull request Dec 25, 2023
@sayakpaul @patrickvonplaten
[Core] add: controlnet support for SDXL (huggingface#4038)
62ed29a 
* add: controlnet sdxl.

* modifications to controlnet.

* run styling.

* add: __init__.pys

* incorporate huggingface#4019 changes.

* run make fix-copies.

* resize the conditioning images.

* remove autocast.

* run styling.

* disable autocast.

* debugging

* device placement.

* back to autocast.

* remove comment.

* save some memory by reusing the vae and unet in the pipeline.

* apply styling.

* Allow low precision sd xl

* finish

* finish

* changes to accommodate the improved VAE.

* modifications to how we handle vae encoding in the training.

* make style

* make existing controlnet fast tests pass.

* change vae checkpoint cli arg.

* fix: vae pretrained paths.

* fix: steps in get_scheduler().

* debugging.

* debugging./

* fix: weight conversion.

* add: docs.

* add: limited tests./

* add: datasets to the requirements.

* update docstrings and incorporate the usage of watermarking.

* incorporate fix from huggingface#4083

* fix watermarking dependency handling.

* run make-fix-copies.

* Empty-Commit

* Update requirements_sdxl.txt

* remove vae upcasting part.

* Apply suggestions from code review

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

* run make style

* run make fix-copies.

* disable suppot for multicontrolnet.

* Apply suggestions from code review

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

* run make fix-copies.

* dtyle/.

* fix-copies.

---------

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>
AmericanPresidentJimmyCarter pushed a commit to AmericanPresidentJimmyCarter/diffusers that referenced this pull request Apr 26, 2024
@patrickvonplaten
Allow low precision vae sd xl (huggingface#4083)
805b0ab 
* Allow low precision sd xl

* finish

* finish

* make style
AmericanPresidentJimmyCarter pushed a commit to AmericanPresidentJimmyCarter/diffusers that referenced this pull request Apr 26, 2024
@sayakpaul @patrickvonplaten
[Core] add: controlnet support for SDXL (huggingface#4038)
6054868 
* add: controlnet sdxl.

* modifications to controlnet.

* run styling.

* add: __init__.pys

* incorporate huggingface#4019 changes.

* run make fix-copies.

* resize the conditioning images.

* remove autocast.

* run styling.

* disable autocast.

* debugging

* device placement.

* back to autocast.

* remove comment.

* save some memory by reusing the vae and unet in the pipeline.

* apply styling.

* Allow low precision sd xl

* finish

* finish

* changes to accommodate the improved VAE.

* modifications to how we handle vae encoding in the training.

* make style

* make existing controlnet fast tests pass.

* change vae checkpoint cli arg.

* fix: vae pretrained paths.

* fix: steps in get_scheduler().

* debugging.

* debugging./

* fix: weight conversion.

* add: docs.

* add: limited tests./

* add: datasets to the requirements.

* update docstrings and incorporate the usage of watermarking.

* incorporate fix from huggingface#4083

* fix watermarking dependency handling.

* run make-fix-copies.

* Empty-Commit

* Update requirements_sdxl.txt

* remove vae upcasting part.

* Apply suggestions from code review

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

* run make style

* run make fix-copies.

* disable suppot for multicontrolnet.

* Apply suggestions from code review

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

* run make fix-copies.

* dtyle/.

* fix-copies.

---------

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@sayakpaul sayakpaul sayakpaul left review comments

@williamberman williamberman Awaiting requested review from williamberman

@yiyixuxu yiyixuxu Awaiting requested review from yiyixuxu

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@patrickvonplaten @HuggingFaceDocBuilderDev @sayakpaul @AmericanPresidentJimmyCarter @wilfrediscoming