eessppeelllloo/stablediffusion
1
0
Fork 0
mirror of https://github.com/Stability-AI/stablediffusion.git synced 2025-12-05 14:30:01 +01:00
Code Issues Packages Projects Releases Wiki Activity

* Use Cutlass ops when possible to +15% speed

Browse Source
This commit is contained in:
Dango233
2022-12-07 18:01:50 +08:00
parent d7440ac160
commit f07c5ec5dc
2 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
ldm/modules/attention.py
View File

@@ -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)

2
ldm/modules/diffusionmodules/model.py
View File

@@ -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