[TRTLLM-14135][feat] Add Qwen-Image-Edit-2511 support by yibinl-nvidia · Pull Request #16095 · NVIDIA/TensorRT-LLM

yibinl-nvidia · 2026-07-08T02:17:34Z

Summary by CodeRabbit

New Features
- Added support for the Hugging Face Qwen-Image-Edit-2511 model for image editing.
- Added a command-line example for editing images with prompts, generation settings, and optional configuration files.
- Added BF16 and FP4 single-GPU example configurations.
- Added support for image-editing pipeline registration and automatic checkpoint detection.
Bug Fixes
- Improved CUDA graph output handling to prevent incorrect tensor reuse during repeated captures.
Documentation
- Updated supported-model lists, feature matrices, and usage examples for Qwen-Image-Edit.

Description

Add Qwen-Image-Edit-2511 support to the TRT-LLM visual generation stack using the existing Qwen-Image transformer, VAE, scheduler, processor, and visual-gen executor path. The implementation adds an image-edit example entrypoint/config, routes Qwen-Image-Edit checkpoints to the edit pipeline, and keeps the architecture aligned with the existing Qwen-Image building blocks.

Supported Features

Feature	Status	Notes
FP8	Supported	E2E image-edit generation was validated against the BF16 baseline.
FP4	Supported with selective quantization	Recommended path uses the FP4 config with sensitive modules excluded from FP4 dynamic quantization; full FP4 quantization showed visible quality regression.
CUDA Graph	Supported	CUDA graph execution is enabled for Qwen-Image-Edit and handles true-CFG positive/negative graph outputs.
torch.compile	Supported	Works for the Qwen-Image-Edit pipeline and can be used together with CUDA graph and CFG parallelism.
CFG parallelism	Supported	`cfg_size=2` splits positive and negative CFG branches across two ranks. Eager CFG=1 vs CFG=2 matched exactly in local validation; with `torch.compile`, LPIPS drift was observed around 0.096 and considered acceptable for this image generation path.

See visuals

Test Coverage

Added Qwen-Image-Edit registry/unit coverage.
Added Qwen-Image-Edit visual-gen example coverage.
Ran local E2E Qwen-Image-Edit generation for BF16, FP8, FP4, CUDA graph, torch.compile, and CFG-parallel configurations.

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.
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.

coderabbitai · 2026-07-18T03:14:26Z

📝 Walkthrough

Walkthrough

Changes

Qwen Image Edit support

Layer / File(s)	Summary
Pipeline registration and contracts `tensorrt_llm/_torch/visual_gen/models/...`	Registers and exports `QwenImageEditPlusPipeline` with generation and warmup defaults.
Image conditioning and latent preparation `tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py`	Adds image loading, prompt embedding, VAE encoding, latent packing, and request adaptation.
Denoising and runtime integration `tensorrt_llm/_torch/visual_gen/models/qwen_image/...`, `tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py`	Implements denoising, CFG handling, output decoding, CUDA graph output retention, and shape-aware graph keys.
Example CLI and engine configurations `examples/visual_gen/models/qwen_image_edit.py`, `examples/visual_gen/configs/qwen-image-edit`, `examples/visual_gen/README.md`	Adds the image-edit CLI, BF16/NVFP4 configurations, and usage commands.
Documentation and validation `docs/source/models/.md`, `tests/unittest/_torch/visual_gen/`, `tests/integration/defs/examples/visual_gen/test_visual_gen.py`	Documents model capabilities and validates registration, compilation, and end-to-end output generation.

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

Sequence Diagram(s)

sequenceDiagram
  participant CLI as qwen_image_edit.py
  participant VisualGen
  participant Pipeline as QwenImageEditPlusPipeline
  participant Transformer as QwenImageTransformer2DModel
  participant Scheduler
  participant VAE
  CLI->>VisualGen: generate image-edit request
  VisualGen->>Pipeline: infer request
  Pipeline->>Transformer: denoise conditioned latents
  Transformer-->>Pipeline: noise prediction
  Pipeline->>Scheduler: update latents
  Pipeline->>VAE: decode output latents
  VAE-->>Pipeline: edited image
  Pipeline-->>VisualGen: PipelineOutput

Suggested reviewers: qijune

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 32.43% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
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.
Title check	✅ Passed	The title clearly names the main change: adding Qwen-Image-Edit-2511 support.
Description check	✅ Passed	The description covers what changed, supported features, test coverage, and checklist completion.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

coderabbitai

Actionable comments posted: 4

🧹 Nitpick comments (1)

tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py (1)
14-14: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Apply the repository’s Python 3.10 typing conventions.

tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py#L14-L14: replace legacy typing aliases with built-in generics and |.

tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py#L338-L338: type req and the PipelineOutput return.

tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py#L980-L985: type the runner and nested key callback.

As per coding guidelines, “Annotate every function” and “prefer built-in generic types and |.”
🤖 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/pipeline_qwen_image_edit.py`
at line 14, Update
tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py:14
to remove legacy typing aliases and use Python 3.10 built-in generics and |
unions; annotate req and the PipelineOutput return in the relevant pipeline
function at line 338. Also update
tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py:980-985
to annotate the runner and nested key callback, following the repository’s
“Annotate every function” convention. These changes are required at all three
listed sites.
Source: Coding guidelines

🤖 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 `@examples/visual_gen/configs/qwen-image-edit-2511-fp4-1gpu.yaml`:
- Around line 18-19: Remove the “trtllm-serve” sharing claim from the comment
near the Qwen-Image-Edit FP4 configuration in
examples/visual_gen/configs/qwen-image-edit-2511-fp4-1gpu.yaml:18-19, rewording
it to describe only the supported offline visual-generation usage. The
references in docs/source/models/supported-models.md:182 and
docs/source/models/visual-generation.md:57 require no direct change unless
needed to keep documentation consistent.

In
`@tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py`:
- Around line 524-547: In the do_true_cfg path, clone the conditional noise_pred
prefix immediately after slicing it and before the second self.transformer call,
so the negative CUDA-graph replay cannot overwrite or alias the conditional
prediction used by CFG. Keep the existing negative prediction, combination, and
norm-rescaling logic unchanged.
- Around line 338-365: Update infer so num_images_per_prompt is not expanded
into the prompts list, which currently causes forward to reject the batch. Pass
the count separately through infer/forward and use it to produce the effective
latent and output batch for each prompt, while preserving prompt batching;
alternatively, reject values greater than one during request validation if the
pipeline cannot support them.

In `@tests/integration/defs/examples/visual_gen/test_visual_gen.py`:
- Around line 1846-1898: Add type annotations throughout the new tests: in
tests/integration/defs/examples/visual_gen/test_visual_gen.py:1846-1898,
annotate fixture parameters and add -> None to test_qwen_image_edit_example; in
tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py:231-272, add
the helper return type, -> None on tests, and record_call(self) -> None; in
tests/unittest/_torch/visual_gen/test_qwen_image_registry.py:84-96, annotate
fixture parameters and add -> None to the new tests.

---

Nitpick comments:
In
`@tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py`:
- Line 14: Update
tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py:14
to remove legacy typing aliases and use Python 3.10 built-in generics and |
unions; annotate req and the PipelineOutput return in the relevant pipeline
function at line 338. Also update
tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py:980-985
to annotate the runner and nested key callback, following the repository’s
“Annotate every function” convention. These changes are required at all three
listed sites.

🪄 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: 5fdd62ff-54ce-4cd4-bfbd-abc1d6bf36dd

📥 Commits

Reviewing files that changed from the base of the PR and between 29bd103 and bbf4efa.

📒 Files selected for processing (14)

docs/source/models/supported-models.md
docs/source/models/visual-generation.md
examples/visual_gen/README.md
examples/visual_gen/configs/qwen-image-edit-2511-bf16-1gpu.yaml
examples/visual_gen/configs/qwen-image-edit-2511-fp4-1gpu.yaml
examples/visual_gen/models/qwen_image_edit.py
tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py
tensorrt_llm/_torch/visual_gen/models/__init__.py
tensorrt_llm/_torch/visual_gen/models/qwen_image/__init__.py
tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py
tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py
tests/integration/defs/examples/visual_gen/test_visual_gen.py
tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py
tests/unittest/_torch/visual_gen/test_qwen_image_registry.py

yibinl-nvidia · 2026-07-20T22:34:25Z

/bot run

tensorrt-cicd · 2026-07-20T22:39:57Z

PR_Github #60479 [ run ] triggered by Bot. Commit: 35e7fd7 Link to invocation

tensorrt-cicd · 2026-07-21T03:01:45Z

PR_Github #60479 [ run ] completed with state SUCCESS. Commit: 35e7fd7
/LLM/main/L0_MergeRequest_PR pipeline #48810 completed with status: 'SUCCESS'

CI Report

Link to invocation

yibinl-nvidia · 2026-07-21T22:10:13Z

/bot run

tensorrt-cicd · 2026-07-21T22:16:40Z

PR_Github #60799 [ run ] triggered by Bot. Commit: 663012b Link to invocation

tensorrt-cicd · 2026-07-21T23:45:53Z

PR_Github #60799 [ run ] completed with state FAILURE. Commit: 663012b
/LLM/main/L0_MergeRequest_PR pipeline #49076 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

zhenhuaw-me · 2026-07-22T02:46:16Z

 | `Lightricks/LTX-2` | Text-to-Video (with Audio), Image-to-Video (with Audio) |
 | `Qwen/Qwen-Image` | Text-to-Image |
 | `Qwen/Qwen-Image-2512` | Text-to-Image |
+| `Qwen/Qwen-Image-Edit-2511` | Image Editing |


Not sure if we should name like text+image-to-image, and also some other models such as Flux also supports text+image-to-image.

Maybe Image Editing (text+image-to-image)?

zhenhuaw-me · 2026-07-22T02:46:46Z

 | `Lightricks/LTX-2` | Text-to-Video (with Audio), Image-to-Video (with Audio) |
 | `Qwen/Qwen-Image` | Text-to-Image |
 | `Qwen/Qwen-Image-2512` | Text-to-Image |
+| `Qwen/Qwen-Image-Edit-2511` | Image Editing |


Not sure if we should name like text+image-to-image, and also some other models such as Flux also supports text+image-to-image.

Maybe Image Editing (text+image-to-image)?

zhenhuaw-me · 2026-07-22T02:50:45Z

+quant_config:
+  quant_algo: NVFP4
+  dynamic: true
+  ignore:
+    - "transformer_blocks.*.attn.to_q"
+    - "transformer_blocks.*.attn.to_k"
+    - "transformer_blocks.*.attn.to_v"
+    - "transformer_blocks.*.attn.add_q_proj"
+    - "transformer_blocks.*.attn.add_k_proj"
+    - "transformer_blocks.*.attn.add_v_proj"


Why these settings? When disclose such, we should have some docs to clarify the methodology.

I briefly mentioned in the comment above that

Q/K/V projections remain in BF16 because they preserve accuracy

Or we could use FP8 quant.

modified to use FP8 quant config instead.

zhenhuaw-me · 2026-07-22T02:51:33Z

+        "--image",
+        nargs="+",


This support is similar to Flux Kontext cc @karljang

yibinl-nvidia · 2026-07-22T21:14:17Z

/bot run

tensorrt-cicd · 2026-07-22T21:23:54Z

PR_Github #61097 [ run ] triggered by Bot. Commit: 51c15b9 Link to invocation

tensorrt-cicd · 2026-07-22T22:30:00Z

PR_Github #61097 [ run ] completed with state FAILURE. Commit: 51c15b9
/LLM/main/L0_MergeRequest_PR pipeline #49352 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>

yibinl-nvidia · 2026-07-23T17:04:18Z

/bot run

tensorrt-cicd · 2026-07-23T17:09:52Z

PR_Github #61344 [ run ] triggered by Bot. Commit: 9d85816 Link to invocation

tensorrt-cicd · 2026-07-23T17:51:18Z

PR_Github #61344 [ run ] completed with state FAILURE. Commit: 9d85816
/LLM/main/L0_MergeRequest_PR pipeline #49571 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

yibinl-nvidia · 2026-07-23T19:18:47Z

/bot run

tensorrt-cicd · 2026-07-23T19:30:18Z

PR_Github #61376 [ run ] triggered by Bot. Commit: 9d85816 Link to invocation

tensorrt-cicd · 2026-07-23T21:43:51Z

PR_Github #61376 [ run ] completed with state FAILURE. Commit: 9d85816
/LLM/main/L0_MergeRequest_PR pipeline #49604 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

chang-l

Does this PR also include trtllm-serve support for the image-edit endpoint?

github-actions Bot assigned yibinl-nvidia Jul 8, 2026

yibinl-nvidia force-pushed the dev-yibinl-qwen-image-edit branch from c9634fc to 3779b15 Compare July 14, 2026 03:23

zhenhuaw-me added the VisualGen label Jul 14, 2026

yibinl-nvidia force-pushed the dev-yibinl-qwen-image-edit branch 2 times, most recently from 6d0ec5f to bbf4efa Compare July 18, 2026 03:02

yibinl-nvidia marked this pull request as ready for review July 18, 2026 03:05

yibinl-nvidia requested review from a team as code owners July 18, 2026 03:05

yibinl-nvidia requested review from EmmaQiaoCh, LarryXFly, QiJune, brnguyen2, karljang, luyiyun1021, mlefeb01, nv-guomingz and yingguo-trt July 18, 2026 03:05

coderabbitai Bot reviewed Jul 18, 2026

View reviewed changes

yingguo-trt reviewed Jul 20, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py Outdated

BowenFu reviewed Jul 20, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py Outdated

yibinl-nvidia force-pushed the dev-yibinl-qwen-image-edit branch 2 times, most recently from 34093c1 to 35e7fd7 Compare July 20, 2026 22:33

yibinl-nvidia removed the request for review from QiJune July 21, 2026 16:46

zhenhuaw-me reviewed Jul 22, 2026

View reviewed changes

yibinl-nvidia force-pushed the dev-yibinl-qwen-image-edit branch from 663012b to 36779f0 Compare July 22, 2026 18:17

karljang reviewed Jul 22, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py Outdated

karljang reviewed Jul 22, 2026

View reviewed changes

Comment thread tests/integration/defs/examples/visual_gen/test_visual_gen.py

yibinl-nvidia force-pushed the dev-yibinl-qwen-image-edit branch from c6ca070 to 51c15b9 Compare July 22, 2026 21:12

yibinl-nvidia added 11 commits July 23, 2026 17:01

Add Qwen-Image-Edit-2511 support

45f3edf

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>

add E2E example test

76ae83f

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>

add cuda graph support

c8f8c88

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>

Add Qwen Image Edit CFG parallelism

67adb9f

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>

Add Qwen Image Edit docs and tests

2458136

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>

Remove Qwen Image Edit quant regression test

8d0aa8f

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>

Enable torch compile with Qwen Image Edit CFG parallelism

160e5a9

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>

Fix Qwen Image Edit pre-commit issues

f6f1272

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>

Address Qwen Image Edit review comments

2ed4aae

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>

Update Qwen Image Edit example docs

a9988a8

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>

Address more review comments

9d85816

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>

yibinl-nvidia force-pushed the dev-yibinl-qwen-image-edit branch from 51c15b9 to 9d85816 Compare July 23, 2026 17:03

chang-l reviewed Jul 23, 2026

View reviewed changes

Uh oh!

Conversation

yibinl-nvidia commented Jul 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Description

Supported Features

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

coderabbitai Bot commented Jul 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yibinl-nvidia commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

yibinl-nvidia commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

yibinl-nvidia commented Jul 22, 2026

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

yibinl-nvidia commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

yibinl-nvidia commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

chang-l left a comment

Choose a reason for hiding this comment

Uh oh!

yibinl-nvidia commented Jul 8, 2026 •

edited

Loading

coderabbitai Bot commented Jul 18, 2026 •

edited

Loading