[discrete diffusion] Add DiffusionGemma pipeline and schedulers

docs/source/en/_toctree.yml

@@ -647,6 +647,8 @@
  title: Z-Image
  title: Image
  - sections:
+ - local: api/pipelines/diffusion_gemma
+ title: DiffusionGemma
  - local: api/pipelines/llada2
  title: LLaDA2
  title: Text

docs/source/en/api/pipelines/diffusion_gemma.md

@@ -0,0 +1,184 @@
+<!--Copyright 2025 The Google and HuggingFace Teams. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+-->
+
+# DiffusionGemma
+
+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`](https://huggingface.co/google/diffusiongemma-26B-A4B-it).
+
+## Usage
+
+```py
+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](#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:
+
+```py
+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.
+
+## Schedulers
+
+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.
+
+```py
+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:
+
+```py
+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).
+
+### Predictor-corrector sampling
+
+`DiscreteDDIMScheduler` supports the leave-one-out predictor-corrector of [Reparameterizing Uniform Diffusion Models](https://huggingface.co/papers/2605.22765). 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.
+
+```py
+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])
+```
+
+## PEFT adapters
+
+The denoiser is a 🤗 Transformers model, so adapters are loaded through its native [PEFT](https://huggingface.co/docs/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:
+
+```py
+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.
+
+## Static cache and compilation
+
+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](#usage). Drop both the `torch.compile` call and
+`cache_implementation="static"` for a simpler dynamic-cache run.
+
+## Adaptive stopping
+
+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`:
+
+```py
+output = pipe(prompt="Why is the sky blue?", gen_length=256, confidence_threshold=None) # disable adaptive stopping
+```
+
+## Callbacks
+
+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.
+
+```py
+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"],
+)
+```
+
+## DiffusionGemmaPipeline
+[[autodoc]] DiffusionGemmaPipeline
+ - all
+ - __call__
+
+## DiffusionGemmaPipelineOutput
+[[autodoc]] pipelines.DiffusionGemmaPipelineOutput

src/diffusers/__init__.py

@@ -396,11 +396,15 @@
  "DDPMScheduler",
  "DDPMWuerstchenScheduler",
  "DEISMultistepScheduler",
+ "DiscreteDDIMScheduler",
+ "DiscreteDDIMSchedulerOutput",
  "DPMSolverMultistepInverseScheduler",
  "DPMSolverMultistepScheduler",
  "DPMSolverSinglestepScheduler",
  "EDMDPMSolverMultistepScheduler",
  "EDMEulerScheduler",
+ "EntropyBoundScheduler",
+ "EntropyBoundSchedulerOutput",
  "EulerAncestralDiscreteScheduler",
  "EulerDiscreteScheduler",
  "FlowMapEulerDiscreteScheduler",
@@ -573,6 +577,8 @@
  "CosmosTextToWorldPipeline",
  "CosmosVideoToWorldPipeline",
  "CycleDiffusionPipeline",
+ "DiffusionGemmaPipeline",
+ "DiffusionGemmaPipelineOutput",
  "DreamLiteMobilePipeline",
  "DreamLitePipeline",
  "DreamLitePipelineOutput",
@@ -1254,11 +1260,15 @@
  DDPMScheduler,
  DDPMWuerstchenScheduler,
  DEISMultistepScheduler,
+ DiscreteDDIMScheduler,
+ DiscreteDDIMSchedulerOutput,
  DPMSolverMultistepInverseScheduler,
  DPMSolverMultistepScheduler,
  DPMSolverSinglestepScheduler,
  EDMDPMSolverMultistepScheduler,
  EDMEulerScheduler,
+ EntropyBoundScheduler,
+ EntropyBoundSchedulerOutput,
  EulerAncestralDiscreteScheduler,
  EulerDiscreteScheduler,
  FlowMapEulerDiscreteScheduler,
@@ -1410,6 +1420,8 @@
  CosmosTextToWorldPipeline,
  CosmosVideoToWorldPipeline,
  CycleDiffusionPipeline,
+ DiffusionGemmaPipeline,
+ DiffusionGemmaPipelineOutput,
  DreamLiteMobilePipeline,
  DreamLitePipeline,
  DreamLitePipelineOutput,

src/diffusers/pipelines/__init__.py

@@ -273,6 +273,7 @@
  "IFPipeline",
  "IFSuperResolutionPipeline",
  ]
+ _import_structure["diffusion_gemma"] = ["DiffusionGemmaPipeline", "DiffusionGemmaPipelineOutput"]
  _import_structure["dreamlite"] = ["DreamLitePipeline", "DreamLiteMobilePipeline", "DreamLitePipelineOutput"]
  _import_structure["easyanimate"] = [
  "EasyAnimatePipeline",
@@ -717,6 +718,7 @@
  WuerstchenDecoderPipeline,
  WuerstchenPriorPipeline,
  )
+ from .diffusion_gemma import DiffusionGemmaPipeline, DiffusionGemmaPipelineOutput
  from .dreamlite import (
  DreamLiteMobilePipeline,
  DreamLitePipeline,

src/diffusers/pipelines/diffusion_gemma/__init__.py

@@ -0,0 +1,49 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+ DIFFUSERS_SLOW_IMPORT,
+ OptionalDependencyNotAvailable,
+ _LazyModule,
+ get_objects_from_module,
+ is_torch_available,
+ is_transformers_available,
+)
+
+
+_dummy_objects = {}
+_import_structure = {}
+
+
+try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from ...utils import dummy_torch_and_transformers_objects # noqa F403
+
+ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+else:
+ _import_structure["pipeline_diffusion_gemma"] = ["DiffusionGemmaPipeline"]
+ _import_structure["pipeline_output"] = ["DiffusionGemmaPipelineOutput"]
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+ try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+
+ except OptionalDependencyNotAvailable:
+ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
+ else:
+ from .pipeline_diffusion_gemma import DiffusionGemmaPipeline
+ from .pipeline_output import DiffusionGemmaPipelineOutput
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(
+ __name__,
+ globals()["__file__"],
+ _import_structure,
+ module_spec=__spec__,
+ )
+
+ for name, value in _dummy_objects.items():
+ setattr(sys.modules[__name__], name, value)