akshan-main/tiny-diffusion-gemma-modular-pipe

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.