README.md

---
library_name: diffusers
tags:
- modular-diffusers
- diffusers
- diffusion-gemma
- text-to-image
- modular-diffusers
- diffusers
- diffusion-gemma
- text-to-image
---
This is a modular diffusion pipeline built with 🧨 Diffusers' modular pipeline framework.

**Pipeline Type**: DiffusionGemmaBlocks

**Description**: Modular blocks for DiffusionGemma block-diffusion text generation.
- `text_encoder` applies the chat template and tokenizes the prompt
- `prepare_generation` sizes the canvases, creates the KV cache and resolves EOS
- `set_timesteps` splits the forward budget and configures the scheduler
- `denoise` generates the text canvas by canvas
- `decode` trims at EOS and decodes the token IDs into text

This pipeline uses a 5-block architecture that can be customized and extended.

## Example Usage

[TODO]

## Pipeline Architecture

This modular pipeline is composed of the following blocks:

1. **text_encoder** (`DiffusionGemmaTextEncoderStep`)
   - Text encoder step that applies the chat template to a `prompt` or a raw `messages` conversation and tokenizes it into the prompt token IDs consumed by the encoder prefill
2. **prepare_generation** (`DiffusionGemmaPrepareGenerationStep`)
   - Prepare step that sizes the generation into canvases, creates the encoder KV cache, and resolves the EOS token used for early stopping and trimming
3. **set_timesteps** (`DiffusionGemmaSetTimestepsStep`)
   - Step that splits the per-canvas forward budget into predictor and corrector steps and configures the scheduler's refinement schedule
4. **denoise** (`DiffusionGemmaDenoiseStep`)
   - Canvas denoise step that iterates over canvases.
5. **decode** (`DiffusionGemmaDecodeStep`)
   - Decode step that trims each generated sequence at its first EOS token and decodes the token IDs into text with the processor 

## Model Components

1. processor (`ProcessorMixin`)
2. model (`DiffusionGemmaForBlockDiffusion`)
3. scheduler (`BlockRefinementScheduler`) 

## Input/Output Specification

**Inputs:**

- `prompt` (`str`, *optional*): Prompt text, wrapped in a chat template and tokenized
- `messages` (`list`, *optional*): A raw chat conversation to encode instead of `prompt`, e.g. `[{"role": "user", "content": "Hello"}]` or a multi-turn / multimodal conversation.
- `image` (`Image | ndarray | Tensor | list | list | list`, *optional*): Image(s) to pair with `prompt` for multimodal generation. For richer layouts, put the image content directly in `messages`.
- `add_generation_prompt` (`bool`, *optional*, defaults to `True`): Whether to add the generation prompt when applying the chat template.
- `gen_length` (`int`, *optional*, defaults to `256`): Number of tokens to generate, rounded up to a multiple of the model's `canvas_length`.
- `cache_implementation` (`str`, *optional*): Set to `"static"` to use a fixed-shape `StaticCache` so the decoder can be compiled.
- `eos_token_id` (`int`, *optional*): EOS token ID for early stopping. Falls back to the processor's tokenizer.
- `num_inference_steps` (`int`, *optional*, defaults to `48`): Number of denoising steps per canvas, i.e. the per-canvas budget of model forwards.
- `eos_early_stop` (`bool`, *optional*, defaults to `True`): Whether to stop generating further canvases once every sequence has emitted EOS.
- `generator` (`Generator`, *optional*): Torch generator for deterministic generation.
- `temperature` (`float`, *optional*, defaults to `0.0`): Sampling temperature (`0.0` is greedy). Other sampling knobs are scheduler config.
- `stability_threshold` (`int`, *optional*, defaults to `1`): Consecutive steps the argmax prediction must be unchanged for a canvas to count as stable. Only used when `confidence_threshold` is set.
- `confidence_threshold` (`float`, *optional*, defaults to `0.005`): Leave the refinement loop early once every example is stable and the mean per-token entropy is below this value. Set to `None` to always run all steps.

**Outputs:**

- `prompt_ids` (`LongTensor`): Tokenized prompt of shape `(batch_size, prompt_length)`.
- `prompt_attention_mask` (`LongTensor`): Attention mask for `prompt_ids`.
- `multimodal_inputs` (`dict`): Image tensors the processor produced for the encoder prefill.
- `canvas_length` (`int`): The model's canvas length, i.e. the number of tokens denoised per block.
- `num_canvases` (`int`): Number of canvases to generate.
- `past_key_values` (`object`): The encoder KV cache reused across canvases and denoising steps.
- `eos_token_id` (`int`): The resolved EOS token ID (user-provided or from the processor's tokenizer).
- `finished` (`Tensor`): Per-example flags marking sequences that already emitted EOS.
- `predictor_steps` (`int`): Predictor steps run per canvas.
- `corrected_steps` (`int`): Number of leading predictor steps that also run corrector sweeps.
- `corrector_steps` (`int`): Corrector sweeps run after each of the first `corrected_steps` predictor steps.
- `decoder_position_ids` (`LongTensor`): Position IDs of the canvas tokens, continuing the running sequence.
- `decoder_attention_mask_mapping` (`object`): The decoder attention mask mapping built over the populated cache plus the canvas.
- `canvas` (`LongTensor`): The noisy canvas of shape `(batch_size, canvas_length)` being denoised.
- `sequences` (`LongTensor`): The generated token IDs of shape `(batch_size, generated_length)`.
- `texts` (`list`): The decoded generated text, one string per prompt.