Allow low precision vae sd xl #4083

patrickvonplaten · 2023-07-13T20:53:51Z

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 · 2023-07-13T21:01:41Z

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

patrickvonplaten · 2023-07-13T21:32:41Z

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

sayakpaul · 2023-07-14T03:51:39Z

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

sayakpaul · 2023-07-14T03:55:04Z

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)


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

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 · 2023-07-14T03:59:24Z

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 · 2023-07-14T08:49:30Z

src/diffusers/models/autoencoder_kl.py

@@ -82,6 +86,7 @@ def __init__(
        norm_num_groups: int = 32,
        sample_size: int = 32,
        scaling_factor: float = 0.18215,
+        force_upcast: float = True,


cc @sayakpaul every config has force_upcast

patrickvonplaten · 2023-07-14T08:49:39Z

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 · 2023-07-14T12:49:02Z

Merging since it's also needed for inpaint

patrickvonplaten · 2023-07-14T12:51:12Z

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

AmericanPresidentJimmyCarter · 2023-07-14T14:17:34Z

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.

* 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 <[email protected]> * run make style * run make fix-copies. * disable suppot for multicontrolnet. * Apply suggestions from code review Co-authored-by: Patrick von Platen <[email protected]> * run make fix-copies. * dtyle/. * fix-copies. --------- Co-authored-by: Patrick von Platen <[email protected]>

wilfrediscoming · 2023-07-25T11:32:48Z

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 · 2023-07-25T11:45:55Z

AFAIK The VAE is specifically suited for SDXL.

* Allow low precision sd xl * finish * finish * make style

* 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 <[email protected]> * run make style * run make fix-copies. * disable suppot for multicontrolnet. * Apply suggestions from code review Co-authored-by: Patrick von Platen <[email protected]> * run make fix-copies. * dtyle/. * fix-copies. --------- Co-authored-by: Patrick von Platen <[email protected]>

* Allow low precision sd xl * finish * finish * make style

* 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 <[email protected]> * run make style * run make fix-copies. * disable suppot for multicontrolnet. * Apply suggestions from code review Co-authored-by: Patrick von Platen <[email protected]> * run make fix-copies. * dtyle/. * fix-copies. --------- Co-authored-by: Patrick von Platen <[email protected]>

* Allow low precision sd xl * finish * finish * make style

* 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 <[email protected]> * run make style * run make fix-copies. * disable suppot for multicontrolnet. * Apply suggestions from code review Co-authored-by: Patrick von Platen <[email protected]> * run make fix-copies. * dtyle/. * fix-copies. --------- Co-authored-by: Patrick von Platen <[email protected]>

* Allow low precision sd xl * finish * finish * make style

* 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 <[email protected]> * run make style * run make fix-copies. * disable suppot for multicontrolnet. * Apply suggestions from code review Co-authored-by: Patrick von Platen <[email protected]> * run make fix-copies. * dtyle/. * fix-copies. --------- Co-authored-by: Patrick von Platen <[email protected]>

* Allow low precision sd xl * finish * finish * make style

* 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 <[email protected]> * run make style * run make fix-copies. * disable suppot for multicontrolnet. * Apply suggestions from code review Co-authored-by: Patrick von Platen <[email protected]> * run make fix-copies. * dtyle/. * fix-copies. --------- Co-authored-by: Patrick von Platen <[email protected]>

patrickvonplaten added 2 commits July 13, 2023 21:53

Allow low precision sd xl

65a5c45

finish

43f842c

patrickvonplaten added 2 commits July 13, 2023 21:24

finish

7171d42

Merge branch 'main' into allow_low_precision_vae_sd_xl

97f69a7

patrickvonplaten requested review from williamberman, yiyixuxu and sayakpaul July 13, 2023 21:30

patrickvonplaten mentioned this pull request Jul 14, 2023

[Core] add: controlnet support for SDXL #4038

Merged

3 tasks

sayakpaul reviewed Jul 14, 2023

View reviewed changes

make style

8ed2ebb

patrickvonplaten commented Jul 14, 2023

View reviewed changes

patrickvonplaten mentioned this pull request Jul 14, 2023

[VAE upscast] Reminder issue #4100

Closed

patrickvonplaten merged commit ad8f985 into main Jul 14, 2023

patrickvonplaten deleted the allow_low_precision_vae_sd_xl branch July 14, 2023 12:51

sayakpaul mentioned this pull request Jul 14, 2023

[SDXL DreamBooth LoRA] add support for text encoder fine-tuning #4097

Merged

sayakpaul added a commit that referenced this pull request Jul 18, 2023

incorporate fix from #4083

695c3ac

orpatashnik pushed a commit to orpatashnik/diffusers that referenced this pull request Aug 1, 2023

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

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

Allow low precision vae sd xl (huggingface#4083)

4fd54e9

* Allow low precision sd xl * finish * finish * make style

yoonseokjin pushed a commit to yoonseokjin/diffusers that referenced this pull request Dec 25, 2023

Allow low precision vae sd xl (huggingface#4083)

beb330b

* Allow low precision sd xl * finish * finish * make style

AmericanPresidentJimmyCarter pushed a commit to AmericanPresidentJimmyCarter/diffusers that referenced this pull request Apr 26, 2024

Allow low precision vae sd xl (huggingface#4083)

805b0ab

* Allow low precision sd xl * finish * finish * make style

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Allow low precision vae sd xl #4083

Allow low precision vae sd xl #4083

Uh oh!

patrickvonplaten commented Jul 13, 2023 •

edited

Loading

Uh oh!

HuggingFaceDocBuilderDev commented Jul 13, 2023 •

edited

Loading

Uh oh!

patrickvonplaten commented Jul 13, 2023

Uh oh!

sayakpaul commented Jul 14, 2023

Uh oh!

sayakpaul Jul 14, 2023

Uh oh!

patrickvonplaten Jul 14, 2023

Uh oh!

sayakpaul commented Jul 14, 2023 •

edited

Loading

Uh oh!

patrickvonplaten Jul 14, 2023

Uh oh!

patrickvonplaten commented Jul 14, 2023

Uh oh!

patrickvonplaten commented Jul 14, 2023

Uh oh!

patrickvonplaten commented Jul 14, 2023

Uh oh!

AmericanPresidentJimmyCarter commented Jul 14, 2023

Uh oh!

wilfrediscoming commented Jul 25, 2023

Uh oh!

sayakpaul commented Jul 25, 2023 •

edited

Loading

Uh oh!

Uh oh!

Allow low precision vae sd xl #4083

Allow low precision vae sd xl #4083

Uh oh!

Conversation

patrickvonplaten commented Jul 13, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Uh oh!

HuggingFaceDocBuilderDev commented Jul 13, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

patrickvonplaten commented Jul 13, 2023

Uh oh!

sayakpaul commented Jul 14, 2023

Uh oh!

sayakpaul Jul 14, 2023

Choose a reason for hiding this comment

Uh oh!

patrickvonplaten Jul 14, 2023

Choose a reason for hiding this comment

Uh oh!

sayakpaul commented Jul 14, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

patrickvonplaten Jul 14, 2023

Choose a reason for hiding this comment

Uh oh!

patrickvonplaten commented Jul 14, 2023

Uh oh!

patrickvonplaten commented Jul 14, 2023

Uh oh!

patrickvonplaten commented Jul 14, 2023

Uh oh!

AmericanPresidentJimmyCarter commented Jul 14, 2023

Uh oh!

wilfrediscoming commented Jul 25, 2023

Uh oh!

sayakpaul commented Jul 25, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

patrickvonplaten commented Jul 13, 2023 •

edited

Loading

HuggingFaceDocBuilderDev commented Jul 13, 2023 •

edited

Loading

sayakpaul commented Jul 14, 2023 •

edited

Loading

sayakpaul commented Jul 25, 2023 •

edited

Loading