diff options
author | unknown <mcgpapu@gmail.com> | 2023-02-05 13:24:57 +0000 |
---|---|---|
committer | unknown <mcgpapu@gmail.com> | 2023-02-05 13:24:57 +0000 |
commit | 501d4e9cf111c2b94805181e3ab2e1b8b5d81f31 (patch) | |
tree | 712d6e405c406e573e8f2b7f44b9e4462432ff61 /modules/devices.py | |
parent | 5e1f4f7464e560a2468812fc9d5cb38659cff86c (diff) | |
parent | ea9bd9fc7409109adcd61b897abc2c8881161256 (diff) | |
download | stable-diffusion-webui-gfx803-501d4e9cf111c2b94805181e3ab2e1b8b5d81f31.tar.gz stable-diffusion-webui-gfx803-501d4e9cf111c2b94805181e3ab2e1b8b5d81f31.tar.bz2 stable-diffusion-webui-gfx803-501d4e9cf111c2b94805181e3ab2e1b8b5d81f31.zip |
Merge branch 'master' of github.com:AUTOMATIC1111/stable-diffusion-webui into gamepad
Diffstat (limited to 'modules/devices.py')
-rw-r--r-- | modules/devices.py | 52 |
1 files changed, 7 insertions, 45 deletions
diff --git a/modules/devices.py b/modules/devices.py index 919048d0..52c3e7cd 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -1,22 +1,17 @@ -import sys, os, shlex +import sys import contextlib import torch from modules import errors -from modules.sd_hijack_utils import CondFunc -from packaging import version + +if sys.platform == "darwin": + from modules import mac_specific -# has_mps is only available in nightly pytorch (for now) and macOS 12.3+. -# check `getattr` and try it for compatibility def has_mps() -> bool: - if not getattr(torch, 'has_mps', False): - return False - try: - torch.zeros(1).to(torch.device("mps")) - return True - except Exception: + if sys.platform != "darwin": return False - + else: + return mac_specific.has_mps def extract_device_id(args, name): for x in range(len(args)): @@ -155,36 +150,3 @@ def test_for_nans(x, where): message += " Use --disable-nan-check commandline argument to disable this check." raise NansException(message) - - -# MPS workaround for https://github.com/pytorch/pytorch/issues/89784 -def cumsum_fix(input, cumsum_func, *args, **kwargs): - if input.device.type == 'mps': - output_dtype = kwargs.get('dtype', input.dtype) - if output_dtype == torch.int64: - return cumsum_func(input.cpu(), *args, **kwargs).to(input.device) - elif cumsum_needs_bool_fix and output_dtype == torch.bool or cumsum_needs_int_fix and (output_dtype == torch.int8 or output_dtype == torch.int16): - return cumsum_func(input.to(torch.int32), *args, **kwargs).to(torch.int64) - return cumsum_func(input, *args, **kwargs) - - -if has_mps(): - 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 - - # MPS workaround for https://github.com/pytorch/pytorch/issues/79383 - CondFunc('torch.Tensor.to', lambda orig_func, self, *args, **kwargs: orig_func(self.contiguous(), *args, **kwargs), - lambda _, self, *args, **kwargs: self.device.type != 'mps' and (args and isinstance(args[0], torch.device) and args[0].type == 'mps' or isinstance(kwargs.get('device'), torch.device) and kwargs['device'].type == 'mps')) - # MPS workaround for https://github.com/pytorch/pytorch/issues/80800 - CondFunc('torch.nn.functional.layer_norm', lambda orig_func, *args, **kwargs: orig_func(*([args[0].contiguous()] + list(args[1:])), **kwargs), - lambda _, *args, **kwargs: args and isinstance(args[0], torch.Tensor) and args[0].device.type == 'mps') - # MPS workaround for https://github.com/pytorch/pytorch/issues/90532 - CondFunc('torch.Tensor.numpy', lambda orig_func, self, *args, **kwargs: orig_func(self.detach(), *args, **kwargs), lambda _, self, *args, **kwargs: self.requires_grad) - elif version.parse(torch.__version__) > version.parse("1.13.1"): - cumsum_needs_int_fix = not torch.Tensor([1,2]).to(torch.device("mps")).equal(torch.ShortTensor([1,1]).to(torch.device("mps")).cumsum(0)) - cumsum_needs_bool_fix = not torch.BoolTensor([True,True]).to(device=torch.device("mps"), dtype=torch.int64).equal(torch.BoolTensor([True,False]).to(torch.device("mps")).cumsum(0)) - cumsum_fix_func = lambda orig_func, input, *args, **kwargs: cumsum_fix(input, orig_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) - |