[None][feat] Enable FA4 + parallel VAE in LTX-2 examples

examples/visual_gen/README.md

@@ -29,7 +29,7 @@ python models/qwen_image_edit.py --image /path/to/source.png --prompt "Make the
 # With engine config (quant, parallelism, etc.)
 python models/wan_t2v.py --visual_gen_args configs/wan2.2-t2v-fp4-1gpu.yaml
 python models/wan_i2v.py --visual_gen_args configs/wan2.2-i2v-fp4-1gpu.yaml --image /path/to/image.png
-python models/ltx2.py --visual_gen_args configs/ltx2-t2v-fp8-1-gpu.yaml
+python models/ltx2.py --visual_gen_args configs/ltx2-1gpu.yaml
 python models/flux1.py --visual_gen_args configs/flux1-dev-fp4-1gpu.yaml
 python models/flux2.py --visual_gen_args configs/flux2-dev-fp4-1gpu.yaml
 python models/cosmos3_ti2v.py --visual_gen_args configs/cosmos3-nano-1gpu.yaml --prompt "A robot arm picks fruit in a grocery store"

examples/visual_gen/configs/ltx2-1gpu.yaml

@@ -0,0 +1,24 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# 1-GPU LTX-2 text-to-video with audio.
+# Shared by offline examples (--visual_gen_args) and trtllm-serve.
+attention_config:
+  backend: FA4
+parallel_config:
+  cfg_size: 1
+  ulysses_size: 1
+cuda_graph_config:
+  enable: false

examples/visual_gen/configs/ltx2-4gpu.yaml

@@ -17,11 +17,12 @@
 # whatever checkpoint is passed via --model_path). Shared by offline
 # examples (--visual_gen_args) and trtllm-serve.
 attention_config:
-  backend: VANILLA
+  backend: FA4
 parallel_config:
   cfg_size: 2
   ulysses_size: 2
   async_ulysses: true
+  parallel_vae_size: 4
 torch_compile_config:
   enable: true
 cuda_graph_config:

examples/visual_gen/configs/ltx2-t2v-fp4-1gpu.yaml → examples/visual_gen/configs/ltx2-fp4-1gpu.yaml

@@ -19,7 +19,7 @@ quant_config:
   quant_algo: NVFP4
   dynamic: true
 attention_config:
-  backend: VANILLA
+  backend: FA4
 parallel_config:
   cfg_size: 1
   ulysses_size: 1

examples/visual_gen/configs/ltx2-t2v-fp8-1gpu.yaml → examples/visual_gen/configs/ltx2-fp8-1gpu.yaml

@@ -19,7 +19,7 @@ quant_config:
   quant_algo: FP8_BLOCK_SCALES
   dynamic: true
 attention_config:
-  backend: VANILLA
+  backend: FA4
 parallel_config:
   cfg_size: 1
   ulysses_size: 1

examples/visual_gen/models/ltx2.py

@@ -17,7 +17,10 @@
 
 Usage:
     python ltx2.py
-    python ltx2.py --visual_gen_args ../configs/ltx2.yaml
+    python ltx2.py --visual_gen_args ../configs/ltx2-1gpu.yaml
+    # Force two-stage on a checkpoint lacking the aux files:
+    python ltx2.py --spatial_upsampler_path <upsampler.safetensors> \
+        --distilled_lora_path <distilled-lora.safetensors>
 """
 
 import argparse
@@ -49,6 +52,26 @@ def main():
             "from --visual_gen_args when set."
         ),
     )
+    parser.add_argument(
+        "--spatial_upsampler_path",
+        type=str,
+        default=None,
+        help=(
+            "Spatial upsampler safetensors path. Setting both this and "
+            "--distilled_lora_path forces two-stage inference. "
+            "Auto-discovered from the checkpoint dir when unset."
+        ),
+    )
+    parser.add_argument(
+        "--distilled_lora_path",
+        type=str,
+        default=None,
+        help=(
+            "Distilled-LoRA safetensors path. Setting both this and "
+            "--spatial_upsampler_path forces two-stage inference. "
+            "Auto-discovered from the checkpoint dir when unset."
+        ),
+    )
     parser.add_argument(
         "--output_path",
         type=str,
@@ -71,6 +94,15 @@ def main():
             **extra_args.pipeline_config,
             "text_encoder_path": text_encoder_path,
         }
+    # Two-stage auto-enables when both aux paths resolve (explicit here, or
+    # auto-discovered from the checkpoint dir); pass both to force it on a
+    # checkpoint that does not bundle them.
+    for key, value in (
+        ("spatial_upsampler_path", args.spatial_upsampler_path),
+        ("distilled_lora_path", args.distilled_lora_path),
+    ):
+        if value is not None:
+            extra_args.pipeline_config = {**extra_args.pipeline_config, key: value}
     visual_gen = VisualGen(model=args.model, args=extra_args)
 
     # --- Model-specific: T2V request construction ---

tests/integration/defs/examples/visual_gen/test_visual_gen_ltx2.py

@@ -430,7 +430,7 @@ def ltx2_two_stage_bf16_video_path(_visual_gen_deps, llm_venv):
 def test_ltx2_example(_visual_gen_deps, llm_root, llm_venv):
     """Run examples/visual_gen/models/ltx2.py with NVFP4 config end-to-end.
 
-    Validates that the LTX-2 example script and ``configs/ltx2-t2v-fp4-1gpu.yaml``
+    Validates that the LTX-2 example script and ``configs/ltx2-fp4-1gpu.yaml``
     work together as documented. The Gemma3 text encoder is passed separately via
     ``--text_encoder_path`` because the shared YAML intentionally omits it to keep
     the config model-path-agnostic.
@@ -445,9 +445,7 @@ def test_ltx2_example(_visual_gen_deps, llm_root, llm_venv):
     output_path = os.path.join(out_dir, "ltx2_output.mp4")
 
     script_path = os.path.join(llm_root, "examples", "visual_gen", "models", "ltx2.py")
-    config_path = os.path.join(
-        llm_root, "examples", "visual_gen", "configs", "ltx2-t2v-fp4-1gpu.yaml"
-    )
+    config_path = os.path.join(llm_root, "examples", "visual_gen", "configs", "ltx2-fp4-1gpu.yaml")
     assert os.path.isfile(script_path), f"Example script not found: {script_path}"
     assert os.path.isfile(config_path), f"Config not found: {config_path}"