Fix for Unet NaNs

2023-04-14 02:22:48 -04:00 · 2023-04-14 02:22:48 -04:00 · 7aab389d6f
commit 7aab389d6f
parent 5ab7f213be
1 changed files with 3 additions and 0 deletions
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@ -256,6 +256,9 @@ def sub_quad_attention_forward(self, x, context=None, mask=None):
    k = k.unflatten(-1, (h, -1)).transpose(1,2).flatten(end_dim=1)
    v = v.unflatten(-1, (h, -1)).transpose(1,2).flatten(end_dim=1)

+    if q.device.type == 'mps':
+        q, k, v = q.contiguous(), k.contiguous(), v.contiguous()
+
    dtype = q.dtype
    if shared.opts.upcast_attn:
        q, k = q.float(), k.float()