[None][refactor] split VisualGen pipeline and model configs by bobboli · Pull Request #14956 · NVIDIA/TensorRT-LLM

bobboli · 2026-06-04T12:15:53Z

Summary

split VisualGen pipeline-level config from per-model DiffusionModelConfig
add BaseDiffusionModel and pass component-specific configs to transformer modules
route multi-component pipelines through model_configs for transformer/transformer_2

Testing

python -m py_compile on changed Python files
git diff --check
pre-commit run --files on changed files

Summary by CodeRabbit

Release Notes

Refactor
- Refactored visual generation pipeline configuration system to support multi-component pipelines through new unified configuration and model base classes, improving pipeline management across supported architectures.

bobboli · 2026-06-04T12:17:20Z

/bot run --disable-fail-fast

bobboli · 2026-06-04T12:19:03Z

/bot kill

bobboli · 2026-06-04T12:23:41Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-04T12:24:16Z

PR_Github #52066 [ run ] triggered by Bot. Commit: 6b1fe47 Link to invocation

coderabbitai · 2026-06-04T12:24:49Z

📝 Walkthrough

Walkthrough

This PR refactors the VisualGen configuration system from a single-component DiffusionModelConfig to a multi-component DiffusionPipelineConfig. A new BaseDiffusionModel base class unifies transformer initialization across all model variants. All pipelines, loaders, and tests are updated to use the new architecture.

Changes

VisualGen Multi-Component Configuration and Model Inheritance

Layer / File(s)	Summary
Configuration Architecture and BaseDiffusionModel `tensorrt_llm/_torch/visual_gen/__init__.py`, `tensorrt_llm/_torch/visual_gen/config.py`, `tensorrt_llm/_torch/visual_gen/models/modeling.py`	`DiffusionModelConfig` now inherits from `_VisualGenConfigBase` with shared runtime fields and accessors (`cache_backend`, `teacache`, `cache_dit`, `get_quant_config`). New `DiffusionPipelineConfig` holds a `model_configs` map of per-component `DiffusionModelConfig` instances, with convenience properties exposing the primary `"transformer"` config. `from_pretrained` now loads all components from `model_index.json` and `config.json` into `component_config_dicts`, then populates `pipeline_config.model_configs` via a new `_make_model_config` helper. New `BaseDiffusionModel` base class stores `model_config` and exposes `component_name` and `pretrained_config` attributes.
BasePipeline and Loader Integration `tensorrt_llm/_torch/visual_gen/pipeline.py`, `tensorrt_llm/_torch/visual_gen/pipeline_loader.py`, `tensorrt_llm/_torch/visual_gen/pipeline_registry.py`, `tensorrt_llm/visual_gen/args.py`	`BasePipeline` now accepts `DiffusionPipelineConfig`, initializes internal `model_configs` from `pipeline_config.model_configs`, and derives `checkpoint_path` from `primary_pretrained_config._name_or_path` for TeaCache variant matching. `PipelineLoader._setup_visual_gen_mapping` propagates both visual-gen and LLM mappings into every component config. `PipelineLoader.load()` constructs `DiffusionPipelineConfig.from_pretrained(...)`. `AutoPipeline.from_config` signature updated. Documentation reflects pipeline-level config.
Transformer Model Base Class Inheritance `tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py`, `tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux.py`, `tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux2.py`, `tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py`, `tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py`, `tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py`	All transformer classes now inherit from `BaseDiffusionModel` instead of `nn.Module`. Each model's `__init__` calls `super().__init__(model_config)` instead of manually storing `self.model_config`.
Pipeline Class Constructor and Component Initialization Updates `tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py`, `tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py`, `tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py`, `tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py`, `tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py`, `tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py`, `tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py`	All pipeline classes now accept `pipeline_config` instead of `model_config` and pass it to `BasePipeline.__init__`. Transformer initialization pulls component-specific config from `self.model_configs[component_name]` (e.g., `"transformer"`, `"transformer_2"` for multi-stage). Variant resolution and TeaCache validation use `pipeline_config.primary_pretrained_config` and `pipeline_config.cache_backend`.
Test Suite Migrations `tests/unittest/_torch/visual_gen/test_*.py`	All test files updated to import and use `DiffusionPipelineConfig` instead of `DiffusionModelConfig`. Test fixtures and assertions updated to use `primary_pretrained_config._name_or_path` and call `DiffusionPipelineConfig.from_pretrained(...)` and `.model_configs["transformer"]` to extract component configs.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

NVIDIA/TensorRT-LLM#13449: Adds Qwen-Image pipeline and transformer support, directly affected by this refactor's wiring changes to DiffusionPipelineConfig/model_configs and BaseDiffusionModel inheritance.

Suggested labels

VisualGen

Suggested reviewers

QiJune
Shixiaowei02
chang-l
zhenhuaw-me

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 46.38% 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
Title check	✅ Passed	The title clearly identifies the main change: refactoring VisualGen configuration by splitting pipeline-level config from per-model config.
Description check	✅ Passed	The PR description covers what was changed (split configs, added BaseDiffusionModel, routed multi-component pipelines) and includes testing performed, but lacks detailed explanation of why these changes were necessary.
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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

coderabbitai

Actionable comments posted: 1

🤖 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 `@tests/unittest/_torch/visual_gen/test_visual_gen_args.py`:
- Line 214: The test unpacks four values from
DiffusionPipelineConfig.load_diffusion_quant_config into qc, _, dwq, daq but
never uses daq; remove the unused unpacked variable by changing the unpacking to
only capture the needed three values (qc, _, dwq) from
load_diffusion_quant_config so daq is not created, or replace daq with a
throwaway underscore if you prefer (qc, _, dwq, _ -> qc, _, dwq), updating the
line that calls DiffusionPipelineConfig.load_diffusion_quant_config accordingly.

🪄 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: e70baaf8-b004-4a89-9044-f9993181245c

📥 Commits

Reviewing files that changed from the base of the PR and between 8361d42 and 387a29b.

📒 Files selected for processing (25)

tensorrt_llm/_torch/visual_gen/__init__.py
tensorrt_llm/_torch/visual_gen/config.py
tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py
tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py
tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py
tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py
tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux.py
tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux2.py
tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py
tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py
tensorrt_llm/_torch/visual_gen/models/modeling.py
tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py
tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py
tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py
tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py
tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py
tensorrt_llm/_torch/visual_gen/pipeline.py
tensorrt_llm/_torch/visual_gen/pipeline_loader.py
tensorrt_llm/_torch/visual_gen/pipeline_registry.py
tensorrt_llm/visual_gen/args.py
tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py
tests/unittest/_torch/visual_gen/test_model_loader.py
tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py
tests/unittest/_torch/visual_gen/test_visual_gen_args.py
tests/unittest/_torch/visual_gen/test_wan_transformer.py

tensorrt-cicd · 2026-06-04T12:27:11Z

PR_Github #52067 [ kill ] triggered by Bot. Commit: 6b1fe47 Link to invocation

bobboli · 2026-06-04T12:29:28Z

/bot kill

tensorrt-cicd · 2026-06-04T12:29:40Z

PR_Github #52068 [ run ] triggered by Bot. Commit: 25ef374 Link to invocation

tensorrt-cicd · 2026-06-04T12:29:44Z

PR_Github #52067 [ kill ] completed with state ABORTED. Commit: 6b1fe47

Link to invocation

tensorrt-cicd · 2026-06-04T12:30:10Z

PR_Github #52066 [ run ] completed with state ABORTED. Commit: 6b1fe47

Link to invocation

bobboli · 2026-06-04T12:33:16Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-04T12:35:01Z

PR_Github #52071 [ kill ] triggered by Bot. Commit: 25ef374 Link to invocation

tensorrt-cicd · 2026-06-04T12:38:27Z

PR_Github #52068 [ run ] completed with state ABORTED. Commit: 25ef374

Link to invocation

tensorrt-cicd · 2026-06-04T12:38:36Z

PR_Github #52071 [ kill ] completed with state SUCCESS. Commit: 25ef374
Successfully killed previous jobs for commit 25ef374

Link to invocation

tensorrt-cicd · 2026-06-04T12:39:31Z

PR_Github #52073 [ run ] triggered by Bot. Commit: 25ef374 Link to invocation

tensorrt-cicd · 2026-06-05T03:46:52Z

PR_Github #52073 [ run ] completed with state SUCCESS. Commit: 25ef374
/LLM/main/L0_MergeRequest_PR pipeline #41404 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

bobboli · 2026-06-05T07:05:50Z

/bot run --disable-fail-fast --reuse-test

tensorrt-cicd · 2026-06-05T07:12:43Z

PR_Github #52288 [ run ] triggered by Bot. Commit: 25ef374 Link to invocation

tensorrt-cicd · 2026-06-05T12:24:10Z

PR_Github #52288 [ run ] completed with state SUCCESS. Commit: 25ef374
/LLM/main/L0_MergeRequest_PR pipeline #41597 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

bobboli · 2026-06-05T13:35:18Z

/bot run --disable-fail-fast --reuse-test

tensorrt-cicd · 2026-06-05T13:41:51Z

PR_Github #52346 [ run ] triggered by Bot. Commit: 25ef374 Link to invocation

tensorrt-cicd · 2026-06-05T18:11:00Z

PR_Github #52346 [ run ] completed with state FAILURE. Commit: 25ef374
/LLM/main/L0_MergeRequest_PR pipeline #41647 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

bobboli · 2026-06-06T04:58:26Z

/bot run --disable-fail-fast --reuse-test

tensorrt-cicd · 2026-06-06T05:03:50Z

PR_Github #52477 [ run ] triggered by Bot. Commit: 25ef374 Link to invocation

zhenhuaw-me

Thank you so much for helping refactoring this settings :)

tensorrt-cicd · 2026-06-06T13:47:44Z

PR_Github #52477 [ run ] completed with state FAILURE. Commit: 25ef374
/LLM/main/L0_MergeRequest_PR pipeline #41769 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

bobboli · 2026-06-06T18:42:05Z

/bot run --reuse-test

tensorrt-cicd · 2026-06-06T18:47:51Z

PR_Github #52524 [ run ] triggered by Bot. Commit: 25ef374 Link to invocation

tensorrt-cicd · 2026-06-07T03:33:56Z

PR_Github #52524 [ run ] completed with state FAILURE. Commit: 25ef374
/LLM/main/L0_MergeRequest_PR pipeline #41812 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

bobboli · 2026-06-07T06:44:08Z

/bot run --reuse-test

tensorrt-cicd · 2026-06-07T06:50:11Z

PR_Github #52558 [ run ] triggered by Bot. Commit: 25ef374 Link to invocation

tensorrt-cicd · 2026-06-07T15:34:36Z

PR_Github #52558 [ run ] completed with state FAILURE. Commit: 25ef374
/LLM/main/L0_MergeRequest_PR pipeline #41843 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

bobboli · 2026-06-07T19:31:42Z

/bot run --reuse-test

tensorrt-cicd · 2026-06-07T19:37:51Z

PR_Github #52598 [ run ] triggered by Bot. Commit: 25ef374 Link to invocation

tensorrt-cicd · 2026-06-08T04:24:41Z

PR_Github #52598 [ run ] completed with state FAILURE. Commit: 25ef374
/LLM/main/L0_MergeRequest_PR pipeline #41879 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: Bo Li <22713281+bobboli@users.noreply.github.com>

tburt-nv · 2026-06-08T17:32:04Z

/bot run --stage-list "GB200-4_GPUs-PyTorch-1, GB200-4_GPUs-PyTorch-2, GB200-4_GPUs-PyTorch-3, GB200-4_GPUs-PyTorch-4"

tburt-nv · 2026-06-08T17:32:45Z

If these stages pass, let's use "/bot reuse-pipeline"

tensorrt-cicd · 2026-06-08T17:38:15Z

PR_Github #52833 [ run ] triggered by Bot. Commit: fde2e27 Link to invocation

tensorrt-cicd · 2026-06-08T21:12:24Z

PR_Github #52833 [ run ] completed with state SUCCESS. Commit: fde2e27
/LLM/main/L0_MergeRequest_PR pipeline #42084 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

bobboli · 2026-06-09T01:34:45Z

/bot reuse-pipeline

tensorrt-cicd · 2026-06-09T01:41:13Z

PR_Github #52908 [ reuse-pipeline ] triggered by Bot. Commit: fde2e27 Link to invocation

tensorrt-cicd · 2026-06-09T01:47:30Z

PR_Github #52908 [ reuse-pipeline ] completed with state SUCCESS. Commit: fde2e27
Reusing PR_Github #52833 (Partly Tested) for commit fde2e27

Link to invocation

…4956) Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>

bobboli requested review from a team as code owners June 4, 2026 12:15

bobboli requested a review from hchings June 4, 2026 12:15

github-actions Bot assigned bobboli Jun 4, 2026

bobboli requested review from chang-l, karljang and zhenhuaw-me June 4, 2026 12:20

bobboli enabled auto-merge (squash) June 4, 2026 12:23

coderabbitai Bot reviewed Jun 4, 2026

View reviewed changes

Comment thread tests/unittest/_torch/visual_gen/test_visual_gen_args.py Outdated

chang-l approved these changes Jun 6, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/visual_gen/config.py

Comment thread tensorrt_llm/_torch/visual_gen/config.py

Comment thread tests/unittest/_torch/visual_gen/test_model_loader.py

zhenhuaw-me approved these changes Jun 6, 2026

View reviewed changes

bobboli added 3 commits June 8, 2026 13:31

[None][refactor] split visual gen pipeline and model configs

1a09ad7

Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>

[None][refactor] remove pipeline model_config alias

9bb57bf

Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>

[None][test] fix visual gen args review lint

fde2e27

Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>

tburt-nv force-pushed the scratch/vg-config-refactor-main branch from 25ef374 to fde2e27 Compare June 8, 2026 17:31

bobboli merged commit bfb4537 into NVIDIA:main Jun 9, 2026
7 checks passed

crazydemo pushed a commit to crazydemo/TensorRT-LLM that referenced this pull request Jun 9, 2026

[None][refactor] split VisualGen pipeline and model configs (NVIDIA#1…

2dac5ec

…4956) Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>

tensorrt-cicd mentioned this pull request Jun 11, 2026

[https://nvbugs/6290965][fix] Rename model_config → pipeline_config on the two Cosmos3 test asserts and… #15263

Closed

2 tasks

Uh oh!

Conversation

bobboli commented Jun 4, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Testing

Summary by CodeRabbit

Release Notes

Uh oh!

bobboli commented Jun 4, 2026

Uh oh!

bobboli commented Jun 4, 2026

Uh oh!

bobboli commented Jun 4, 2026

Uh oh!

tensorrt-cicd commented Jun 4, 2026

Uh oh!

coderabbitai Bot commented Jun 4, 2026

Walkthrough

Changes

Estimated code review effort

Possibly related PRs

Suggested labels

Suggested reviewers

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tensorrt-cicd commented Jun 4, 2026

Uh oh!

bobboli commented Jun 4, 2026

Uh oh!

tensorrt-cicd commented Jun 4, 2026

Uh oh!

tensorrt-cicd commented Jun 4, 2026

Uh oh!

tensorrt-cicd commented Jun 4, 2026

Uh oh!

bobboli commented Jun 4, 2026

Uh oh!

tensorrt-cicd commented Jun 4, 2026

Uh oh!

tensorrt-cicd commented Jun 4, 2026

Uh oh!

tensorrt-cicd commented Jun 4, 2026

Uh oh!

tensorrt-cicd commented Jun 4, 2026

Uh oh!

tensorrt-cicd commented Jun 5, 2026

Uh oh!

bobboli commented Jun 5, 2026

Uh oh!

tensorrt-cicd commented Jun 5, 2026

Uh oh!

tensorrt-cicd commented Jun 5, 2026

Uh oh!

bobboli commented Jun 5, 2026

Uh oh!

tensorrt-cicd commented Jun 5, 2026

Uh oh!

tensorrt-cicd commented Jun 5, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

bobboli commented Jun 6, 2026

Uh oh!

tensorrt-cicd commented Jun 6, 2026

Uh oh!

zhenhuaw-me left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Jun 6, 2026

Uh oh!

bobboli commented Jun 6, 2026

Uh oh!

tensorrt-cicd commented Jun 6, 2026

Uh oh!

bobboli commented Jun 4, 2026 •

edited by coderabbitai Bot

Loading