# BlockRefinementScheduler

`BlockRefinementScheduler` 管理离散 token 扩散中的分块迭代细化。在每一步中，它会确定置信度最高的 token，并在模型以高置信度预测出不同 token 时，可选地修改已确定的 token。

此调度器由 [LLaDA2Pipeline](/docs/diffusers/v0.43.0/zh/api/pipelines/llada2#diffusers.LLaDA2Pipeline) 使用。

## BlockRefinementScheduler[[diffusers.BlockRefinementScheduler]]

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

```python
diffusers.BlockRefinementScheduler(block_length: int = 32, num_inference_steps: int = 32, threshold: float = 0.95, editing_threshold: float | None = None, minimal_topk: int = 1)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/schedulers/scheduling_block_refinement.py#L54)

用于分块迭代优化（按置信度提交）的调度器。

在每一步中，调度器会从模型 logits 中采样候选 token，并提交置信度最高的那些。每步提交的 token 数量通过将块长度均匀分配到各个细化步骤来确定。

可选支持编辑：在所有掩码标记解析完成后，如果模型以高于正向 `editing_threshold` 的置信度预测了不同的标记，则可以替换该标记（`None`、`0.0` 或负值会禁用编辑）。

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

```python
add_noise(original_samples: torch.LongTensor, attention_mask: torch.LongTensor, prompt_length: int, block_length: int, mask_token_id: int, generator: torch.Generator | None = None)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/schedulers/scheduling_block_refinement.py#L456)

**参数：**

original_samples (`torch.LongTensor`，形状为 `(batch_size, seq_len)`)：干净的 token ID。

attention_mask (`torch.LongTensor`，形状为 `(batch_size, seq_len)`)：填充掩码（有效位置为 1，填充位置为 0）。

prompt_length (`int`) : 需要保留且不被掩码的前导提示词元数量。

block_length (`int`) : 用于掩码的块大小。

mask_token_id (`int`) : 用于填充被屏蔽位置的 Token ID。

generator (`torch.Generator`，*可选*)：用于可复现性的随机数生成器。

**返回：** `tuple[torch.LongTensor, torch.LongTensor, torch.BoolTensor, torch.BoolTensor]`

`(noisy, noisy_rev, masked, masked_rev)` — 两个互补的含噪序列及其对应的布尔掩码。

应用前向（加噪）过程，用于半自回归块掩码。

对于提示词之后的每个块，随机选取一部分有效（非填充）token 并用 `mask_token_id` 进行替换。返回两个互补的视图：`noisy` 和 `noisy_rev`，其中一个视图中的掩码位置对应另一个视图中的未掩码位置。

#### 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)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/schedulers/scheduling_block_refinement.py#L412)

**参数：**

step_idx (`int`) : 当前块内的细化步骤索引。

masks_remaining (`bool`) : 块中是否仍包含掩码标记。

editing_enabled (`bool`) : 编辑模式是否处于激活状态。

editing_transfer_index (`torch.BoolTensor`)：本步骤中被编辑的 token。

post_steps (`int`) : 迄今为止已执行的掩码后编辑步骤数。

max_post_steps (`int`) : 允许的最大掩码后编辑步数。

finished (`torch.BoolTensor`)：每个批次的完成标志（来自 EOS 检测）。

**返回：** `bool`

如果应继续细化，则为 `True`；若要中断，则为 `False`。

确定当前块的内部细化循环是否应继续。

#### 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)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/schedulers/scheduling_block_refinement.py#L362)

**参数：**

cur_x (`torch.LongTensor`，形状为 `(batch_size, seq_len)`)：当前完整序列，包含截至当前窗口的所有块。

sampled_tokens (`torch.LongTensor`，形状为 `(batch_size, block_length)`)：调度器在此步骤中采样的 Token。

final_transfer (`torch.BoolTensor`，形状为 `(batch_size, block_length)`)：已提交位置和已编辑位置的组合掩码。

finished (`torch.BoolTensor`，形状为 `(batch_size,)`)：当前每个批次的完成标志。

eos_token_id (`int`) : EOS token 的 ID。

mask_token_id (`int`) : 掩码标记 ID。

prompt_length (`int`) : 序列开头的提示词元数量。

**返回：** `torch.BoolTensor`

更新了完成标志。

在提交 EOS 令牌时，更新每个批次的完成标志。

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

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

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/schedulers/scheduling_block_refinement.py#L102)

将 `block_length` 个 token 提交均匀地分布在 `num_inference_steps` 个步骤中。

#### 步骤[[diffusers.BlockRefinementScheduler.step]]

```python
step(model_output: torch.Tensor, timestep: int | torch.Tensor, sample: torch.LongTensor, 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 | None = None, editing_threshold: float | None = None, minimal_topk: int | None = None, prompt_mask: torch.BoolTensor | None = None, generator: torch.Generator | None = None, return_dict: bool = True)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/schedulers/scheduling_block_refinement.py#L181)

**参数：**

model_output (`torch.Tensor`，形状为 `(batch_size, block_length, vocab_size)`)：模型针对当前块输出的原始 logits。

timestep (`int` 或 `torch.Tensor`)：当前步在块细化计划中的索引。

sample (`torch.LongTensor`，形状为 `(batch_size, block_length)`)：当前块的 token ID（在基于掩码的模式下，未提交位置包含掩码 token）。

mask_token_id (`int`, *optional*) : 用于掩码位置的 Token ID。当值为 `None` 时，调度器将以均匀损坏模式运行：它在内部跟踪已提交的位置（在 `timestep == 0` 时重置），并使用均匀随机 Token 对未提交的位置重新加噪，这与 DiffusionGemma 的块细化采样器相匹配。

temperature (`float`) : 采样温度。

top_p (`float`, *可选*)：核采样截断值。

top_k (`int`，*可选*)：Top-k 采样截断值。

sampling_method (`str`) : 采样方法（`auto`、`greedy`、`multinomial`）。

threshold (`float`, *可选*) : 提交 token 的置信度阈值。默认使用配置值。

editing_threshold (`float`, *可选*) : 编辑非掩码 token 的置信度阈值；必须为正值才能启用编辑。默认使用配置值。

minimal_topk (`int`，*可选*)：每个步骤提交的最小 token 数。默认使用配置值。

prompt_mask (`torch.BoolTensor`, *可选*) : 形状为 `(block_length,)` 的布尔掩码，其中 `True` 标记提示（不可编辑）位置。

generator (`torch.Generator`，*可选*)：用于采样的随机数生成器。

return_dict (`bool`) : 是否返回 `BlockRefinementSchedulerOutput` 或元组。

执行单步精化：从 logits 中采样，提交置信度高的 token，并可选择编辑已有的 token。

## BlockRefinementSchedulerOutput[[diffusers.BlockRefinementSchedulerOutput]]

#### diffusers.BlockRefinementSchedulerOutput[[diffusers.BlockRefinementSchedulerOutput]]

```python
diffusers.BlockRefinementSchedulerOutput(prev_sample: torch.LongTensor, transfer_index: torch.BoolTensor, editing_transfer_index: torch.BoolTensor, sampled_tokens: torch.LongTensor, sampled_probs: torch.Tensor, pred_logits: torch.Tensor)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/schedulers/scheduling_block_refinement.py#L27)

**参数：**

prev_sample (`torch.LongTensor`，形状为 `(batch_size, block_length)`) : 当前细化步骤之后更新的块 token。

transfer_index (`torch.BoolTensor`，形状为 `(batch_size, block_length)`)：布尔掩码，指示哪些 token 已被提交（掩码填充）。

editing_transfer_index (`torch.BoolTensor`，形状为 `(batch_size, block_length)`)：布尔掩码，指示哪些 token 被编辑（非掩码替换）。

sampled_tokens (`torch.LongTensor`，形状为 `(batch_size, block_length)`)：从模型 logits 中采样的 token ID。

sampled_probs (`torch.Tensor`，形状为 `(batch_size, block_length)`)：采样 token 的概率。

pred_logits (`torch.Tensor`，形状为 `(batch_size, block_length, vocab_size)`)：去噪器 logits，用于下一步的自条件化。

用于块细化调度的输出类。