# DiffusionGemma

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

## 用法

```py
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` 编译（参见 [静态缓存与编译](#static-cache-and-compilation)）；若希望进行更简单的动态缓存运行，可去掉这两项。

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

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

```py
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`、熵边界等）在调度器配置中设置：

```py
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](https://huggingface.co/papers/2605.22765) 中提出的留一法预测-校正器。它通过 `corrector_steps` 次 Gibbs 扫描来细化画布，从含噪边缘分布的单坐标条件分布中重新采样置信度最低的位置，这一过程保持该边缘分布不变，且无需额外训练成本即可提升生成质量。它可直接在已发布的检查点上运行：对于均匀扩散，去噪器与留一法后验在闭式解下是可互换的，因此校正器无需任何重新训练即可恢复其所需的留一法量。

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

```py
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](https://huggingface.co/docs/peft) 集成进行加载，而不是通过 diffusers 的 `load_lora_weights` API。由于该集成与适配器类型无关，相同的调用可以加载 LoRA、DoRA 或任何其他 PEFT 适配器（例如 TRL 的 `SFTTrainer` 的输出）。直接在模型组件上管理适配器：

```py
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](#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`：

```py
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。

```py
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]]

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

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py#L53)

**参数：**

model ([DiffusionGemmaForBlockDiffusion](https://huggingface.co/docs/transformers/v5.22.2/en/model_doc/diffusion_gemma#transformers.DiffusionGemmaForBlockDiffusion))：块扩散去噪器（因果编码器 + 双向解码器，权重共享）。

scheduler （[BlockRefinementScheduler](/docs/diffusers/v0.43.0/zh/api/schedulers/block_refinement#diffusers.BlockRefinementScheduler)、[DiscreteDDIMScheduler](/docs/diffusers/v0.43.0/zh/api/schedulers/discrete_ddim#diffusers.DiscreteDDIMScheduler) 或 [EntropyBoundScheduler](/docs/diffusers/v0.43.0/zh/api/schedulers/entropy_bound#diffusers.EntropyBoundScheduler)）：在每次去噪步骤中提交并重新加噪画布 token 的采样器。

processor ([ProcessorMixin](https://huggingface.co/docs/transformers/v5.22.2/en/main_classes/processors#transformers.ProcessorMixin))：用于应用聊天模板并解码生成 token 的处理器。

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

DiffusionGemma 是一种块扩散（block-diffusion）编码器-解码器模型：因果编码器将干净的提示词（以及任何先前生成的块）读取到 KV 缓存中，而双向解码器则通过交叉注意力机制关注该缓存，对固定大小为 `canvas_length` 个 token 的“画布”进行去噪。生成过程交替执行两个循环：外层是对画布的自回归循环，内层是去噪循环。在去噪循环的每一步中，模型会采样候选 token，通过 [BlockRefinementScheduler](/docs/diffusers/v0.43.0/zh/api/schedulers/block_refinement#diffusers.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__]]

```python
__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)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py#L163)

**参数：**

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](/docs/diffusers/v0.43.0/zh/api/pipelines/diffusion_gemma#diffusers.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](/docs/diffusers/v0.43.0/zh/api/pipelines/diffusion_gemma#diffusers.DiffusionGemmaPipelineOutput) 或 `tuple`

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

使用块扩散生成文本。

示例：
```python
>>> 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]]

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

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/diffusion_gemma/pipeline_output.py#L25)

**参数：**

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

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

DiffusionGemma 块扩散生成的输出类。