Support LLaDa-Image - #14815
Support LLaDa-Image#14815lucasruan1618 wants to merge 3 commits into
Conversation
github-actions
Bot
added
documentation
|
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. 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 |
![sergereview[bot]](https://avatars.githubusercontent.com/in/3673521?s=60&v=4){kind=small}
There was a problem hiding this comment.
🤗 Serge says:
Large, well-structured port of LLaDA-Image: four models in transformer_llada_image.py, a pipeline with an output dataclass, docs, dummies and both test layers. The model file follows the Z-Image/Lumina lineage closely (ADALN constant, per-token modulation, pad_sequence batching) and the attention classes correctly use AttentionModuleMixin + dispatch_attention_fn. A few things need attention before merge.
Correctness
- Global monkeypatch of a Transformers registry at import time.
pipeline_llada_image.pymutatestransformers.modeling_rope_utils.ROPE_INIT_FUNCTIONSat module import, so importingdiffusers.pipelines.llada_imagesilently changes behaviour for every Transformers model in the process that resolvesrope_type="default". - Random weight initialization in a load-only model.
LLaDAImageTransformer2DModel.__init__andLLaDAImageQueryAttention.__init__callnn.init.normal_/xavier_uniform_/zeros_. For checkpoint-loaded inference these are pure overhead and can mask a missing key (a layer that fails to load gets plausible random init instead of the usual meta/empty tensor signal). Drop them or move them behind an explicit init hook. t_embeddertimestep embedding is computed in the caller's autocast context, unlike the referenceTimestepEmbedderintransformer_z_image.py, which wraps the frequency computation intorch.amp.autocast(..., enabled=False). Confirm numerics under bf16 autocast.- Undocumented
LLaDAImageSigVQModelconfig args. NoArgs:section at all while sibling models document each config value;[[autodoc]]renders an empty parameter list for a model with 13 config knobs. generate_vq_tokensdepends on remote-code API and magic constants.self.text_encoder.generate_bd_image_logic(...),image_token_offset = 157184,<|reserved_token_N|>names andblock_length=32, steps=8, cfg_scale=2.0are hard-coded against one remote checkpoint with nohasattrguard — a different text encoder gives anAttributeErrordeep in the call. Add an up-front check plus a comment naming where the offset comes from.
Docs
- Both new pages (
api/models/llada_image_transformer2d.md,api/pipelines/llada_image.md) omit the Apache license header that every other page indocs/source/en/api/pipelines/carries (comparelongcat_image.md). docs/source/en/_toctree.ymltitles the model entryLLaDA-Imagewhile surrounding entries use the class name (LatteTransformer3DModel,LongCatImageTransformer2DModel).
Unrelated change
- The
pipeline_utils.pydownload()hunk broadens custom-component allow patterns from{component}/{module}.pyto{component}/*.py— a behaviour change for every pipeline with custom components. It belongs in its own PR with its own test, or at minimum needs justification in the description (currently unmentioned).
Tests
- Good coverage overall (model + pipeline layers, memory mixin, no slow/LoRA tests, per project convention). Concerns:
test_transformers_5_default_rope_is_materializedhand-builds aSimpleNamespaceconfig and a barenn.Modulerotary_emb, then asserts the pipeline wrote the values it itself computed — verifying the pipeline against itself rather than the real Transformers rotary-embedding contract (seetesting.mdon mocks / call-level test doubles).test_vq_conditionedmonkeypatchesgenerate_bd_image_logiconto the text encoder, so the VQ path is exercised only against a stub returning exactly what the pipeline expects. Given the magic157184offset and the token-count validation, that is the part most worth testing against real behaviour.- Six
@pytest.mark.skips on the training/compile mixins are all justified by the list-valued input/output signature, which is fine — but the number of shared tests opted out of is worth a maintainer look, since it stems fromforwardtakinglist[torch.Tensor]rather than a batched tensor.
serge v0.1.0 · model: claude-opus-5 · 21 LLM turns · 27 tool calls · 731.0s · 2153181 in / 57702 out tokens
|
|
||
| # The official LLaDA2 remote model uses this Transformers 4 compatibility entry. Transformers 5 no longer | ||
| # registers it, so make it available before DiffusionPipeline loads the custom text encoder. | ||
| if "default" not in ROPE_INIT_FUNCTIONS: |
There was a problem hiding this comment.
This mutates a global Transformers registry at import time; since diffusers/pipelines/__init__.py lazily exposes llada_image, any import path changes RoPE initialization process-wide for every Transformers model resolving rope_type="default", not just the LLaDA2 encoder loaded here. Please scope it: register the entry inside __init__ (or just before the remote model loads) and restore it afterwards, or pass the init function directly to the rotary module materialized below — you already hold the _default_rope_parameters reference, so line 117 does not need the registry at all.
| self.text_encoder.to(execution_device) | ||
|
|
||
| prompts = [prompt] if isinstance(prompt, str) else prompt | ||
| image_token_offset = 157184 |
There was a problem hiding this comment.
image_token_offset = 157184 is an unexplained magic constant tied to one checkpoint's tokenizer, and generate_bd_image_logic (line 283) is a remote-code method only the official LLaDA2 encoder exposes — any other PreTrainedModel gives an AttributeError mid-loop. Please (a) comment where 157184 comes from (e.g. the id of the first image VQ token in the LLaDA2 vocabulary) so it can be re-derived, and (b) fail early with a clear message, e.g.
if not hasattr(self.text_encoder, "generate_bd_image_logic"):
raise ValueError(
"`generation_mode='vq'` requires the LLaDA2 text encoder, which exposes `generate_bd_image_logic`."
)Same for hard-coded block_length=32, steps=8, cfg_scale=2.0 — surface them as arguments or note they are the published reference defaults.
| nn.Linear(semantic_feat_dim, dim, bias=True), | ||
| ) | ||
|
|
||
| nn.init.normal_(self.semantic_embedder[1].weight, mean=0.0, std=0.02) |
There was a problem hiding this comment.
These four nn.init calls (and nn.init.normal_(self.sigvq_pad_token, ...) on line 430) run on every instantiation including from_pretrained, where the checkpoint immediately overwrites them; beyond wasted work on a 3840-dim model, a random init masks a missing/renamed key as plausible-looking noise instead of an obviously uninitialized tensor. The comparable Z-Image transformer does not init in __init__ — recommend dropping these and the xavier_uniform_/zeros_ pair in LLaDAImageQueryAttention.__init__ (lines 1248-1249), unless training-from-scratch parity requires them.
|
|
||
| def forward(self, timestep: torch.Tensor, hidden_dtype: torch.dtype) -> torch.Tensor: | ||
| half_dim = self.frequency_embedding_dim // 2 | ||
| frequencies = torch.exp( |
There was a problem hiding this comment.
The reference TimestepEmbedder this is ported from (transformer_z_image.py, timestep_embedding) computes the sinusoidal frequencies inside torch.amp.autocast(..., enabled=False). Here cos/sin are float32 explicitly but the surrounding autocast is not disabled, so under bf16 autocast the self.mlp(...) call and downstream fusion can differ from the reference. Confirm the numerics match the original under bf16, or mirror the explicit autocast-disable.
|
|
||
|
|
||
| class LLaDAImageSigVQModel(ModelMixin, ConfigMixin, AttentionMixin): | ||
| r""" |
There was a problem hiding this comment.
Unlike LLaDAImageTransformer2DModel, LLaDAImageQueryFormerModel and LLaDAImageTextProjectionModel, this docstring has no Args: section, so the [[autodoc]] LLaDAImageSigVQModel block in docs/source/en/api/models/llada_image_transformer2d.md renders 13 config parameters with no documentation. Please document image_size, patch_size, codebook_size, codebook_embed_dim, semantic_embed_dim and friends in the sibling classes' style.
| # add custom component files | ||
| allow_patterns += [f"{k}/{f}.py" for k, f in custom_components.items()] | ||
| # Add custom component modules and their local Python dependencies. | ||
| allow_patterns += [f"{folder_name}/*.py" for folder_name in custom_components] |
There was a problem hiding this comment.
This widens the download allow-list for all pipelines with custom components from {component}/{module}.py to {component}/*.py, changing what gets fetched for every existing custom-component repo, not just LLaDA-Image — and the PR description does not mention it. It also discards the mapping's values: _get_custom_components_and_folders returns {component: module_name} precisely so only the declared module is allowed, and pulling every .py in the folder is a broader trust surface for remote-code repos. Please split this into its own PR with a test covering a custom component whose module imports a local sibling, or narrow it to the specific dependency-resolution case needed here.
| @@ -0,0 +1,57 @@ | |||
| # LLaDA-Image | |||
There was a problem hiding this comment.
Missing the Apache license header comment that every other page under docs/source/en/api/pipelines/ starts with (see longcat_image.md). Same for docs/source/en/api/models/llada_image_transformer2d.md.
| # LLaDA-Image | |
| <!--Copyright 2026 The HuggingFace Team. 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. | |
| --> | |
| # LLaDA-Image |
| - local: api/models/latte_transformer3d | ||
| title: LatteTransformer3DModel | ||
| - local: api/models/llada_image_transformer2d | ||
| title: LLaDA-Image |
There was a problem hiding this comment.
Entries in the models section are titled by class name (LatteTransformer3DModel, LongCatImageTransformer2DModel, Krea2Transformer2DModel). LLaDA-Image here also collides with the identically titled pipeline entry at line 623, making the two hard to tell apart in the sidebar.
| title: LLaDA-Image | |
| title: LLaDAImageTransformer2DModel |
| pytest.skip("This regression only affects Transformers 5 and later.") | ||
|
|
||
| components = self.get_dummy_components() | ||
| rotary_emb = torch.nn.Module() |
There was a problem hiding this comment.
Both sides of the assertion come from the code under test: a torch.nn.Module with a SimpleNamespace config stands in for the real rotary embedding, and the test asserts the pipeline wrote back exactly what _default_rope_parameters computes from that same namespace. It stays green if Transformers renames original_inv_freq, changes attention_scaling semantics, or stops using rope_theta — precisely the contract this workaround depends on. Per testing.md, exercise the real component: instantiate an actual Transformers rotary module (or the real text encoder under the version guard) and assert its inv_freq is no longer a meta/zero tensor after pipeline construction.
| def test_vq_conditioned(self): | ||
| pipe = self.get_pipeline().to(torch_device) | ||
|
|
||
| def generate_bd_image_logic(text_encoder, data, block_length, steps, gen_length, cfg_scale): |
There was a problem hiding this comment.
This stub returns input_ids concatenated with arange(gen_length) + 157184, i.e. exactly the layout generate_vq_tokens slices and offsets, so the test cannot fail on a wrong offset, a wrong slice window, or a mismatched token count — the three things the validation at pipeline_llada_image.py:295-299 exists to catch. At minimum add cases driving those error paths (a stub returning too few tokens, and one returning an out-of-codebook id) so the ValueErrors are covered.
This branch has not been deployed
2 participants
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
LLaDAImagePipelineregisters the eight components already described by the publishedmodel_index.json:text_encoderandtokenizerqueryformertext_projectionsigvqtransformervaeschedulerThe pipeline is loaded with the usual
DiffusionPipeline.from_pretrainedpath. The published LLaDA2 text encoder is custom remote code, so users must passtrust_remote_code=Truewhen loading the official checkpoint.text_encoderremains 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, whenguidance_scale > 1, an empty or supplied negative prompt. It then denoises random Flux2 latent patches using classifier-free guidance.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 editing
generation_mode="editing"requires an inputimage. 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.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.compiletraces the model, avoiding module-state mutation during export while retaining the normal eager-mode cache.Diffusers integration details
from_pretrainedoverride.defaultRoPE registry entry required by the checkpoint's custom text encoder when Transformers 5 omits it.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, andgit diff --check: passed.llada_image_text_to_image.png.llada_image_vq_transformers_5.png.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.compiletests 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
self-reviewskill on the diff?Who can review?