[TRTLLM-11343][feat] LTX-2 Two Stage pipeline support

[yibinl-nvidia]

Summary by CodeRabbit

• New Features

  • Added two-stage LTX-2 pipeline with spatial upsampling and optional distilled LoRA refinement support
  • New CLI arguments --spatial_upsampler_path and --distilled_lora_path to enable advanced generation workflows

• Bug Fixes

  • Removed duplicate CLI argument aliases
  • Improved resolution validation for two-stage versus one-stage generation modes

• Chores

  • Added FP8 quantization utilities for handling compressed model weights

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.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR introduces a two-stage LTX2 text-to-video generation pipeline with supporting infrastructure. Changes include a latent upsampler module, a new pipeline implementation that performs Stage 1 denoising followed by upsampling and Stage 2 refinement, FP8 quantization helpers for LoRA weight handling, and automatic pipeline detection logic.

Changes

Cohort / File(s)	Summary
CLI and Configuration examples/visual_gen/visual_gen_ltx2.py, tensorrt_llm/_torch/visual_gen/config.py	Removed duplicate CLI aliases for common arguments and added --spatial_upsampler_path and --distilled_lora_path flags with validation. Extended VisualGenArgs with two new string fields (spatial_upsampler_path, distilled_lora_path) defaulting to empty strings.
Two-Stage Pipeline Core tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/upsampler.py, tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py	Added LatentUpsampler module with pixel shuffle, blur downsampling, residual blocks, and rational resampling support for 2D/3D latents. Implemented new LTX2TwoStagesPipeline class orchestrating Stage 1 inference, learned latent upsampling, and Stage 2 refinement with distilled LoRA support, quantization handling (FP8/NVFP4), and final decoding.
Pipeline Infrastructure tensorrt_llm/_torch/visual_gen/models/ltx2/__init__.py, tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py, tensorrt_llm/_torch/visual_gen/pipeline_loader.py, tensorrt_llm/_torch/visual_gen/pipeline_registry.py	Updated __all__ exports to include LTX2TwoStagesPipeline. Enhanced _assert_resolution() with optional is_two_stage parameter to enforce appropriate divisibility factors (64 for two-stage, 32 for one-stage). Modified PipelineLoader to pass optional upsampler and LoRA paths to component loading. Updated AutoPipeline.from_config() to conditionally switch to LTX2TwoStagesPipeline when both upsampler and LoRA paths are present.
FP8 Quantization Utilities tensorrt_llm/quantization/utils/fp8_utils.py	Added unpack_col_major_tma_aligned_packed_tensor() to reverse TMA-aligned packing with scale reconstruction. Added inverse_transform_sf() to recover block-scale grids from packed FP8 scale layouts with configurable block size.
Tests tests/unittest/_torch/thop/parallel/test_fp8_quantize.py	Added round-trip unit tests for FP8 scale layout transformations: pack→unpack roundtrip for aligned tensors and transform→inverse transform roundtrip for block-scale grids, both verifying exact reconstruction.

Sequence Diagram(s)

sequenceDiagram
    actor User
    participant LTX2TwoStagesPipeline
    participant LTX2Pipeline as Stage1<br/>(LTX2Pipeline)
    participant LatentUpsampler
    participant Stage2 as Stage2<br/>(Refinement)
    participant Decoder
    
    User->>LTX2TwoStagesPipeline: forward(prompt, height=512, width=768, ...)
    LTX2TwoStagesPipeline->>LTX2Pipeline: forward(..., height/2, width/2)<br/>Stage 1 at half resolution
    LTX2Pipeline-->>LTX2TwoStagesPipeline: video_latent, audio_latent
    
    LTX2TwoStagesPipeline->>LatentUpsampler: forward(video_latent)<br/>2x spatial upsampling
    LatentUpsampler-->>LTX2TwoStagesPipeline: upsampled_video_latent
    
    LTX2TwoStagesPipeline->>Stage2: _refinement_denoise(...)<br/>3-step distilled schedule
    Note over Stage2: Apply distilled LoRA<br/>Patchify → Encode prompt<br/>Blend image cond → Euler steps<br/>Unpatchify
    Stage2-->>LTX2TwoStagesPipeline: refined_video_latent
    
    LTX2TwoStagesPipeline->>Decoder: decode(video_latent, audio_latent)<br/>Tiled video + audio decode
    Decoder-->>LTX2TwoStagesPipeline: video_frames, audio
    
    LTX2TwoStagesPipeline-->>User: output (video, audio)

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

🚥 Pre-merge checks | ✅ 1 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 51.11% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Description check	⚠️ Warning	The PR description is entirely blank/empty, containing only the template structure without any substantive content in the required sections.	Provide a clear description explaining the feature implementation, test coverage details, and confirm all checklist items have been properly reviewed and completed.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly describes the main change: adding LTX-2 two-stage pipeline support. It is concise, specific, and directly reflects the primary feature added across multiple files.

_{✏️ 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: 2

🧹 Nitpick comments (2)

tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py (2)

504-506: Prefix unused variable with underscore.

The unexpected variable from load_state_dict is never used. Prefix it with an underscore to indicate it's intentionally ignored and silence the RUF059 lint warning.

🔧 Proposed fix

-            missing, unexpected = self.spatial_upsampler.load_state_dict(
+            missing, _unexpected = self.spatial_upsampler.load_state_dict(
                 state_dict, strict=False,
             )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py`
around lines 504 - 506, The tuple unpacking from
self.spatial_upsampler.load_state_dict currently assigns (missing, unexpected)
but only missing is used; rename the unused variable to _unexpected (e.g.,
(missing, _unexpected) = self.spatial_upsampler.load_state_dict(...)) to signal
it is intentionally ignored and silence the RUF059 lint warning while leaving
spatial_upsampler.load_state_dict unchanged.

---

487-494: Avoid bare except: — catch specific exceptions or log the error.

The current try-except-pass silently swallows all exceptions when parsing upsampler config metadata. This makes debugging difficult if the config format is malformed or if other unexpected errors occur. As per coding guidelines, limit the except to the smallest set of errors possible.

♻️ Proposed fix to catch specific exceptions

             try:
                 with safetensors.torch.safe_open(sft_paths[0], framework="pt") as f:
                     meta = f.metadata()
                     if meta and "config" in meta:
                         import json
                         config = json.loads(meta["config"])
-            except Exception:
-                pass
+            except (KeyError, json.JSONDecodeError, TypeError) as e:
+                logger.debug(f"Could not parse upsampler config from metadata: {e}")

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py`
around lines 487 - 494, Replace the bare "except Exception: pass" around the
safetensors metadata parsing in pipeline_ltx2_two_stages.py with targeted
exception handling: catch file/IO errors (e.g., FileNotFoundError, OSError),
safetensors-specific errors from safetensors.torch.safe_open, and
json.JSONDecodeError when calling json.loads, and log the caught error (or
re-raise if unrecoverable) so failures parsing the upsampler config are visible;
update the try/except that surrounds the safe_open/meta/json.loads sequence (the
block that opens sft_paths[0], reads f.metadata(), checks "config", and calls
json.loads) to handle those specific exceptions and use the module/class logger
to record the exception details.

🤖 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/_torch/visual_gen/models/ltx2/ltx2_core/upsampler.py`:
- Line 1: Replace the ambiguous EN DASH character in the copyright header at the
top of the file (the file-level comment line containing "Copyright (c) 2025–2026
Lightricks Ltd.") with a standard HYPHEN-MINUS so the range reads "2025-2026";
edit that header string in
tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/upsampler.py to use '-'
instead of '–' and commit the change so the pre-commit hooks pass.

In `@tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py`:
- Line 1: Replace the ambiguous EN DASH character in the SPDX copyright header
line that starts with "SPDX-FileCopyrightText" so the year range uses a standard
HYPHEN-MINUS (`-`) instead of `–`; update that header string in
tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py to use
"2025-2026" (hyphen-minus) to satisfy the linter and pre-commit checks.

---

Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py`:
- Around line 504-506: The tuple unpacking from
self.spatial_upsampler.load_state_dict currently assigns (missing, unexpected)
but only missing is used; rename the unused variable to _unexpected (e.g.,
(missing, _unexpected) = self.spatial_upsampler.load_state_dict(...)) to signal
it is intentionally ignored and silence the RUF059 lint warning while leaving
spatial_upsampler.load_state_dict unchanged.
- Around line 487-494: Replace the bare "except Exception: pass" around the
safetensors metadata parsing in pipeline_ltx2_two_stages.py with targeted
exception handling: catch file/IO errors (e.g., FileNotFoundError, OSError),
safetensors-specific errors from safetensors.torch.safe_open, and
json.JSONDecodeError when calling json.loads, and log the caught error (or
re-raise if unrecoverable) so failures parsing the upsampler config are visible;
update the try/except that surrounds the safe_open/meta/json.loads sequence (the
block that opens sft_paths[0], reads f.metadata(), checks "config", and calls
json.loads) to handle those specific exceptions and use the module/class logger
to record the exception details.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: b435a581-f5bc-4c83-a405-002eb3d43d43

📥 Commits

Reviewing files that changed from the base of the PR and between 46f4ca3 and b2bd565.

📒 Files selected for processing (10)

• examples/visual_gen/visual_gen_ltx2.py
• tensorrt_llm/_torch/visual_gen/config.py
• tensorrt_llm/_torch/visual_gen/models/ltx2/__init__.py
• tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/upsampler.py
• tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py
• tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py
• tensorrt_llm/_torch/visual_gen/pipeline_loader.py
• tensorrt_llm/_torch/visual_gen/pipeline_registry.py
• tensorrt_llm/quantization/utils/fp8_utils.py
• tests/unittest/_torch/thop/parallel/test_fp8_quantize.py

[yibinl-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #40378 [ run ] triggered by Bot. Commit: 050710f Link to invocation

[tensorrt-cicd]

PR_Github #40378 [ run ] completed with state FAILURE. Commit: 050710f
/LLM/main/L0_MergeRequest_PR pipeline #31476 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

Link to invocation

[yibinl-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #40740 [ run ] triggered by Bot. Commit: 7e23017 Link to invocation

[tensorrt-cicd]

PR_Github #40740 [ run ] completed with state SUCCESS. Commit: 7e23017
/LLM/main/L0_MergeRequest_PR pipeline #31760 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

Link to invocation

[yibinl-nvidia]

/bot run

[yibinl-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #41200 [ run ] triggered by Bot. Commit: 21a93d1 Link to invocation

[tensorrt-cicd]

PR_Github #41200 [ run ] completed with state FAILURE. Commit: 21a93d1
/LLM/main/L0_MergeRequest_PR pipeline #32161 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

Link to invocation

[yibinl-nvidia]

/bot run --disable-fail-fast

[yibinl-nvidia]

@zhenhuaw-me @chang-l Now that tests are added, comments are addressed, could you please review this PR again?

[tensorrt-cicd]

PR_Github #41236 [ run ] triggered by Bot. Commit: d6ae648 Link to invocation

[yibinl-nvidia]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #41846 [ run ] triggered by Bot. Commit: 4325b88 Link to invocation

[tensorrt-cicd]

PR_Github #41846 [ run ] completed with state SUCCESS. Commit: 4325b88
/LLM/main/L0_MergeRequest_PR pipeline #32715 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

Link to invocation

[yibinl-nvidia]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #41910 [ run ] triggered by Bot. Commit: a18fce1 Link to invocation

[tensorrt-cicd]

PR_Github #41910 [ run ] completed with state SUCCESS. Commit: a18fce1
/LLM/main/L0_MergeRequest_PR pipeline #32769 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

Link to invocation

[yibinl-nvidia]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #41931 [ run ] triggered by Bot. Commit: a18fce1 Link to invocation

[tensorrt-cicd]

PR_Github #41931 [ run ] completed with state SUCCESS. Commit: a18fce1
/LLM/main/L0_MergeRequest_PR pipeline #32788 completed with status: 'SUCCESS'

CI Report

Link to invocation

[Shixiaowei02]

Approve for the doc changes.