diff --git a/examples/visual_gen/README.md b/examples/visual_gen/README.md index 51e59f0a2eb4..1eb0ac57593a 100644 --- a/examples/visual_gen/README.md +++ b/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" diff --git a/examples/visual_gen/configs/ltx2-1gpu.yaml b/examples/visual_gen/configs/ltx2-1gpu.yaml new file mode 100644 index 000000000000..89b72030817b --- /dev/null +++ b/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 diff --git a/examples/visual_gen/configs/ltx2-4gpu.yaml b/examples/visual_gen/configs/ltx2-4gpu.yaml index 1cd6b97c9876..90e8f8efe822 100644 --- a/examples/visual_gen/configs/ltx2-4gpu.yaml +++ b/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: diff --git a/examples/visual_gen/configs/ltx2-t2v-fp4-1gpu.yaml b/examples/visual_gen/configs/ltx2-fp4-1gpu.yaml similarity index 98% rename from examples/visual_gen/configs/ltx2-t2v-fp4-1gpu.yaml rename to examples/visual_gen/configs/ltx2-fp4-1gpu.yaml index 862d8a4fa640..ded67bdc6236 100644 --- a/examples/visual_gen/configs/ltx2-t2v-fp4-1gpu.yaml +++ b/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 diff --git a/examples/visual_gen/configs/ltx2-t2v-fp8-1gpu.yaml b/examples/visual_gen/configs/ltx2-fp8-1gpu.yaml similarity index 98% rename from examples/visual_gen/configs/ltx2-t2v-fp8-1gpu.yaml rename to examples/visual_gen/configs/ltx2-fp8-1gpu.yaml index 0f4afcdb4384..cf5463ba881a 100644 --- a/examples/visual_gen/configs/ltx2-t2v-fp8-1gpu.yaml +++ b/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 diff --git a/examples/visual_gen/models/ltx2.py b/examples/visual_gen/models/ltx2.py index 54ba38e8db6c..bbbd556ddeca 100644 --- a/examples/visual_gen/models/ltx2.py +++ b/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 \ + --distilled_lora_path """ 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 --- diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_ltx2.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_ltx2.py index 07cbadfb806c..a853709a1479 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_ltx2.py +++ b/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}"