hoelee/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2025-08-08 13:19:54 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge pull request #4918 from brkirch/pytorch-fixes

Browse Source 
Fixes for PyTorch 1.12.1 when using MPS
This commit is contained in:
AUTOMATIC1111
2022-11-27 13:47:01 +03:00
committed by GitHub
parent 10923f9b3a e247b7400a
commit cc90dcc933
4 changed files with 25 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
modules/devices.py
View File

@@ -2,9 +2,10 @@ import sys, os, shlex
import contextlib
import torch
from modules import errors
from packaging import version
# has_mps is only available in nightly pytorch (for now) and MasOS 12.3+.
# 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):
@@ -99,9 +100,25 @@ def autocast(disable=False):
# MPS workaround for https://github.com/pytorch/pytorch/issues/79383
def mps_contiguous(input_tensor, device):
return input_tensor.contiguous() if device.type == 'mps' else input_tensor
orig_tensor_to = torch.Tensor.to
def tensor_to_fix(self, *args, **kwargs):
if self.device.type != 'mps' and \
((len(args) > 0 and isinstance(args[0], torch.device) and args[0].type == 'mps') or \
(isinstance(kwargs.get('device'), torch.device) and kwargs['device'].type == 'mps')):
self = self.contiguous()
return orig_tensor_to(self, *args, **kwargs)
def mps_contiguous_to(input_tensor, device):
return mps_contiguous(input_tensor, device).to(device)
# MPS workaround for https://github.com/pytorch/pytorch/issues/80800
orig_layer_norm = torch.nn.functional.layer_norm
def layer_norm_fix(*args, **kwargs):
if len(args) > 0 and isinstance(args[0], torch.Tensor) and args[0].device.type == 'mps':
args = list(args)
args[0] = args[0].contiguous()
return orig_layer_norm(*args, **kwargs)
# PyTorch 1.13 doesn't need these fixes but unfortunately is slower and has regressions that prevent training from working
if has_mps() and version.parse(torch.__version__) < version.parse("1.13"):
torch.Tensor.to = tensor_to_fix
torch.nn.functional.layer_norm = layer_norm_fix