[None][feat] Add Qwen image support - #13449
Merged
Merged
Conversation
svc-trtllm-gh-bot
added
the
Community want to contribute
Contributor
📝 WalkthroughWalkthroughThis pull request adds support for the Qwen-Image text-to-image model to TensorRT-LLM's VisualGen stack, including a transformer implementation, pipeline integration, serving configuration, example scripts, registry updates, and comprehensive parity and registry tests. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant QwenImagePipeline
participant Qwen25VL as Qwen2.5-VL<br/>(Prompt Encoder)
participant Transformer as Qwen Image<br/>Transformer2DModel
participant VAE as AutoencoderKL<br/>Qwen Image
participant Scheduler as FlowMatch Euler<br/>Discrete Scheduler
participant Output as Output<br/>MediaOutput
User->>QwenImagePipeline: infer(prompt, negative_prompt, steps, cfg_scale)
QwenImagePipeline->>Qwen25VL: encode(prompt)
Qwen25VL-->>QwenImagePipeline: text_embeddings, text_mask
alt Classifier-Free Guidance
QwenImagePipeline->>Qwen25VL: encode(negative_prompt)
Qwen25VL-->>QwenImagePipeline: uncond_embeddings, uncond_mask
end
QwenImagePipeline->>QwenImagePipeline: initialize_latents()
loop for each timestep
QwenImagePipeline->>Transformer: forward(latents, text_embeddings, timestep)
Transformer-->>QwenImagePipeline: model_pred
alt with CFG
QwenImagePipeline->>Transformer: forward(latents, uncond_embeddings, timestep)
Transformer-->>QwenImagePipeline: uncond_pred
QwenImagePipeline->>QwenImagePipeline: apply_cfg(model_pred, uncond_pred, cfg_scale)
end
QwenImagePipeline->>Scheduler: step(pred, timestep, latents)
Scheduler-->>QwenImagePipeline: updated_latents
end
QwenImagePipeline->>VAE: decode(latents)
VAE-->>QwenImagePipeline: image (uint8)
QwenImagePipeline->>Output: wrap image
Output-->>User: MediaOutput(image)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 11
🧹 Nitpick comments (2)
tests/unittest/_torch/visual_gen/test_qwen_image_parity.py (2)
23-31: Avoid checking in a developer-local checkpoint default.Baking
/home/scratch.asteiner/...into the repo means everyone else silently skips this suite unless they overrideQWEN_IMAGE_CKPT. Prefer env-var-only discovery or a shared repo convention for model roots.Proposed fix
-_DEFAULT_CKPT = "/home/scratch.asteiner/trtllm-qwen-image/models/qwen-image" +_DEFAULT_CKPT = "" _CKPT_ENV = "QWEN_IMAGE_CKPT" @@ def _ckpt_path() -> Path | None: - path = Path(os.environ.get(_CKPT_ENV, _DEFAULT_CKPT)) + ckpt = os.environ.get(_CKPT_ENV, _DEFAULT_CKPT) + if not ckpt: + return None + path = Path(ckpt) if not (path / "transformer" / "config.json").is_file(): return None return path🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/visual_gen/test_qwen_image_parity.py` around lines 23 - 31, Remove the developer-local hardcoded default in _DEFAULT_CKPT and make _ckpt_path() rely only on the environment variable _CKPT_ENV (or return None when unset/invalid); specifically, replace the current behavior where _DEFAULT_CKPT points to "/home/scratch.asteiner/..." so that _ckpt_path() reads os.environ.get(_CKPT_ENV) and validates that path (checking for "transformer/config.json"), returning None if the env var is not set or the file is missing; update references to _DEFAULT_CKPT to remove or deprecate it so tests only run when QWEN_IMAGE_CKPT is explicitly provided.
411-466: Please verify Qwen-Image also lands in scheduled perf coverage.This slow parity test is good correctness coverage, but it will not catch latency/throughput regressions unless the new model path was also added to the perf sanity definitions and the
test-db/ QA perf lists elsewhere in the PR.As per coding guidelines "If the PR touches performance-sensitive paths ... check whether a perf test entry is present or updated in: (a) tests/integration/test_lists/test-db/l0_perf.yml ... and (b) tests/integration/test_lists/qa/llm_perf_*.yml ...".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/visual_gen/test_qwen_image_parity.py` around lines 411 - 466, The PR added a slow parity test (test_qwen_image_transformer_full_parity / QwenImageTransformer2DModel) but did not add the new model path to scheduled performance coverage; add entries for the Qwen-Image path to the perf test lists so latency/throughput regressions are caught: update tests/integration/test_lists/test-db/l0_perf.yml to include a perf entry referencing test_qwen_image_transformer_full_parity or QwenImageTransformer2DModel and also add corresponding entries in tests/integration/test_lists/qa/llm_perf_*.yml (matching naming and tags used by existing perf jobs), ensuring the test is included in the test-db and QA perf suites and any required markers (e.g., slow or device requirements) are preserved.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@examples/visual_gen/README.md`:
- Around line 279-303: The markdown table has mismatched column counts and a
misplaced footnote that breaks the table; update each argument row (e.g.,
`--model_path`, `--linear_type`, `--true_cfg_scale`, `--negative_prompt`) so
every row has seven pipe-separated cells to match the header, ensure empty
values are represented with a clear placeholder (e.g., "—") rather than dropping
a cell, and move the footnote definition `[^qi]` outside/after the table block
so it does not terminate the table prematurely.
In `@examples/visual_gen/visual_gen_qwen_image.py`:
- Around line 180-181: When load_prompts(args.prompts_file, args.num_prompts)
returns an empty list, subsequent batch processing divides by len(times) and
crashes; add a guard right after prompts = load_prompts(...) to detect if
prompts is empty and either raise a clear error (reject the file) or early-exit,
or ensure avg_time calculation is protected by checking len(times) > 0 before
dividing (e.g., set avg_time = 0 or skip averaging when times is empty). Update
the code paths that use prompts, times, and avg_time (references: load_prompts,
prompts, times, avg_time) so empty prompt files are handled safely and do not
cause a ZeroDivisionError.
- Around line 139-144: In load_prompts, strip each line before checking for
emptiness or comments so indented comment lines (e.g., " # ...") are ignored;
change the comprehension to call line.strip() once into a variable and then
filter with that stripped value (refer to function load_prompts and the variable
prompts) so you test stripped_line and stripped_line.startswith("#") before
adding to prompts.
In `@tensorrt_llm/_torch/visual_gen/models/__init__.py`:
- Line 24: This module lacks the required NVIDIA source header; add the standard
NVIDIA copyright/header block at the top of
tensorrt_llm/_torch/visual_gen/models/__init__.py (the file that exports
QwenImagePipeline) and ensure the header reflects the current year for a
modified file; place the header before any imports and preserve existing license
formatting used across other Python files in the repo.
In `@tensorrt_llm/_torch/visual_gen/models/qwen_image/__init__.py`:
- Around line 4-12: The module docstring in the Qwen-Image package still claims
the transformer is a stub raising NotImplementedError and references
PHASE1_PLAN.md and registry-only scaffolding; update it to accurately describe
the current implementation by removing the "stub" wording and
NotImplementedError claim, noting that the implemented transformer stack and
real pipeline are now re-exported (mentioning the pipeline and transformer stack
by name), and adjust or remove the AutoPipeline/phase-plan references so the
docstring reflects that the real pipeline and sidecars are loaded and
functional.
In `@tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py`:
- Around line 396-403: The negative prompt list is not being expanded to match
the fully expanded prompts batch, causing a CFG-size mismatch with latents in
infer(); fix by expanding/broadcasting negative prompts exactly the same way as
prompts are expanded: read negative_prompt into negatives (singleton or list),
then build negative = [n for n in negatives for _ in range(num_per)] but applied
across the expanded prompt list (i.e., use the same comprehension/replication
logic you used to turn req.prompt into prompts), ensuring the resulting negative
list length equals len(prompts) * num_per (or simply matches len(prompts) after
your expansion) so the CFG batch and latents sizes align in
pipeline_qwen_image.py (use the same variables num_per, prompts,
negative/negatives to locate and update the code).
- Around line 535-536: The log guard dereferences self.rank which may not exist
(BasePipeline sets self.mapping, not rank); change the check to safely handle
missing attribute, e.g. replace the `if self.rank == 0:` guard with `if
getattr(self, "rank", 0) == 0:` or an equivalent `hasattr(self, "rank") and
self.rank == 0` so logger.info("Pipeline total: ...") cannot raise
AttributeError; update the check near the logger.info call in
pipeline_qwen_image.py referencing self.rank/self.mapping/BasePipeline to locate
the change.
In `@tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py`:
- Around line 829-852: The method currently accepts optional timestep and
img_shapes but then uses them unconditionally (timestep.to(...) and
self.pos_embed(...)), which will raise obscure errors; either change the
signature to remove Optional for timestep and img_shapes or add explicit early
validation: check if timestep is None or img_shapes is None at the start of the
function (before timestep.to and self.pos_embed) and raise a clear ValueError
identifying the missing parameter(s). Reference the parameters timestep and
img_shapes and the calls self.time_text_embed(...) and self.pos_embed(... ) in
transformer_qwen_image.py when adding the guard or updating the type hints.
- Around line 51-60: The code uses an assert in get_timestep_embedding to
validate timesteps shape which can be stripped under python -O; replace the
assert with raising a ValueError with the same descriptive message (e.g.,
"Timesteps should be a 1d-array") so invalid inputs fail fast and consistently;
update the validation in get_timestep_embedding to raise ValueError when
len(timesteps.shape) != 1.
In `@tensorrt_llm/_torch/visual_gen/pipeline_registry.py`:
- Line 5: The file pipeline_registry.py is missing the required NVIDIA copyright
header; add the standard multi-line NVIDIA header block at the very top of the
module (before any imports or code) and update the copyright year to the current
year. Ensure the header appears in pipeline_registry.py and precedes the
existing decorator/registration code such as the `@register_pipeline` usages so
the file now complies with the /*** NVIDIA ... ***/ style header required by the
repository guidelines.
In `@tests/unittest/_torch/visual_gen/test_qwen_image_registry.py`:
- Line 95: The pytest.raises call uses a normal string for a regex pattern;
change the match argument to a raw string to make the regex explicit and
consistent with other tests (update the pytest.raises(...) call that currently
uses "Missing keys|Unexpected keys" to use r"Missing keys|Unexpected keys").
Ensure only the match string is modified to a raw string in the test function
located in test_qwen_image_registry.py.
---
Nitpick comments:
In `@tests/unittest/_torch/visual_gen/test_qwen_image_parity.py`:
- Around line 23-31: Remove the developer-local hardcoded default in
_DEFAULT_CKPT and make _ckpt_path() rely only on the environment variable
_CKPT_ENV (or return None when unset/invalid); specifically, replace the current
behavior where _DEFAULT_CKPT points to "/home/scratch.asteiner/..." so that
_ckpt_path() reads os.environ.get(_CKPT_ENV) and validates that path (checking
for "transformer/config.json"), returning None if the env var is not set or the
file is missing; update references to _DEFAULT_CKPT to remove or deprecate it so
tests only run when QWEN_IMAGE_CKPT is explicitly provided.
- Around line 411-466: The PR added a slow parity test
(test_qwen_image_transformer_full_parity / QwenImageTransformer2DModel) but did
not add the new model path to scheduled performance coverage; add entries for
the Qwen-Image path to the perf test lists so latency/throughput regressions are
caught: update tests/integration/test_lists/test-db/l0_perf.yml to include a
perf entry referencing test_qwen_image_transformer_full_parity or
QwenImageTransformer2DModel and also add corresponding entries in
tests/integration/test_lists/qa/llm_perf_*.yml (matching naming and tags used by
existing perf jobs), ensuring the test is included in the test-db and QA perf
suites and any required markers (e.g., slow or device requirements) are
preserved.
🪄 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: 05d6a6f4-29c7-4cb8-ac9f-d04e3352d236
📒 Files selected for processing (12)
docs/source/models/visual-generation.mdexamples/visual_gen/README.mdexamples/visual_gen/serve/README.mdexamples/visual_gen/serve/configs/qwen_image.ymlexamples/visual_gen/visual_gen_qwen_image.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/transformer_qwen_image.pytensorrt_llm/_torch/visual_gen/pipeline_registry.pytests/unittest/_torch/visual_gen/test_qwen_image_parity.pytests/unittest/_torch/visual_gen/test_qwen_image_registry.py
chang-l
added
VisualGen
and removed
Community want to contribute
svc-trtllm-gh-bot
added
the
Community want to contribute
chang-l
reviewed
Apr 24, 2026
chang-l
left a comment
chang-l
left a comment
Collaborator
There was a problem hiding this comment.
Does this PR also support Qwen-Image-2512 : https://huggingface.co/Qwen/Qwen-Image-2512?
Contributor
Author
|
Pushed Summary:
On Qwen-Image-2512: its Hugging Face |
Contributor
Author
|
Update on Qwen-Image-2512 validation: I downloaded Command shape: QWEN_IMAGE_CKPT=/models/qwen-image-2512 python3 -m pytest /tmp/qwen_tests/test_qwen_image_registry.py /tmp/qwen_tests/test_qwen_image_parity.py -qResult: So the transformer/registry parity path is now validated for both |
chang-l
reviewed
Apr 29, 2026
chang-l
left a comment
chang-l
left a comment
Collaborator
There was a problem hiding this comment.
I dropped a few other comments and thanks for your patience!
Also adding @zhenhuaw-me to help review as well
Contributor
Author
|
Pushed Summary:
Validation:
|
yibinl-nvidia
force-pushed
the
add-qwen-image-support
branch
4 times, most recently
from
May 22, 2026 03:53
413b6f9 to
66e8292
Compare
yibinl-nvidia
removed
the
Community want to contribute
Collaborator
|
/bot run |
Collaborator
|
PR_Github #49842 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49842 [ run ] completed with state
|
Collaborator
|
/bot run |
Collaborator
|
PR_Github #49846 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49846 [ run ] completed with state
|
Collaborator
|
/bot run |
Collaborator
|
PR_Github #51180 [ run ] completed with state |
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Collaborator
|
/bot reuse-pipeline |
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Collaborator
|
PR_Github #51313 [ reuse-pipeline ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #51313 [ reuse-pipeline ] completed with state |
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
chang-l
reviewed
Jun 1, 2026
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Collaborator
|
/bot run |
Collaborator
|
PR_Github #51506 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #51506 [ run ] completed with state
|
Collaborator
|
/bot run |
Collaborator
|
PR_Github #51555 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #51555 [ run ] completed with state
|
Collaborator
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #51588 [ run ] triggered by Bot. Commit: |
Collaborator
|
/bot kill |
Collaborator
|
PR_Github #51588 [ run ] completed with state
|
Collaborator
|
/bot run |
chang-l
approved these changes
Jun 2, 2026
Collaborator
|
/bot run |
Collaborator
|
PR_Github #51680 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #51680 [ run ] completed with state |
This was referenced Jun 4, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Validation
Summary by CodeRabbit
New Features
Documentation
Tests