Merge pull request #8866 from brkirch/mps-torch-2-0-nn-linear-workarounds

Add PyTorch 2.0 support for macOS, fix image generation on macOS 13.2.X

1 files changed, 8 insertions, 1 deletions

diff --git a/modules/mac_specific.py b/modules/mac_specific.py
index 18e6ff72..6fe8dea0 100644
--- a/modules/mac_specific.py
+++ b/modules/mac_specific.py
@@ -1,4 +1,5 @@
 import torch
+import platform
 from modules import paths
 from modules.sd_hijack_utils import CondFunc
 from packaging import version
@@ -32,6 +33,10 @@ if has_mps:
     # MPS fix for randn in torchsde
     CondFunc('torchsde._brownian.brownian_interval._randn', lambda _, size, dtype, device, seed: torch.randn(size, dtype=dtype, device=torch.device("cpu"), generator=torch.Generator(torch.device("cpu")).manual_seed(int(seed))).to(device), lambda _, size, dtype, device, seed: device.type == 'mps')
 
+    if platform.mac_ver()[0].startswith("13.2."):
+        # MPS workaround for https://github.com/pytorch/pytorch/issues/95188, thanks to danieldk (https://github.com/explosion/curated-transformers/pull/124)
+        CondFunc('torch.nn.functional.linear', lambda _, input, weight, bias: (torch.matmul(input, weight.t()) + bias) if bias is not None else torch.matmul(input, weight.t()), lambda _, input, weight, bias: input.numel() > 10485760)
+
     if version.parse(torch.__version__) < version.parse("1.13"):
         # PyTorch 1.13 doesn't need these fixes but unfortunately is slower and has regressions that prevent training from working
 
@@ -49,4 +54,6 @@ if has_mps:
         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)