Skip to content

[None][feat] Add GlmImage text-to-image pipeline support to VisualGen#16018

Open
jloftin-nv wants to merge 3 commits into
NVIDIA:mainfrom
jloftin-nv:dev-jloftin-glm-image
Open

[None][feat] Add GlmImage text-to-image pipeline support to VisualGen#16018
jloftin-nv wants to merge 3 commits into
NVIDIA:mainfrom
jloftin-nv:dev-jloftin-glm-image

Conversation

@jloftin-nv

@jloftin-nv jloftin-nv commented Jul 7, 2026

Copy link
Copy Markdown

Summary by CodeRabbit

  • New Features

    • Added support for the GLM-Image visual generation model, including a new end-to-end generation pipeline and model components.
    • Added a runnable example, configuration files, and serving settings for single-GPU text-to-image generation.
    • Extended the visual generation docs and README with GLM-Image availability and usage examples.
  • Tests

    • Added coverage for generation quality, batching, supported behavior, and quantization accuracy.
    • Added model-level checks comparing results against a reference implementation.

Description

Added support for GLMImage T2I pipeline along with unit tests and example to VisualGen. Wanted to get this merged then follow up with cache, parallelism, and maybe sage attention support. Note: this does not contain a port of the AR portion, if we want that I can close this and add or follow up with second MR.

Test Coverage

Added test_glm_image_transformer.py and test_glm_image_pipeline.py

PR Checklist

Please review the following before submitting your PR:

  • PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

  • PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

  • Test cases are provided for new code paths (see test instructions)

  • If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

  • Any new dependencies have been scanned for license and vulnerabilities

  • CODEOWNERS updated if ownership changes

  • Documentation updated as needed

  • Update tava architecture diagram if there is a significant design change in PR.

  • The reviewers assigned automatically/manually are appropriate for the PR.

  • [ X] Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

@jloftin-nv
jloftin-nv requested review from a team as code owners July 7, 2026 00:04
@jloftin-nv
jloftin-nv requested review from QiJune and kaiyux July 7, 2026 00:04
@coderabbitai

coderabbitai Bot commented Jul 7, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

📝 Walkthrough

Walkthrough

This PR adds support for the zai-org/GLM-Image text-to-image model in TensorRT-LLM's visual generation module. It introduces GlmImagePipeline and GlmImageTransformer2DModel, registers them in the pipeline/model registries, adds a new VISION_LANGUAGE_ENCODER pipeline component, provides example scripts/configs, updates documentation, and adds unit/integration tests.

Changes

GLM-Image model integration

Layer / File(s) Summary
Documentation and example usage
docs/source/models/visual-generation.md, examples/visual_gen/README.md, examples/visual_gen/models/glm_image.py, examples/visual_gen/configs/glm-image-fp8-1gpu.yaml, examples/visual_gen/serve/configs/glm_image.yml
Docs add GLM-Image to supported models/feature matrix with a footnote on limitations; README and a new CLI script/configs demonstrate default and FP8 1-GPU usage.
Pipeline registration and exports
tensorrt_llm/_torch/visual_gen/models/__init__.py, tensorrt_llm/_torch/visual_gen/models/glm_image/__init__.py, tensorrt_llm/_torch/visual_gen/pipeline_registry.py
Registers GlmImagePipeline in the models package __all__/imports, adds the glm_image package initializer re-exporting pipeline/transformer/attention classes, and adds a VISION_LANGUAGE_ENCODER enum value to PipelineComponent.
GlmImageTransformer2DModel implementation
tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py
Implements quantization-aware linear helpers, feed-forward/embedding/AdaLN modules, fused-QKV attention, transformer blocks, the main model with forward, and weight loading/post-load logic.
GlmImagePipeline implementation
tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py
Implements latent/timestep helpers, image validation, AR prior-token generation, prompt/glyph encoding, component loading, forward/infer orchestration with denoising loop, latent decoding, CFG embedding alignment, and transformer weight loading/initialization.
Unit and integration tests
tests/unittest/_torch/visual_gen/test_glm_image_transformer.py, tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py, tests/integration/test_lists/test-db/l0_b200.yml
Adds transformer structure/forward/HF-comparison tests, pipeline correctness/generation/quantization tests against HuggingFace references, and registers both test modules in the B200 test list.

Estimated code review effort: 4 (Complex) | ~75 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Client
  participant Pipeline as GlmImagePipeline
  participant VLEncoder as Vision-Language Encoder
  participant Transformer as GlmImageTransformer2DModel
  participant VAE

  Client->>Pipeline: forward(prompt, height, width, ...)
  Pipeline->>VLEncoder: generate_prior_tokens(prompt)
  VLEncoder-->>Pipeline: prior_token_ids
  Pipeline->>Pipeline: encode_prompt / align CFG embeds
  Pipeline->>Pipeline: retrieve_timesteps / calculate_shift
  loop denoising steps
    Pipeline->>Transformer: forward_fn(latents, embeds, timestep)
    Transformer-->>Pipeline: noise prediction
  end
  Pipeline->>VAE: decode latents
  VAE-->>Pipeline: image tensor
  Pipeline-->>Client: PipelineOutput
Loading

Possibly related PRs

Suggested reviewers: laikhtewari, zhenhuaw-me, lfr-0531

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title follows the required [None][feat] format and clearly states the main change: GLMImage text-to-image pipeline support in VisualGen.
Description check ✅ Passed The description includes the required Description, Test Coverage, and PR Checklist sections and is specific enough for the PR template.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 6

🧹 Nitpick comments (1)
tensorrt_llm/_torch/visual_gen/models/glm_image/__init__.py (1)

19-19: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Sort __all__ per ruff RUF022.

Static analysis flags this list as unsorted.

♻️ Proposed fix
-__all__ = ["GlmImagePipeline", "GlmImageTransformer2DModel", "GlmImageAttention"]
+__all__ = ["GlmImageAttention", "GlmImagePipeline", "GlmImageTransformer2DModel"]
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/visual_gen/models/glm_image/__init__.py` at line 19, The
__all__ export list in the glm_image module is not sorted and triggers ruff
RUF022. Update the __all__ assignment in the module so the exported symbols are
in alphabetical order, keeping the same names and using the existing __all__
symbol as the place to fix it.

Source: Linters/SAST tools

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py`:
- Around line 1096-1099: The initialization log in _init_transformer uses the
wrong copied model name and should be updated to match the actual GlmImage
pipeline. Change the logger.info message in GlmImagePipeline's _init_transformer
method so it references GlmImage rather than HunyuanVideo1.5, keeping the rest
of the transformer initialization unchanged.
- Line 516: The prompt parameter in the GLM image pipeline is using an implicit
optional type, which Ruff flags. Update the prompt annotation in the relevant
pipeline function to use explicit Python 3.10+ union syntax with None included,
matching the existing Union[str, List[str]] style only if needed elsewhere and
keeping the signature in line with static typing expectations.

In `@tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py`:
- Line 141: The Ruff warnings in the GLM image transformer need cleanup: in
transformer_glm_image, make the out_dim parameter explicitly Optional instead of
a bare None default, rename any unpacked values that are not used so they don’t
trigger unused-local errors, and remove enumerate() in loops where idx is never
referenced. Apply the same typing/linter fixes in the affected methods and
blocks around the listed locations, especially the relevant transformer methods
in transformer_glm_image.
- Around line 436-448: The QKV path in the transformer block is currently
bypassing the auxiliary text-stream projections, so the text stream never uses
the intended HF-mapped weights. Update the forward logic around self.get_qkv in
transformer_glm_image.py to apply add_q_proj, add_k_proj, and add_v_proj to the
encoder/text hidden states, and run norm_added_q and norm_added_k where
appropriate before combining with the image stream. Keep the existing
hidden_states concatenation and downstream QKV flow, but ensure the
text-specific projections are actually used instead of feeding both streams
through the same fused QKV path.

In `@tests/integration/test_lists/test-db/l0_b200.yml`:
- Around line 248-249: Keep
unittest/_torch/visual_gen/test_glm_image_transformer.py in the l0 list for
smoke/HF-parity coverage, but remove
unittest/_torch/visual_gen/test_glm_image_pipeline.py from l0_b200 because it is
too broad for pre-merge CI. Either replace it with a smaller pipeline smoke test
in the same test list or move the full pipeline suite to a heavier/post-merge
list so the l0 tier does not trigger the full generation, HF comparison,
quantization, memory, and E2E coverage.

In `@tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py`:
- Around line 21-39: Set and restore TLLM_DISABLE_MPI around the TensorRT-LLM
imports in test_glm_image_pipeline so the MPI-disabling side effects take effect
before Linear and other tensorrt_llm modules are loaded. Move the env var setup
ahead of the first TensorRT-LLM import, and update the _cleanup_mpi_env fixture
to save the prior value and restore it afterward instead of always removing the
key. Keep the change localized to the module-level import/setup section and the
_cleanup_mpi_env fixture.

---

Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/models/glm_image/__init__.py`:
- Line 19: The __all__ export list in the glm_image module is not sorted and
triggers ruff RUF022. Update the __all__ assignment in the module so the
exported symbols are in alphabetical order, keeping the same names and using the
existing __all__ symbol as the place to fix it.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 3c70e717-9640-41da-9fa8-d574e58d07dc

📥 Commits

Reviewing files that changed from the base of the PR and between a0c406f and 4225964.

📒 Files selected for processing (13)
  • docs/source/models/visual-generation.md
  • examples/visual_gen/README.md
  • examples/visual_gen/configs/glm-image-fp8-1gpu.yaml
  • examples/visual_gen/models/glm_image.py
  • examples/visual_gen/serve/configs/glm_image.yml
  • tensorrt_llm/_torch/visual_gen/models/__init__.py
  • tensorrt_llm/_torch/visual_gen/models/glm_image/__init__.py
  • tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py
  • tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py
  • tensorrt_llm/_torch/visual_gen/pipeline_registry.py
  • tests/integration/test_lists/test-db/l0_b200.yml
  • tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py
  • tests/unittest/_torch/visual_gen/test_glm_image_transformer.py

Comment thread tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py Outdated
Comment thread tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py Outdated
Comment thread tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py Outdated
Comment thread tests/integration/test_lists/test-db/l0_b200.yml
Comment thread tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py Outdated
Comment thread tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py Outdated
Comment thread tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py Outdated
)


class GlmImageAdaLayerNormContinuous(torch.nn.Module):

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not a blocker, cross-cutting across both your VG ports: this AdaLayerNormContinuous / timestep-embedding / FeedForward stack is a near-duplicate of what Qwen-Image already has and what the HunyuanVideo 1.5 PR (#15562) adds again. visual_gen/modules/ has no home for these yet — could the most uniform ones (timestep embeddings first) move into a shared module instead of a third copy, or is there a port-fidelity reason to keep per-model copies? Worth a team-sync call either way.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Mainly just following precedent with each model having repeat/similar AdaNorms. Is it ok if I resolve this comment here and we can do a separate MR unifying timestep and AdaNorms?

Comment thread examples/visual_gen/configs/glm-image-fp8-1gpu.yaml Outdated
Signed-off-by: Joseph Loftin <jloftin@nvidia.com>
Signed-off-by: Joseph Loftin <jloftin@nvidia.com>
Signed-off-by: Joseph Loftin <jloftin@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants