Fix batched DiffusionGemma adaptive stopping

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.

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] = []