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

I know this is closed and these things are in docs, but just wanted to say that if you’re running into this issue to install the following:

pip install xformers
pip install triton==2.0.0.dev20221120

And to add this to your pipeline:

import torch
from diffusers import DiffusionPipeline
from xformers.ops import MemoryEfficientAttentionFlashAttentionOp

pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16)
pipe = pipe.to("cuda")
pipe.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp)
# Workaround for not accepting attention shape using VAE for Flash Attention
pipe.vae.enable_xformers_memory_efficient_attention(attention_op=None)

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

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

I implemented probably the most simplistic form of tiling possible here: https://github.com/carson-katri/dream-textures/blob/aa0132b42dd14ddbf9491c13a7a46a01da2c880a/generator_process/actions/upscale.py

I’m sure there are much better approaches that would limit seams. Perhaps just tiling the latent decoding process? Not entirely sure. Looking forward to seeing the improvements that will be made in this pipeline!

I’m working on a tile-based solution that runs the upscale model on small, overlapping patches of a larger source image and then merges them back into the full sized result. Much of the code is borrowed from realesrgan upscaler which supports this. Will try and publish code as soon as it’s working

@carson-katri correct you can try

pipeline.enable_attention_slicing()

and that should reduce some memory in exchange for a small speed decrease and enable larger inputs. With xformers installed it should be less as you point out!

The model may support xformers and attention slicing, which could help I assume.

Thank you for the explanation, I thought this may be the case. Resizing the image to 128x128 would produce a 512x512 image, correct?

Thank you for reporting this. The reason this happens is that your initial image gets bigger e.g. 512x512 the latent representations end up being 512 (latent dim) x 512 (H) x 512 (W) in the decoding attention block which gets reshaped into batch = 1, seq = 512 x 512 (H xW) and channel = 512 and that obviously will not work as vanilla attention is quadratic mem/compute in the seq length. Thus as you increase your initial image dims, the more mem. it will use.

A solution for now in the above code is to downsize the initial image to something manageable e.g.:

low_res_img = low_res_img.resize((128, 128))

As noted below you can also install xformers or try with attention slicing:

pipeline.enable_attention_slicing()
# pipeline.enable_xformers_memory_efficient_attention()