Blank black image output with stable diffusion 2.1 using autocast

[fralumz]

Describe the bug

Using stable diffusion pipeline with torch.autocast and the stabilityai/stable-diffusion-2-1 model, the images generate are all blank black images.

Reproduction

import torch
from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler

model_id = "stabilityai/stable-diffusion-2-1"

scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler, revision="fp16", torch_dtype=torch.float16)
pipe = pipe.to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"

image = pipe(prompt, height=768, width=768).images[0]

image.save("astronaut_rides_horse.png") # works fine

with torch.autocast("cuda"):
    image = pipe(prompt, height=768, width=768).images[0] # generates blank image
    image.save("astronaut_rides_horse_autocast.png") 

Logs

Python 3.10.8 | packaged by conda-forge | (main, Nov 24 2022, 14:07:00) [MSC v.1916 64 bit (AMD64)] on win32
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
>>>
>>> model_id = "stabilityai/stable-diffusion-2-1"
>>>
>>> scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
>>> pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler, revision="fp16", torch_dtype=torch.float16)
Fetching 12 files: 100%|█████████████████████████████████████████████████████████████| 12/12 [00:00<00:00, 6015.50it/s]
>>> pipe = pipe.to("cuda")
>>>
>>> prompt = "a photo of an astronaut riding a horse on mars"
>>>
>>> image = pipe(prompt, height=768, width=768).images[0]
100%|██████████████████████████████████████████████████████████████████████████████████| 50/50 [00:14<00:00,  3.57it/s]
>>>
>>> image.save("astronaut_rides_horse.png") # works fine
>>>
>>> with torch.autocast("cuda"):
...     image = pipe(prompt, height=768, width=768).images[0] # generates blank image
...     image.save("astronaut_rides_horse_autocast.png")
...
100%|██████████████████████████████████████████████████████████████████████████████████| 50/50 [00:12<00:00,  3.95it/s]

System Info

• diffusers version: 0.10.0.dev0
• Platform: Windows-10-10.0.19044-SP0
• Python version: 3.10.8
• PyTorch version (GPU?): 1.13.0 (True)
• Huggingface_hub version: 0.11.1
• Transformers version: 4.25.1
• Using GPU in script?: yes
• Using distributed or parallel set-up in script?: no

[patrickvonplaten]

Hey @fralumz,

It is not recommended to use autocast - it is both slower then pure fp16 and more instable/controllabel

[HughPH]

After removing with torch.autocast("cuda"): I get images again with 0.10.1 but it's back to black output with 0.10.2 or above, including 0.12.0dev0

[patrickvonplaten]

Hey @HughPH,

Could you please post a reproducible code snippet including your current diffusers, pytorch, etc.. versions (you can get them with diffusers-cli env)

[HughPH]

Yes, when I finish work I'll duplicate my environment, get a minimal script together, and then make sure it breaks with 0.10.2

[patrickvonplaten]

Since a couple of people seem to be following this issue, I'd like to clarify a couple of things regarding the 2.1 model.

The 2.1 model was trained with xformers flash attention. Upon release, it was noticed that without using xformers, the model produces black images for torch.autocast and pure fp16 as written by the authors here: Stability-AI/stablediffusion@c12d960
The reason is that when using xformers flash attention as explained here: https://deepai.org/publication/flashattention-fast-and-memory-efficient-exact-attention-with-io-awareness
values are much less prone to overflow since there are less additions.
=> So we strongly recommend to always use xformers if possible when using SD 2.1. (It's a pity that there are no clean xformers wheels yet, but this should change very soon with PyTorch 2.0).

If you cannot use xformers, then there is a chance that you run into precision problems in half precision.
To fix this, we found that upcasting the query and key states just before the attention allows the model to be run
in inference in fp16. See PR here: #1590

Note: This fixes inference only for pure fp16, not for autocast. Autocast is much harder to control as it will automatically downcast weights to fp16 before the softmax computation which will then lead to black images.
Also, 2.1 cannot be fine-tuned in mixed-precision because of this the layer is prone to overflow.

Conclusion:

• Please do not use torch.autocast(...) with `diffusers SD 2.1 (actually don't use it at all, we only see disadvantages compared to pure fp16)
• Please use xformers if you can.

Also trying to make this a bit clearer in the docs: #2021

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