Fix MPS on PyTorch 2.0.1, Intel Macs
diff --git a/modules/mac_specific.py b/modules/mac_specific.py
index 6fe8dea0726..40ce2101764 100644
--- a/modules/mac_specific.py
+++ b/modules/mac_specific.py
@@ -54,6 +54,11 @@ def cumsum_fix(input, cumsum_func, *args, **kwargs):
         CondFunc('torch.cumsum', cumsum_fix_func, None)
         CondFunc('torch.Tensor.cumsum', cumsum_fix_func, None)
         CondFunc('torch.narrow', lambda orig_func, *args, **kwargs: orig_func(*args, **kwargs).clone(), None)
-    if version.parse(torch.__version__) == version.parse("2.0"):
+
         # MPS workaround for https://github.com/pytorch/pytorch/issues/96113
-        CondFunc('torch.nn.functional.layer_norm', lambda orig_func, x, normalized_shape, weight, bias, eps, **kwargs: orig_func(x.float(), normalized_shape, weight.float() if weight is not None else None, bias.float() if bias is not None else bias, eps).to(x.dtype), lambda *args, **kwargs: len(args) == 6)
+        CondFunc('torch.nn.functional.layer_norm', lambda orig_func, x, normalized_shape, weight, bias, eps, **kwargs: orig_func(x.float(), normalized_shape, weight.float() if weight is not None else None, bias.float() if bias is not None else bias, eps).to(x.dtype), lambda _, input, *args, **kwargs: len(args) == 4 and input.device.type == 'mps')
+
+        # MPS workaround for https://github.com/pytorch/pytorch/issues/92311
+        if platform.processor() == 'i386':
+            for funcName in ['torch.argmax', 'torch.Tensor.argmax']:
+                CondFunc(funcName, lambda _, input, *args, **kwargs: torch.max(input.float() if input.dtype == torch.int64 else input, *args, **kwargs)[1], lambda _, input, *args, **kwargs: input.device.type == 'mps')
\ No newline at end of file
diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 372555ffaf4..f10865cd1e7 100644
--- 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()

**Describe what this pull request is trying to achieve.**
- Fix NaNs occurring in sub-quadratic attention by making q, k, and v contiguous when using MPS
- Fix crash that occurs with MPS on PyTorch 2.0.1 due to LayerNorm still not accepting float16 inputs
- Fix generation failing on Intel Macs with k-diffusion and UniPC samplers

**Environment this was tested in**
 - OS: macOS 13.3.1
 - Browser: Safari
 - Graphics card: M1 Max 64 GB, Radeon Pro Vega 20 (4 GB)

Fixes #8555.