AutoencoderKL encoder outputs NaN for large images

[antoine-scenario]

Describe the bug

AutoEncoderKL encoder loaded from runwayml/stable-diffusion-v1-5 outputs NaN for large images. I observe this behavior for image sizes starting from around 1500x1500 with vae_tiling disabled. I tried with float32, float16, with and without xFormers. Is it an expected behavior ?

I would have liked to use vae_tiling but it produces tiles artefacts as reported in #1441.

Reproduction

from PIL import Image
import torch
import urllib.request
import os

from diffusers import StableDiffusionImg2ImgPipeline
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import preprocess


image_file = urllib.request.urlopen("https://upload.wikimedia.org/wikipedia/commons/3/32/A_photograph_of_an_astronaut_riding_a_horse_2022-08-28.png")
init_image = Image.open(image_file)

up_size = (2048, 2048)
upsampled_image = init_image.resize(up_size, Image.Resampling.BILINEAR)
pipe = StableDiffusionImg2ImgPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float32)
pipe.enable_xformers_memory_efficient_attention()
device = "cuda"

with torch.no_grad():
    vae = pipe.vae.to(device)
    vae.disable_tiling()
    preprocessed = preprocess(upsampled_image).to(torch.float32).to(device)
    latents = vae.encode(preprocessed).latent_dist.sample()
    print(latents)

Logs

tensor([[[[nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          ...,
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan]],

         [[nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          ...,
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan]],

         [[nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          ...,
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan]],

         [[nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          ...,
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan]]]], device='cuda:0')

System Info

• diffusers version: 0.14.0
• Platform: Linux-5.15.0-1028-aws-x86_64-with-glibc2.35
• Python version: 3.10.6
• PyTorch version (GPU?): 2.0.0+cu117 (True)
• Huggingface_hub version: 0.13.3
• Transformers version: 4.27.4
• Accelerate version: 0.15.0
• xFormers version: 0.0.18
• Using GPU in script?:
• Using distributed or parallel set-up in script?:

[patrickvonplaten]

This is likely related to: facebookresearch/xformers#719

[patrickvonplaten]

Not that after 425192f you should be able to just use torch 2.0 and it should work

[antoine-scenario]

Thanks @patrickvonplaten !
Unfortunately, the issue persists when I run the script after having set diffusers to 425192f

[patrickvonplaten]

Hey @antoine-scenario,

I cannot reproduce the issue. When running your code above, I get:

tensor([[[[ 3.6682,  4.5076,  5.6030,  ...,  1.8190,  4.3212,  0.5521],
          [ 7.0356,  5.2478,  2.6387,  ...,  5.9366,  5.4313,  6.6636],
          [ 4.4871,  4.4803,  7.0959,  ...,  3.4377,  3.7891,  1.4332],
          ...,
          [ 5.5094,  6.2821,  8.7655,  ...,  4.3388,  6.7097,  5.6651],
          [ 5.1794,  4.6551,  8.2015,  ...,  6.7687,  6.4646,  2.5777],
          [ 2.5079,  3.8722,  2.6401,  ...,  2.3823,  3.1924,  6.6117]],

         [[ 2.8880,  3.2955,  0.7796,  ..., -0.1046, -0.6432,  1.8139],
          [-0.0606,  0.2110, -0.2602,  ..., -1.6427, -2.5984,  0.3116],
          [ 1.3094,  1.1901, -0.1676,  ..., -0.9184, -2.4686,  0.2785],
          ...,
          [-0.9043, -0.9740,  1.6073,  ..., -0.2104,  2.2977,  0.7033],
          [-5.3002, -0.0696, -1.5857,  ...,  1.2968, -0.6559,  0.6579],
          [-1.0852,  0.9874,  1.1869,  ...,  1.9165, -2.8102,  1.7200]],

         [[ 2.6312,  4.7043,  5.8275,  ...,  6.8181,  9.0096,  7.6515],
          [ 4.9425, 10.7664,  7.3586,  ...,  8.3967,  9.2255,  8.6791],
          [ 6.4644,  6.5642,  6.7298,  ...,  9.8467,  9.7069,  8.5922],
          ...,
          [-3.1227, -2.9758, -0.5441,  ..., -2.6374, -3.4337, -1.2992],
          [-3.3511, -2.4731, -2.5276,  ..., -2.7271, -1.3782, -1.5970],
          [-2.0939, -0.9700, -0.7706,  ..., -1.6660, -2.3485, -0.9619]],

         [[-3.1845, -0.5830,  0.7922,  ..., -1.2312, -0.8559, -2.7735],
          [-1.2046,  0.1267,  0.3111,  ..., -0.0862,  2.2621, -0.1901],
          [-0.4298, -0.2186, -0.3073,  ...,  1.3177,  1.2746, -2.1658],
          ...,
          [ 7.0325,  6.8949,  8.4607,  ...,  6.7027,  6.1686,  5.1244],
          [ 8.5836,  8.9068,  7.3452,  ...,  5.9527,  6.5639,  6.5665],
          [ 5.3573,  7.1136,  6.8975,  ...,  6.4285,  4.7341,  5.9127]]]],
       device='cuda:0')

Here are my lib versions:

- `diffusers` version: 0.16.0
- Platform: Linux-5.15.0-69-generic-x86_64-with-glibc2.35
- Python version: 3.10.6
- PyTorch version (GPU?): 2.0.0+cu117 (True)
- Huggingface_hub version: 0.13.3
- Transformers version: 4.28.0.dev0
- Accelerate version: 0.18.0
- xFormers version: 0.0.19.dev516
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

[patrickvonplaten]

Ah you probably have to upgrade xformers to 0.0.19.dev516

[antoine-scenario]

Upgrading diffusers to 0.16.0 and xformers to 0.0.19.dev516 solved the issue. Thanks @patrickvonplaten !

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