From c7697e7a8d740dabc0bae851bfcf0d37b97f033c Mon Sep 17 00:00:00 2001 From: patil-suraj Date: Thu, 3 Nov 2022 17:57:53 +0100 Subject: [PATCH 1/2] handle device for randn in euler step --- .../schedulers/scheduling_euler_ancestral_discrete.py | 11 ++++++++++- src/diffusers/schedulers/scheduling_euler_discrete.py | 11 ++++++++++- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py index 3fe52a498044..409db0985604 100644 --- a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py @@ -217,7 +217,16 @@ def step( prev_sample = sample + derivative * dt device = model_output.device if torch.is_tensor(model_output) else "cpu" - noise = torch.randn(model_output.shape, dtype=model_output.dtype, generator=generator).to(device) + if device == "mps": + # randn does not work reproducibly on mps + noise = torch.randn(model_output.shape, dtype=model_output.dtype, device="cpu", generator=generator).to( + device + ) + else: + noise = torch.randn(model_output.shape, dtype=model_output.dtype, device=device, generator=generator).to( + device + ) + prev_sample = prev_sample + noise * sigma_up if not return_dict: diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py index 93aeb8cc3865..5d3ebcf8d786 100644 --- a/src/diffusers/schedulers/scheduling_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_discrete.py @@ -214,7 +214,16 @@ def step( gamma = min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigma <= s_tmax else 0.0 device = model_output.device if torch.is_tensor(model_output) else "cpu" - noise = torch.randn(model_output.shape, dtype=model_output.dtype, generator=generator).to(device) + if device == "mps": + # randn does not work reproducibly on mps + noise = torch.randn(model_output.shape, dtype=model_output.dtype, device="cpu", generator=generator).to( + device + ) + else: + noise = torch.randn(model_output.shape, dtype=model_output.dtype, device=device, generator=generator).to( + device + ) + eps = noise * s_noise sigma_hat = sigma * (gamma + 1) From 465fa3ef44c90972f4efa37c175706841466dbc5 Mon Sep 17 00:00:00 2001 From: patil-suraj Date: Thu, 3 Nov 2022 18:12:01 +0100 Subject: [PATCH 2/2] convert device to str --- src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py | 2 +- src/diffusers/schedulers/scheduling_euler_discrete.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py index 409db0985604..134b45a73b08 100644 --- a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py @@ -217,7 +217,7 @@ def step( prev_sample = sample + derivative * dt device = model_output.device if torch.is_tensor(model_output) else "cpu" - if device == "mps": + if str(device) == "mps": # randn does not work reproducibly on mps noise = torch.randn(model_output.shape, dtype=model_output.dtype, device="cpu", generator=generator).to( device diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py index 5d3ebcf8d786..6425072ac3c9 100644 --- a/src/diffusers/schedulers/scheduling_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_discrete.py @@ -214,7 +214,7 @@ def step( gamma = min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigma <= s_tmax else 0.0 device = model_output.device if torch.is_tensor(model_output) else "cpu" - if device == "mps": + if str(device) == "mps": # randn does not work reproducibly on mps noise = torch.randn(model_output.shape, dtype=model_output.dtype, device="cpu", generator=generator).to( device