Diffusers documentation
BlockRefinementScheduler
BlockRefinementScheduler
BlockRefinementScheduler 管理离散 token 扩散中的分块迭代细化。在每一步中，它会确定置信度最高的 token，并可选择在模型以高置信度预测出不同 token 时，对已确定的 token 进行编辑。
此调度器由 LLaDA2Pipeline 使用。
BlockRefinementScheduler
class diffusers.BlockRefinementScheduler
< source >( block_length: int = 32num_inference_steps: int = 32threshold: float = 0.95editing_threshold: float | None = Noneminimal_topk: int = 1 )
用于分块迭代优化（按置信度提交）的调度器。
在每一步中，调度器会从模型 logits 中采样候选 token，并提交置信度最高的那些。每步提交的 token 数量通过将块长度均匀分配到各个细化步骤来确定。
可选支持编辑：在所有掩码标记解析完成后，如果模型以高于正向 editing_threshold 的置信度预测了不同的标记，则可以替换该标记（None、0.0 或负值会禁用编辑）。
add_noise
< source >( original_samples: torch.LongTensorattention_mask: torch.LongTensorprompt_length: intblock_length: intmask_token_id: intgenerator: torch.Generator | None = None ) → tuple[torch.LongTensor, torch.LongTensor, torch.BoolTensor, torch.BoolTensor]
参数
- original_samples（形状为
(batch_size, seq_len)的torch.LongTensor） — 干净的 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
< source >( step_idx: intmasks_remaining: boolediting_enabled: boolediting_transfer_index: torch.BoolTensorpost_steps: intmax_post_steps: intfinished: torch.BoolTensor ) → bool
参数
- step_idx (
int) — 当前块内的细化步骤索引。 - masks_remaining (
bool) — 块中是否仍包含掩码标记。 - editing_enabled (
bool) — 是否启用编辑模式。 - editing_transfer_index (
torch.BoolTensor) — 本步骤中被编辑的 token。 - post_steps (
int) — 目前已执行的 post-mask 编辑步骤数。 - max_post_steps (
int) — 允许的最大后掩码编辑步数。 - finished (
torch.BoolTensor) — 每个批次的完成标志（来自 EOS 检测）。
返回值
bool
如果应继续细化，则为 True；若要中断，则为 False。
确定当前块的内部细化循环是否应继续。
check_eos_finished
< source >( cur_x: torch.LongTensorsampled_tokens: torch.LongTensorfinal_transfer: torch.BoolTensorfinished: torch.BoolTensoreos_token_id: intmask_token_id: intprompt_length: int ) → torch.BoolTensor
参数
- cur_x（形状为
(batch_size, seq_len)的torch.LongTensor）— 包含截至当前窗口为止所有块的当前完整序列。 - sampled_tokens（形状为
(batch_size, block_length)的torch.LongTensor）— 调度器在此步骤中采样的 Token。 - final_transfer（形状为
(batch_size, block_length)的torch.BoolTensor）— 已提交位置和已编辑位置的组合掩码。 - finished（形状为
(batch_size,)的torch.BoolTensor）— 当前每个批次的完成标志。 - eos_token_id (
int) — EOS token 的 ID。 - mask_token_id (
int) — 掩码 token 的 ID。 - prompt_length (
int) — 序列开头的提示词 token 数量。
返回值
torch.BoolTensor
更新了完成标志。
在提交 EOS 令牌时，更新每个批次的完成标志。
将 block_length 个 token 提交均匀地分布在 num_inference_steps 个步骤中。
步骤
< source >( model_output: torch.Tensortimestep: int | torch.Tensorsample: torch.LongTensormask_token_id: int | None = Nonetemperature: float = 0.0top_p: float | None = Nonetop_k: int | None = Nonesampling_method: str = 'auto'threshold: float | None = Noneediting_threshold: float | None = Noneminimal_topk: int | None = Noneprompt_mask: torch.BoolTensor | None = Nonegenerator: torch.Generator | None = Nonereturn_dict: bool = True )
参数
- model_output（形状为
(batch_size, block_length, vocab_size)的torch.Tensor）— 模型针对当前块输出的原始 logits。 - timestep (
int或torch.Tensor) — 当前步在块细化计划中的索引。 - sample（形状为
(batch_size, block_length)的torch.LongTensor）— 当前块的 token ID（在基于掩码的模式下，包含未提交位置的掩码 token）。 - mask_token_id (
int, 可选) — 用于掩码位置的 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 的置信度阈值；必须为正值才能启用编辑。默认使用 config 中的值。 - minimal_topk (
int, 可选) — 每个步骤中需要提交的最小 token 数量。默认使用配置值。 - prompt_mask (
torch.BoolTensor, 可选) — 形状为(block_length,)的布尔掩码，其中True标记提示（不可编辑）位置。 - generator (
torch.Generator，可选) — 用于采样的随机数生成器。 - return_dict (
bool) — 是否返回BlockRefinementSchedulerOutput或元组。
执行单步精化：从 logits 中采样，提交置信度高的 token，并可选择编辑已有的 token。
BlockRefinementSchedulerOutput
class diffusers.BlockRefinementSchedulerOutput
< source >( prev_sample: torch.LongTensortransfer_index: torch.BoolTensorediting_transfer_index: torch.BoolTensorsampled_tokens: torch.LongTensorsampled_probs: torch.Tensorpred_logits: torch.Tensor )
参数
- prev_sample（形状为
(batch_size, block_length)的torch.LongTensor）— 当前细化步骤之后更新的块 token。 - transfer_index（形状为
(batch_size, block_length)的torch.BoolTensor）— 布尔掩码，指示哪些 token 已被提交（掩码填充）。 - editing_transfer_index（形状为
(batch_size, block_length)的torch.BoolTensor）— 布尔掩码，指示哪些 token 被编辑（非掩码替换）。 - sampled_tokens（形状为
(batch_size, block_length)的torch.LongTensor）— 从模型 logits 中采样的 token ID。 - sampled_probs（形状为
(batch_size, block_length)的torch.Tensor）— 采样 token 的概率。 - pred_logits（形状为
(batch_size, block_length, vocab_size)的torch.Tensor）— 去噪器 logits，用于下一步的自条件化。
用于块细化调度的输出类。