Make a Hugging Face "optimized" pipeline example

[exo-pla-net]

Currently, using Stable Diffusion with huggingface/diffusers often fails due to VRAM shortage, even when using a 16GB GPU, and even when trying to make fairly humble generations, like an image with a 16:9 aspect ratio (which is important for making videos).

Basujinal's fork uses drastically less VRAM, allowing larger images to be generated.

However, his code is nigh impenetrable.

Thus, a developer either has to choose between two flawed options:

• Opting for easy development using huggingface's philosophy of clean and simple code, but being plagued by CUDA out-of-memory failures.
• Opting for being free of VRAM worries, but be reduced to glacially slow development, due to having to rely on buggy and esoteric code.

This sucks! If someone were to replicate basujinal's "optimized" SD modifications, using Hugging Face's Contribution philosophy, we could have the best of both worlds: clean and easy development and no VRAM worries.

I'd put my money on VRAM shortages being the primary reason why devs choose not to use huggingface/diffusers. This is tragic, since I think clean code led by Hugging Face is the future. We need a Hugging Face "optimized" pipeline.

[Inkorak]

I was able to upgrade the pipeline quite easily to implement efficient attention from this repository.

from diffusers.models.attention import CrossAttention
import math
import torch
from torch import einsum
from einops import rearrange
import types

def forward(self, x, context=None, mask=None):
    batch_size, sequence_length, dim = x.shape

    h = self.heads

    q = self.to_q(x)
    context = context if context is not None else x
    k = self.to_k(context)
    v = self.to_v(context)
    del context, x

    q = self.reshape_heads_to_batch_dim(q)
    k = self.reshape_heads_to_batch_dim(k)
    v = self.reshape_heads_to_batch_dim(v)

    r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device)

    stats = torch.cuda.memory_stats(q.device)
    mem_total = torch.cuda.get_device_properties(0).total_memory
    mem_active = stats['active_bytes.all.current']
    mem_free = mem_total - mem_active

    mem_required = q.shape[0] * q.shape[1] * k.shape[1] * 4 * 2.5
    steps = 1

    if mem_required > mem_free:
        steps = 2**(math.ceil(math.log(mem_required / mem_free, 2)))

    slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]
    for i in range(0, q.shape[1], slice_size):
        end = i + slice_size
        s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k)
        s1 *= self.scale

        s2 = s1.softmax(dim=-1)
        del s1

        r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v)
        del s2

    del q, k, v

    r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h)
    del r1

    return self.to_out(r2)

def optimize_attention(model):
    for module in model.modules():
        if isinstance(module, CrossAttention):
            module.forward = types.MethodType(forward, module)

from diffusers import StableDiffusionPipeline
from torch import autocast

pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=True)
pipe = pipe.to("cuda")

optimize_attention(pipe.unet)

prompt = "a photo of an astronaut riding a horse on mars"
with autocast("cuda"):
    image = pipe(prompt).images[0]  

[sradc]

Thanks @Inkorak. Looking at watch -n1 nvidia-smi, this looked like it saved about ~500mb, for a 512x512 image, on an 8gb card (though not sure how accurate that estimate is).

To avoid OOM, I made the change:

 pipe = StableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4", 
    use_auth_token=True, 
    revision="fp16", 
    torch_dtype=torch.float16)

And to get the image, with diffusers '0.2.4', :

with autocast("cuda"):
    image = pipe(prompt)["sample"][0]

[anton-l]

Update: the VRAM optimizations are planned for this week's release, feel free to comment on this PR: #366 🤗

[Inkorak]

@sradc This implementation looks at how much VRAM is free and adapts accordingly. In terms of memory consumption, it becomes more and more efficient if there is not enough VRAM, at the expense of speed. Try increasing the resolution beyond what your graphics card can handle with a regular pipeline.

[patrickvonplaten]

Here we have a very optimized stable diffusion example:

pipe = StableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4", 
    revision="fp16", 
    torch_dtype=torch.float16,
    use_auth_token=True
)
pipe = pipe.to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
pipe.enable_attention_slicing()
with autocast("cuda"):
    image = pipe(prompt).images[0] 

=> 3.3 GB VRAM only

You can now find this example on the main readme: https://github.com/huggingface/diffusers/blob/main/README.md#text-to-image-generation-with-stable-diffusion

[exo-pla-net]

This problem is largely solved, so I'm closing this issue.

However, Inkorak's dynamic VRAM+attention adjustment is interesting and would be useful. If I'm understanding it correctly, it would allow us to make arbitrarily large images or batches. If we decide we want to implement such a feature, we should open a new issue.

There's some more discussion here.