[None][feat] Add Qwen-Image-Layered baseline support - #15096
Conversation
📝 WalkthroughWalkthroughThis PR introduces the Qwen-Image-Layered pipeline variant for TensorRT-LLM, enabling per-layer frame-axis positioning with 3D RoPE, TRTLLM masked attention support, and Cache-DiT acceleration. It extends the existing Qwen-Image transformer with configuration flags and adds a complete layered pipeline implementation supporting VAE encoding, scheduler-driven denoising, true CFG with optional normalization, and layered latent packing. ChangesQwen-Image-Layered Pipeline Implementation and Integration
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/_torch/visual_gen/test_qwen_image_layered_registry.py (1)
227-274:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winCoverage gap: add a negative test for layered latent frame-count validation.
Please add one test in
tests/unittest/_torch/visual_gen/test_qwen_image_layered_registry.pythat passes latent input shaped(B, C, F, H, W)withF != 1and assertsValueError. This closes the edge-case contract for latent-image inputs.As per coding guidelines, tests under
tests/**should explicitly call out whether coverage is sufficient, and this path needs follow-up coverage.🤖 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 `@tests/unittest/_torch/visual_gen/test_qwen_image_layered_registry.py` around lines 227 - 274, Add a new unit test that constructs the same QwenImageTransformer2DModel (use the same model creation pattern from test_layered_transformer_forward_sanity, referencing QwenImageTransformer2DModel and model_config) but passes hidden_states with shape (B, C, F, H, W) where F != 1 (e.g. F=2) and assert that calling the model raises ValueError; create a test function name like test_layered_transformer_rejects_multi_frame_latent, set device/dtype the same as the existing test, wrap the model call in a context that expects ValueError, and include any needed encoder_hidden_states/encoder_hidden_states_mask similar to the existing test for parity.Source: Coding guidelines
🧹 Nitpick comments (1)
tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py (1)
467-467: 💤 Low valueRuff B019:
lru_cacheon methods can cause memory leaks.The
@functools.lru_cachedecorator on_compute_condition_freqsretains strong references toself, preventing garbage collection of the entire module instance. This is inherited from the parent class pattern (line 381_compute_video_freqsalso useslru_cache), so it's consistent with the existing design. However, for long-running processes, eachQwenEmbedLayer3DRopeinstance will remain in memory as long as any cached result exists.If you expect many pipeline instances to be created/destroyed, consider using
functools.cache(Python 3.9+) or a WeakMethod-based cache pattern.🤖 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/qwen_image/transformer_qwen_image.py` at line 467, The method decorator `@functools.lru_cache` on _compute_condition_freqs holds strong references to self and can leak QwenEmbedLayer3DRope instances; replace the decorator with a non-leaking alternative (either use functools.cache if you require a global unbounded cache and run on Python 3.9+, or implement a WeakMethod-based cache keyed by weakrefs to self) and mirror the same change for the sibling method _compute_video_freqs; ensure the cache is defined at function-level (or via a WeakKeyDictionary/weakref.WeakMethod helper) so instances can be garbage collected and add a simple compatibility guard if functools.cache is used only when available.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/qwen_image/pipeline_qwen_image_layered.py`:
- Around line 530-537: In the is_latent_image branch in
pipeline_qwen_image_layered (where image is converted to dtype/device and
calculated_width is set using self.vae_scale_factor) add an explicit validation
that rejects multi-frame latent inputs: detect 5D latent tensors (B, C, F, H, W)
or when image.shape[-3] (the frames dimension) exists and ensure F == 1; if F !=
1 raise a clear ValueError like "Layered latent inputs must have exactly one
frame (F==1) for packing/layers=1, got F={F}." Apply the same check to the other
latent-image handling block earlier in this file (the other branch that
currently checks spatial evenness) so both latent paths fail fast with a clear
error instead of triggering an internal reshape failure.
- Around line 438-441: The trimming code that builds generated_ids_trimmed
assumes model_inputs.input_ids and generated_ids have the same batch length;
change the zip call to use zip(model_inputs.input_ids, generated_ids,
strict=True) so mismatched batch sizes raise immediately instead of silently
truncating; update the expression that computes generated_ids_trimmed
(referenced by the variable name generated_ids_trimmed and the operands
model_inputs.input_ids and generated_ids) to use strict zip and run tests to
ensure any mismatch surfaces as an exception.
In `@tests/unittest/_torch/visual_gen/test_qwen_image_registry.py`:
- Around line 21-37: The import block fails ruff-format because names/group
order isn't lint-sorted; reorder the import statements so imported names are
alphabetically sorted inside each from ... import (...) group and the import
groups themselves follow the linter's expected grouping; specifically sort the
first block containing AttentionConfig, DiffusionModelConfig,
create_attention_metadata_state (alphabetize those names), reorder
PIPELINE_REGISTRY and AutoPipeline so AutoPipeline comes first, and ensure the
remaining single-name imports (QwenImagePipeline, QwenImageTransformer2DModel,
QuantConfig, QuantAlgo, QuantAttentionConfig) are placed in the correct group
order per ruff/ruff-format, then rerun pre-commit.
---
Outside diff comments:
In `@tests/unittest/_torch/visual_gen/test_qwen_image_layered_registry.py`:
- Around line 227-274: Add a new unit test that constructs the same
QwenImageTransformer2DModel (use the same model creation pattern from
test_layered_transformer_forward_sanity, referencing QwenImageTransformer2DModel
and model_config) but passes hidden_states with shape (B, C, F, H, W) where F !=
1 (e.g. F=2) and assert that calling the model raises ValueError; create a test
function name like test_layered_transformer_rejects_multi_frame_latent, set
device/dtype the same as the existing test, wrap the model call in a context
that expects ValueError, and include any needed
encoder_hidden_states/encoder_hidden_states_mask similar to the existing test
for parity.
---
Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py`:
- Line 467: The method decorator `@functools.lru_cache` on
_compute_condition_freqs holds strong references to self and can leak
QwenEmbedLayer3DRope instances; replace the decorator with a non-leaking
alternative (either use functools.cache if you require a global unbounded cache
and run on Python 3.9+, or implement a WeakMethod-based cache keyed by weakrefs
to self) and mirror the same change for the sibling method _compute_video_freqs;
ensure the cache is defined at function-level (or via a
WeakKeyDictionary/weakref.WeakMethod helper) so instances can be garbage
collected and add a simple compatibility guard if functools.cache is used only
when available.
🪄 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: 61e476dc-3b53-4b6b-ab47-641ec635bd2d
📒 Files selected for processing (13)
tensorrt_llm/_torch/visual_gen/cache/cache_dit_enablers.pytensorrt_llm/_torch/visual_gen/models/__init__.pytensorrt_llm/_torch/visual_gen/models/qwen_image/__init__.pytensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.pytensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_layered.pytensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.pytensorrt_llm/_torch/visual_gen/modules/attention.pytensorrt_llm/_torch/visual_gen/pipeline_registry.pytests/unittest/_torch/visual_gen/test_attention_integration.pytests/unittest/_torch/visual_gen/test_qwen_image_layered_parity.pytests/unittest/_torch/visual_gen/test_qwen_image_layered_pipeline_config.pytests/unittest/_torch/visual_gen/test_qwen_image_layered_registry.pytests/unittest/_torch/visual_gen/test_qwen_image_registry.py
yumin066
changed the title
87db44d to
b2c327e
Compare
|
/bot run |
|
PR_Github #52931 [ run ] triggered by Bot. Commit: |
|
PR_Github #52931 [ run ] completed with state |
b2c327e to
07fb4de
Compare
125488b to
1e0cf6b
Compare
|
/bot run |
|
PR_Github #54470 [ run ] triggered by Bot. Commit: |
|
PR_Github #54470 [ run ] completed with state
|
1e0cf6b to
9edb365
Compare
|
/bot run |
|
PR_Github #54599 [ run ] triggered by Bot. Commit: |
|
PR_Github #54599 [ run ] completed with state
|
|
/bot help |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break. |
|
/bot reuse-pipeline |
|
PR_Github #54701 [ reuse-pipeline ] triggered by Bot. Commit: |
|
PR_Github #54701 [ reuse-pipeline ] completed with state |
b575bd1 to
52f0e5b
Compare
|
/bot run |
|
PR_Github #61208 [ run ] triggered by Bot. Commit: |
52f0e5b to
054ebd6
Compare
|
/bot run |
|
PR_Github #61219 [ run ] triggered by Bot. Commit: |
|
PR_Github #61208 [ run ] completed with state |
054ebd6 to
1af0acb
Compare
Signed-off-by: Min Yu <minyu@nvidia.com>
|
/bot run |
1af0acb to
8818aa9
Compare
|
PR_Github #61232 [ run ] triggered by Bot. Commit: |
|
PR_Github #61219 [ run ] completed with state |
|
PR_Github #61232 [ run ] completed with state |
yibinl-nvidia
left a comment
yibinl-nvidia
left a comment
There was a problem hiding this comment.
LGTM, left one question.
|
✅ LFS objects already in storage (1 file) — no sync needed. These LFS-tracked files are already present in this repository's LFS storage:
|
9 participants
@coderabbitai summary
Description
This PR adds baseline Qwen-Image-Layered support to VisualGen.
Main changes:
QwenImageLayeredPipeline.Scope note: FP8 SageAttention, Cache-DiT integration, FP8 quantization presets, and parallelization follow-ups were intentionally split out of this PR per review feedback.
Test Coverage
pre-commit run --files $(git diff --name-only upstream/main...HEAD)pytestwas not run locally because the local Python environment does not havepytestinstalled.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.
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.