[TRTLLM-12133][fix] refactor VisualGen examples - #13298
Conversation
|
/bot run --disable-fail-fast |
|
PR_Github #44817 [ run ] triggered by Bot. Commit: |
|
PR_Github #44817 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #45001 [ run ] triggered by Bot. Commit: |
|
PR_Github #45001 [ run ] completed with state |
59c4125 to
9b08bd5
Compare
zhenhuaw-me
left a comment
zhenhuaw-me
left a comment
There was a problem hiding this comment.
I would prefer standalone per-model examples, i.e., no shared utils across examples, but since we are switching to new style example. Should be fine to have this fix in short term.
|
|
||
| """Shared argparse helpers for visual-generation example scripts.""" | ||
|
|
||
| from tensorrt_llm._torch.visual_gen.config import CacheDiTConfig, TeaCacheConfig |
There was a problem hiding this comment.
This is one of the reasons that why I want to switch to the new style example - keep CLI interface and dict/yaml config synced is painful...
|
/bot run --disable-fail-fast |
|
PR_Github #45454 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughA refactoring that centralizes CLI argument registration and configuration building for cache, quantization, attention backend, and optimization settings across visual generation examples into new shared utilities, removing duplicate code from four example scripts. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 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}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tensorrt_llm/_torch/visual_gen/example_utils.py (1)
159-169: Consider a more specific return type annotation.The return type could be more precise to aid type checking.
Suggested improvement
-def build_cache_config(args) -> dict: +def build_cache_config(args) -> dict[str, CacheDiTConfig | TeaCacheConfig]:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/visual_gen/example_utils.py` around lines 159 - 169, The function build_cache_config currently uses a broad return type of dict; change its signature to a more specific type such as Dict[str, Any] (or Mapping[str, Any]) to improve type checking and IDE hints, and update the import from typing accordingly; ensure the signature for build_cache_config(args) -> Dict[str, Any] (or -> Mapping[str, Any]) and keep the returned keys compatible with VisualGenArgs by referencing the same "cache" key produced by build_cache_dit_config and build_teacache_config.tensorrt_llm/_torch/visual_gen/utils.py (1)
38-45: Add return type annotation.Per coding guidelines, Python functions should always have return type annotations.
Suggested fix
-def linear_type_to_quant_config(linear_type: str): +def linear_type_to_quant_config(linear_type: str) -> dict | None: """Map a --linear_type CLI shortcut to a quant_config dict for VisualGenArgs."""🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/visual_gen/utils.py` around lines 38 - 45, The function linear_type_to_quant_config lacks a return type annotation; update its signature to include an explicit return type (e.g. -> Optional[Dict[str, Any]]) and add the required typing imports (Optional, Dict, Any) if not already present so the mapping return value or None is correctly typed; keep the function body unchanged and ensure the annotation references the same function name linear_type_to_quant_config.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/example_utils.py`:
- Around line 159-169: The function build_cache_config currently uses a broad
return type of dict; change its signature to a more specific type such as
Dict[str, Any] (or Mapping[str, Any]) to improve type checking and IDE hints,
and update the import from typing accordingly; ensure the signature for
build_cache_config(args) -> Dict[str, Any] (or -> Mapping[str, Any]) and keep
the returned keys compatible with VisualGenArgs by referencing the same "cache"
key produced by build_cache_dit_config and build_teacache_config.
In `@tensorrt_llm/_torch/visual_gen/utils.py`:
- Around line 38-45: The function linear_type_to_quant_config lacks a return
type annotation; update its signature to include an explicit return type (e.g.
-> Optional[Dict[str, Any]]) and add the required typing imports (Optional,
Dict, Any) if not already present so the mapping return value or None is
correctly typed; keep the function body unchanged and ensure the annotation
references the same function name linear_type_to_quant_config.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: baf7511f-bb55-4b60-8896-b43d2a50c8d9
📒 Files selected for processing (7)
examples/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/_torch/visual_gen/example_utils.pytensorrt_llm/_torch/visual_gen/utils.pytests/integration/defs/examples/test_visual_gen.py
|
PR_Github #45454 [ run ] completed with state
|
49521dc to
5bdbec3
Compare
|
/bot run --disable-fail-fast --add-multi-gpu-test |
|
PR_Github #46685 [ run ] triggered by Bot. Commit: |
|
PR_Github #46685 [ run ] completed with state |
Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
1d6c144 to
def1d59
Compare
|
/bot run --disable-fail-fast --add-multi-gpu-test |
|
PR_Github #48054 [ run ] triggered by Bot. Commit: |
|
PR_Github #48054 [ run ] completed with state
|
Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
Signed-off-by: o-stoner <245287810+o-stoner@users.noreply.github.com>
|
/bot run --disable-fail-fast --add-multi-gpu-test |
|
PR_Github #48252 [ run ] triggered by Bot. Commit: |
|
PR_Github #48252 [ run ] completed with state
|
|
These examples scripts being refactored have been updated/removed in #14632; closing. |
3 participants
Summary by CodeRabbit
Description
Cache-DiT integration resulted in many redundant functions across pipeline example files (Wan, LTX-2, FLUX.2). For example, the _cache_dit_config_from_args() helper is repeated across all example scripts. In this PR, we extract these redundant example functions found into a shared example utility. Additionally, we enable a missing FA4 attention backend for LTX-2 in the example scripts.
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.