Skip to content

Commit 889e572

Browse files
RyanJDickpsychedelicious
authored andcommitted
Add a workaround for broken sliced attention on MPS with torch 2.4.1.
1 parent cc83ded commit 889e572

File tree

1 file changed

+9
-0
lines changed

1 file changed

+9
-0
lines changed

invokeai/backend/stable_diffusion/diffusers_pipeline.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -204,6 +204,15 @@ def _adjust_memory_efficient_attention(self, latents: torch.Tensor):
204204
else:
205205
raise Exception("torch-sdp attention slicing not available")
206206

207+
# See https://github.com/invoke-ai/InvokeAI/issues/7049 for context.
208+
# Bumping torch from 2.2.2 to 2.4.1 caused the sliced attention implementation to produce incorrect results.
209+
# For now, if a user is on an MPS device and has not explicitly set the attention_type, then we select the
210+
# non-sliced torch-sdp implementation. This keeps things working on MPS at the cost of increased peak memory
211+
# utilization.
212+
if torch.backends.mps.is_available():
213+
assert hasattr(torch.nn.functional, "scaled_dot_product_attention")
214+
return
215+
207216
# the remainder if this code is called when attention_type=='auto'
208217
if self.unet.device.type == "cuda":
209218
if is_xformers_available() and prefer_xformers:

0 commit comments

Comments
 (0)