From 4489725e1c209a2230d88665d3ca2e3b9b12ce7a Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Mon, 29 Jun 2026 20:58:50 +0200 Subject: [PATCH 1/2] add doc pages for the discrete ddim and entropy bound schedulers Co-Authored-By: Claude Opus 4.8 --- docs/source/en/_toctree.yml | 4 +++ .../en/api/pipelines/diffusion_gemma.md | 6 ++--- .../source/en/api/schedulers/discrete_ddim.md | 27 +++++++++++++++++++ .../source/en/api/schedulers/entropy_bound.md | 26 ++++++++++++++++++ 4 files changed, 60 insertions(+), 3 deletions(-) create mode 100644 docs/source/en/api/schedulers/discrete_ddim.md create mode 100644 docs/source/en/api/schedulers/entropy_bound.md diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index d256ebdeffe7..83bc19670d96 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -718,6 +718,8 @@ title: DDPMScheduler - local: api/schedulers/deis title: DEISMultistepScheduler + - local: api/schedulers/discrete_ddim + title: DiscreteDDIMScheduler - local: api/schedulers/multistep_dpm_solver_inverse title: DPMSolverMultistepInverse - local: api/schedulers/multistep_dpm_solver @@ -730,6 +732,8 @@ title: EDMDPMSolverMultistepScheduler - local: api/schedulers/edm_euler title: EDMEulerScheduler + - local: api/schedulers/entropy_bound + title: EntropyBoundScheduler - local: api/schedulers/euler_ancestral title: EulerAncestralDiscreteScheduler - local: api/schedulers/euler diff --git a/docs/source/en/api/pipelines/diffusion_gemma.md b/docs/source/en/api/pipelines/diffusion_gemma.md index b83afc491c3a..ea24672a4ebe 100644 --- a/docs/source/en/api/pipelines/diffusion_gemma.md +++ b/docs/source/en/api/pipelines/diffusion_gemma.md @@ -76,11 +76,11 @@ the model's image inputs automatically. The scheduler is the sampler that denoises each canvas, and it is interchangeable: swap it to change the sampling strategy without touching anything else. Three schedulers are available: -- `BlockRefinementScheduler` (default): commits the most confident tokens each step (above `threshold`, plus an even +- [`BlockRefinementScheduler`] (default): commits the most confident tokens each step (above `threshold`, plus an even per-step quota) and renoises the rest. `editing_threshold` additionally lets it re-edit already committed tokens. -- `DiscreteDDIMScheduler`: samples each position from the exact discrete posterior of the uniform corruption process +- [`DiscreteDDIMScheduler`]: samples each position from the exact discrete posterior of the uniform corruption process (D3PM). It is parameter free, and the final step deterministically commits the predicted tokens. -- `EntropyBoundScheduler`: commits the lowest-entropy positions whose joint entropy stays under `entropy_bound`, so +- [`EntropyBoundScheduler`]: commits the lowest-entropy positions whose joint entropy stays under `entropy_bound`, so roughly independent tokens are accepted together. It anneals its sampling temperature from `t_max` (`0.8`) on the first step down to `t_min` (`0.4`) on the last, matching the released checkpoint's sampler. diff --git a/docs/source/en/api/schedulers/discrete_ddim.md b/docs/source/en/api/schedulers/discrete_ddim.md new file mode 100644 index 000000000000..a83f70764cc1 --- /dev/null +++ b/docs/source/en/api/schedulers/discrete_ddim.md @@ -0,0 +1,27 @@ + + +# DiscreteDDIMScheduler + +The `DiscreteDDIMScheduler` samples each canvas position from the exact discrete posterior of the uniform corruption +process (D3PM), following [Structured Denoising Diffusion Models in Discrete State-Spaces](https://huggingface.co/papers/2107.03006). +It is parameter free, and the final step deterministically commits the predicted tokens. An optional predictor-corrector +mode adds the leave-one-out Gibbs sweeps of [Reparameterizing Uniform Diffusion Models](https://huggingface.co/papers/2605.22765) +through `corrector_steps`. + +This scheduler is used by [`DiffusionGemmaPipeline`]. + +## DiscreteDDIMScheduler +[[autodoc]] DiscreteDDIMScheduler + +## DiscreteDDIMSchedulerOutput +[[autodoc]] schedulers.scheduling_discrete_ddim.DiscreteDDIMSchedulerOutput diff --git a/docs/source/en/api/schedulers/entropy_bound.md b/docs/source/en/api/schedulers/entropy_bound.md new file mode 100644 index 000000000000..60d03b396cd9 --- /dev/null +++ b/docs/source/en/api/schedulers/entropy_bound.md @@ -0,0 +1,26 @@ + + +# EntropyBoundScheduler + +The `EntropyBoundScheduler` commits the lowest-entropy positions whose joint entropy stays under `entropy_bound`, so +roughly independent tokens are accepted together and the rest are renoised. It anneals its sampling temperature from +`t_max` on the first step down to `t_min` on the last, matching the released checkpoint's sampler. Proposed in +[Beyond Next-Token Prediction](https://huggingface.co/papers/2505.24857). + +This scheduler is used by [`DiffusionGemmaPipeline`]. + +## EntropyBoundScheduler +[[autodoc]] EntropyBoundScheduler + +## EntropyBoundSchedulerOutput +[[autodoc]] schedulers.scheduling_entropy_bound.EntropyBoundSchedulerOutput From c24e5e1b2e01cf64a328771c147fb4114808786b Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Mon, 29 Jun 2026 21:08:06 +0200 Subject: [PATCH 2/2] fix the cited paper titles to match the linked papers Co-Authored-By: Claude Opus 4.8 --- docs/source/en/api/pipelines/diffusion_gemma.md | 2 +- docs/source/en/api/schedulers/discrete_ddim.md | 2 +- docs/source/en/api/schedulers/entropy_bound.md | 2 +- src/diffusers/schedulers/scheduling_discrete_ddim.py | 11 ++++++----- src/diffusers/schedulers/scheduling_entropy_bound.py | 3 ++- 5 files changed, 11 insertions(+), 9 deletions(-) diff --git a/docs/source/en/api/pipelines/diffusion_gemma.md b/docs/source/en/api/pipelines/diffusion_gemma.md index ea24672a4ebe..bb3adfe7b514 100644 --- a/docs/source/en/api/pipelines/diffusion_gemma.md +++ b/docs/source/en/api/pipelines/diffusion_gemma.md @@ -108,7 +108,7 @@ greedy). ### Predictor-corrector sampling -`DiscreteDDIMScheduler` supports the leave-one-out predictor-corrector of [Reparameterizing Uniform Diffusion Models](https://huggingface.co/papers/2605.22765). It refines the canvas with `corrector_steps` Gibbs sweeps that resample the least-confident positions from the one-coordinate conditional of the noisy marginal, which leaves that marginal invariant and improves generation at no extra training cost. It works directly on the released checkpoint: for uniform diffusion the denoiser and the leave-one-out posterior are interchangeable in closed form, so the corrector recovers the leave-one-out quantities it needs without any retraining. +`DiscreteDDIMScheduler` supports the leave-one-out predictor-corrector of [Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation](https://huggingface.co/papers/2605.22765). It refines the canvas with `corrector_steps` Gibbs sweeps that resample the least-confident positions from the one-coordinate conditional of the noisy marginal, which leaves that marginal invariant and improves generation at no extra training cost. It works directly on the released checkpoint: for uniform diffusion the denoiser and the leave-one-out posterior are interchangeable in closed form, so the corrector recovers the leave-one-out quantities it needs without any retraining. The corrector sweeps are folded into the `num_inference_steps` budget rather than added on top: the pipeline runs fewer predictor steps and spends the freed forwards on correctors, so the total number of model forwards stays `num_inference_steps` and the predictor-corrector costs the same as plain ancestral sampling. diff --git a/docs/source/en/api/schedulers/discrete_ddim.md b/docs/source/en/api/schedulers/discrete_ddim.md index a83f70764cc1..01d5afecf073 100644 --- a/docs/source/en/api/schedulers/discrete_ddim.md +++ b/docs/source/en/api/schedulers/discrete_ddim.md @@ -15,7 +15,7 @@ specific language governing permissions and limitations under the License. The `DiscreteDDIMScheduler` samples each canvas position from the exact discrete posterior of the uniform corruption process (D3PM), following [Structured Denoising Diffusion Models in Discrete State-Spaces](https://huggingface.co/papers/2107.03006). It is parameter free, and the final step deterministically commits the predicted tokens. An optional predictor-corrector -mode adds the leave-one-out Gibbs sweeps of [Reparameterizing Uniform Diffusion Models](https://huggingface.co/papers/2605.22765) +mode adds the leave-one-out Gibbs sweeps of [Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation](https://huggingface.co/papers/2605.22765) through `corrector_steps`. This scheduler is used by [`DiffusionGemmaPipeline`]. diff --git a/docs/source/en/api/schedulers/entropy_bound.md b/docs/source/en/api/schedulers/entropy_bound.md index 60d03b396cd9..cd770d42e900 100644 --- a/docs/source/en/api/schedulers/entropy_bound.md +++ b/docs/source/en/api/schedulers/entropy_bound.md @@ -15,7 +15,7 @@ specific language governing permissions and limitations under the License. The `EntropyBoundScheduler` commits the lowest-entropy positions whose joint entropy stays under `entropy_bound`, so roughly independent tokens are accepted together and the rest are renoised. It anneals its sampling temperature from `t_max` on the first step down to `t_min` on the last, matching the released checkpoint's sampler. Proposed in -[Beyond Next-Token Prediction](https://huggingface.co/papers/2505.24857). +[Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking](https://huggingface.co/papers/2505.24857). This scheduler is used by [`DiffusionGemmaPipeline`]. diff --git a/src/diffusers/schedulers/scheduling_discrete_ddim.py b/src/diffusers/schedulers/scheduling_discrete_ddim.py index 3e0a61273ae9..890a75e694f6 100644 --- a/src/diffusers/schedulers/scheduling_discrete_ddim.py +++ b/src/diffusers/schedulers/scheduling_discrete_ddim.py @@ -57,11 +57,12 @@ class DiscreteDDIMScheduler(SchedulerMixin, ConfigMixin): or jump to a uniformly random token. Unlike masked diffusion, there is no mask token; uncommitted positions carry random tokens. - An optional predictor-corrector mode follows "Reparameterizing Uniform Diffusion Models" via the leave-one-out - (LOO) denoiser (https://huggingface.co/papers/2605.22765). When `corrector_steps > 0`, the pipeline runs that many - Gibbs corrector sweeps after each predictor step (see [`~DiscreteDDIMScheduler.step_correct`]), resampling the - least-confident positions from the one-coordinate conditional `Cat(alpha_s * x0_loo + (1 - alpha_s) / K)` while - holding the rest fixed, which leaves the marginal `p_s` invariant and improves generation at no training cost. + An optional predictor-corrector mode follows "Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and + Absorbing State Reformulation" via the leave-one-out (LOO) denoiser (https://huggingface.co/papers/2605.22765). + When `corrector_steps > 0`, the pipeline runs that many Gibbs corrector sweeps after each predictor step (see + [`~DiscreteDDIMScheduler.step_correct`]), resampling the least-confident positions from the one-coordinate + conditional `Cat(alpha_s * x0_loo + (1 - alpha_s) / K)` while holding the rest fixed, which leaves the marginal + `p_s` invariant and improves generation at no training cost. Args: num_inference_steps (`int`, defaults to 32): diff --git a/src/diffusers/schedulers/scheduling_entropy_bound.py b/src/diffusers/schedulers/scheduling_entropy_bound.py index a31c31d6152e..cd17641cecd0 100644 --- a/src/diffusers/schedulers/scheduling_entropy_bound.py +++ b/src/diffusers/schedulers/scheduling_entropy_bound.py @@ -57,7 +57,8 @@ class EntropyBoundScheduler(SchedulerMixin, ConfigMixin): joint mutual information between the accepted tokens, so they are approximately independent. Accepted positions keep their sampled token; the rest are renoised with uniformly random tokens (there is no mask token). - Proposed in "Beyond Next-Token Prediction" (https://huggingface.co/papers/2505.24857). + Proposed in "Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking" + (https://huggingface.co/papers/2505.24857). The sampling temperature is annealed from `t_max` on the first step down to `t_min` on the last, matching the released checkpoint's sampler (sharper sampling as denoising advances). It is applied to the logits before both the