diffusers/v0.43.0/zh/api/pipelines/diffusion_gemma.md

DiffusionGemma

DiffusionGemma 是一种块扩散（block-diffusion）编码器-解码器语言模型。因果编码器将干净的提示词（以及任何先前生成的块）读取到 KV 缓存中，而双向解码器则通过交叉注意力机制对该缓存进行关注，从而对一个固定大小为 canvas_length 个 token 的“画布”进行去噪。生成过程交替执行两个循环：外层是针对画布的自回归循环，内层是去噪循环。在去噪循环的每一步中，模型会采样候选 token，通过 BlockRefinementScheduler 在均匀损坏模式下提交置信度最高的 token，并对其余 token 重新加噪。该模型本身位于 transformers 库中，类名为 DiffusionGemmaForBlockDiffusion；发布的检查点为 google/diffusiongemma-26B-A4B-it。

用法

import torch
from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion

from diffusers import BlockRefinementScheduler, DiffusionGemmaPipeline

model_id = "google/diffusiongemma-26B-A4B-it"
model = DiffusionGemmaForBlockDiffusion.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto")
processor = AutoProcessor.from_pretrained(model_id)
scheduler = BlockRefinementScheduler()

pipe = DiffusionGemmaPipeline(model=model, scheduler=scheduler, processor=processor)
pipe.model.model.decoder = torch.compile(pipe.model.model.decoder, mode="reduce-overhead", fullgraph=True)
output = pipe(
    prompt="Why is the sky blue?",
    gen_length=256,
    num_inference_steps=48,
    cache_implementation="static",
)
print(output.texts[0])

num_inference_steps 是每个画布的降噪步数（48 与发布的检查点一致）；步数越少速度越快，但质量越低。cache_implementation="static" 允许解码器使用 cudagraphs 进行 torch.compile 编译（参见 静态缓存与编译）；若希望进行更简单的动态缓存运行，可去掉这两项。

对于多轮或多模态输入，请传入原始的 messages 对话，而不是 prompt。它是一个由 {"role", "content"} 字典组成的列表，采用常见的聊天格式，处理器会将其通过其聊天模板进行处理：

messages = [
    {"role": "user", "content": "Why is the sky blue?"},
]
# or with an image:
messages = [
    {
        "role": "user",
        "content": [
            {"type": "image", "image": image},
            {"type": "text", "text": "Describe this image."},
        ],
    },
]
output = pipe(messages=messages, gen_length=256)

对于单个用户轮次，您可以跳过 messages，并在 prompt 旁边传入 image；处理器会自动将其转换为模型的图像输入。

调度器

调度器（scheduler）是用于对每个画布进行去噪的采样器，并且可以互换：只需替换调度器即可更改采样策略，而无需修改其他任何部分。目前提供三种调度器：

• BlockRefinementScheduler（默认）：在每一步中提交置信度最高的 token（高于 threshold 的 token，外加一个均匀的每步配额），并对其余 token 重新加噪。editing_threshold 还允许其重新编辑已提交的 token。
• DiscreteDDIMScheduler：从均匀腐蚀过程（D3PM）的精确离散后验分布中为每个位置采样。该调度器无需参数，且最后一步会确定性地提交预测的 token。
• EntropyBoundScheduler：提交联合熵保持在 entropy_bound 以下的最低熵位置，从而将大致独立的 token 一起接受。其采样温度从第一步的 t_max（0.8）退火至最后一步的 t_min（0.4），与已发布检查点的采样器保持一致。

from diffusers import DiscreteDDIMScheduler, EntropyBoundScheduler

pipe.scheduler = DiscreteDDIMScheduler()
# or: pipe.scheduler = EntropyBoundScheduler(entropy_bound=0.1)
output = pipe(prompt="Why is the sky blue?", gen_length=256, num_inference_steps=48)
print(output.texts[0])

调度器特定的采样参数（块细化threshold/top_k、熵边界等）在调度器配置中设置：

from diffusers import BlockRefinementScheduler

pipe.scheduler = BlockRefinementScheduler.from_config(pipe.scheduler.config, threshold=0.9)

EntropyBoundScheduler 会在去噪步骤内部对其采样温度（t_max/t_min）进行退火；DiscreteDDIMScheduler 和 BlockRefinementScheduler 则使用传入流水线的固定 temperature（贪心策略下为 0.0）。

预测器-校正器采样

DiscreteDDIMScheduler 支持 Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation 中提出的留一法预测-校正器。它通过 corrector_steps 次 Gibbs 扫描来细化画布，从含噪边缘分布的单坐标条件分布中重新采样置信度最低的位置，这一过程保持该边缘分布不变，且无需额外训练成本即可提升生成质量。它可直接在已发布的检查点上运行：对于均匀扩散，去噪器与留一法后验在闭式解下是可互换的，因此校正器无需任何重新训练即可恢复其所需的留一法量。

校正器扫描被纳入 num_inference_steps 的预算中，而非额外叠加：流水线会减少预测器步数，并将释放出的前向计算用于校正器，因此模型前向调用的总次数仍为 num_inference_steps，预测器-校正器的开销与普通祖先采样相同。

from diffusers import DiscreteDDIMScheduler

pipe.scheduler = DiscreteDDIMScheduler(corrector_steps=2, corrector_k=12)
output = pipe(prompt="Why is the sky blue?", gen_length=256, num_inference_steps=48)
print(output.texts[0])

PEFT 适配器

去噪器是一个 🤗 Transformers 模型，因此适配器通过其原生的 PEFT 集成进行加载，而不是通过 diffusers 的 load_lora_weights API。由于该集成与适配器类型无关，相同的调用可以加载 LoRA、DoRA 或任何其他 PEFT 适配器（例如 TRL 的 SFTTrainer 的输出）。直接在模型组件上管理适配器：

pipe.model.load_adapter("path/to/adapter", adapter_name="sft")  # LoRA, DoRA, ...
pipe.model.set_adapter("sft")
output = pipe(prompt="Why is the sky blue?", gen_length=256)

pipe.model.disable_adapters()  # run the base model
pipe.model.delete_adapter("sft")

适配器保持激活且未合并：DiffusionGemma 将编码器和解码器的基础权重绑定在一起，因此将适配器融合到其中会破坏这两个分支。

静态缓存与编译

Pipeline 会为每个 block 将编码器预填充一次，并写入可复用的缓存（默认为 DynamicCache）。传入 cache_implementation="static" 则改用固定形状的 StaticCache，其形状允许你使用 cudagraphs 对解码器进行 torch.compile 以进一步提速（pipeline 会为每个 step 做标记并克隆 logits，以避免 cudagraph 内存被覆盖）；这正是 Usage 中展示的配置。若希望进行更简单的动态缓存运行，可同时去掉 torch.compile 调用和 cache_implementation="static"。

自适应停止

一个 block 通常会在消耗完所有 num_inference_steps 之前收敛，因此默认情况下，当某个 batch 样本的 argmax 预测在 stability_threshold 步内保持稳定，且其平均逐 token 熵低于 confidence_threshold（0.005，即发布 checkpoint 所使用的值）时，pipeline 会冻结该样本。一旦所有样本均被冻结，去噪循环即结束。在质量相当的情况下，这大约将 decoder 前向传播次数减半，是提升吞吐量最大的单一手段。传入 confidence_threshold=None 可始终执行完整的 num_inference_steps：

output = pipe(prompt="Why is the sky blue?", gen_length=256, confidence_threshold=None)  # disable adaptive stopping

回调

回调函数会在每个去噪步骤之后运行。通过传入 callback_on_step_end_tensor_inputs 来选择包含在 callback_kwargs 中的张量；canvas（当前块 token）和 logits 均可用。从回调函数中返回 {"canvas": ...} 即可替换 canvas。

def on_step_end(pipe, step, timestep, callback_kwargs):
    canvas = callback_kwargs["canvas"]
    # Inspect or modify `canvas` here.
    return {"canvas": canvas}

out = pipe(
    prompt="Why is the sky blue?",
    callback_on_step_end=on_step_end,
    callback_on_step_end_tensor_inputs=["canvas"],
)

DiffusionGemmaPipeline[[diffusers.DiffusionGemmaPipeline]]

diffusers.DiffusionGemmaPipeline[[diffusers.DiffusionGemmaPipeline]]

diffusers.DiffusionGemmaPipeline(model: Any, scheduler: BlockRefinementScheduler | DiscreteDDIMScheduler | EntropyBoundScheduler, processor: Any)

源代码

参数：

model (DiffusionGemmaForBlockDiffusion)：块扩散去噪器（因果编码器 + 双向解码器，权重共享）。

scheduler （BlockRefinementScheduler、DiscreteDDIMScheduler 或 EntropyBoundScheduler）：在每次去噪步骤中提交并重新加噪画布 token 的采样器。

processor (ProcessorMixin)：用于应用聊天模板并解码生成 token 的处理器。

用于 DiffusionGemma 块扩散文本生成的 Pipeline。

DiffusionGemma 是一种块扩散（block-diffusion）编码器-解码器模型：因果编码器将干净的提示词（以及任何先前生成的块）读取到 KV 缓存中，而双向解码器则通过交叉注意力机制关注该缓存，对固定大小为 canvas_length 个 token 的“画布”进行去噪。生成过程交替执行两个循环：外层是对画布的自回归循环，内层是去噪循环。在去噪循环的每一步中，模型会采样候选 token，通过 BlockRefinementScheduler（采用均匀损坏模式，mask_token_id=None）提交置信度最高的 token，并对其余 token 重新加噪。

该模型应为 DiffusionGemmaForBlockDiffusion 的实例，需暴露 forward(input_ids, decoder_input_ids=..., self_conditioning_logits=..., ...) 方法，并返回形状为 [batch, canvas_length, vocab_size] 的 logits，覆盖整个 canvas。请参阅 https://huggingface.co/google/diffusiongemma-26B-A4B-it 上的模型卡片。

调用[[diffusers.DiffusionGemmaPipeline.call]]

__call__(prompt: str | list[str] | None = None, messages: list[dict] | None = None, image: Any | list[Any] | None = None, add_generation_prompt: bool = True, gen_length: int = 256, num_inference_steps: int = 48, temperature: float = 0.0, cache_implementation: str | None = None, eos_early_stop: bool = True, eos_token_id: int | None = None, stability_threshold: int = 1, confidence_threshold: float | None = 0.005, generator: torch.Generator | None = None, output_type: str = 'text', return_dict: bool = True, callback_on_step_end: Callable[[Any, int, int, dict], dict] | PipelineCallback | MultiPipelineCallbacks | None = None, callback_on_step_end_tensor_inputs: list[str] | None = None)

源代码

参数：

prompt (str 或 List[str]，可选)：提示文本，会经过聊天模板包装并由处理器进行分词。请提供此项或 messages 其中之一。

messages (List[Dict], 可选) : 要编码的原始聊天对话，例如 [{"role": "user", "content": "Hello"}] 或多轮/多模态对话。对于超出单个用户轮次的任何情况，请使用此参数代替 prompt。

image (PIL.Image.Image 或 List，可选)：用于与 prompt 配对以进行多模态生成的图像；处理器会将其转换为模型的图像输入。如需更丰富的布局，请将图像内容直接放入 messages 中。

add_generation_prompt (bool，默认为 True)：在应用聊天模板时，是否添加生成提示。

gen_length (int，默认值为 256)：要生成的 token 数量，向上取整为模型 canvas_length 的整数倍。

num_inference_steps (int，默认为 48)：每个画布的降噪步数。

temperature (float，默认为 0.0)：DiscreteDDIMScheduler/BlockRefinementScheduler 的采样温度（0.0 表示贪心采样）；EntropyBoundScheduler 会忽略该参数，并自行退火其温度。其他采样参数（例如 top_k、threshold、t_min/t_max）属于调度器配置，请在调度器上设置，例如 pipe.scheduler = BlockRefinementScheduler.from_config(pipe.scheduler.config, top_k=...)。

cache_implementation (str, 可选) : 设置为 "static" 可针对每个块将编码器预填充一次到持久的 StaticCache 中，并使用固定形状运行解码器，而不是在每一步都重新编码完整序列。固定形状还允许你编译解码器，例如 pipe.model.model.decoder = torch.compile(pipe.model.model.decoder, fullgraph=True)。

eos_early_stop (bool，默认为 True)：当所有序列都输出 EOS 后，是否停止生成更多的画布。

eos_token_id (int, 可选) : 用于提前停止的 EOS token ID。如果未指定，则回退到处理器的分词器。

stability_threshold (int，默认为 1)：示例的 argmax 预测必须连续保持不变的步数，达到该步数后该示例才被计为稳定。仅在设置了 confidence_threshold 时使用。

confidence_threshold (float，可选，默认值为 0.005)：当某个示例已稳定（参见 stability_threshold）且其调度器形状预测 logits 的平均逐 token 熵低于此值时，冻结该示例。一旦所有示例均被冻结，该块的降噪循环即结束。在保持质量一致的前提下可加速生成过程；默认值与发布的检查点相匹配。设置为 None 可始终运行全部 num_inference_steps。

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

output_type (str，默认为 "text") : "text" 将序列解码为字符串（需要处理器）；"seq" 仅返回 token ID。

return_dict (bool，默认为 True)：是否返回 DiffusionGemmaPipelineOutput 而不是元组。

callback_on_step_end (Callable 或 PipelineCallback，可选)：在每次去噪步骤之后运行的回调，其签名为 callback_on_step_end(self, step, timestep, callback_kwargs)。允许的张量键：canvas、logits。

callback_on_step_end_tensor_inputs (List[str], 可选) : 传递给回调的张量键。

返回： DiffusionGemmaPipelineOutput 或 tuple

生成的 token ID（sequences），以及当 output_type="text" 时，解码后的 texts。

使用块扩散生成文本。

示例：

>>> import torch
>>> from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion
>>> from diffusers import BlockRefinementScheduler, DiffusionGemmaPipeline

>>> model_id = "google/diffusiongemma-26B-A4B-it"
>>> model = DiffusionGemmaForBlockDiffusion.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto")
>>> processor = AutoProcessor.from_pretrained(model_id)
>>> scheduler = BlockRefinementScheduler()

>>> pipe = DiffusionGemmaPipeline(model=model, scheduler=scheduler, processor=processor)
>>> output = pipe(prompt="Why is the sky blue?", gen_length=256)
>>> print(output.texts[0])

DiffusionGemmaPipelineOutput[[diffusers.DiffusionGemmaPipelineOutput]]

diffusers.DiffusionGemmaPipelineOutput[[diffusers.DiffusionGemmaPipelineOutput]]

diffusers.DiffusionGemmaPipelineOutput(sequences: torch.LongTensor, texts: list[str] | None = None)

源代码

参数：

sequences (torch.LongTensor，形状为 (batch_size, gen_length))：生成的 token ID（已去除提示部分）。

texts (list[str]，可选)：解码后的文本，每个序列对应一个字符串。仅在 output_type="text" 时设置。

DiffusionGemma 块扩散生成的输出类。