From f07c5ec5dc60b15a5342fad21a0cb3b5d7998bc0 Mon Sep 17 00:00:00 2001 From: Dango233 Date: Wed, 7 Dec 2022 18:01:50 +0800 Subject: [PATCH] * Use Cutlass ops when possible to +15% speed --- ldm/modules/attention.py | 2 +- ldm/modules/diffusionmodules/model.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/ldm/modules/attention.py b/ldm/modules/attention.py index d504d93..574951c 100644 --- a/ldm/modules/attention.py +++ b/ldm/modules/attention.py @@ -201,7 +201,7 @@ class MemoryEfficientCrossAttention(nn.Module): self.to_v = nn.Linear(context_dim, inner_dim, bias=False) self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout)) - self.attention_op: Optional[Any] = None + self.attention_op: Optional[Any] = xformers.ops.MemoryEfficientAttentionCutlassOp if hasattr(xformers.ops, "MemoryEfficientAttentionCutlassOp") else None def forward(self, x, context=None, mask=None): q = self.to_q(x) diff --git a/ldm/modules/diffusionmodules/model.py b/ldm/modules/diffusionmodules/model.py index b089eeb..af5bdb3 100644 --- a/ldm/modules/diffusionmodules/model.py +++ b/ldm/modules/diffusionmodules/model.py @@ -234,7 +234,7 @@ class MemoryEfficientAttnBlock(nn.Module): kernel_size=1, stride=1, padding=0) - self.attention_op: Optional[Any] = None + self.attention_op: Optional[Any] = xformers.ops.MemoryEfficientAttentionCutlassOp if hasattr(xformers.ops, "MemoryEfficientAttentionCutlassOp") else None def forward(self, x): h_ = x