diff --git a/src/diffusers/schedulers/scheduling_deis_multistep.py b/src/diffusers/schedulers/scheduling_deis_multistep.py index c22d2b5b439d..ea745dd42c52 100644 --- a/src/diffusers/schedulers/scheduling_deis_multistep.py +++ b/src/diffusers/schedulers/scheduling_deis_multistep.py @@ -16,6 +16,7 @@ # The codebase is modified based on https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py import math +import warnings from typing import Literal import numpy as np @@ -1015,6 +1016,12 @@ def add_noise( `torch.Tensor`: The noisy samples. """ + if self.num_inference_steps is None: + warnings.warn( + "`add_noise` was called before `set_timesteps`, so the noise level may be incorrect. " + "Call `set_timesteps` before `add_noise`." + ) + # Make sure sigmas and timesteps have the same device and dtype as original_samples sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype) if original_samples.device.type == "mps" and torch.is_floating_point(timesteps): diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py index edcfa4576e53..67a0cf287bd1 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py @@ -15,6 +15,7 @@ # DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver import math +import warnings from typing import Literal import numpy as np @@ -1321,6 +1322,12 @@ def add_noise( `torch.Tensor`: The noisy samples. """ + if self.num_inference_steps is None: + warnings.warn( + "`add_noise` was called before `set_timesteps`, so the noise level may be incorrect. " + "Call `set_timesteps` before `add_noise`." + ) + # Make sure sigmas and timesteps have the same device and dtype as original_samples sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype) if original_samples.device.type == "mps" and torch.is_floating_point(timesteps): diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py b/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py index 7a63d0b36392..909de11429f0 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py @@ -15,6 +15,7 @@ # DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver import math +import warnings from typing import Literal import numpy as np @@ -1295,6 +1296,12 @@ def add_noise( `torch.Tensor`: The noisy samples. """ + if self.num_inference_steps is None: + warnings.warn( + "`add_noise` was called before `set_timesteps`, so the noise level may be incorrect. " + "Call `set_timesteps` before `add_noise`." + ) + # Make sure sigmas and timesteps have the same device and dtype as original_samples sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype) if original_samples.device.type == "mps" and torch.is_floating_point(timesteps): diff --git a/src/diffusers/schedulers/scheduling_unipc_multistep.py b/src/diffusers/schedulers/scheduling_unipc_multistep.py index 5c2cbcc13ff1..28a958484594 100644 --- a/src/diffusers/schedulers/scheduling_unipc_multistep.py +++ b/src/diffusers/schedulers/scheduling_unipc_multistep.py @@ -16,6 +16,7 @@ # The codebase is modified based on https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py import math +import warnings from typing import Literal import numpy as np @@ -1268,6 +1269,12 @@ def add_noise( `torch.Tensor`: The noisy samples. """ + if self.num_inference_steps is None: + warnings.warn( + "`add_noise` was called before `set_timesteps`, so the noise level may be incorrect. " + "Call `set_timesteps` before `add_noise`." + ) + # Make sure sigmas and timesteps have the same device and dtype as original_samples sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype) if original_samples.device.type == "mps" and torch.is_floating_point(timesteps): diff --git a/tests/schedulers/test_scheduler_dpm_multi.py b/tests/schedulers/test_scheduler_dpm_multi.py index 28c354709dc9..c55a4baaac69 100644 --- a/tests/schedulers/test_scheduler_dpm_multi.py +++ b/tests/schedulers/test_scheduler_dpm_multi.py @@ -1,6 +1,8 @@ import tempfile import unittest +import warnings +import pytest import torch from diffusers import ( @@ -366,3 +368,21 @@ def test_beta_sigmas(self): def test_exponential_sigmas(self): self.check_over_configs(use_exponential_sigmas=True) + + +@pytest.mark.parametrize( + "scheduler_class", + [DPMSolverMultistepScheduler, DPMSolverSinglestepScheduler, UniPCMultistepScheduler, DEISMultistepScheduler], +) +def test_add_noise_before_set_timesteps_warns(scheduler_class): + scheduler = scheduler_class() + sample = torch.ones(1) + noise = torch.zeros(1) + + with pytest.warns(UserWarning, match="`add_noise` was called before `set_timesteps`"): + scheduler.add_noise(sample, noise, torch.tensor([200])) + + scheduler.set_timesteps(10) + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + scheduler.add_noise(sample, noise, scheduler.timesteps[:1])