docs/source/en/api/pipelines/diffusion_gemma.md
DiffusionGemma is a block-diffusion encoder-decoder language model. A causal encoder reads the clean prompt (and any
previously generated blocks) into a KV cache, and a bidirectional decoder denoises a fixed-size "canvas" of
canvas_length tokens by cross-attending to that cache. Generation alternates an outer autoregressive loop over
canvases with an inner denoising loop, where each step samples candidate tokens, commits the most confident ones via
[BlockRefinementScheduler] in uniform corruption mode, and renoises the rest. The model itself lives in
transformers as DiffusionGemmaForBlockDiffusion; the released checkpoint is
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 is the number of denoising steps per canvas (48 matches the released checkpoint); fewer steps are
faster but lower quality. cache_implementation="static" lets the decoder be torch.compile-d with cudagraphs (see
Static cache and compilation); drop both for a simpler dynamic-cache run.
For multi-turn or multimodal inputs, pass a raw messages conversation instead of prompt. It is a list of
{"role", "content"} dicts in the usual chat format, which the processor runs through its chat template:
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)
For a single user turn you can skip messages and pass an image alongside the prompt; the processor turns it into
the model's image inputs automatically.
The scheduler is the sampler that denoises each canvas, and it is interchangeable: swap it to change the sampling strategy without touching anything else. Three schedulers are available:
BlockRefinementScheduler] (default): commits the most confident tokens each step (above threshold, plus an even
per-step quota) and renoises the rest. editing_threshold additionally lets it re-edit already committed tokens.DiscreteDDIMScheduler]: samples each position from the exact discrete posterior of the uniform corruption process
(D3PM). It is parameter free, and the final step deterministically commits the predicted tokens.EntropyBoundScheduler]: commits the lowest-entropy positions whose joint entropy stays under entropy_bound, so
roughly independent tokens are accepted together. It anneals its sampling temperature from t_max (0.8) on the
first step down to t_min (0.4) on the last, matching the released checkpoint's sampler.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])
Scheduler-specific sampling knobs (the block-refinement threshold/top_k, the entropy bound, ...) are set on the
scheduler config:
from diffusers import BlockRefinementScheduler
pipe.scheduler = BlockRefinementScheduler.from_config(pipe.scheduler.config, threshold=0.9)
EntropyBoundScheduler anneals its sampling temperature (t_max/t_min) internally over the denoising steps;
DiscreteDDIMScheduler and BlockRefinementScheduler use the flat temperature passed to the pipeline (0.0 for
greedy).
DiscreteDDIMScheduler supports the leave-one-out predictor-corrector of Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation. It refines the canvas with corrector_steps Gibbs sweeps that resample the least-confident positions from the one-coordinate conditional of the noisy marginal, which leaves that marginal invariant and improves generation at no extra training cost. It works directly on the released checkpoint: for uniform diffusion the denoiser and the leave-one-out posterior are interchangeable in closed form, so the corrector recovers the leave-one-out quantities it needs without any retraining.
The corrector sweeps are folded into the num_inference_steps budget rather than added on top: the pipeline runs fewer predictor steps and spends the freed forwards on correctors, so the total number of model forwards stays num_inference_steps and the predictor-corrector costs the same as plain ancestral sampling.
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])
The denoiser is a 🤗 Transformers model, so adapters are loaded through its native PEFT integration rather than the diffusers load_lora_weights API. Because that integration is adapter-type-agnostic, the same calls load LoRA, DoRA, or any other PEFT adapter (e.g. the output of TRL's SFTTrainer). Manage adapters on the model component directly:
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")
Adapters stay active and unmerged: DiffusionGemma ties the encoder and decoder base weights, so fusing an adapter into them would corrupt both branches.
The pipeline prefills the encoder once per block into a reusable cache (a DynamicCache by default). Passing
cache_implementation="static" uses a fixed-shape StaticCache instead, whose shapes let you torch.compile the
decoder with cudagraphs for a further speedup (the pipeline marks each step and clones the logits so cudagraph memory
is not overwritten); this is the setup shown in Usage. Drop both the torch.compile call and
cache_implementation="static" for a simpler dynamic-cache run.
A block usually converges before all num_inference_steps are spent, so by default the pipeline leaves a block's
denoising loop early once every example's argmax prediction is stable for stability_threshold steps and the mean
per-token entropy falls below confidence_threshold (0.005, the value used by the released checkpoint). This roughly
halves the number of decoder forwards at matched quality and is the largest single throughput lever. Pass
confidence_threshold=None to always run the full num_inference_steps:
output = pipe(prompt="Why is the sky blue?", gen_length=256, confidence_threshold=None) # disable adaptive stopping
Callbacks run after each denoising step. Pass callback_on_step_end_tensor_inputs to select which tensors are
included in callback_kwargs; canvas (the current block tokens) and logits are available. Return {"canvas": ...}
from the callback to replace the 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"],
)
[[autodoc]] DiffusionGemmaPipeline - all - call
[[autodoc]] pipelines.DiffusionGemmaPipelineOutput