From 51e73fc05ec75f1c4d3f1cdbf1201a1b6dc99cac Mon Sep 17 00:00:00 2001 From: Mihail Dumitrescu Date: Sun, 11 Sep 2022 01:52:17 +0300 Subject: [PATCH] Keep using float16 in ldm.modules.attention; ~12% speed up and less max VRAM Tested on nvidia eGPU setup so YMMV with the default half precision math. Speed from 1.69it/s to 1.89it/s and max VRAM from 4.44G to 3.37G for generating 512x512 images. Measured after applying a separate PR #484 Move model.half() before model.to(device) --- ldm/modules/attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ldm/modules/attention.py b/ldm/modules/attention.py index 24aef292793..37b2c6b5149 100644 --- a/ldm/modules/attention.py +++ b/ldm/modules/attention.py @@ -193,7 +193,7 @@ def forward(self, x, context=None, mask=None): sim.masked_fill_(~mask, max_neg_value) # attention, what we cannot get enough of - attn = sim.softmax(dim=-1) + attn = sim.softmax(dim=-1, dtype=sim.dtype) out = einsum('b i j, b j d -> b i d', attn, v) out = rearrange(out, '(b h) n d -> b n (h d)', h=h)