diffusers/pr_14741/en/api/schedulers/block_refinement.md

# BlockRefinementScheduler

The `BlockRefinementScheduler` manages block-wise iterative refinement for discrete token diffusion. At each step it
commits the most confident tokens and optionally edits already-committed tokens when the model predicts a different
token with high confidence.

This scheduler is used by [LLaDA2Pipeline](/docs/diffusers/pr_14741/en/api/pipelines/llada2#diffusers.LLaDA2Pipeline).

For the uniform corruption process, which has no mask token, use [UniformRefinementScheduler](/docs/diffusers/pr_14741/en/api/schedulers/uniform_refinement#diffusers.UniformRefinementScheduler) instead.

This scheduler follows the shared [discrete diffusion scheduler](overview#discrete-diffusion-schedulers) contract: a decreasing
`float` corruption level in `(0, 1]`, `step(model_output, timestep, sample)`, sampling knobs on the config, and a
[DiscreteSchedulerOutput](/docs/diffusers/pr_14741/en/api/schedulers/entropy_bound#diffusers.DiscreteSchedulerOutput) return.

## BlockRefinementScheduler[[diffusers.BlockRefinementScheduler]]

#### diffusers.BlockRefinementScheduler[[diffusers.BlockRefinementScheduler]]

```python
diffusers.BlockRefinementScheduler(block_length: int = 32, num_inference_steps: int = 32, mask_token_id: int | None = None, temperature: float = 0.0, top_p: float | None = None, top_k: int | None = None, sampling_method: str = 'auto', threshold: float = 0.95, editing_threshold: float | None = None)
```

[Source](https://github.com/huggingface/diffusers/blob/vr_14741/src/diffusers/schedulers/scheduling_block_refinement.py#L56)

**Parameters:**

block_length (`int`, defaults to 32) : The block size this scheduler is configured for. Pipelines read it as their default block size; the commit quota itself is taken from the width of `sample`.

num_inference_steps (`int`, defaults to 32) : The number of refinement steps the commit quota is spread across.

mask_token_id (`int`, *optional*) : Token ID marking an undecided position. Required by [step()](/docs/diffusers/pr_14741/en/api/schedulers/block_refinement#diffusers.BlockRefinementScheduler.step); it lives in the config because it is a property of the tokenizer, matching `AmusedScheduler`.

temperature (`float`, defaults to 0.0) : Sampling temperature applied to the logits when drawing candidates. `0.0` takes the argmax. The confidence driving the quota is measured on the unscaled distribution, so `threshold` and `editing_threshold` do not move with this value.

top_p (`float`, *optional*) : Nucleus sampling cutoff.

top_k (`int`, *optional*) : Top-k sampling cutoff.

sampling_method (`str`, defaults to `"auto"`) : One of `"auto"`, `"greedy"`, `"multinomial"`. `"auto"` draws multinomially when `temperature != 0`.

threshold (`float`, defaults to 0.95) : Confidence above which a masked position commits even if the quota is already met.

editing_threshold (`float`, *optional*) : Confidence above which an already-resolved position is overwritten with a different predicted token. Must be positive to enable editing; `None`, `0.0`, or negative disables it.

Scheduler for block-wise iterative refinement (commit-by-confidence).

At each step, the scheduler samples candidate tokens from model logits and commits those with the highest
confidence. The number of tokens to commit per step is determined by evenly distributing the block length across
the number of refinement steps.

Optionally supports editing: after all mask tokens are resolved, tokens can be replaced if the model predicts a
different token with confidence above a positive `editing_threshold` (`None`, `0.0`, or negative disables editing).

This scheduler models the absorbing (masked) corruption process. For the uniform process, where every position
always holds a real token and there is no mask token, use [UniformRefinementScheduler](/docs/diffusers/pr_14741/en/api/schedulers/uniform_refinement#diffusers.UniformRefinementScheduler).

#### add_noise[[diffusers.BlockRefinementScheduler.add_noise]]

```python
add_noise(original_samples: torch.LongTensor, timesteps: float | torch.Tensor, generator: torch.Generator | None = None, **kwargs)
```

[Source](https://github.com/huggingface/diffusers/blob/vr_14741/src/diffusers/schedulers/scheduling_block_refinement.py#L654)

**Parameters:**

original_samples (`torch.LongTensor` of shape `(batch_size, seq_len)`) : Clean token IDs.

timesteps (`float` or `torch.Tensor`) : Masking probability. A scalar applies to the whole batch; a tensor of shape `(batch_size,)` or `(batch_size, 1)` gives a per-example rate.

generator (`torch.Generator`, *optional*) : RNG for reproducibility.

**Returns:** `tuple[torch.LongTensor, torch.BoolTensor]`

the noisy tokens and the boolean mask of noised positions.

Apply the forward (noising) process: replace each position with `mask_token_id` with probability `timesteps`.

`timesteps` is the corruption level in `[0, 1]`, so it *is* the expected masking fraction — `1` masks
everything, `0` masks nothing. The caller chooses it, matching every other `add_noise` in the library.

#### check_block_should_continue[[diffusers.BlockRefinementScheduler.check_block_should_continue]]

```python
check_block_should_continue(step_idx: int, masks_remaining: bool, editing_enabled: bool, editing_transfer_index: torch.BoolTensor, post_steps: int, max_post_steps: int, finished: torch.BoolTensor)
```

[Source](https://github.com/huggingface/diffusers/blob/vr_14741/src/diffusers/schedulers/scheduling_block_refinement.py#L604)

**Parameters:**

step_idx (`int`) : Current refinement step index within this block.

masks_remaining (`bool`) : Whether any mask tokens remain in the block.

editing_enabled (`bool`) : Whether editing mode is active.

editing_transfer_index (`torch.BoolTensor`) : Which tokens were edited in this step.

post_steps (`int`) : Number of post-mask editing steps taken so far.

max_post_steps (`int`) : Maximum allowed post-mask editing steps.

finished (`torch.BoolTensor`) : Per-batch finished flags (from EOS detection).

**Returns:** `bool`

`True` if refinement should continue, `False` to break.

Determine whether the inner refinement loop should continue for the current block.

#### check_eos_finished[[diffusers.BlockRefinementScheduler.check_eos_finished]]

```python
check_eos_finished(cur_x: torch.LongTensor, sampled_tokens: torch.LongTensor, final_transfer: torch.BoolTensor, finished: torch.BoolTensor, eos_token_id: int, mask_token_id: int, prompt_length: int)
```

[Source](https://github.com/huggingface/diffusers/blob/vr_14741/src/diffusers/schedulers/scheduling_block_refinement.py#L548)

**Parameters:**

cur_x (`torch.LongTensor` of shape `(batch_size, seq_len)`) : Current full sequence including all blocks up to the current window.

sampled_tokens (`torch.LongTensor` of shape `(batch_size, block_length)`) : Tokens sampled by the scheduler in this step.

final_transfer (`torch.BoolTensor` of shape `(batch_size, block_length)`) : Combined mask of committed and edited positions.

finished (`torch.BoolTensor` of shape `(batch_size,)`) : Current per-batch finished flags.

eos_token_id (`int`) : EOS token ID.

mask_token_id (`int`) : Mask token ID.

prompt_length (`int`) : Number of prompt tokens at the start of the sequence.

**Returns:** `torch.BoolTensor`

Updated finished flags.

Update per-batch finished flags when EOS tokens are committed.

#### get_num_transfer_tokens[[diffusers.BlockRefinementScheduler.get_num_transfer_tokens]]

```python
get_num_transfer_tokens(block_length: int, num_inference_steps: int)
```

[Source](https://github.com/huggingface/diffusers/blob/vr_14741/src/diffusers/schedulers/scheduling_block_refinement.py#L219)

Evenly distribute `block_length` token commits across `num_inference_steps` steps.

Deprecated: the per-step quota is now computed inline in [step()](/docs/diffusers/pr_14741/en/api/schedulers/block_refinement#diffusers.BlockRefinementScheduler.step) from `step_index`,
so there is no schedule tensor to build.

#### index_for_timestep[[diffusers.BlockRefinementScheduler.index_for_timestep]]

```python
index_for_timestep(timestep: Union[float, torch.FloatTensor], schedule_timesteps: Optional[torch.FloatTensor] = None)
```

[Source](https://github.com/huggingface/diffusers/blob/vr_14741/src/diffusers/schedulers/scheduling_block_refinement.py#L179)

**Parameters:**

timestep (`float` or `torch.FloatTensor`) : The timestep to find the index for.

schedule_timesteps (`torch.FloatTensor`, *optional*) : The schedule timesteps to validate against. If `None`, the scheduler's timesteps are used.

**Returns:** `int`

The index of the timestep.

Get the index for the given timestep.

#### set_begin_index[[diffusers.BlockRefinementScheduler.set_begin_index]]

```python
set_begin_index(begin_index: int = 0)
```

[Source](https://github.com/huggingface/diffusers/blob/vr_14741/src/diffusers/schedulers/scheduling_block_refinement.py#L136)

**Parameters:**

begin_index (`int`, defaults to `0`) : The begin index for the scheduler.

Sets the begin index for the scheduler. This function should be run from pipeline before the inference.

#### set_timesteps[[diffusers.BlockRefinementScheduler.set_timesteps]]

```python
set_timesteps(num_inference_steps: int, device: str | torch.device | None = None, **kwargs)
```

[Source](https://github.com/huggingface/diffusers/blob/vr_14741/src/diffusers/schedulers/scheduling_block_refinement.py#L146)

**Parameters:**

num_inference_steps (`int`) : The number of refinement steps.

device (`str` or `torch.device`, *optional*) : The device the timesteps should be moved to.

Set the discrete timestep grid, as the decreasing corruption level `t` in `[0, 1]`.

The grid matches [set_timesteps()](/docs/diffusers/pr_14741/en/api/schedulers/discrete_ddim#diffusers.DiscreteDDIMScheduler.set_timesteps) — `1.0` down to `1 / num_inference_steps` — so the
discrete schedulers are interchangeable in a pipeline loop. `timesteps` is the public loop variable; the commit
quota is derived from the integer `step_index` so it stays exact for any `num_inference_steps`.

#### step[[diffusers.BlockRefinementScheduler.step]]

```python
step(model_output: torch.Tensor, timestep: float | torch.Tensor, sample: torch.LongTensor, generator: torch.Generator | None = None, return_dict: bool = True, **kwargs)
```

[Source](https://github.com/huggingface/diffusers/blob/vr_14741/src/diffusers/schedulers/scheduling_block_refinement.py#L390)

**Parameters:**

model_output (`torch.Tensor` of shape `(batch_size, block_length, vocab_size)`) : Raw logits from the model for the current block.

timestep (`float` or `torch.Tensor`) : The current corruption level, one entry of `~BlockRefinementScheduler.timesteps`.

sample (`torch.LongTensor` of shape `(batch_size, block_length)`) : Current block token IDs, with `mask_token_id` at the positions still undecided.

generator (`torch.Generator`, *optional*) : RNG for sampling.

return_dict (`bool`) : Whether to return a [DiscreteSchedulerOutput](/docs/diffusers/pr_14741/en/api/schedulers/entropy_bound#diffusers.DiscreteSchedulerOutput) or a plain tuple.

Perform a single refinement step: sample from logits, commit confident masked positions, and optionally edit
already-resolved ones.

#### step_edit[[diffusers.BlockRefinementScheduler.step_edit]]

```python
step_edit(model_output: torch.Tensor, sample: torch.LongTensor, generator: torch.Generator | None = None, return_dict: bool = True)
```

[Source](https://github.com/huggingface/diffusers/blob/vr_14741/src/diffusers/schedulers/scheduling_block_refinement.py#L490)

**Parameters:**

model_output (`torch.Tensor` of shape `(batch_size, block_length, vocab_size)`) : Raw logits from the model for the current block.

sample (`torch.LongTensor` of shape `(batch_size, block_length)`) : Current block token IDs, with every position resolved.

generator (`torch.Generator`, *optional*) : RNG for sampling.

return_dict (`bool`) : Whether to return a [DiscreteSchedulerOutput](/docs/diffusers/pr_14741/en/api/schedulers/entropy_bound#diffusers.DiscreteSchedulerOutput) or a plain tuple.

Overwrite already-resolved positions whose prediction is both different and confident.

This is the post-mask refinement phase: once no `mask_token_id` remains there is nothing left to unmask, so the
step is a confidence-thresholded overwrite rather than a diffusion step and takes **no** `timestep`. It also
does not advance `step_index`, so a pipeline can run as many sweeps as it likes after exhausting the schedule.

## DiscreteSchedulerOutput[[diffusers.DiscreteSchedulerOutput]]

#### diffusers.DiscreteSchedulerOutput[[diffusers.DiscreteSchedulerOutput]]

```python
diffusers.DiscreteSchedulerOutput(prev_sample: LongTensor, pred_original_sample: LongTensor, sampled_probs: Tensor, pred_logits: Tensor, committed_mask: BoolTensor, edited_mask: typing.Optional[torch.BoolTensor] = None)
```

[Source](https://github.com/huggingface/diffusers/blob/vr_14741/src/diffusers/schedulers/scheduling_utils.py#L80)

**Parameters:**

prev_sample (`torch.LongTensor` of shape `(batch_size, sequence_length)`) : Computed sample `(x_{t-1})` of the previous timestep, as token IDs. `prev_sample` should be used as the next model input in the denoising loop.

pred_original_sample (`torch.LongTensor` of shape `(batch_size, sequence_length)`) : The predicted clean token at each position, i.e. the discrete analog of `x0`. Sampled from the model distribution rather than taken as its argmax, unless the sampling configuration is greedy.

sampled_probs (`torch.Tensor` of shape `(batch_size, sequence_length)`) : Probability of each token in `pred_original_sample` under the **unmodified** denoiser distribution, so that confidence thresholds mean the same thing across schedulers and across sampling temperatures.

pred_logits (`torch.Tensor` of shape `(batch_size, sequence_length, vocab_size)`) : The distribution the tokens were actually drawn from, after any shaping the scheduler applies (temperature, top-k, top-p, or a schedule-dependent annealing). Returned because the scheduler owns that shaping, so it is the only holder of this tensor; pipelines that self-condition the denoiser on its own prediction need exactly this and not the raw logits they passed in.

committed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length)`) : Positions that adopted their predicted token in this step. For schedulers that resample every position each step, this is every position.

edited_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length)`, *optional*) : Subset of positions that overwrote an already-committed token. `None` for schedulers with no editing concept.

Base class for the output of a discrete diffusion scheduler's `step` function.

Discrete diffusion operates on token IDs rather than continuous latents, so a step yields the token predictions and
the per-position bookkeeping of which positions were decided, alongside the `prev_sample` hand-off common to every
scheduler.