Add workaround for using MPS with torchsde

1 files changed, 14 insertions, 0 deletions

diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index 5fefb227..8b11f569 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -6,6 +6,7 @@ import tqdm
 from PIL import Image

 import inspect

 import k_diffusion.sampling

+import torchsde._brownian.brownian_interval

 import ldm.models.diffusion.ddim

 import ldm.models.diffusion.plms

 from modules import prompt_parser, devices, processing, images

@@ -367,6 +368,19 @@ class TorchHijack:
         return torch.randn_like(x)

 

 

+# MPS fix for randn in torchsde

+def torchsde_randn(size, dtype, device, seed):

+    if device.type == 'mps':

+        generator = torch.Generator(devices.cpu).manual_seed(int(seed))

+        return torch.randn(size, dtype=dtype, device=devices.cpu, generator=generator).to(device)

+    else:

+        generator = torch.Generator(device).manual_seed(int(seed))

+        return torch.randn(size, dtype=dtype, device=device, generator=generator)

+

+

+torchsde._brownian.brownian_interval._randn = torchsde_randn

+

+

 class KDiffusionSampler:

     def __init__(self, funcname, sd_model):

         denoiser = k_diffusion.external.CompVisVDenoiser if sd_model.parameterization == "v" else k_diffusion.external.CompVisDenoiser