From 504092aaebdd18d7eac7ce7e728fca9e7d97f3b4 Mon Sep 17 00:00:00 2001 From: "DESKTOP-E51EATJ\\user" Date: Tue, 11 Aug 2026 21:09:43 +0800 Subject: [PATCH 1/2] Allow UniPC to set the customed timestep. --- .../schedulers/scheduling_unipc_multistep.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_unipc_multistep.py b/src/diffusers/schedulers/scheduling_unipc_multistep.py index 5c2cbcc13ff1..8ce72b12a802 100644 --- a/src/diffusers/schedulers/scheduling_unipc_multistep.py +++ b/src/diffusers/schedulers/scheduling_unipc_multistep.py @@ -25,7 +25,6 @@ from ..utils import deprecate, is_scipy_available from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput - if is_scipy_available(): import scipy.stats @@ -321,6 +320,7 @@ def set_timesteps( device: str | torch.device = None, sigmas: list[float] | None = None, mu: bool | None = None, + timesteps: list[int] | None = None, ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -338,6 +338,19 @@ def set_timesteps( """ if self.config.use_dynamic_shifting and mu is None: raise ValueError("`mu` must be passed when `use_dynamic_shifting` is set to be `True`") + if num_inference_steps is None and timesteps is None and sigmas is None: + raise ValueError("Must pass exactly one of `num_inference_steps`, `timesteps`, or `sigmas`.") + if num_inference_steps is not None and timesteps is not None: + raise ValueError("Can only pass one of `num_inference_steps` or `custom_timesteps`.") + if timesteps is not None and self.config.use_karras_sigmas: + raise ValueError("Cannot use `timesteps` with `config.use_karras_sigmas = True`") + if timesteps is not None and self.config.use_exponential_sigmas: + raise ValueError("Cannot set `timesteps` with `config.use_exponential_sigmas = True`.") + if timesteps is not None and self.config.use_beta_sigmas: + raise ValueError("Cannot set `timesteps` with `config.use_beta_sigmas = True`.") + if timesteps is not None and sigmas is not None: + raise ValueError("Cannot set `timesteps` and `sigmas` at the same time.") + if sigmas is not None: if not self.config.use_flow_sigmas: @@ -348,7 +361,9 @@ def set_timesteps( num_inference_steps = len(sigmas) # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if self.config.timestep_spacing == "linspace": + if timesteps is not None: + timesteps = np.array(timesteps).astype(np.int64) + elif self.config.timestep_spacing == "linspace": timesteps = ( np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps + 1) .round()[::-1][:-1] From 218850a6aa51a83e6260701b890124be1b7f147d Mon Sep 17 00:00:00 2001 From: "DESKTOP-E51EATJ\\user" Date: Tue, 11 Aug 2026 21:21:13 +0800 Subject: [PATCH 2/2] update parameters explaination. --- src/diffusers/schedulers/scheduling_unipc_multistep.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/diffusers/schedulers/scheduling_unipc_multistep.py b/src/diffusers/schedulers/scheduling_unipc_multistep.py index 8ce72b12a802..615ecd6a49f3 100644 --- a/src/diffusers/schedulers/scheduling_unipc_multistep.py +++ b/src/diffusers/schedulers/scheduling_unipc_multistep.py @@ -335,6 +335,10 @@ def set_timesteps( automatically. mu (`float`, *optional*): Optional mu parameter for dynamic shifting when using exponential time shift type. + timesteps (`list[int]`, *optional*): + Custom timesteps used to support arbitrary timesteps schedule. If `None`, timesteps will be generated + based on the `timestep_spacing` attribute. If `timesteps` is passed, `num_inference_steps` and `sigmas` + must be `None`, and `timestep_spacing` attribute will be ignored. """ if self.config.use_dynamic_shifting and mu is None: raise ValueError("`mu` must be passed when `use_dynamic_shifting` is set to be `True`")