[OOM] Memory blows out when trying to upscale images larger than 128x128 using StableDiffusionUpscalePipeline

Describe the bug

When trying to upscale images larger than 128x128 the progress goes to 100% and then crashes with CUDA OOM.

With 512x512 images it’s trying to allocate 256.00 GiB!

Reproduction

import requests
from PIL import Image
from io import BytesIO
from diffusers import StableDiffusionUpscalePipeline
import torch

model_id = "stabilityai/stable-diffusion-x4-upscaler"
pipeline = StableDiffusionUpscalePipeline.from_pretrained(model_id, revision="fp16", torch_dtype=torch.float16)
pipeline = pipeline.to("cuda")

url = "https://www.freepnglogos.com/uploads/512x512-logo/512x512-transparent-circle-instagram-media-network-social-logo-new-16.png"
response = requests.get(url)
low_res_img = Image.open(BytesIO(response.content)).convert("RGB")
prompt=""
upscaled_image = pipeline(prompt=prompt, image=low_res_img).images[0]
display(upscaled_image)

Logs

RuntimeError: CUDA out of memory. Tried to allocate 256.00 GiB (GPU 0; 14.76 GiB total capacity; 4.77 GiB already allocated; 8.28 GiB free; 5.18 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

System Info

• diffusers version: 0.9.0
• Platform: Linux-5.10.133±x86_64-with-Ubuntu-18.04-bionic
• Python version: 3.7.15
• PyTorch version (GPU?): 1.12.1+cu113 (True)
• Huggingface_hub version: 0.11.0
• Transformers version: 4.24.0

[blefaudeux]

you can upscale a 512x image with a ~20GB GPU (I didn´t try with less), with the linked PR & using xformers in the attentions in the VAE (when properly picked up by the enablement, hence another PR). I’ve this running just fine on a private fork, it looks like all the missing pieces are arriving here (see this PR) else I can PR the required missing bits

[kashif]

I tried using it with xformers, i believe, and I think I got the same issue… i can re-run it… But the issue occurs in the creating of this empty tensor in the default attention block:

https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention.py#L331-L337