diff --git a/src/diffusers/schedulers/scheduling_heun_discrete.py b/src/diffusers/schedulers/scheduling_heun_discrete.py index f2aaa738233b..cb6cb9e79565 100644 --- a/src/diffusers/schedulers/scheduling_heun_discrete.py +++ b/src/diffusers/schedulers/scheduling_heun_discrete.py @@ -342,7 +342,7 @@ def set_timesteps( timesteps = torch.from_numpy(timesteps) timesteps = torch.cat([timesteps[:1], timesteps[1:].repeat_interleave(2)]) - self.timesteps = timesteps.to(device=device) + self.timesteps = timesteps.to(device=device, dtype=torch.float32) # empty dt and derivative self.prev_derivative = None diff --git a/src/diffusers/schedulers/scheduling_lms_discrete.py b/src/diffusers/schedulers/scheduling_lms_discrete.py index 3d4a794c62e8..bcf9d9b59e11 100644 --- a/src/diffusers/schedulers/scheduling_lms_discrete.py +++ b/src/diffusers/schedulers/scheduling_lms_discrete.py @@ -311,7 +311,7 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32) self.sigmas = torch.from_numpy(sigmas).to(device=device) - self.timesteps = torch.from_numpy(timesteps).to(device=device) + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.float32) self._step_index = None self._begin_index = None self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication diff --git a/tests/schedulers/test_scheduler_lcm.py b/tests/schedulers/test_scheduler_lcm.py index c2c6530faa11..f3f6e9ba5837 100644 --- a/tests/schedulers/test_scheduler_lcm.py +++ b/tests/schedulers/test_scheduler_lcm.py @@ -99,7 +99,7 @@ def test_add_noise_device(self, num_inference_steps=10): scaled_sample = scheduler.scale_model_input(sample, 0.0) self.assertEqual(sample.shape, scaled_sample.shape) - noise = torch.randn_like(scaled_sample).to(torch_device) + noise = torch.randn(scaled_sample.shape).to(torch_device) t = scheduler.timesteps[5][None] noised = scheduler.add_noise(scaled_sample, noise, t) self.assertEqual(noised.shape, scaled_sample.shape) diff --git a/tests/schedulers/test_schedulers.py b/tests/schedulers/test_schedulers.py index fc7f22d2a8e5..42ca1bc54155 100755 --- a/tests/schedulers/test_schedulers.py +++ b/tests/schedulers/test_schedulers.py @@ -361,7 +361,7 @@ def model(sample, t, *args): if isinstance(t, torch.Tensor): num_dims = len(sample.shape) # pad t with 1s to match num_dims - t = t.reshape(-1, *(1,) * (num_dims - 1)).to(sample.device).to(sample.dtype) + t = t.reshape(-1, *(1,) * (num_dims - 1)).to(sample.device, dtype=sample.dtype) return sample * t / (t + 1) @@ -722,7 +722,7 @@ def test_add_noise_device(self): scaled_sample = scheduler.scale_model_input(sample, 0.0) self.assertEqual(sample.shape, scaled_sample.shape) - noise = torch.randn_like(scaled_sample).to(torch_device) + noise = torch.randn(scaled_sample.shape).to(torch_device) t = scheduler.timesteps[5][None] noised = scheduler.add_noise(scaled_sample, noise, t) self.assertEqual(noised.shape, scaled_sample.shape)