Warn when add_noise is called before set_timesteps

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):

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):

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):

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):

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])