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 freezes each batch
example once its argmax prediction is stable for stability_threshold steps and its mean per-token entropy falls below
confidence_threshold (0.005, the value used by the released checkpoint). The denoising loop ends once every example
is frozen. 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