From 249ee767d9d1ec05af588e98071ed133904e9149 Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Tue, 4 Aug 2026 19:31:23 +0200 Subject: [PATCH 1/3] Fix DiffusionGemma stopping --- .../pipeline_diffusion_gemma.py | 39 ++++++++++++------- 1 file changed, 25 insertions(+), 14 deletions(-) diff --git a/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py b/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py index 5222ead8813b..06453a925cb5 100644 --- a/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py +++ b/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py @@ -223,8 +223,9 @@ def __call__( Only used when `confidence_threshold` is set. confidence_threshold (`float`, *optional*, defaults to `0.005`): Leave a block's denoising loop early once every example is stable (see `stability_threshold`) and the - mean per-token entropy of the prediction is below this value. Speeds up generation at matched quality; - the default matches the released checkpoint. Set to `None` to always run all `num_inference_steps`. + mean per-token entropy of the scheduler-shaped prediction logits is below this value. Speeds up + generation at matched quality; the default matches the released checkpoint. Set to `None` to always + run all `num_inference_steps`. generator (`torch.Generator`, *optional*): RNG for sampling. output_type (`str`, defaults to `"text"`): @@ -347,6 +348,8 @@ def __call__( 0, text_config.vocab_size, (batch_size, canvas_length), device=device, generator=generator ) self_conditioning_logits = None + finished_denoising = torch.zeros(batch_size, dtype=torch.bool, device=device) + argmax_canvas = canvas # Adaptive stopping history: the last `stability_threshold` argmax predictions of this block's canvas. argmax_history = torch.full( (max(stability_threshold, 1), batch_size, canvas_length), -1, dtype=torch.long, device=device @@ -380,7 +383,11 @@ def __call__( canvas = scheduler_output.prev_sample # Self-condition on the logits the scheduler sampled from: temperature-shaped for the reference # EntropyBound sampler, the raw denoiser logits for the others. - self_conditioning_logits = scheduler_output.pred_logits + pred_logits = scheduler_output.pred_logits + if finished_denoising.any(): + canvas = torch.where(finished_denoising[:, None], argmax_canvas, canvas) + pred_logits = torch.where(finished_denoising[:, None, None], self_conditioning_logits, pred_logits) + self_conditioning_logits = pred_logits # Predictor-corrector (https://huggingface.co/papers/2605.22765): a scheduler exposing `corrector_steps` # + `step_correct` refines the canvas with extra Gibbs sweeps on the first `corrected_steps` predictor @@ -408,21 +415,25 @@ def __call__( global_step += 1 progress_bar.update() - # Adaptive stopping: leave this block early once every example's argmax prediction is stable across - # `stability_threshold` steps and confident (mean per-token entropy below `confidence_threshold`). + # Adaptive stopping: freeze each example once its scheduler-shaped prediction is stable across + # `stability_threshold` steps and confident (mean per-token entropy below `confidence_threshold`), + # then leave the block once every example is finished. if confidence_threshold is not None: - argmax_canvas = logits.argmax(dim=-1) - stable = (argmax_history == argmax_canvas[None]).all(dim=-1).all(dim=0) + next_argmax_canvas = pred_logits.argmax(dim=-1) + next_argmax_canvas = torch.where(finished_denoising[:, None], argmax_canvas, next_argmax_canvas) + stable = (argmax_history == next_argmax_canvas[None]).all(dim=-1).all(dim=0) argmax_history = torch.roll(argmax_history, shifts=-1, dims=0) - argmax_history[-1] = argmax_canvas - confident = torch.distributions.Categorical(logits=logits.float()).entropy().mean(-1) < ( + argmax_history[-1] = next_argmax_canvas + confident = torch.distributions.Categorical(logits=pred_logits.float()).entropy().mean(-1) < ( confidence_threshold ) - if bool((stable & confident).all()): - # Commit the converged prediction. Ancestral schedulers (e.g. DiscreteDDIM) only clean the - # canvas on their final step, so the in-progress canvas may still hold noise tokens; the - # denoiser argmax is the converged answer (and equals the canvas for commit-style schedulers). - canvas = argmax_canvas + finished_denoising = finished_denoising | (stable & confident) + argmax_canvas = next_argmax_canvas + # Commit each converged prediction. Ancestral schedulers (e.g. DiscreteDDIM) only clean the canvas + # on their final step, so the in-progress canvas may still hold noise tokens; the denoiser argmax + # is the converged answer (and equals the canvas for commit-style schedulers). + canvas = torch.where(finished_denoising[:, None], argmax_canvas, canvas) + if bool(finished_denoising.all()): break # Append the denoised canvas and extend the context for the next block. From 4ab3ef72f9e4c8893e24bdd629d091777d888dd2 Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Tue, 4 Aug 2026 19:31:25 +0200 Subject: [PATCH 2/3] Test batched stopping --- .../diffusion_gemma/test_diffusion_gemma.py | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/tests/pipelines/diffusion_gemma/test_diffusion_gemma.py b/tests/pipelines/diffusion_gemma/test_diffusion_gemma.py index c01b7adbc81f..b2737a6a7772 100644 --- a/tests/pipelines/diffusion_gemma/test_diffusion_gemma.py +++ b/tests/pipelines/diffusion_gemma/test_diffusion_gemma.py @@ -1,4 +1,5 @@ import unittest +from types import SimpleNamespace import torch @@ -103,6 +104,46 @@ def test_generate(self): self.assertEqual(sequences.shape, (1, self.canvas_length)) self.assertEqual(len(texts), 1) + def test_adaptive_stopping_freezes_finished_rows_and_uses_scheduler_logits(self): + vocab_size = 8 + self.pipe.model.config.get_text_config(decoder=True).vocab_size = vocab_size + forward_calls = 0 + + def forward(decoder_input_ids, **kwargs): + nonlocal forward_calls + forward_calls += 1 + batch_size, canvas_length = decoder_input_ids.shape + return SimpleNamespace(logits=torch.zeros(batch_size, canvas_length, vocab_size)) + + class DeterministicScheduler: + config = SimpleNamespace(corrector_steps=0) + + def set_timesteps(self, num_inference_steps, device=None): + self.timesteps = torch.arange(num_inference_steps, device=device) + + def step(self, model_output, timestep, sample, return_dict=True): + del model_output, return_dict + token_ids = ([1, 3], [1, 4], [2, 5], [2, 5], [2, 6])[timestep] + tokens = torch.tensor(token_ids, device=sample.device)[:, None].expand_as(sample) + pred_logits = torch.full((*sample.shape, vocab_size), -100.0, device=sample.device) + pred_logits.scatter_(-1, tokens[..., None], 100.0) + return SimpleNamespace(prev_sample=tokens, pred_logits=pred_logits) + + self.pipe.model.forward = forward + self.pipe.scheduler = DeterministicScheduler() + output = self.pipe( + prompt=["Short prompt.", "A somewhat longer prompt for the second batch row."], + gen_length=self.canvas_length, + num_inference_steps=5, + confidence_threshold=0.005, + eos_early_stop=False, + output_type="seq", + ) + + self.assertEqual(forward_calls, 4) + self.assertTrue((output.sequences[0] == 1).all()) + self.assertTrue((output.sequences[1] == 5).all()) + def test_callback_receives_advertised_keys(self): observed: list[str] = [] From 98ac49eeb3fd573c88da2d21fd40d29ddb4f343a Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Tue, 4 Aug 2026 19:33:30 +0200 Subject: [PATCH 3/3] Style DiffusionGemma docs --- .../pipelines/diffusion_gemma/pipeline_diffusion_gemma.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py b/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py index 06453a925cb5..42847678c758 100644 --- a/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py +++ b/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py @@ -224,8 +224,8 @@ def __call__( confidence_threshold (`float`, *optional*, defaults to `0.005`): Leave a block's denoising loop early once every example is stable (see `stability_threshold`) and the mean per-token entropy of the scheduler-shaped prediction logits is below this value. Speeds up - generation at matched quality; the default matches the released checkpoint. Set to `None` to always - run all `num_inference_steps`. + generation at matched quality; the default matches the released checkpoint. Set to `None` to always run + all `num_inference_steps`. generator (`torch.Generator`, *optional*): RNG for sampling. output_type (`str`, defaults to `"text"`):