feat: add BlockRefinementScheduler for commit-by-confidence scheduling

Extract the confidence-based token commit logic from BlockRefinementPipeline
into a dedicated BlockRefinementScheduler, following diffusers conventions.

The scheduler owns:
- Transfer schedule computation (get_num_transfer_tokens)
- Timestep management (set_timesteps)
- Step logic: confidence-based mask-filling and optional token editing

The pipeline now delegates scheduling to self.scheduler.step() and accepts
a scheduler parameter in __init__.

Author: kashif

docs/source/en/api/pipelines/block_refinement.md

@@ -21,19 +21,16 @@ You can set default sampling parameters when creating the pipeline. Passing `Non
 falls back to `pipe.config`.
 
 ```py
-from diffusers import BlockRefinementPipeline
+from diffusers import BlockRefinementPipeline, BlockRefinementScheduler
 
+scheduler = BlockRefinementScheduler()
 pipe = BlockRefinementPipeline(
     model=model,
+    scheduler=scheduler,
     tokenizer=tokenizer,
-    gen_length=256,
-    block_length=32,
-    steps=16,
-    temperature=0.8,
-    sampling_method="multinomial",
 )
 
-out = pipe(prompt="Explain gradient descent.")
+out = pipe(prompt="Explain gradient descent.", gen_length=256, block_length=32, steps=16, temperature=0.8)
 print(out.texts[0])
 ```
 
@@ -61,3 +58,9 @@ out = pipe(
 
 ## BlockRefinementPipelineOutput
 [[autodoc]] pipelines.BlockRefinementPipelineOutput
+
+## BlockRefinementScheduler
+[[autodoc]] BlockRefinementScheduler
+
+## BlockRefinementSchedulerOutput
+[[autodoc]] schedulers.scheduling_block_refinement.BlockRefinementSchedulerOutput

examples/discrete_diffusion/sample_block_refinement.py

@@ -5,7 +5,7 @@
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-from diffusers import BlockRefinementPipeline
+from diffusers import BlockRefinementPipeline, BlockRefinementScheduler
 
 
 def main():
@@ -38,7 +38,8 @@ def main():
     if tokenizer.mask_token_id is None:
         raise ValueError("Tokenizer must have `mask_token_id` for block refinement sampling.")
 
-    pipe = BlockRefinementPipeline(model=model, tokenizer=tokenizer).to(args.device)
+    scheduler = BlockRefinementScheduler()
+    pipe = BlockRefinementPipeline(model=model, scheduler=scheduler, tokenizer=tokenizer).to(args.device)
     gen = torch.Generator(device=args.device).manual_seed(args.seed)
 
     prompt_ids = tokenizer(args.prompt, return_tensors="pt")["input_ids"].to(args.device)

examples/discrete_diffusion/sample_llada2.py

@@ -29,7 +29,7 @@
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-from diffusers import LLaDA2Pipeline
+from diffusers import BlockRefinementScheduler, LLaDA2Pipeline
 from diffusers.hooks import apply_group_offloading
 
 
@@ -207,7 +207,8 @@ def main():
     model.eval()
 
     # Create pipeline
-    pipe = LLaDA2Pipeline(model=model, tokenizer=tokenizer)
+    scheduler = BlockRefinementScheduler()
+    pipe = LLaDA2Pipeline(model=model, scheduler=scheduler, tokenizer=tokenizer)
 
     # Apply sequential CPU offload if requested
     if args.offload == "sequential":

examples/discrete_diffusion/train_block_refinement_cap.py

@@ -26,7 +26,7 @@
 from accelerate.utils import ProjectConfiguration, set_seed
 from torch.utils.data import DataLoader, Dataset
 
-from diffusers import BlockRefinementPipeline
+from diffusers import BlockRefinementPipeline, BlockRefinementScheduler
 from diffusers.training_utils import compute_confidence_aware_loss
 
 
@@ -249,7 +249,8 @@ def main():
     dataloader = DataLoader(dataset, batch_size=cfg.per_device_train_batch_size, shuffle=True, drop_last=True)
 
     model = TinyBlockRefinementLM(vocab_size=cfg.vocab_size)
-    pipe = BlockRefinementPipeline(model=model, tokenizer=None)
+    scheduler = BlockRefinementScheduler()
+    pipe = BlockRefinementPipeline(model=model, scheduler=scheduler, tokenizer=None)
 
     optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.learning_rate, weight_decay=cfg.weight_decay)
 

src/diffusers/__init__.py

@@ -347,6 +347,8 @@
     _import_structure["schedulers"].extend(
         [
             "AmusedScheduler",
+            "BlockRefinementScheduler",
+            "BlockRefinementSchedulerOutput",
             "CMStochasticIterativeScheduler",
             "CogVideoXDDIMScheduler",
             "CogVideoXDPMScheduler",
@@ -1121,6 +1123,8 @@
         from .quantizers import DiffusersQuantizer
         from .schedulers import (
             AmusedScheduler,
+            BlockRefinementScheduler,
+            BlockRefinementSchedulerOutput,
             CMStochasticIterativeScheduler,
             CogVideoXDDIMScheduler,
             CogVideoXDPMScheduler,

src/diffusers/pipelines/block_refinement/pipeline_block_refinement.py

@@ -20,6 +20,7 @@
 import torch
 
 from ...callbacks import MultiPipelineCallbacks, PipelineCallback
+from ...schedulers import BlockRefinementScheduler
 from ...utils import BaseOutput
 from ..pipeline_utils import DiffusionPipeline, DiscreteDiffusionPipelineMixin
 
@@ -30,16 +31,6 @@ class BlockRefinementPipelineOutput(BaseOutput):
     texts: Optional[List[str]] = None
 
 
-def _get_num_transfer_tokens(block_length: int, steps: int) -> torch.LongTensor:
-    if steps <= 0:
-        return torch.zeros((0,), dtype=torch.long)
-    base = int(block_length) // int(steps)
-    remainder = int(block_length) % int(steps)
-    out = torch.full((int(steps),), base, dtype=torch.long)
-    out[:remainder] += 1
-    return out
-
-
 class BlockRefinementPipeline(DiffusionPipeline, DiscreteDiffusionPipelineMixin):
     """
     Block-wise iterative refinement pipeline for token generation.
@@ -52,17 +43,19 @@ class BlockRefinementPipeline(DiffusionPipeline, DiscreteDiffusionPipelineMixin)
     """
 
     model: Any
+    scheduler: BlockRefinementScheduler
     tokenizer: Any
 
     _callback_tensor_inputs = ["cur_x", "x0", "x0_p", "transfer_index", "confidence", "active_block"]
 
     def __init__(
         self,
         model: Any,
+        scheduler: BlockRefinementScheduler,
         tokenizer: Optional[Any] = None,
     ):
         super().__init__()
-        self.register_modules(model=model, tokenizer=tokenizer)
+        self.register_modules(model=model, scheduler=scheduler, tokenizer=tokenizer)
 
     @property
     def num_timesteps(self):
@@ -310,6 +303,8 @@ def __call__(
 
         steps = min(int(steps), int(gen_length) // int(minimal_topk))
 
+        self.scheduler.set_timesteps(steps, device=model_device)
+
         num_blocks = (prompt_length + int(gen_length) + int(block_length) - 1) // int(block_length)
         total_length = int(num_blocks) * int(block_length)
 
@@ -333,7 +328,6 @@ def __call__(
 
         prefill_blocks = prompt_length // int(block_length)
         self._num_timesteps = int(steps) * max(int(num_blocks) - int(prefill_blocks), 0)
-        transfer_schedule = _get_num_transfer_tokens(int(block_length), int(steps)).to(device=model_device)
 
         finished = torch.zeros((batch_size,), device=model_device, dtype=torch.bool)
         resolved_attention_mode: str = str(attention_mask_mode)
@@ -362,8 +356,9 @@ def __call__(
                 if finished.all():
                     break
 
-                active_block = cur_x[:, -int(block_length) :] == int(mask_token_id)
-                masks_remaining = active_block.sum() > 0
+                block_tokens = cur_x[:, -int(block_length) :]
+                active_block = block_tokens == int(mask_token_id)
+                masks_remaining = active_block.any()
 
                 if not masks_remaining and not editing_enabled:
                     break
@@ -390,47 +385,26 @@ def __call__(
                     use_multinomial=use_multinomial,
                 )
 
-                # --- Mask-filling transfer ---
-                transfer_index = torch.zeros_like(x0, dtype=torch.bool)
-                if masks_remaining and step_idx < int(steps):
-                    clamped_step = min(step_idx, len(transfer_schedule) - 1)
-                    num_to_transfer = int(transfer_schedule[clamped_step].item())
-
-                    confidence = torch.where(
-                        active_block,
-                        x0_p.to(dtype=torch.float32),
-                        torch.full_like(x0_p, -torch.inf, dtype=torch.float32),
-                    )
-
-                    for b in range(batch_size):
-                        if finished[b]:
-                            continue
-                        high_conf = confidence[b] > float(threshold)
-                        if high_conf.sum().item() >= num_to_transfer:
-                            transfer_index[b] = high_conf
-                        else:
-                            k = min(num_to_transfer, int(active_block[b].sum().item()))
-                            if k > 0:
-                                _, idx = torch.topk(confidence[b], k=k)
-                                transfer_index[b, idx] = True
-
-                # --- Editing transfer (non-mask, non-prompt positions) ---
-                editing_transfer_index = torch.zeros_like(x0, dtype=torch.bool)
-                if editing_enabled:
-                    old_block_tokens = cur_x[:, -int(block_length) :]
-                    editable = (~active_block) & (~prompt_mask_in_block.unsqueeze(0))
-                    editing_conf = torch.where(
-                        editable, x0_p.to(dtype=torch.float32), torch.full_like(x0_p, -torch.inf, dtype=torch.float32)
-                    )
-                    high_conf_edit = editing_conf > float(editing_threshold)
-                    token_changed = x0 != old_block_tokens
-                    editing_transfer_index = high_conf_edit & token_changed & editable
+                scheduler_output = self.scheduler.step(
+                    sampled_tokens=x0,
+                    sampled_probs=x0_p,
+                    timestep=step_idx,
+                    sample=block_tokens,
+                    mask_token_id=int(mask_token_id),
+                    threshold=float(threshold),
+                    editing_threshold=editing_threshold,
+                    minimal_topk=int(minimal_topk),
+                    prompt_mask=prompt_mask_in_block,
+                    generator=generator,
+                    return_dict=True,
+                )
 
+                transfer_index = scheduler_output.transfer_index
+                editing_transfer_index = scheduler_output.editing_transfer_index
                 final_transfer = transfer_index | editing_transfer_index
+
                 if final_transfer.any():
-                    updated = cur_x[:, -int(block_length) :].clone()
-                    updated[final_transfer] = x0[final_transfer]
-                    cur_x[:, -int(block_length) :] = updated
+                    cur_x[:, -int(block_length) :] = scheduler_output.prev_sample
 
                 # Break if no masks remain and no edits were made.
                 if not masks_remaining and not editing_transfer_index.any():

src/diffusers/pipelines/llada2/pipeline_llada2.py

@@ -33,12 +33,15 @@
         >>> from transformers import AutoModelForCausalLM, AutoTokenizer
         >>> from diffusers import LLaDA2Pipeline
 
+        >>> from diffusers import BlockRefinementScheduler
+
         >>> model_id = "inclusionAI/LLaDA2.0-mini"
         >>> model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, torch_dtype=torch.bfloat16)
         >>> tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
         >>> model = model.to("cuda")
+        >>> scheduler = BlockRefinementScheduler()
 
-        >>> pipe = LLaDA2Pipeline(model=model, tokenizer=tokenizer)
+        >>> pipe = LLaDA2Pipeline(model=model, scheduler=scheduler, tokenizer=tokenizer)
         >>> output = pipe(prompt="What is the meaning of life?", gen_length=256)
         >>> print(output.texts[0])
         ```

src/diffusers/schedulers/__init__.py

@@ -40,6 +40,7 @@
 else:
     _import_structure["deprecated"] = ["KarrasVeScheduler", "ScoreSdeVpScheduler"]
     _import_structure["scheduling_amused"] = ["AmusedScheduler"]
+    _import_structure["scheduling_block_refinement"] = ["BlockRefinementScheduler", "BlockRefinementSchedulerOutput"]
     _import_structure["scheduling_consistency_decoder"] = ["ConsistencyDecoderScheduler"]
     _import_structure["scheduling_consistency_models"] = ["CMStochasticIterativeScheduler"]
     _import_structure["scheduling_ddim"] = ["DDIMScheduler"]
@@ -145,6 +146,7 @@
     else:
         from .deprecated import KarrasVeScheduler, ScoreSdeVpScheduler
         from .scheduling_amused import AmusedScheduler
+        from .scheduling_block_refinement import BlockRefinementScheduler, BlockRefinementSchedulerOutput
         from .scheduling_consistency_decoder import ConsistencyDecoderScheduler
         from .scheduling_consistency_models import CMStochasticIterativeScheduler
         from .scheduling_ddim import DDIMScheduler