Mirror/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2024-06-07 21:20:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Rework MPS randn fix, add randn_like fix

Browse Source 
torch.manual_seed() already sets a CPU generator, so there is no reason to create a CPU generator manually. torch.randn_like also needs a MPS fix for k-diffusion, but a torch hijack with randn_like already exists so it can also be used for that.
This commit is contained in:
brkirch 2022-11-30 08:02:39 -05:00
parent 4d5f1691dd
commit 0fddb4a1c0
2 changed files with 8 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
modules/devices.py
View File

@ -66,24 +66,15 @@ dtype_vae = torch.float16
def randn(seed, shape): def randn(seed, shape):
# Pytorch currently doesn't handle setting randomness correctly when the metal backend is used.
if device.type == 'mps':
generator = torch.Generator(device=cpu)
generator.manual_seed(seed)
noise = torch.randn(shape, generator=generator, device=cpu).to(device)
return noise
torch.manual_seed(seed) torch.manual_seed(seed)
if device.type == 'mps':
return torch.randn(shape, device=cpu).to(device)
return torch.randn(shape, device=device) return torch.randn(shape, device=device)
def randn_without_seed(shape): def randn_without_seed(shape):
# Pytorch currently doesn't handle setting randomness correctly when the metal backend is used.
if device.type == 'mps': if device.type == 'mps':
generator = torch.Generator(device=cpu) return torch.randn(shape, device=cpu).to(device)
noise = torch.randn(shape, generator=generator, device=cpu).to(device)
return noise
return torch.randn(shape, device=device) return torch.randn(shape, device=device)

6
modules/sd_samplers.py
View File

@ -365,6 +365,9 @@ class TorchHijack:
if noise.shape == x.shape: if noise.shape == x.shape:
return noise return noise
if x.device.type == 'mps':
return torch.randn_like(x, device=devices.cpu).to(x.device)
else:
return torch.randn_like(x) return torch.randn_like(x)
@ -429,8 +432,7 @@ class KDiffusionSampler:
self.model_wrap.step = 0 self.model_wrap.step = 0
self.eta = p.eta or opts.eta_ancestral self.eta = p.eta or opts.eta_ancestral
if self.sampler_noises is not None: k_diffusion.sampling.torch = TorchHijack(self.sampler_noises if self.sampler_noises is not None else [])
k_diffusion.sampling.torch = TorchHijack(self.sampler_noises)
extra_params_kwargs = {} extra_params_kwargs = {}
for param_name in self.extra_params: for param_name in self.extra_params: