[TRTLLM-11804][feat] Mechanical refactoring VisualGen API#12807
Open
zhenhuaw-me wants to merge 1 commit intoNVIDIA:mainfrom
Open
[TRTLLM-11804][feat] Mechanical refactoring VisualGen API#12807zhenhuaw-me wants to merge 1 commit intoNVIDIA:mainfrom
zhenhuaw-me wants to merge 1 commit intoNVIDIA:mainfrom
Conversation
zhenhuaw-me
commented
Apr 7, 2026
| # from tensorrt_llm.visual_gen.args import VisualGenArgs | ||
| # from tensorrt_llm.visual_gen import VisualGenArgs | ||
| # from tensorrt_llm import VisualGenArgs | ||
| from tensorrt_llm._torch.visual_gen.config import VisualGenArgs |
Member
Author
There was a problem hiding this comment.
Note for reviewers: VisualGenArgs needs more work to refactor, so not moved here in this PR.
Contributor
📝 WalkthroughWalkthroughThe pull request refactors the VisualGen API by renaming constructor parameters ( Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (1 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: 2
🧹 Nitpick comments (2)
tensorrt_llm/visual_gen/visual_gen.py (1)
36-36: Sort__all__for consistency with linter rules.Static analysis flags that
__all__is not sorted. Consider applying isort-style sorting:-__all__ = ["VisualGen", "VisualGenParams", "MediaOutput", "VisualGenError", "VisualGenResult"] +__all__ = ["MediaOutput", "VisualGen", "VisualGenError", "VisualGenParams", "VisualGenResult"]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/visual_gen/visual_gen.py` at line 36, The __all__ export list is unsorted; update the __all__ declaration so its entries are in sorted (isort-style) alphabetical order to satisfy the linter. Locate the __all__ assignment (containing "VisualGen", "VisualGenParams", "MediaOutput", "VisualGenError", "VisualGenResult") and reorder the items alphabetically (by name) while preserving the same symbols and string formatting.tensorrt_llm/visual_gen/__init__.py (1)
18-25: Sort__all__for consistency.Static analysis flags the unsorted
__all__. Consider alphabetical ordering:♻️ Proposed fix
__all__ = [ - "VisualGen", - "VisualGenArgs", - "VisualGenError", - "VisualGenParams", - "VisualGenResult", "MediaOutput", + "VisualGen", + "VisualGenArgs", + "VisualGenError", + "VisualGenParams", + "VisualGenResult", ]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/visual_gen/__init__.py` around lines 18 - 25, The __all__ list is unsorted; update it to be alphabetically ordered to satisfy static analysis. Locate the __all__ definition in tensorrt_llm.visual_gen.__init__ and reorder the exported symbols (VisualGen, VisualGenArgs, VisualGenError, VisualGenParams, VisualGenResult, MediaOutput) into alphabetical order so the list is stable and passes linting.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tensorrt_llm/__init__.py`:
- Around line 135-136: The import line importing VisualGen, VisualGenArgs,
VisualGenError, VisualGenParams, VisualGenResult from .visual_gen incorrectly
includes a stale "# noqa: F811 (VisualGenArgs already imported above)" comment;
remove that inline noqa comment so the import reads normally without the F811
suppression. Locate the import statement that references
VisualGen/VisualGenArgs/VisualGenError/VisualGenParams/VisualGenResult and
delete only the trailing noqa comment text.
In `@tensorrt_llm/visual_gen/visual_gen.py`:
- Around line 428-434: result_sync calls result() via
asyncio.run_coroutine_threadsafe but if timeout expires await_responses can
return None causing result() to access response.error_msg and raise
AttributeError; update result_sync/result to check for a None response from
await_responses (or from result()) and handle it explicitly (e.g., raise a
TimeoutError or return a clear error object/string) instead of dereferencing
response.error_msg—look at the result_sync, result, and await_responses paths
and add a None-check and appropriate error handling/raise when response is None.
---
Nitpick comments:
In `@tensorrt_llm/visual_gen/__init__.py`:
- Around line 18-25: The __all__ list is unsorted; update it to be
alphabetically ordered to satisfy static analysis. Locate the __all__ definition
in tensorrt_llm.visual_gen.__init__ and reorder the exported symbols (VisualGen,
VisualGenArgs, VisualGenError, VisualGenParams, VisualGenResult, MediaOutput)
into alphabetical order so the list is stable and passes linting.
In `@tensorrt_llm/visual_gen/visual_gen.py`:
- Line 36: The __all__ export list is unsorted; update the __all__ declaration
so its entries are in sorted (isort-style) alphabetical order to satisfy the
linter. Locate the __all__ assignment (containing "VisualGen",
"VisualGenParams", "MediaOutput", "VisualGenError", "VisualGenResult") and
reorder the items alphabetically (by name) while preserving the same symbols and
string formatting.
🪄 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: Pro
Run ID: 623e7db0-2bac-4882-8242-b86aa7664351
📒 Files selected for processing (14)
examples/visual_gen/quickstart_example.pyexamples/visual_gen/visual_gen_flux.pyexamples/visual_gen/visual_gen_ltx2.pyexamples/visual_gen/visual_gen_wan_i2v.pyexamples/visual_gen/visual_gen_wan_t2v.pytensorrt_llm/__init__.pytensorrt_llm/_torch/visual_gen/config.pytensorrt_llm/bench/benchmark/visual_gen.pytensorrt_llm/commands/serve.pytensorrt_llm/visual_gen/__init__.pytensorrt_llm/visual_gen/args.pytensorrt_llm/visual_gen/visual_gen.pytests/integration/defs/examples/test_visual_gen.pytests/unittest/_torch/visual_gen/test_visual_gen_args.py
💤 Files with no reviewable changes (1)
- tensorrt_llm/_torch/visual_gen/config.py
zhenhuaw-me
force-pushed
the
update-api-step1
branch
from
718eca3 to
8ddd412
Compare
Member
Author
|
/bot run |
Collaborator
|
PR_Github #42245 [ run ] triggered by Bot. Commit: |
- Move VisualGenArgs to public visual_gen/args.py; remove to_dict()/from_dict() - Rename VisualGen constructor params: model_path→model, diffusion_args→args - Add VisualGenError exception class; replace bare RuntimeError raises - Rename DiffusionGenerationResult→VisualGenResult; add done property and result_sync() - Fix req_counter thread safety with itertools.count() - Export VisualGenArgs, VisualGenError, VisualGenResult from public __init__.py files - Update all callers: serve, bench, examples, tests Signed-off-by: Zhenhua Wang <zhenhuaw@nvidia.com>
zhenhuaw-me
force-pushed
the
update-api-step1
branch
from
ff3ab43 to
7bb6074
Compare
Member
Author
|
/bot run |
Collaborator
|
PR_Github #42281 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #42281 [ run ] completed with state
|
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
3 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 by CodeRabbit
Breaking Changes
VisualGenconstructor parameters: usemodelandargsinstead ofmodel_pathanddiffusion_args.New Features
VisualGenErrorexception for improved error handling in visual generation tasks.VisualGenResultwith a newresult_sync()method for synchronous result retrieval and adoneproperty for checking completion status.Description
Test Coverage
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)
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.