Rework MPS randn fix, add randn_like fix (0fddb4a1) · Commits · github_fork / Stable Diffusion Webui

modules/devices.py

+3 −12

Original line number	Diff line number	Diff line
		@@ -66,24 +66,15 @@ dtype_vae = torch.float16


		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)
		if device.type == 'mps':
		return torch.randn(shape, device=cpu).to(device)
		return torch.randn(shape, device=device)


		def randn_without_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)
		noise = torch.randn(shape, generator=generator, device=cpu).to(device)
		return noise

		return torch.randn(shape, device=cpu).to(device)
		return torch.randn(shape, device=device)

+5 −3

Original line number	Diff line number	Diff line
		@@ -365,6 +365,9 @@ class TorchHijack:
		if noise.shape == x.shape:
		return noise

		if x.device.type == 'mps':
		return torch.randn_like(x, device=devices.cpu).to(x.device)
		else:
		return torch.randn_like(x)


		@@ -429,8 +432,7 @@ class KDiffusionSampler:
		self.model_wrap.step = 0
		self.eta = p.eta or opts.eta_ancestral

		if self.sampler_noises is not None:
		k_diffusion.sampling.torch = TorchHijack(self.sampler_noises)
		k_diffusion.sampling.torch = TorchHijack(self.sampler_noises if self.sampler_noises is not None else [])

		extra_params_kwargs = {}
		for param_name in self.extra_params: