Support LLaDa-Image

[lucasruan1618]

What does this PR do?

This PR adds native LLaDA-Image support to Diffusers. It introduces LLaDAImagePipeline, a single pipeline for text-to-image generation, VQ-conditioned generation, and instruction-guided image editing.

The implementation ports the published LLaDA-Image inference architecture while using Diffusers' standard component registration, serialization, device placement, CPU offloading, group offloading, attention processor, and pipeline loading interfaces.

Pipeline architecture

LLaDAImagePipeline registers the eight components already described by the published model_index.json:

Component	Role in inference
text_encoder and tokenizer	The custom LLaDA2 language model produces prompt representations and, in VQ mode, autoregressively generates image-token IDs.
queryformer	Appends learned image-generation queries to the tokenized prompt before the LLaDA2 backbone runs.
text_projection	Maps the LLaDA2 hidden states into the caption-feature dimension expected by the denoising transformer.
sigvq	Encodes a reference image into semantic VQ features for editing, or maps MLLM-generated VQ IDs into equivalent features for VQ-conditioned generation.
transformer	Runs flow-matching denoising over Flux2 latent patches, conditioned on text and optional SigVQ features/source latents.
vae	Encodes the source image for editing and decodes final Flux2 latents into images.
scheduler	Provides the published flow-matching timestep schedule.

The pipeline is loaded with the usual DiffusionPipeline.from_pretrained path. The published LLaDA2 text encoder is custom remote code, so users must pass trust_remote_code=True when loading the official checkpoint.

import torch

from diffusers import LLaDAImagePipeline


pipe = LLaDAImagePipeline.from_pretrained(
    "inclusionAI/LLaDA-Image",
    dtype=torch.bfloat16,
    trust_remote_code=True,
)
pipe.enable_model_cpu_offload()

text_encoder remains resident because the pipeline directly calls its embedding layer and language backbone; the QueryFormer output must be inserted between those calls. The remaining model components participate in the normal offloading sequence: QueryFormer, text projection, SigVQ, denoising transformer, and VAE.

Supported inference modes

Text-to-image

generation_mode="text" is the default path. The pipeline encodes the positive prompt and, when guidance_scale > 1, an empty or supplied negative prompt. It then denoises random Flux2 latent patches using classifier-free guidance.

image = pipe(
    prompt="A red fox walking through fresh snow, cinematic photography",
    height=1024,
    width=1024,
    num_inference_steps=50,
    guidance_scale=5.0,
    generator=torch.Generator("cuda").manual_seed(42),
).images[0]

VQ-conditioned generation

generation_mode="vq" asks the LLaDA2 image-generation head for VQ token IDs. The pipeline converts those IDs to SigVQ semantic features and supplies them to the denoising transformer alongside prompt features. The frontend VQ grid is capped at 512 pixels on its longest side, matching the reference implementation.

image = pipe(
    prompt="A friendly robot tending a rooftop garden, colorful editorial illustration",
    generation_mode="vq",
    height=1024,
    width=1024,
    num_inference_steps=50,
    guidance_scale=5.0,
).images[0]

Image editing

generation_mode="editing" requires an input image. The pipeline normalizes and encodes that image twice: SigVQ produces semantic image features, and the Flux2 VAE produces source latents. The denoising transformer receives both forms of conditioning with the text instruction.

from diffusers.utils import load_image


source = load_image(
    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png"
)
image = pipe(
    prompt="Turn it into a watercolor painting with a soft blue background",
    image=source,
    generation_mode="editing",
    height=1024,
    width=1024,
    num_inference_steps=50,
    guidance_scale=5.0,
).images[0]

The pipeline validates mode-specific inputs before inference. Text and VQ modes reject image; editing requires it; VQ dimensions must be divisible by 16; and all output dimensions must match the Flux2 VAE and latent-patch scaling requirements.

New public model components

The PR adds and exports four serializable Diffusers models:

• LLaDAImageTransformer2DModel, the variable-resolution denoising transformer.
• LLaDAImageQueryFormerModel, which refines learned generation queries against token embeddings.
• LLaDAImageTextProjectionModel, which connects LLaDA2 hidden states to transformer caption features.
• LLaDAImageSigVQModel, which supports both image-to-VQ encoding and VQ-token-to-semantic-feature lookup.

The transformer preserves the reference model's list-valued output so a batch may contain samples with different spatial shapes. Its RoPE cache is bypassed only while torch.compile traces the model, avoiding module-state mutation during export while retaining the normal eager-mode cache.

Diffusers integration details

• Adds lazy imports and dependency dummy objects for the pipeline and all four models.
• Uses the standard pipeline loader without a checkpoint-specific from_pretrained override.
• Extends the generic custom-component downloader to include sibling Python files, allowing custom components such as LLaDA2 to import their local configuration and fused-MoE modules.
• Restores the standard default RoPE registry entry required by the checkpoint's custom text encoder when Transformers 5 omits it.
• Re-materializes LLaDA2's non-persistent RoPE frequency buffer after Transformers 5 meta-device loading. Without this repair, the buffer contains invalid values and VQ sampling emits text/control tokens outside the SigVQ codebook.
• Keeps the published model-index component names and configuration fields compatible with the official checkpoint.
• Adds model and pipeline API documentation, plus entries in the documentation table of contents.
• Supports save_pretrained / from_pretrained, dtype loading, device maps, CPU/disk offload, model CPU offload, group offload, callbacks, batching, and supported image output types.

Tests

• tests/models/transformers/test_models_transformer_llada_image.py: model serialization, deterministic outputs, dtype loading, CPU/disk/group offload, gradient checkpointing, attention processor behavior, compilation, and direct QueryFormer, projection, and SigVQ forward coverage.
• tests/pipelines/llada_image/test_pipeline_llada_image.py: shared pipeline contracts for loading, batching, callbacks, serialization, dtype handling, accelerator integration, and offloading; plus focused VQ and image-editing tests.

Validation performed:

• pytest tests/models/transformers/test_models_transformer_llada_image.py -q: 43 passed, 12 skipped.
• pytest tests/pipelines/llada_image/test_pipeline_llada_image.py -q: 38 passed, 1 skipped.
• make quality, utils/check_dummies.py, utils/check_copies.py, and git diff --check: passed.
• Manual full-checkpoint runs on an NVIDIA A100 80 GB, bfloat16, model CPU offload, and 50 steps:
  • Text-to-image at 1024×1024, CFG 5.0, seed 42: llada_image_text_to_image.png.
  • VQ-conditioned generation at 512×512, CFG 5.0, seed 43, using Transformers 5.15.1: llada_image_vq_transformers_5.png.
  • Image editing at 512×512, CFG 5.0, seed 44, using the generated fox image as the source: llada_image_edit.png.

The skips cover generic test utilities that cannot operate on the transformer's intentionally list-valued input/output interface, plus AOT package loading, which currently cannot deserialize list-valued inputs. The standard eager, dynamic-shape, and repeated-block torch.compile tests pass.

Self-review

Verdict: READY

No blocking or non-blocking issues remain. The Transformers 5.15.1 VQ failure was traced to the official remote text encoder's non-persistent RoPE buffer being initialized on the meta device during sharded loading. The corrected buffer reproduces the Transformers 4.57.6 reference tokens exactly in the reduced comparison, and the full 512×512 Transformers 5 output is byte-identical to the reference-runtime output.

No likely-dead inference paths were found. The text, VQ, and editing paths are all traced from LLaDAImagePipeline.__call__ and covered by tests. The official configuration schema was checked against every new model constructor, and the port preserves upstream model math apart from the compile-safe RoPE cache guard and Diffusers device/offloading integration.

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Did you read our philosophy doc?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link if applicable.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?
• Are you the author (or part of the team) of the model/pipeline?

Who can review?

• Pipelines and models: @yiyixuxu and @asomoza

[github-actions]

Hi @lucasruan1618, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice.

Once the PR links an issue (or gets the no-issue-needed label), you can ignore this message — it stays here as a comment, but it no longer applies.