[None][feat] Add GlmImage text-to-image pipeline support to VisualGen#16018
[None][feat] Add GlmImage text-to-image pipeline support to VisualGen#16018jloftin-nv wants to merge 3 commits into
Conversation
📝 WalkthroughWalkthroughThis PR adds support for the ChangesGLM-Image model integration
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
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
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 valueSort
__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
📒 Files selected for processing (13)
docs/source/models/visual-generation.mdexamples/visual_gen/README.mdexamples/visual_gen/configs/glm-image-fp8-1gpu.yamlexamples/visual_gen/models/glm_image.pyexamples/visual_gen/serve/configs/glm_image.ymltensorrt_llm/_torch/visual_gen/models/__init__.pytensorrt_llm/_torch/visual_gen/models/glm_image/__init__.pytensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.pytensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.pytensorrt_llm/_torch/visual_gen/pipeline_registry.pytests/integration/test_lists/test-db/l0_b200.ymltests/unittest/_torch/visual_gen/test_glm_image_pipeline.pytests/unittest/_torch/visual_gen/test_glm_image_transformer.py
070072c to
c86415b
Compare
| ) | ||
|
|
||
|
|
||
| class GlmImageAdaLayerNormContinuous(torch.nn.Module): |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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?
Signed-off-by: Joseph Loftin <jloftin@nvidia.com>
Signed-off-by: Joseph Loftin <jloftin@nvidia.com>
c86415b to
5a4e454
Compare
Summary by CodeRabbit
New Features
Tests
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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.