docs/source/en/model_doc/diffusion_gemma.md
This model was contributed to Hugging Face Transformers on 2026-06-10.
DiffusionGemma is engineered to reduce the sequential bottlenecks of standard causal language models. It employs an encoder-decoder architecture specifically optimized for inference speed.
The encoder operates in a prefill capacity, processing the initial prompt and generating the KV cache. The decoder then utilizes bidirectional attention to process an input block (a 'canvas') of tokens, accessing the cached context via cross-attention.
During inference, DiffusionGemma leverages multi-canvas sampling. Rather than generating one token at a time, the model iteratively denoises a full block of tokens using a diffusion sampler. Once a canvas is fully denoised, it is processed by the encoder and appended to the KV cache, after which the model generates the next canvas. This block-autoregressive approach facilitates text generation at higher speeds.
You can find the model card and checkpoint here.
Despite it being a text diffusion model and having a custom generation loop, most of the interface is shared with other model that can generate text with .generate(). If you're using another transformers model in your app, you should be able to directly replace it with this model.
Common caveats:
use_cache. It always uses a KV cache;top_k won't be available at release day, but will be added over time if they are compatible with text diffusion.from transformers import DiffusionGemmaForBlockDiffusion, AutoProcessor
model = DiffusionGemmaForBlockDiffusion.from_pretrained(
"google/diffusiongemma-26B-A4B-it", device_map="auto",
)
processor = AutoProcessor.from_pretrained("google/diffusiongemma-26B-A4B-it")
messages = [
{
"role": "user", "content": [
{"type": "image", "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"},
{"type": "text", "text": "What is shown in this image?"},
]
},
]
inputs = processor.apply_chat_template(
messages,
tokenize=True,
return_dict=True,
return_tensors="pt",
add_generation_prompt=True,
).to(model.device)
input_len = inputs["input_ids"].shape[-1]
# Set `cache_implementation="static"` in `generate` to trigger `torch.compile`.
# Compilation is much faster, after warming up!
output = model.generate(**inputs, max_new_tokens=256)
print(processor.decode(output.sequences[0][input_len:], skip_special_tokens=True))
Like other models that can generate text, you can set a streamer class to stream text. Unlike other models, DiffusionGemma generates intermediate drafts before the final text. You can visualize them with TextDiffusionStreamer
from transformers import TextDiffusionStreamer
# (... copy from the example above, up to the `generate` call)
streamer = TextDiffusionStreamer(tokenizer=processor.tokenizer)
model.generate(**inputs, max_new_tokens=256, streamer=streamer)
[[autodoc]] DiffusionGemmaTextConfig
[[autodoc]] DiffusionGemmaConfig
[[autodoc]] DiffusionGemmaGenerationOutput
[[autodoc]] DiffusionGemmaGenerationMixin
[[autodoc]] DiffusionGemmaGenerationConfig
[[autodoc]] EntropyBoundSamplerConfig
[[autodoc]] EntropyBoundSampler
[[autodoc]] StableAndConfidentStoppingCriteria
[[autodoc]] LinearTemperatureScheduleLogitsProcessor
[[autodoc]] DiffusionGemmaPreTrainedModel - forward
[[autodoc]] DiffusionGemmaModel - forward
[[autodoc]] DiffusionGemmaEncoderModel - forward
[[autodoc]] DiffusionGemmaEncoderTextModel - forward
[[autodoc]] DiffusionGemmaDecoderModel - forward
[[autodoc]] DiffusionGemmaForBlockDiffusion - forward