From 3f4b8cc981e24dc3affb01cc8d8616f62f8bb8e1 Mon Sep 17 00:00:00 2001 From: Min Yu <171526537+yumin066@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:31:16 +0800 Subject: [PATCH] [None][feat] Add LTX-2.3 audio-video and BF16 retake pipelines Add LTX-2.3 generation and retake workflows with reusable LTX-2 components, examples, configuration, documentation and smoke/LPIPS tests. Include subsequent compatibility fixes and native pipeline/component cleanup. Includes the LTX-2.3 implementation originally contributed by rajeshwarid179 in NVIDIA/TensorRT-LLM#17132. Signed-off-by: Min Yu <171526537+yumin066@users.noreply.github.com> --- docs/source/models/supported-models.md | 2 + docs/source/models/visual-generation.md | 4 +- examples/visual_gen/README.md | 7 + .../configs/ltx23-retake-bf16-1gpu.yaml | 27 + .../configs/ltx23-retake-bf16-2gpu.yaml | 28 + .../configs/ltx23-retake-fp4-1gpu.yaml | 32 + .../configs/ltx23-retake-fp8-1gpu.yaml | 32 + .../configs/ltx23-t2v-bf16-1gpu.yaml | 25 + .../configs/ltx23-t2v-bf16-2gpu.yaml | 26 + .../configs/ltx23-t2v-fp4-1gpu.yaml | 32 + .../configs/ltx23-t2v-fp8-1gpu.yaml | 32 + examples/visual_gen/models/ltx23.py | 126 +++ examples/visual_gen/models/ltx23_retake.py | 135 +++ .../attention/backends/flashinfer.py | 153 +++- tensorrt_llm/_torch/visual_gen/config.py | 29 +- .../_torch/visual_gen/models/__init__.py | 3 + .../ltx2/ltx2_core/video_vae/video_vae.py | 6 +- .../visual_gen/models/ltx2/pipeline_ltx2.py | 20 +- .../models/ltx2/transformer_ltx2.py | 11 +- .../_torch/visual_gen/models/ltx23/NOTICE | 6 + .../visual_gen/models/ltx23/__init__.py | 11 + .../visual_gen/models/ltx23/checkpoint.py | 45 + .../models/ltx23/ltx23_core/__init__.py | 4 + .../ltx23/ltx23_core/audio_vae/__init__.py | 15 + .../ltx23/ltx23_core/audio_vae/activations.py | 152 ++++ .../ltx23/ltx23_core/audio_vae/amp_block.py | 46 + .../ltx23/ltx23_core/audio_vae/audio_vae.py | 178 ++++ .../models/ltx23/ltx23_core/audio_vae/bw.py | 95 ++ .../ltx23/ltx23_core/audio_vae/downsample.py | 97 ++ .../audio_vae/model_configurator.py | 53 ++ .../ltx23/ltx23_core/audio_vae/vocoder.py | 152 ++++ .../models/ltx23/ltx23_core/connector.py | 106 +++ .../models/ltx23/ltx23_core/modality.py | 26 + .../models/ltx23/ltx23_core/video_vae.py | 295 +++++++ .../visual_gen/models/ltx23/media_io.py | 211 +++++ .../visual_gen/models/ltx23/pipeline_ltx23.py | 666 ++++++++++++++ .../models/ltx23/pipeline_ltx23_retake.py | 832 ++++++++++++++++++ .../models/ltx23/transformer_ltx23.py | 493 +++++++++++ .../_torch/visual_gen/pipeline_registry.py | 9 +- .../ltx23_lpips_golden_video.json | 33 + .../ltx23_retake_lpips_golden_video.json | 34 + .../visual_gen_lpips_golden_media_ltx23.zip | 3 + .../visual_gen/test_minimax_h3_e2e.py | 35 +- .../visual_gen/test_visual_gen_ltx23.py | 458 ++++++++++ .../visual_gen/visual_gen_test_utils.py | 73 +- .../test_lists/test-db/l0_b200.yml | 6 + .../integration/test_lists/test-db/l0_cpu.yml | 3 + .../test_lists/test-db/l0_dgx_b200.yml | 2 + .../visual_gen/test_ltx23_pipeline_config.py | 150 ++++ .../_torch/visual_gen/test_ltx23_registry.py | 64 ++ .../test_ltx23_retake_initialization.py | 100 +++ 51 files changed, 5123 insertions(+), 60 deletions(-) create mode 100644 examples/visual_gen/configs/ltx23-retake-bf16-1gpu.yaml create mode 100644 examples/visual_gen/configs/ltx23-retake-bf16-2gpu.yaml create mode 100644 examples/visual_gen/configs/ltx23-retake-fp4-1gpu.yaml create mode 100644 examples/visual_gen/configs/ltx23-retake-fp8-1gpu.yaml create mode 100644 examples/visual_gen/configs/ltx23-t2v-bf16-1gpu.yaml create mode 100644 examples/visual_gen/configs/ltx23-t2v-bf16-2gpu.yaml create mode 100644 examples/visual_gen/configs/ltx23-t2v-fp4-1gpu.yaml create mode 100644 examples/visual_gen/configs/ltx23-t2v-fp8-1gpu.yaml create mode 100644 examples/visual_gen/models/ltx23.py create mode 100644 examples/visual_gen/models/ltx23_retake.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/ltx23/NOTICE create mode 100644 tensorrt_llm/_torch/visual_gen/models/ltx23/__init__.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/ltx23/checkpoint.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/__init__.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/__init__.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/activations.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/amp_block.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/audio_vae.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/bw.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/downsample.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/model_configurator.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/vocoder.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/connector.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/modality.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/video_vae.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/ltx23/media_io.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/ltx23/pipeline_ltx23.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/ltx23/pipeline_ltx23_retake.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/ltx23/transformer_ltx23.py create mode 100644 tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/ltx23_lpips_golden_video.json create mode 100644 tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/ltx23_retake_lpips_golden_video.json create mode 100644 tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media_ltx23.zip create mode 100644 tests/integration/defs/examples/visual_gen/test_visual_gen_ltx23.py create mode 100644 tests/unittest/_torch/visual_gen/test_ltx23_pipeline_config.py create mode 100644 tests/unittest/_torch/visual_gen/test_ltx23_registry.py create mode 100644 tests/unittest/_torch/visual_gen/test_ltx23_retake_initialization.py diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index ee61a8d5d012..53ddd835f8b4 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -218,6 +218,7 @@ For full documentation, see the [Visual Generation](./visual-generation.md) page | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | Image-to-Video | | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | Text-to-Video, Image-to-Video | | `Lightricks/LTX-2` | Text-to-Video (with Audio), Image-to-Video (with Audio) | +| `Lightricks/LTX-2.3` | Text-to-Video (with Audio), Video Retake (with Audio) | | `MiniMaxAI/MiniMax-H3` | Text-to-Video (with Audio), First/Last-Frame-to-Video (with Audio) | | `Qwen/Qwen-Image` | Text-to-Image | | `Qwen/Qwen-Image-2512` | Text-to-Image | @@ -238,6 +239,7 @@ For full documentation, see the [Visual Generation](./visual-generation.md) page | **Wan 2.1** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | | **Wan 2.2** | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | | **LTX-2** | Yes | Yes | No | Yes | Yes | No | No | Yes | Yes | Yes | Yes | No | +| **LTX-2.3** | Yes | Yes | No | No | Yes | Yes | Yes | Yes | Yes | No | No | No | | **MiniMax-H3** | Yes | Yes | No | No | Yes | Yes | No | Yes | Yes | No | No | No | | **Qwen-Image** | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No | | **Qwen-Image-Layered** | Yes | No | No | No | No | No | Yes | Yes | Yes | No | No | No | diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index c25b43e6a258..1d2b21aa0530 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -37,6 +37,7 @@ TensorRT-LLM **VisualGen** provides a unified inference stack for diffusion mode | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | Text-to-Video, Image-to-Video | | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | Text-to-Video (3-step distilled) | | `Lightricks/LTX-2` | Text-to-Video (with Audio), Image-to-Video (with Audio) | +| `Lightricks/LTX-2.3` | Text-to-Video (with Audio), Video Retake (with Audio) | | `Qwen/Qwen-Image` | Text-to-Image | | `Qwen/Qwen-Image-2512` | Text-to-Image | | `Qwen/Qwen-Image-Layered` | Image-to-Image | @@ -65,6 +66,7 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models | **Wan 2.2** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | | **FastWan 2.2** | Yes | Yes | No | No | No | No | No | No | Yes | Yes | Yes | No | No | No | No | | **LTX-2** | Yes | Yes | Yes | Yes | No | Yes | Yes | No | No | Yes | Yes | Yes | Yes | No | Yes | +| **LTX-2.3** | Yes | Yes | No | No | No | No | Yes | Yes | Yes | Yes | Yes | No | No | No | No | | **MiniMax-H3** | Yes | Yes | No | No | No | No | Yes | Yes | No | Yes | Yes | No | No | No | No | | **Qwen-Image** | Yes | Yes | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No | No | | **Qwen-Image-Layered** | No | No | No | No | No | No | No | No | Yes | Yes | Yes | No | No | No | No | @@ -392,7 +394,7 @@ Configured under `VisualGenArgs.parallel_config`. Modes can be combined: - **CFG Parallelism** (`cfg_size: 2`): Splits positive/negative guidance prompts across GPUs. FLUX uses embedded guidance without a separate negative prompt path; CFG parallelism is not applicable to FLUX or the distilled FastWan 2.2 model. - **Ulysses Parallelism** (`ulysses_size: N`): Splits the sequence dimension across GPUs for longer sequences. - **Async Ulysses A2A pipeline** (`async_ulysses: true` in `parallel_config`): Overlaps per-rank V/Q/K projection compute with the cross-rank all-to-all on a dedicated side stream. Requires `ulysses_size > 1` and an NVLink-connected GPU domain (uses PyTorch `_SymmetricMemory` with CUDA IPC for peer pushes; not currently supported across nodes without MNNVL). Currently wired for WAN and LTX-2 self-attention. -- **Parallel VAE** (`parallel_vae_size: N`): Parallelizes the final VAE decode (constraint: `parallel_vae_size ≤ world_size`). WAN/Cosmos3 shard along a spatial axis; LTX-2 and MiniMax-H3 distribute independent tiles. +- **Parallel VAE** (`parallel_vae_size: N`): Parallelizes the final VAE decode (constraint: `parallel_vae_size ≤ world_size`). WAN/Cosmos3 shard along a spatial axis; LTX-2, LTX-2.3, and MiniMax-H3 distribute independent tiles. - **Context Parallel (CP)** — Partitions the sequence into shards so that each rank computes partial attention. Requires an LSE-capable attention backend (`FA4` or `CUTEDSL`). CP can be composed with Ulysses, giving a total sequence-parallel (SP) degree = `cp_size · ulysses_size`. The CP degree depends on the implementation below: - **Attention2D** (`attn2d_size: [N, M]`): Shards the sequence axis across an `N × M` device mesh (CP degree = `N · M`; total SP degree = `N · M · ulysses_size`). - **Ring Attention** (`ring_size: N`): Shards the sequence axis across a 1D ring of `N` ranks, streaming K/V blocks (CP degree = `N`; total SP degree = `N · ulysses_size`; mutually exclusive with Attention2D). diff --git a/examples/visual_gen/README.md b/examples/visual_gen/README.md index c7cf8d994b89..ba7bf89c8819 100644 --- a/examples/visual_gen/README.md +++ b/examples/visual_gen/README.md @@ -16,6 +16,7 @@ for feature details. ```bash python quickstart_example.py +python models/ltx23.py python models/glm_image.py # With engine config (quant, parallelism, etc.) @@ -23,6 +24,12 @@ python models/glm_image.py 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-1gpu.yaml +python models/ltx23.py --visual_gen_args configs/ltx23-t2v-bf16-1gpu.yaml +python models/ltx23.py --visual_gen_args configs/ltx23-t2v-fp8-1gpu.yaml +python models/ltx23.py --visual_gen_args configs/ltx23-t2v-fp4-1gpu.yaml +python models/ltx23_retake.py --visual_gen_args configs/ltx23-retake-bf16-1gpu.yaml --source /path/to/retake_input.mp4 +python models/ltx23_retake.py --visual_gen_args configs/ltx23-retake-fp8-1gpu.yaml --source /path/to/retake_input.mp4 +python models/ltx23_retake.py --visual_gen_args configs/ltx23-retake-fp4-1gpu.yaml --source /path/to/retake_input.mp4 # FP8 blockwise config: configs/minimax-h3-fp8-blockwise-1gpu.yaml. # 8-GPU Ulysses config: configs/minimax-h3-bf16-8gpu.yaml. python models/minimax_h3.py --model --visual_gen_args configs/minimax-h3-bf16-1gpu.yaml diff --git a/examples/visual_gen/configs/ltx23-retake-bf16-1gpu.yaml b/examples/visual_gen/configs/ltx23-retake-bf16-1gpu.yaml new file mode 100644 index 000000000000..edb68ded4489 --- /dev/null +++ b/examples/visual_gen/configs/ltx23-retake-bf16-1gpu.yaml @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 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 BF16 LTX-2.3 retake. +pipeline_config: + workflow: retake +attention_config: + backend: VANILLA +parallel_config: + cfg_size: 1 + ulysses_size: 1 +torch_compile_config: + enable: true +cuda_graph_config: + enable: true diff --git a/examples/visual_gen/configs/ltx23-retake-bf16-2gpu.yaml b/examples/visual_gen/configs/ltx23-retake-bf16-2gpu.yaml new file mode 100644 index 000000000000..9e038216aa7f --- /dev/null +++ b/examples/visual_gen/configs/ltx23-retake-bf16-2gpu.yaml @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 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. + +# 2-GPU BF16 LTX-2.3 retake. +pipeline_config: + workflow: retake +attention_config: + backend: VANILLA +parallel_config: + cfg_size: 1 + ulysses_size: 2 + parallel_vae_size: 2 +torch_compile_config: + enable: true +cuda_graph_config: + enable: true diff --git a/examples/visual_gen/configs/ltx23-retake-fp4-1gpu.yaml b/examples/visual_gen/configs/ltx23-retake-fp4-1gpu.yaml new file mode 100644 index 000000000000..0664904fbf6f --- /dev/null +++ b/examples/visual_gen/configs/ltx23-retake-fp4-1gpu.yaml @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 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.3 retake with dynamic NVFP4 transformer quantization at every step +# and FlashInfer NVFP4 Q/K + FP8 V video self-attention on SM100/SM103 GPUs. +pipeline_config: + workflow: retake +quant_config: + quant_algo: NVFP4 + dynamic: true +attention_config: + backend: FLASHINFER + quant_attention_config: + qk_dtype: nvfp4 + v_dtype: fp8 +parallel_config: + cfg_size: 1 + ulysses_size: 1 +cuda_graph_config: + enable: true diff --git a/examples/visual_gen/configs/ltx23-retake-fp8-1gpu.yaml b/examples/visual_gen/configs/ltx23-retake-fp8-1gpu.yaml new file mode 100644 index 000000000000..dec5ec1c7e42 --- /dev/null +++ b/examples/visual_gen/configs/ltx23-retake-fp8-1gpu.yaml @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 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.3 retake with dynamic FP8 block-scale transformer quantization +# and FlashInfer MXFP8 Q/K + FP8 V video self-attention on SM100/SM103 GPUs. +pipeline_config: + workflow: retake +quant_config: + quant_algo: FP8_BLOCK_SCALES + dynamic: true +attention_config: + backend: FLASHINFER + quant_attention_config: + qk_dtype: mxfp8 + v_dtype: fp8 +parallel_config: + cfg_size: 1 + ulysses_size: 1 +cuda_graph_config: + enable: true diff --git a/examples/visual_gen/configs/ltx23-t2v-bf16-1gpu.yaml b/examples/visual_gen/configs/ltx23-t2v-bf16-1gpu.yaml new file mode 100644 index 000000000000..e90782660e5c --- /dev/null +++ b/examples/visual_gen/configs/ltx23-t2v-bf16-1gpu.yaml @@ -0,0 +1,25 @@ +# 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.3 text-to-video with audio, BF16. +attention_config: + backend: VANILLA +parallel_config: + cfg_size: 1 + ulysses_size: 1 +torch_compile_config: + enable: true +cuda_graph_config: + enable: true diff --git a/examples/visual_gen/configs/ltx23-t2v-bf16-2gpu.yaml b/examples/visual_gen/configs/ltx23-t2v-bf16-2gpu.yaml new file mode 100644 index 000000000000..1c804e4258ef --- /dev/null +++ b/examples/visual_gen/configs/ltx23-t2v-bf16-2gpu.yaml @@ -0,0 +1,26 @@ +# 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. + +# 2-GPU LTX-2.3 text-to-video with audio, BF16. +attention_config: + backend: VANILLA +parallel_config: + cfg_size: 1 + ulysses_size: 2 + parallel_vae_size: 2 +torch_compile_config: + enable: true +cuda_graph_config: + enable: true diff --git a/examples/visual_gen/configs/ltx23-t2v-fp4-1gpu.yaml b/examples/visual_gen/configs/ltx23-t2v-fp4-1gpu.yaml new file mode 100644 index 000000000000..dd127dd145ba --- /dev/null +++ b/examples/visual_gen/configs/ltx23-t2v-fp4-1gpu.yaml @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 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.3 text-to-video with audio, dynamic NVFP4 transformer weights +# and FlashInfer NVFP4 Q/K + FP8 V video self-attention on SM100/SM103 GPUs. +quant_config: + quant_algo: NVFP4 + dynamic: true +attention_config: + backend: FLASHINFER + quant_attention_config: + qk_dtype: nvfp4 + v_dtype: fp8 +parallel_config: + cfg_size: 1 + ulysses_size: 1 +torch_compile_config: + enable: true +cuda_graph_config: + enable: true diff --git a/examples/visual_gen/configs/ltx23-t2v-fp8-1gpu.yaml b/examples/visual_gen/configs/ltx23-t2v-fp8-1gpu.yaml new file mode 100644 index 000000000000..e22e48bc5532 --- /dev/null +++ b/examples/visual_gen/configs/ltx23-t2v-fp8-1gpu.yaml @@ -0,0 +1,32 @@ +# 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.3 text-to-video with audio, dynamic FP8 transformer weights +# and FlashInfer MXFP8 Q/K + FP8 V video self-attention on SM100/SM103 GPUs. +quant_config: + quant_algo: FP8_BLOCK_SCALES + dynamic: true +attention_config: + backend: FLASHINFER + quant_attention_config: + qk_dtype: mxfp8 + v_dtype: fp8 +parallel_config: + cfg_size: 1 + ulysses_size: 1 +torch_compile_config: + enable: true +cuda_graph_config: + enable: true diff --git a/examples/visual_gen/models/ltx23.py b/examples/visual_gen/models/ltx23.py new file mode 100644 index 000000000000..4b7ed42e302e --- /dev/null +++ b/examples/visual_gen/models/ltx23.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python3 +# 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. +"""LTX-2.3 Text-to-Video generation with audio. + +Usage: + python ltx23.py + python ltx23.py --visual_gen_args ../configs/ltx23-t2v-bf16-1gpu.yaml + # Optional: override the default HuggingFace model with a local checkpoint. + python ltx23.py --model /path/to/ltx-2.3-22b-dev.safetensors +""" + +import argparse +from dataclasses import replace + +import torch.nn.functional as F + +from tensorrt_llm import VisualGen, VisualGenArgs +from tensorrt_llm._torch.visual_gen.models.ltx23.checkpoint import ( + LTX23_MODEL_ID, + resolve_ltx23_checkpoint, +) + + +def main() -> None: + parser = argparse.ArgumentParser(description="LTX-2.3 Text-to-Video example") + parser.add_argument( + "--model", + type=str, + default=LTX23_MODEL_ID, + help="Checkpoint file, directory, or Hub ID. Directories/Hub IDs select the dev checkpoint.", + ) + parser.add_argument( + "--visual_gen_args", + dest="visual_gen_args", + type=str, + default=None, + help="Path to YAML config (same as trtllm-serve --visual_gen_args)", + ) + parser.add_argument( + "--text_encoder_path", + type=str, + default=None, + help=( + "Gemma3 text encoder path. Overrides pipeline_config.text_encoder_path " + "from --visual_gen_args when set." + ), + ) + parser.add_argument( + "--output_path", + type=str, + default="ltx23_t2v_output.mp4", + help="Path to save the output video", + ) + args = parser.parse_args() + + extra_args = ( + VisualGenArgs.from_yaml(args.visual_gen_args) if args.visual_gen_args else VisualGenArgs() + ) + text_encoder_path = args.text_encoder_path + if text_encoder_path is not None: + extra_args.pipeline_config = { + **extra_args.pipeline_config, + "text_encoder_path": text_encoder_path, + } + extra_args.compilation_config.resolutions = [(512, 768)] + extra_args.compilation_config.num_frames = [49] + model = resolve_ltx23_checkpoint( + args.model, workflow="generation", revision=extra_args.revision + ) + visual_gen = VisualGen(model=model, args=extra_args) + + params = visual_gen.default_params + params.height = 512 + params.width = 768 + params.num_frames = 49 + params.frame_rate = 24.0 + params.num_inference_steps = 8 + params.guidance_scale = 5.0 + params.seed = 42 + params.negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted" + + output = visual_gen.generate( + inputs=( + "A woman with long brown hair and light skin smiles at the camera while standing in a " + "sunlit park, her hair gently blowing in the breeze as she tilts her head slightly to the side." + ), + params=params, + ) + + metrics = output.metrics + if metrics is not None: + print( + f"latency: generation={metrics.generation:.3f}s " + f"pre_denoise={metrics.pre_denoise:.3f}s " + f"denoise={metrics.denoise:.3f}s " + f"post_denoise={metrics.post_denoise:.3f}s" + ) + + output_to_save = output + if output.video is not None and output.audio is not None: + if output.frame_rate is None or output.audio_sample_rate is None: + raise ValueError("LTX-2.3 audio/video output must carry frame and sample rates.") + audio_samples = round(output.video.shape[-4] * output.audio_sample_rate / output.frame_rate) + padding = max(0, audio_samples - output.audio.shape[-1]) + if padding: + # Pad only the saved copy; preserve the generated waveform for scoring. + output_to_save = replace(output, audio=F.pad(output.audio, (0, padding))) + output_to_save.save(args.output_path) + print(f"Saved: {args.output_path}") + + +if __name__ == "__main__": + main() diff --git a/examples/visual_gen/models/ltx23_retake.py b/examples/visual_gen/models/ltx23_retake.py new file mode 100644 index 000000000000..b12a5b8634e5 --- /dev/null +++ b/examples/visual_gen/models/ltx23_retake.py @@ -0,0 +1,135 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 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. +"""Run one-GPU native LTX-2.3 retake with the public VisualGen API. + +Defaults to the LTX-2.3 distilled 1.1 checkpoint. +Use ``--model`` to select a checkpoint file, directory, or HuggingFace Hub ID. +Use ``--text_encoder_path`` to override the Gemma text encoder, or pass +precomputed prompt conditioning with ``--prompt_conditioning_path``. Source +decoding requires PyAV, and audio conditioning requires a ``torchaudio`` build +compatible with the installed PyTorch. +Both video and audio are regenerated in the selected window by default. +Use ``--no-regenerate_video`` or ``--no-regenerate_audio`` to freeze a modality. +Frozen streams still pass through the native VAE; they are not file copies. +An input without an audio track generates audio from scratch, even when +``--no-regenerate_audio`` is set, matching the upstream retake pipeline. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from tensorrt_llm import VisualGen, VisualGenArgs +from tensorrt_llm._torch.visual_gen.models.ltx23.checkpoint import ( + LTX23_MODEL_ID, + resolve_ltx23_checkpoint, +) + +_DEFAULT_CONFIG = Path(__file__).resolve().parents[1] / "configs" / "ltx23-retake-bf16-1gpu.yaml" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--model", + type=str, + default=LTX23_MODEL_ID, + help="Checkpoint file, directory, or Hub ID. Directories/Hub IDs select distilled 1.1.", + ) + parser.add_argument("--source", required=True, help="Source MP4 to retake.") + parser.add_argument( + "--start", + type=float, + default=None, + help="Retake start time in seconds. Defaults to the beginning of the video.", + ) + parser.add_argument( + "--end", + type=float, + default=None, + help="Retake end time in seconds. Defaults to the end of the video.", + ) + parser.add_argument("--output_path", default="retake_output.mp4", help="Output MP4 path.") + parser.add_argument("--prompt", default="", help="Retake text prompt.") + parser.add_argument( + "--regenerate_audio", + action=argparse.BooleanOptionalAction, + default=True, + help="Regenerate source audio in the window (default: true).", + ) + parser.add_argument( + "--regenerate_video", + action=argparse.BooleanOptionalAction, + default=True, + help="Regenerate video in the window (default: true).", + ) + parser.add_argument( + "--visual_gen_args", + default=str(_DEFAULT_CONFIG), + help="VisualGen YAML configuration.", + ) + parser.add_argument( + "--text_encoder_path", + default=None, + help="Optional Gemma text encoder path override.", + ) + parser.add_argument( + "--prompt_conditioning_path", + default=None, + help="Precomputed prompt conditioning safetensors path.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + visual_gen_args = VisualGenArgs.from_yaml(args.visual_gen_args) + pipeline_config = dict(visual_gen_args.pipeline_config) + if args.prompt_conditioning_path is not None: + if args.prompt: + raise ValueError("--prompt and --prompt_conditioning_path are mutually exclusive.") + if args.text_encoder_path is not None: + raise ValueError( + "--text_encoder_path and --prompt_conditioning_path are mutually exclusive." + ) + pipeline_config["text_encoder_path"] = None + elif args.text_encoder_path is not None: + pipeline_config["text_encoder_path"] = args.text_encoder_path + visual_gen_args = visual_gen_args.model_copy(update={"pipeline_config": pipeline_config}) + + model = resolve_ltx23_checkpoint( + args.model, workflow="retake", revision=visual_gen_args.revision + ) + with VisualGen(model=model, args=visual_gen_args) as visual_gen: + params = visual_gen.default_params + params.extra_params["retake_video_path"] = args.source + params.extra_params["retake_regenerate_video"] = args.regenerate_video + params.extra_params["retake_regenerate_audio"] = args.regenerate_audio + if args.prompt_conditioning_path is not None: + params.extra_params["retake_prompt_conditioning_path"] = args.prompt_conditioning_path + if args.start is not None: + params.extra_params["retake_start_time"] = args.start + if args.end is not None: + params.extra_params["retake_end_time"] = args.end + output = visual_gen.generate(inputs=args.prompt, params=params) + output.save(args.output_path) + + print(f"Saved: {args.output_path}") + + +if __name__ == "__main__": + main() diff --git a/tensorrt_llm/_torch/visual_gen/attention/backends/flashinfer.py b/tensorrt_llm/_torch/visual_gen/attention/backends/flashinfer.py index 168f36d64f64..d7e7d63ee909 100644 --- a/tensorrt_llm/_torch/visual_gen/attention/backends/flashinfer.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/flashinfer.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """FlashInfer attention backend for visual generation models.""" +import hashlib import math from typing import Any, Literal @@ -48,6 +49,7 @@ def __init__( self._metadata_state = ( attention_metadata_state if attention_metadata_state is not None else {} ) + self._static_global_scale = self._metadata_state.get("_ltx23_static_global_scale") == 1.0 self._single_prefill = flashinfer.single_prefill_with_kv_cache self._batch_prefill_cls = flashinfer.BatchPrefillWithRaggedKVCacheWrapper @@ -143,6 +145,10 @@ def _run_batch_prefill( is_causal: bool, key_padding_mask: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor]: + if self._static_global_scale: + return self._run_capture_safe_batch_prefill( + q, k, v, is_causal=is_causal, key_padding_mask=key_padding_mask + ) batch_size, query_length = q.shape[:2] key_value_length = k.shape[1] stream = torch.cuda.current_stream(q.device) @@ -240,6 +246,142 @@ def _run_batch_prefill( # Batch prefill reports log2 LSE; VisualGen exposes natural-log LSE. return output, lse * math.log(2.0) + def _prefill_mask_digest(self, mask: torch.Tensor | None, *, capturing: bool) -> str | None: + if mask is None: + return None + version = None if torch.is_inference(mask) else mask._version + identity = (id(mask), mask.data_ptr(), version) + identities = self._metadata_state.setdefault("flashinfer_mask_identities", {}) + if capturing: + if identity not in identities: + raise RuntimeError("FlashInfer padding mask must be warmed before CUDA capture.") + return identities[identity][1] + digest = hashlib.sha256(mask.detach().to("cpu", torch.bool).numpy().tobytes()).hexdigest() + # Retain the tensor so its identity cannot be recycled while a plan uses it. + identities[identity] = (mask, digest) + return digest + + def _run_capture_safe_batch_prefill( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + is_causal: bool, + key_padding_mask: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Reuse immutable mask-content plans across warmup and capture streams. + + The component's shared metadata state requires serialized forwards. + Planning and mask readback occur only outside capture; a changed mask + owns a new wrapper rather than mutating a plan retained by an old graph. + """ + capturing = torch.cuda.is_current_stream_capturing() + mask_digest = self._prefill_mask_digest(key_padding_mask, capturing=capturing) + batch_size, query_length = q.shape[:2] + key_value_length = k.shape[1] + plan_key = ( + q.device, + batch_size, + query_length, + key_value_length, + self.num_heads, + self.num_kv_heads, + self.head_dim, + q.dtype, + is_causal, + mask_digest, + ) + plans = self._metadata_state.setdefault("flashinfer_immutable_prefill_plans", {}) + entry = plans.get(plan_key) + if entry is None: + if capturing: + raise RuntimeError("FlashInfer batched prefill must be warmed before CUDA capture.") + workspaces = self._metadata_state.setdefault("flashinfer_immutable_workspaces", {}) + workspace = workspaces.get(q.device) + if workspace is None: + workspace = torch.empty(_WORKSPACE_BYTES, dtype=torch.uint8, device=q.device) + workspaces[q.device] = workspace + wrapper = self._batch_prefill_cls(workspace, "NHD") + query_indptr = ( + torch.arange(batch_size + 1, dtype=torch.int32, device=q.device) * query_length + ) + key_value_indptr = ( + torch.arange(batch_size + 1, dtype=torch.int32, device=q.device) * key_value_length + ) + custom_mask = None + if key_padding_mask is not None: + custom_mask = ( + key_padding_mask.to(device=q.device, dtype=torch.bool) + .unsqueeze(1) + .expand(-1, query_length, -1) + .reshape(-1) + .clone() + ) + torch.cuda.current_stream(q.device).synchronize() + wrapper.plan( + query_indptr, + key_value_indptr, + self.num_heads, + self.num_kv_heads, + self.head_dim, + custom_mask=custom_mask, + causal=is_causal, + sm_scale=self.scale, + q_data_type=q.dtype, + kv_data_type=k.dtype, + o_data_type=q.dtype, + ) + entry = (wrapper, query_indptr, key_value_indptr, custom_mask) + plans[plan_key] = entry + output, lse = entry[0].run( + q.contiguous().view(-1, self.num_heads, self.head_dim), + k.contiguous().view(-1, self.num_kv_heads, self.head_dim), + v.contiguous().view(-1, self.num_kv_heads, self.head_dim), + return_lse=True, + ) + return ( + output.view(batch_size, query_length, self.num_heads, self.head_dim), + lse.view(batch_size, query_length, self.num_heads).transpose(1, 2) * math.log(2.0), + ) + + @staticmethod + def _quantize_qk_static_global( + tensor: torch.Tensor, mode: _BlockScaledQKMode + ) -> tuple[torch.Tensor, torch.Tensor]: + """Encode NHD Q/K with global one and native, input-dependent block scales.""" + from flashinfer.quantization import fp4_quantize, mxfp8_quantize + + batch, sequence, heads, dimension = tensor.shape + padded_sequence = ((sequence + 127) // 128) * 128 + data = tensor.transpose(1, 2).contiguous() + if padded_sequence != sequence: + data = torch.nn.functional.pad(data, (0, 0, 0, padded_sequence - sequence)) + matrix = data.reshape(batch * heads * padded_sequence, dimension) + if mode == "mxfp8": + encoded, factors = mxfp8_quantize( + matrix, is_sf_swizzled_layout=True, alignment=32, backend="cute-dsl" + ) + encoded = encoded.view(batch, heads, padded_sequence, dimension)[:, :, :sequence] + factors = factors.view(torch.float8_e8m0fnu) + elif mode == "nvfp4": + # fp4_quantize takes the inverse global dequantization scale. Both + # encode and FMHA decode use one; only native block factors vary. + inverse_global = torch.ones(1, device=tensor.device, dtype=torch.float32) + encoded, factors = fp4_quantize( + matrix, + inverse_global, + sf_vec_size=16, + is_sf_swizzled_layout=True, + backend="cute-dsl", + ) + encoded = encoded.view(torch.float4_e2m1fn_x2) + encoded = encoded.view(batch, heads, padded_sequence, dimension // 2)[:, :, :sequence] + factors = factors.view(torch.float8_e4m3fn) + else: + raise ValueError(f"Unsupported static-global Q/K mode: {mode!r}.") + return encoded.transpose(1, 2).contiguous(), factors + @staticmethod def _quantize_fp8_per_tensor(tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: fp8_dtype = torch.float8_e4m3fn @@ -265,8 +407,14 @@ def _run_blockscaled_sm10x( from flashinfer.attention.cute_dsl.fmha_blockscaled import cute_dsl_fmha_blockscaled_prefill from flashinfer.cute_dsl.attention.fmha.quantize import quantize_blockscaled_qk - q_store, k_store, q_sf, k_sf, q_scale, k_scale = quantize_blockscaled_qk(q, k, qk_mode) - v_store, v_scale = self._quantize_fp8_per_tensor(v) + if self._static_global_scale: + q_store, q_sf = self._quantize_qk_static_global(q, qk_mode) + k_store, k_sf = self._quantize_qk_static_global(k, qk_mode) + v_store = v.clamp(-448.0, 448.0).to(torch.float8_e4m3fn).contiguous() + q_scale = k_scale = v_scale = 1.0 + else: + q_store, k_store, q_sf, k_sf, q_scale, k_scale = quantize_blockscaled_qk(q, k, qk_mode) + v_store, v_scale = self._quantize_fp8_per_tensor(v) output = torch.empty_like(q) lse = torch.empty( (q.shape[0], q.shape[1], self.num_heads), @@ -287,6 +435,7 @@ def _run_blockscaled_sm10x( scale_q=q_scale, scale_k=k_scale, scale_v=v_scale, + scale_o=1.0, ) # CuTe DSL reports log2 LSE; VisualGen exposes natural-log LSE. return output, lse.transpose(1, 2).contiguous() * math.log(2.0) diff --git a/tensorrt_llm/_torch/visual_gen/config.py b/tensorrt_llm/_torch/visual_gen/config.py index ea7cbc861762..8c0b8aab1936 100644 --- a/tensorrt_llm/_torch/visual_gen/config.py +++ b/tensorrt_llm/_torch/visual_gen/config.py @@ -635,7 +635,20 @@ def from_pretrained( resolved_pipeline_config = dict(args.pipeline_config) if args else {} extra_attrs.update(resolved_pipeline_config) - h3_workflow = resolved_pipeline_config.get("workflow") + model_index_path = checkpoint_path / "model_index.json" + if not model_index_path.exists(): + model_index_path = checkpoint_path / "modular_model_index.json" + model_index = {} + if model_index_path.exists(): + with open(model_index_path) as f: + model_index = json.load(f) + + h3_workflow = None + if model_index.get("_class_name") == "MiniMaxH3ModularPipeline" or ( + model_index.get("_class_name") is None + and model_index.get("transformer") == ["diffusers", "MiniMaxH3Transformer3DModel"] + ): + h3_workflow = resolved_pipeline_config.get("workflow") if h3_workflow is not None: if h3_workflow not in ("fl2va", "ref2va"): raise ValueError("MiniMax-H3 workflow must be 'fl2va' or 'ref2va'.") @@ -670,12 +683,7 @@ def from_pretrained( if not getattr(pretrained_config, "_name_or_path", None): pretrained_config._name_or_path = str(checkpoint_path) - model_index_path = checkpoint_path / "model_index.json" - if not model_index_path.exists(): - model_index_path = checkpoint_path / "modular_model_index.json" if model_index_path.exists(): - with open(model_index_path) as f: - model_index = json.load(f) if "boundary_ratio" in model_index and "transformer_2" in model_index: transformer_2_spec = model_index.get("transformer_2") if transformer_2_spec and transformer_2_spec[0] is not None: @@ -687,6 +695,15 @@ def from_pretrained( native_config = cls._try_load_safetensors_config(checkpoint_path) if native_config is not None: + if "transformer" in native_config and ( + "vae" in native_config or "audio_vae" in native_config + ): + from .models.ltx23.pipeline_ltx23 import detect_native_ltx_pipeline + + if detect_native_ltx_pipeline(native_config) == "LTX23Pipeline": + ltx23_workflow = resolved_pipeline_config.get("workflow", "generation") + if ltx23_workflow not in ("generation", "retake"): + raise ValueError("LTX-2.3 workflow must be 'generation' or 'retake'.") transformer_dict = native_config.get("transformer", {}) component_config_dicts["transformer"] = transformer_dict transformer_2_dict = native_config.get("transformer_2") diff --git a/tensorrt_llm/_torch/visual_gen/models/__init__.py b/tensorrt_llm/_torch/visual_gen/models/__init__.py index 29495a37c540..94b99fadf1e5 100644 --- a/tensorrt_llm/_torch/visual_gen/models/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/models/__init__.py @@ -38,6 +38,7 @@ from .glm_image import GlmImagePipeline from .hunyuan_video1_5 import HunyuanVideo15Pipeline from .ltx2 import LTX2Pipeline # noqa: F401 +from .ltx23 import LTX23Pipeline, LTX23RetakePipeline from .minimax_h3 import MiniMaxH3Pipeline from .qwen_image import QwenImageEditPlusPipeline, QwenImagePipeline from .qwen_image_layered import QwenImageLayeredPipeline @@ -57,6 +58,8 @@ "WanDMDPipeline", "WanImageToVideoPipeline", "Cosmos3OmniMoTPipeline", + "LTX23Pipeline", + "LTX23RetakePipeline", "register_pipeline", "HunyuanVideo15Pipeline", ] diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/video_vae/video_vae.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/video_vae/video_vae.py index 49ffda1fb215..80cf165fafc5 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/video_vae/video_vae.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/video_vae/video_vae.py @@ -326,6 +326,8 @@ def _make_decoder_block( class VideoDecoder(nn.Module): _DEFAULT_NORM_NUM_GROUPS = 32 + _channel_reducing_blocks = ("compress_all",) + _make_decoder_block = staticmethod(_make_decoder_block) def __init__( self, @@ -354,7 +356,7 @@ def __init__( block_config = block_params if isinstance(block_params, dict) else {} if block_name == "res_x_y": feature_channels = feature_channels * block_config.get("multiplier", 2) - if block_name == "compress_all": + if block_name in self._channel_reducing_blocks: feature_channels = feature_channels * block_config.get("multiplier", 1) self.conv_in = make_conv_nd( dims=convolution_dimensions, @@ -371,7 +373,7 @@ def __init__( block_config = ( {"num_layers": block_params} if isinstance(block_params, int) else block_params ) - block, feature_channels = _make_decoder_block( + block, feature_channels = self._make_decoder_block( block_name=block_name, block_config=block_config, in_channels=feature_channels, diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py index 2b8d1685bc9b..b4e3caaacf44 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py @@ -7,6 +7,7 @@ import json import os import time +from dataclasses import replace from io import BytesIO from pathlib import Path from typing import Any, Dict, List, Optional, Set, Tuple, Union @@ -421,8 +422,9 @@ def _clone_value(v): if isinstance(v, torch.Tensor): return v.clone() if isinstance(v, Modality): - # NOTE: must be kept in sync with Modality dataclass fields. - return Modality( + # Preserve modality subclasses when cloning the shared tensor fields. + return replace( + v, latent=v.latent.clone(), timesteps=v.timesteps.clone(), positions=v.positions.clone(), @@ -636,6 +638,10 @@ class LTX2Pipeline(BasePipeline): ``transformers`` library. """ + _transformer_class = LTXModel + _cuda_graph_runner_class = _LTX2CUDAGraphRunner + _supports_quantized_attention = False + def __init__(self, model_config): _sa_cfg = model_config.attention.sparse_attention_config if _sa_cfg is not None and getattr(_sa_cfg, "algorithm", None) == "vsa": @@ -761,7 +767,11 @@ def _init_transformer(self) -> None: fall back to the same defaults the reference uses. """ attn_cfg = getattr(self.pipeline_config, "attention", None) - if attn_cfg is not None and getattr(attn_cfg, "quant_attention_config", None) is not None: + if ( + attn_cfg is not None + and getattr(attn_cfg, "quant_attention_config", None) is not None + and not self._supports_quantized_attention + ): raise NotImplementedError( "Quantized attention is not yet supported for the LTX-2 pipeline." ) @@ -780,7 +790,7 @@ def _init_transformer(self) -> None: f"apply_gated_attention={apply_gated_attention}" ) - self.transformer = LTXModel( + self.transformer = self._transformer_class( model_type=LTXModelType.AudioVideo, num_attention_heads=getattr(cfg, "num_attention_heads", 32), attention_head_dim=getattr(cfg, "attention_head_dim", 128), @@ -843,7 +853,7 @@ def _setup_cuda_graphs(self): if not self.pipeline_config.cuda_graph.enable: return - runner = _LTX2CUDAGraphRunner( + runner = self._cuda_graph_runner_class( CUDAGraphRunnerConfig(use_cuda_graph=True), ) self.transformer.register_cuda_graph_extra_key_fns(runner) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py index 7be3b9f16506..6ef0ed669ccc 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py @@ -1522,6 +1522,9 @@ class LTXModel(BaseDiffusionModel): follows the Lightricks reference implementation. """ + _transformer_block_class = BasicAVTransformerBlock + _cache_text_kv = True + def __init__( self, *, @@ -1915,7 +1918,7 @@ def _init_transformer_blocks( else None ) blocks: list[nn.Module] = [ - BasicAVTransformerBlock( + self._transformer_block_class( idx=idx, video=video_config, audio=audio_config, @@ -2179,13 +2182,15 @@ def prepare_text_cache( v_ctx, v_mask, v_pe, v_cross_pe = self.video_args_preprocessor.prepare_text_cache( video_context, video_context_mask, video_positions, dtype ) - v_kv = [block.attn2.project_kv(v_ctx) for block in self.transformer_blocks] + if self._cache_text_kv: + v_kv = [block.attn2.project_kv(v_ctx) for block in self.transformer_blocks] if audio_context is not None: a_ctx, a_mask, a_pe, a_cross_pe = self.audio_args_preprocessor.prepare_text_cache( audio_context, audio_context_mask, audio_positions, dtype ) - a_kv = [block.audio_attn2.project_kv(a_ctx) for block in self.transformer_blocks] + if self._cache_text_kv: + a_kv = [block.audio_attn2.project_kv(a_ctx) for block in self.transformer_blocks] # cos/sin are token-major: token axis is dim 1 for both SPLIT and # INTERLEAVED rope, matching `_make_pe_local`'s `cos[:, s:e]` shard. if self._audio_pad > 0: diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx23/NOTICE b/tensorrt_llm/_torch/visual_gen/models/ltx23/NOTICE new file mode 100644 index 000000000000..d4ab26aa7361 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/ltx23/NOTICE @@ -0,0 +1,6 @@ +Files in this directory marked SPDX-License-Identifier: LicenseRef-LTX-2 are +covered by the LTX-2 Community License Agreement. Files marked +SPDX-License-Identifier: Apache-2.0 are licensed under Apache 2.0. +Please refer to the file headers and the root LICENSE file for license information. + +LTX-2 repo commit id: 63fd9a4f8659989748623f635f3ed66b7c6ac076 diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx23/__init__.py b/tensorrt_llm/_torch/visual_gen/models/ltx23/__init__.py new file mode 100644 index 000000000000..9453cfad6739 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/ltx23/__init__.py @@ -0,0 +1,11 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 Lightricks Ltd. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: LicenseRef-LTX-2 + +from .pipeline_ltx23 import LTX23Pipeline +from .pipeline_ltx23_retake import LTX23RetakePipeline + +__all__ = [ + "LTX23Pipeline", + "LTX23RetakePipeline", +] diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx23/checkpoint.py b/tensorrt_llm/_torch/visual_gen/models/ltx23/checkpoint.py new file mode 100644 index 000000000000..8a50c186726a --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/ltx23/checkpoint.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Resolve a single native LTX-2.3 checkpoint before pipeline configuration/loading.""" + +from pathlib import Path + +from tensorrt_llm.llmapi.utils import download_hf_partial + +LTX23_MODEL_ID = "Lightricks/LTX-2.3" + +_WORKFLOW_CHECKPOINTS = { + "generation": "ltx-2.3-22b-dev.safetensors", + "retake": "ltx-2.3-22b-distilled-1.1.safetensors", +} + + +def resolve_ltx23_checkpoint(model: str, *, workflow: str, revision: str | None = None) -> str: + """Resolve a local file, directory, or Hub ID to one workflow's checkpoint. + + Explicit safetensors files are preserved. Directories and Hub repositories + must contain the canonical filename for the requested workflow; no fallback + to another checkpoint variant is allowed. Hub downloads use the shared + file-lock-protected utility, including its offline-cache behavior. + """ + if workflow not in _WORKFLOW_CHECKPOINTS: + raise ValueError(f"Unsupported LTX-2.3 workflow: {workflow!r}") + filename = _WORKFLOW_CHECKPOINTS[workflow] + path = Path(model) + if path.is_file(): + if path.suffix != ".safetensors": + raise ValueError(f"LTX-2.3 checkpoint must be a .safetensors file: {model}") + return str(path) + if path.is_dir(): + checkpoint = path / filename + elif path.suffix == ".safetensors" or path.is_absolute() or model.startswith((".", "~")): + raise FileNotFoundError(f"LTX-2.3 checkpoint path not found: {model}") + else: + checkpoint = ( + Path(download_hf_partial(model, allow_patterns=[filename], revision=revision)) + / filename + ) + if not checkpoint.is_file(): + raise FileNotFoundError(f"LTX-2.3 {workflow} checkpoint not found: {checkpoint}") + return str(checkpoint) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/__init__.py b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/__init__.py new file mode 100644 index 000000000000..79cf2c262f54 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 Lightricks Ltd. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: LicenseRef-LTX-2 +"""LTX-2.3-only core components.""" diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/__init__.py b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/__init__.py new file mode 100644 index 000000000000..9bea267d98cb --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/__init__.py @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 Lightricks Ltd. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: LicenseRef-LTX-2 + +from .audio_vae import encode_audio +from .bw import VocoderWithBWE +from .model_configurator import AudioEncoderConfigurator +from .vocoder import LTX23VocoderConfigurator + +__all__ = [ + "AudioEncoderConfigurator", + "LTX23VocoderConfigurator", + "VocoderWithBWE", + "encode_audio", +] diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/activations.py b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/activations.py new file mode 100644 index 000000000000..d5884323237d --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/activations.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 Lightricks Ltd. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: LicenseRef-LTX-2 + +import math + +import torch +import torch.nn.functional as F +from torch import nn + + +def _sinc(x: torch.Tensor) -> torch.Tensor: + return torch.where( + x == 0, + torch.tensor(1.0, device=x.device, dtype=x.dtype), + torch.sin(math.pi * x) / math.pi / x, + ) + + +def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor: + even = kernel_size % 2 == 0 + half_size = kernel_size // 2 + delta_f = 4 * half_width + amplitude = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 + if amplitude > 50.0: + beta = 0.1102 * (amplitude - 8.7) + elif amplitude >= 21.0: + beta = 0.5842 * (amplitude - 21) ** 0.4 + 0.07886 * (amplitude - 21.0) + else: + beta = 0.0 + window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) + time = ( + torch.arange(-half_size, half_size) + 0.5 if even else torch.arange(kernel_size) - half_size + ) + if cutoff == 0: + filter_ = torch.zeros_like(time) + else: + filter_ = 2 * cutoff * window * _sinc(2 * cutoff * time) + filter_ /= filter_.sum() + return filter_.view(1, 1, kernel_size) + + +class LowPassFilter1d(nn.Module): + def __init__( + self, + cutoff: float = 0.5, + half_width: float = 0.6, + stride: int = 1, + kernel_size: int = 12, + ) -> None: + super().__init__() + if cutoff < 0: + raise ValueError("Minimum cutoff must be larger than zero.") + if cutoff > 0.5: + raise ValueError("A cutoff above 0.5 does not make sense.") + self.pad_left = kernel_size // 2 - int(kernel_size % 2 == 0) + self.pad_right = kernel_size // 2 + self.stride = stride + self.register_buffer("filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + _, n_channels, _ = x.shape + x = F.pad(x, (self.pad_left, self.pad_right), mode="replicate") + return F.conv1d( + x, + self.filter.expand(n_channels, -1, -1), + stride=self.stride, + groups=n_channels, + ) + + +class UpSample1d(nn.Module): + def __init__( + self, + ratio: int = 2, + persistent: bool = True, + window_type: str = "kaiser", + ) -> None: + super().__init__() + self.ratio = ratio + + if window_type == "hann": + rolloff = 0.99 + lowpass_filter_width = 6 + width = math.ceil(lowpass_filter_width / rolloff) + kernel_size = 2 * width * ratio + 1 + self.pad = width + self.pad_left = 2 * width * ratio + self.pad_right = kernel_size - ratio + time_axis = (torch.arange(kernel_size) / ratio - width) * rolloff + time_clamped = time_axis.clamp(-lowpass_filter_width, lowpass_filter_width) + window = torch.cos(time_clamped * math.pi / lowpass_filter_width / 2) ** 2 + sinc_filter = (torch.sinc(time_axis) * window * rolloff / ratio).view(1, 1, -1) + else: + kernel_size = int(6 * ratio // 2) * 2 + self.pad = kernel_size // ratio - 1 + self.pad_left = self.pad * ratio + (kernel_size - ratio) // 2 + self.pad_right = self.pad * ratio + (kernel_size - ratio + 1) // 2 + sinc_filter = kaiser_sinc_filter1d( + cutoff=0.5 / ratio, + half_width=0.6 / ratio, + kernel_size=kernel_size, + ) + + self.register_buffer("filter", sinc_filter, persistent=persistent) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + _, n_channels, _ = x.shape + x = F.pad(x, (self.pad, self.pad), mode="replicate") + filt = self.filter.to(dtype=x.dtype, device=x.device).expand(n_channels, -1, -1) + x = self.ratio * F.conv_transpose1d(x, filt, stride=self.ratio, groups=n_channels) + return x[..., self.pad_left : -self.pad_right] + + +class DownSample1d(nn.Module): + """Fixed 2x anti-alias downsampler for AMP activations.""" + + def __init__(self) -> None: + super().__init__() + self.lowpass = LowPassFilter1d( + cutoff=0.25, + half_width=0.3, + stride=2, + kernel_size=12, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.lowpass(x) + + +class Activation1d(nn.Module): + def __init__(self, activation: nn.Module) -> None: + super().__init__() + self.act = activation + self.upsample = UpSample1d() + self.downsample = DownSample1d() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.downsample(self.act(self.upsample(x))) + + +class SnakeBeta(nn.Module): + def __init__(self, in_features: int) -> None: + super().__init__() + self.alpha = nn.Parameter(torch.zeros(in_features)) + self.beta = nn.Parameter(torch.zeros(in_features)) + self.eps = 1e-9 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + alpha = self.alpha.exp().view(1, -1, 1) + beta = self.beta.exp().view(1, -1, 1) + return x + (1.0 / (beta + self.eps)) * torch.sin(x * alpha).pow(2) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/amp_block.py b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/amp_block.py new file mode 100644 index 000000000000..56367065817f --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/amp_block.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 Lightricks Ltd. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: LicenseRef-LTX-2 + +import torch +from torch import nn + +from .activations import Activation1d, SnakeBeta + + +class AMPBlock1(nn.Module): + """BigVGAN anti-aliased multi-periodicity residual block.""" + + def __init__( + self, + channels: int, + kernel_size: int = 3, + dilation: tuple[int, ...] = (1, 3, 5), + ) -> None: + super().__init__() + + def convs(dilations: tuple[int, ...]) -> nn.ModuleList: + return nn.ModuleList( + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=d, + padding=(kernel_size * d - d) // 2, + ) + for d in dilations + ) + + self.convs1 = convs(dilation) + self.convs2 = convs((1,) * len(dilation)) + self.acts1 = nn.ModuleList([Activation1d(SnakeBeta(channels)) for _ in self.convs1]) + self.acts2 = nn.ModuleList([Activation1d(SnakeBeta(channels)) for _ in self.convs2]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for conv1, conv2, act1, act2 in zip( + self.convs1, self.convs2, self.acts1, self.acts2, strict=True + ): + residual = conv2(act2(conv1(act1(x)))) + x = x + residual + return x diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/audio_vae.py b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/audio_vae.py new file mode 100644 index 000000000000..1a019cbf04c4 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/audio_vae.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 Lightricks Ltd. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: LicenseRef-LTX-2 + +import torch + +from tensorrt_llm.inputs.multimodal_data import AudioData + +from ....ltx2.ltx2_core.audio_vae.audio_vae import ( + LATENT_DOWNSAMPLE_FACTOR, + build_mid_block, + run_mid_block, +) +from ....ltx2.ltx2_core.audio_vae.causal_conv_2d import make_conv2d +from ....ltx2.ltx2_core.audio_vae.causality_axis import CausalityAxis +from ....ltx2.ltx2_core.audio_vae.ops import PerChannelStatistics +from ....ltx2.ltx2_core.normalization import NormType, build_normalization_layer +from ....ltx2.ltx2_core.patchifier import AudioPatchifier +from ....ltx2.ltx2_core.types import AudioLatentShape +from .downsample import build_downsampling_path + + +class AudioEncoder(torch.nn.Module): + """Compress a log-mel spectrogram into normalized audio latents.""" + + def __init__( + self, + *, + ch: int, + ch_mult: tuple[int, ...] = (1, 2, 4, 8), + num_res_blocks: int, + attn_resolutions: set[int], + dropout: float = 0.0, + resamp_with_conv: bool = True, + in_channels: int, + resolution: int, + z_channels: int, + double_z: bool = True, + mid_block_add_attention: bool = True, + norm_type: NormType = NormType.GROUP, + causality_axis: CausalityAxis = CausalityAxis.WIDTH, + sample_rate: int = 16000, + mel_hop_length: int = 160, + n_fft: int = 1024, + is_causal: bool = True, + mel_bins: int = 64, + ) -> None: + super().__init__() + + self.per_channel_statistics = PerChannelStatistics(latent_channels=ch) + self.sample_rate = sample_rate + self.mel_hop_length = mel_hop_length + self.n_fft = n_fft + self.mel_bins = mel_bins + self.patchifier = AudioPatchifier( + patch_size=1, + audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR, + sample_rate=sample_rate, + hop_length=mel_hop_length, + is_causal=is_causal, + ) + self.double_z = double_z + + self.conv_in = make_conv2d( + in_channels, + ch, + kernel_size=3, + stride=1, + causality_axis=causality_axis, + ) + self.non_linearity = torch.nn.SiLU() + self.down, block_in = build_downsampling_path( + ch=ch, + ch_mult=ch_mult, + num_res_blocks=num_res_blocks, + resolution=resolution, + dropout=dropout, + norm_type=norm_type, + causality_axis=causality_axis, + attn_resolutions=attn_resolutions, + resamp_with_conv=resamp_with_conv, + ) + self.mid = build_mid_block( + channels=block_in, + temb_channels=0, + dropout=dropout, + norm_type=norm_type, + causality_axis=causality_axis, + add_attention=mid_block_add_attention, + ) + self.norm_out = build_normalization_layer(block_in, normtype=norm_type) + self.conv_out = make_conv2d( + block_in, + 2 * z_channels if double_z else z_channels, + kernel_size=3, + stride=1, + causality_axis=causality_axis, + ) + + def forward(self, spectrogram: torch.Tensor) -> torch.Tensor: + """Encode a ``(B, C, time, mel_bins)`` spectrogram into ``(B, C, frames, mel_bins)``.""" + h = self.conv_in(spectrogram) + for level, stage in enumerate(self.down): + for block_idx, block in enumerate(stage.block): + h = block(h, temb=None) + if stage.attn: + h = stage.attn[block_idx](h) + if level < len(self.down) - 1: + h = stage.downsample(h) + h = run_mid_block(self.mid, h) + h = self.conv_out(self.non_linearity(self.norm_out(h))) + # Use the posterior mean; discard the log-variance half. + means = torch.chunk(h, 2, dim=1)[0] if self.double_z else h + latent_shape = AudioLatentShape.from_torch_shape(means.shape) + latent_patched = self.patchifier.patchify(means) + latent_normalized = self.per_channel_statistics.normalize(latent_patched) + return self.patchifier.unpatchify(latent_normalized, latent_shape) + + +def _waveform_to_mel( + samples: torch.Tensor, + source_sample_rate: int, + audio_encoder: AudioEncoder, +) -> torch.Tensor: + try: + import torchaudio + except ImportError as exc: # pragma: no cover - environment dependent + raise ImportError( + "The LTX-2.3 native audio encoder needs `torchaudio` for waveform " + "resampling and mel preprocessing. Install a torchaudio build matching " + "the installed torch in the VisualGen runtime environment." + ) from exc + + if source_sample_rate != audio_encoder.sample_rate: + samples = torchaudio.functional.resample( + samples, + source_sample_rate, + audio_encoder.sample_rate, + ).to(samples) + + mel_transform = torchaudio.transforms.MelSpectrogram( + sample_rate=audio_encoder.sample_rate, + n_fft=audio_encoder.n_fft, + win_length=audio_encoder.n_fft, + hop_length=audio_encoder.mel_hop_length, + f_min=0.0, + f_max=audio_encoder.sample_rate / 2.0, + n_mels=audio_encoder.mel_bins, + window_fn=torch.hann_window, + center=True, + pad_mode="reflect", + power=1.0, + mel_scale="slaney", + norm="slaney", + ).to(device=samples.device) + mel = torch.log(torch.clamp(mel_transform(samples), min=1e-5)) + return mel.to(dtype=samples.dtype).permute(0, 1, 3, 2).contiguous() + + +def encode_audio( + audio: AudioData, + audio_encoder: AudioEncoder, +) -> torch.Tensor: + """Encode an audio waveform into audio VAE latents.""" + parameter = next(audio_encoder.parameters()) + + samples = torch.as_tensor(audio.samples, device=parameter.device) + if samples.ndim != 3: + raise ValueError(f"Expected audio samples with shape (B, C, T); got {samples.shape}.") + input_channels = audio_encoder.conv_in.conv.in_channels + if samples.shape[1] == 1 and input_channels == 2: + samples = samples.expand(-1, 2, -1).contiguous() + elif samples.shape[1] != input_channels: + raise ValueError( + f"Audio encoder expects {input_channels} channel(s), got {samples.shape[1]}." + ) + mel_spectrogram = _waveform_to_mel(samples, audio.sample_rate, audio_encoder) + return audio_encoder(mel_spectrogram.to(dtype=parameter.dtype)) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/bw.py b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/bw.py new file mode 100644 index 000000000000..ff6ddb53baca --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/bw.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 Lightricks Ltd. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: LicenseRef-LTX-2 + +from typing import TYPE_CHECKING + +import torch +import torch.nn.functional as F +from torch import nn + +from .activations import UpSample1d + +if TYPE_CHECKING: + from .vocoder import Vocoder + + +class _STFTFn(nn.Module): + def __init__(self, filter_length: int, hop_length: int, win_length: int) -> None: + super().__init__() + self.hop_length = hop_length + self.left_pad = max(0, win_length - hop_length) + n_freqs = filter_length // 2 + 1 + self.register_buffer("forward_basis", torch.zeros(n_freqs * 2, 1, filter_length)) + + def forward(self, y: torch.Tensor) -> torch.Tensor: + if y.dim() == 2: + y = y.unsqueeze(1) + y = F.pad(y, (self.left_pad, 0)) + spec = F.conv1d(y, self.forward_basis, stride=self.hop_length) + n_freqs = spec.shape[1] // 2 + real, imag = spec[:, :n_freqs], spec[:, n_freqs:] + return torch.sqrt(real**2 + imag**2) + + +class MelSTFT(nn.Module): + def __init__( + self, + filter_length: int, + hop_length: int, + win_length: int, + n_mel_channels: int, + ) -> None: + super().__init__() + self.stft_fn = _STFTFn(filter_length, hop_length, win_length) + n_freqs = filter_length // 2 + 1 + self.register_buffer("mel_basis", torch.zeros(n_mel_channels, n_freqs)) + + def mel_spectrogram(self, y: torch.Tensor) -> torch.Tensor: + magnitude = self.stft_fn(y) + mel = torch.matmul(self.mel_basis.to(magnitude.dtype), magnitude) + return torch.log(torch.clamp(mel, min=1e-5)) + + +class VocoderWithBWE(nn.Module): + def __init__( + self, + vocoder: "Vocoder", + bwe_generator: "Vocoder", + mel_stft: MelSTFT, + input_sampling_rate: int, + output_sampling_rate: int, + hop_length: int, + ) -> None: + super().__init__() + self.vocoder = vocoder + self.bwe_generator = bwe_generator + self.mel_stft = mel_stft + self.input_sampling_rate = input_sampling_rate + self.output_sampling_rate = output_sampling_rate + self.hop_length = hop_length + with torch.device("cpu"): + self.resampler = UpSample1d( + ratio=output_sampling_rate // input_sampling_rate, + persistent=False, + window_type="hann", + ) + + def forward(self, mel_spec: torch.Tensor) -> torch.Tensor: + input_dtype = mel_spec.dtype + with torch.autocast(device_type=mel_spec.device.type, enabled=False): + x = self.vocoder(mel_spec.float()) + length_low_rate = x.shape[-1] + output_length = length_low_rate * self.output_sampling_rate // self.input_sampling_rate + + remainder = length_low_rate % self.hop_length + if remainder: + x = F.pad(x, (0, self.hop_length - remainder)) + + batch, n_channels, _ = x.shape + mel = self.mel_stft.mel_spectrogram(x.flatten(0, 1)) + mel = mel.reshape(batch, n_channels, mel.shape[1], mel.shape[2]) + residual = self.bwe_generator(mel.transpose(2, 3)) + skip = self.resampler(x) + assert residual.shape == skip.shape, f"residual {residual.shape} != skip {skip.shape}" + return torch.clamp(residual + skip, -1, 1)[..., :output_length].to(input_dtype) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/downsample.py b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/downsample.py new file mode 100644 index 000000000000..b40dad82d284 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/downsample.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 Lightricks Ltd. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: LicenseRef-LTX-2 + +import torch + +from ....ltx2.ltx2_core.audio_vae.attention import AttnBlock +from ....ltx2.ltx2_core.audio_vae.causality_axis import CausalityAxis +from ....ltx2.ltx2_core.audio_vae.resnet import ResnetBlock +from ....ltx2.ltx2_core.normalization import NormType + + +class Downsample(torch.nn.Module): + """Downsample by two using an asymmetrically padded convolution or average pooling.""" + + def __init__( + self, + in_channels: int, + with_conv: bool, + causality_axis: CausalityAxis = CausalityAxis.HEIGHT, + ) -> None: + super().__init__() + self.with_conv = with_conv + + if causality_axis != CausalityAxis.NONE and not self.with_conv: + raise ValueError("causality is only supported when `with_conv=True`.") + + if self.with_conv: + # F.pad tuple order is (left, right, top, bottom). + match causality_axis: + case CausalityAxis.NONE: + self.padding = (0, 1, 0, 1) + case CausalityAxis.WIDTH: + self.padding = (2, 0, 0, 1) + case CausalityAxis.HEIGHT: + self.padding = (0, 1, 2, 0) + case CausalityAxis.WIDTH_COMPATIBILITY: + self.padding = (1, 0, 0, 1) + case _: + raise ValueError(f"Invalid causality_axis: {causality_axis}") + self.conv = torch.nn.Conv2d( + in_channels, in_channels, kernel_size=3, stride=2, padding=0 + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if not self.with_conv: + return torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2) + + return self.conv(torch.nn.functional.pad(x, self.padding)) + + +def build_downsampling_path( + *, + ch: int, + ch_mult: tuple[int, ...], + num_res_blocks: int, + resolution: int, + dropout: float, + norm_type: NormType, + causality_axis: CausalityAxis, + attn_resolutions: set[int], + resamp_with_conv: bool, +) -> tuple[torch.nn.ModuleList, int]: + """Build the encoder downsampling path; mirror of :func:`build_upsampling_path`.""" + down_modules = torch.nn.ModuleList() + curr_res = resolution + block_in = ch + + for level, multiplier in enumerate(ch_mult): + stage = torch.nn.Module() + stage.block = torch.nn.ModuleList() + stage.attn = torch.nn.ModuleList() + block_in = ch * (ch_mult[level - 1] if level else 1) + block_out = ch * multiplier + + for _ in range(num_res_blocks): + stage.block.append( + ResnetBlock( + in_channels=block_in, + out_channels=block_out, + temb_channels=0, + dropout=dropout, + norm_type=norm_type, + causality_axis=causality_axis, + ) + ) + block_in = block_out + if curr_res in attn_resolutions: + stage.attn.append(AttnBlock(block_in, norm_type=norm_type)) + + if level < len(ch_mult) - 1: + stage.downsample = Downsample(block_in, resamp_with_conv, causality_axis=causality_axis) + curr_res //= 2 + + down_modules.append(stage) + + return down_modules, block_in diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/model_configurator.py b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/model_configurator.py new file mode 100644 index 000000000000..82592c58b976 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/model_configurator.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 Lightricks Ltd. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: LicenseRef-LTX-2 + +from ....ltx2.ltx2_core.audio_vae.causality_axis import CausalityAxis +from ....ltx2.ltx2_core.normalization import NormType +from .audio_vae import AudioEncoder + + +class AudioEncoderConfigurator: + """Create an AudioEncoder from the LTX-2.3 native config dict.""" + + @classmethod + def from_config(cls, config: dict) -> AudioEncoder: + audio_vae_cfg = config.get("audio_vae", {}) + model_params = audio_vae_cfg.get("model", {}).get("params", {}) + ddconfig = model_params.get("ddconfig", {}) + preprocessing_cfg = audio_vae_cfg.get("preprocessing", {}) + stft_cfg = preprocessing_cfg.get("stft", {}) + mel_bins = ( + ddconfig.get("mel_bins") + or preprocessing_cfg.get("mel", {}).get("n_mel_channels") + or audio_vae_cfg.get("variables", {}).get("mel_bins") + ) + if mel_bins is None: + raise ValueError("LTX-2.3 audio VAE config does not define the mel-bin count.") + + attn_type = ddconfig.get("attn_type", "vanilla") + if attn_type != "vanilla": + raise ValueError( + f"LTX-2.3 native audio encoder supports attn_type='vanilla'; got {attn_type!r}." + ) + + return AudioEncoder( + ch=ddconfig.get("ch", 128), + ch_mult=tuple(ddconfig.get("ch_mult", (1, 2, 4))), + num_res_blocks=ddconfig.get("num_res_blocks", 2), + attn_resolutions=set(ddconfig.get("attn_resolutions", (8, 16, 32))), + resolution=ddconfig.get("resolution", 256), + z_channels=ddconfig.get("z_channels", 8), + double_z=ddconfig.get("double_z", True), + dropout=ddconfig.get("dropout", 0.0), + resamp_with_conv=ddconfig.get("resamp_with_conv", True), + in_channels=ddconfig.get("in_channels", 2), + mid_block_add_attention=ddconfig.get("mid_block_add_attention", True), + norm_type=NormType(ddconfig.get("norm_type", "pixel")), + causality_axis=CausalityAxis(ddconfig.get("causality_axis", "height")), + sample_rate=model_params.get("sampling_rate", 16000), + mel_hop_length=stft_cfg.get("hop_length", 160), + n_fft=stft_cfg.get("filter_length", 1024), + is_causal=stft_cfg.get("causal", True), + mel_bins=mel_bins, + ) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/vocoder.py b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/vocoder.py new file mode 100644 index 000000000000..1cc3e7b175cc --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/audio_vae/vocoder.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 Lightricks Ltd. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: LicenseRef-LTX-2 + +from collections.abc import Sequence + +import einops +import torch +from torch import nn + +from ....ltx2.ltx2_core.audio_vae.vocoder import Vocoder as LTX2Vocoder +from .activations import Activation1d, SnakeBeta +from .amp_block import AMPBlock1 +from .bw import MelSTFT, VocoderWithBWE + + +class Vocoder(LTX2Vocoder): + """LTX-2.3 BigVGAN AMP1 mel-spectrogram generator.""" + + def __init__( + self, + resblock_kernel_sizes: Sequence[int] = (3, 7, 11), + upsample_rates: Sequence[int] = (6, 5, 2, 2, 2), + upsample_kernel_sizes: Sequence[int] = (16, 15, 8, 4, 4), + resblock_dilation_sizes: Sequence[Sequence[int]] = ((1, 3, 5),) * 3, + upsample_initial_channel: int = 1024, + output_sampling_rate: int = 24000, + use_tanh_at_final: bool = True, + apply_final_activation: bool = True, + use_bias_at_final: bool = True, + ) -> None: + # Initialize shared convolution/upsampling layers without LTX-2 residual blocks. + super().__init__( + resblock_kernel_sizes=[], + upsample_rates=list(upsample_rates), + upsample_kernel_sizes=list(upsample_kernel_sizes), + resblock_dilation_sizes=[], + upsample_initial_channel=upsample_initial_channel, + output_sample_rate=output_sampling_rate, + ) + self.output_sampling_rate = output_sampling_rate + self.num_kernels = len(resblock_kernel_sizes) + self.use_tanh_at_final = use_tanh_at_final + self.apply_final_activation = apply_final_activation + + for upsample in self.ups: + channels = upsample.out_channels + for kernel_size, dilations in zip( + resblock_kernel_sizes, resblock_dilation_sizes, strict=True + ): + self.resblocks.append(AMPBlock1(channels, kernel_size, tuple(dilations))) + + final_channels = self.conv_post.in_channels + self.act_post = Activation1d(SnakeBeta(final_channels)) + if not use_bias_at_final: + self.conv_post = nn.Conv1d(final_channels, 2, 7, 1, padding=3, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x.transpose(2, 3) + if x.dim() == 4: + assert x.shape[1] == 2, "Input must have 2 channels for stereo" + x = einops.rearrange(x, "b s c t -> b (s c) t") + + x = self.conv_pre(x) + for i, upsample in enumerate(self.ups): + x = upsample(x) + start = i * self.num_kernels + outputs = torch.stack( + [self.resblocks[index](x) for index in range(start, start + self.num_kernels)] + ) + x = outputs.mean(dim=0) + + x = self.conv_post(self.act_post(x)) + if self.apply_final_activation: + x = torch.tanh(x) if self.use_tanh_at_final else torch.clamp(x, -1, 1) + return x + + +def _vocoder_from_config( + cfg: dict, + apply_final_activation: bool = True, + output_sampling_rate: int | None = None, +) -> Vocoder: + keys = ( + "resblock_kernel_sizes", + "upsample_rates", + "upsample_kernel_sizes", + "resblock_dilation_sizes", + "upsample_initial_channel", + "output_sampling_rate", + "use_tanh_at_final", + "use_bias_at_final", + ) + kwargs = {key: cfg[key] for key in keys if cfg.get(key) is not None} + if output_sampling_rate is not None: + kwargs["output_sampling_rate"] = output_sampling_rate + return Vocoder( + **kwargs, + apply_final_activation=apply_final_activation, + ) + + +class LTX23VocoderConfigurator: + """Build an LTX-2.3 vocoder from checkpoint configuration.""" + + @classmethod + def from_config(cls, config: dict) -> VocoderWithBWE | LTX2Vocoder: + def check_config_value(cfg: dict, key: str, expected: str | bool) -> None: + actual = cfg.get(key) + if actual != expected: + raise ValueError( + f"LTX-2.3 vocoder config mismatch: expected {key}={expected!r}, got {actual!r}" + ) + + cfg = config.get("vocoder", {}) + if "bwe" not in cfg: + check_config_value(cfg, "resblock", "1") + check_config_value(cfg, "stereo", True) + from ....ltx2.ltx2_core.audio_vae.model_configurator import VocoderConfigurator + + return VocoderConfigurator.from_config(config) + + vocoder_cfg = cfg.get("vocoder", {}) + bwe_cfg = cfg["bwe"] + for section in (vocoder_cfg, bwe_cfg): + check_config_value(section, "resblock", "AMP1") + check_config_value(section, "stereo", True) + check_config_value(section, "activation", "snakebeta") + + vocoder = _vocoder_from_config( + vocoder_cfg, + output_sampling_rate=bwe_cfg["input_sampling_rate"], + ) + bwe_generator = _vocoder_from_config( + bwe_cfg, + apply_final_activation=False, + output_sampling_rate=bwe_cfg["output_sampling_rate"], + ) + mel_stft = MelSTFT( + filter_length=bwe_cfg["n_fft"], + hop_length=bwe_cfg["hop_length"], + win_length=bwe_cfg["n_fft"], + n_mel_channels=bwe_cfg["num_mels"], + ) + return VocoderWithBWE( + vocoder=vocoder, + bwe_generator=bwe_generator, + mel_stft=mel_stft, + input_sampling_rate=bwe_cfg["input_sampling_rate"], + output_sampling_rate=bwe_cfg["output_sampling_rate"], + hop_length=bwe_cfg["hop_length"], + ) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/connector.py b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/connector.py new file mode 100644 index 000000000000..c7f607b0dd65 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/connector.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 Lightricks Ltd. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: LicenseRef-LTX-2 +"""LTX-2.3 per-modality feature projections and LTX-2 connector configuration.""" + +import math + +import torch + +from ...ltx2.ltx2_core.connector import Embeddings1DConnector +from ...ltx2.ltx2_core.rope import LTXRopeType + +# Gemma-3-12b-it exposes 49 hidden states (48 layers plus embeddings). +_NUM_GEMMA_HIDDEN_STATES = 49 + + +class LTX23GemmaFeaturesExtractor(torch.nn.Module): + """Split Gemma feature extractor (video + audio), pre-connector.""" + + def __init__( + self, + caption_channels: int = 3840, + video_dim: int = 4096, + audio_dim: int = 2048, + num_hidden_states: int = _NUM_GEMMA_HIDDEN_STATES, + ) -> None: + super().__init__() + in_features = caption_channels * num_hidden_states + # Rescale uses the Gemma hidden width, not the flattened 3840*49 size. + self.embedding_dim = caption_channels + self.video_aggregate_embed = torch.nn.Linear(in_features, video_dim, bias=True) + self.audio_aggregate_embed = torch.nn.Linear(in_features, audio_dim, bias=True) + + def _project(self, linear: torch.nn.Linear, x: torch.Tensor) -> torch.Tensor: + return linear(x * math.sqrt(linear.out_features / self.embedding_dim)) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return self._project(self.video_aggregate_embed, x), self._project( + self.audio_aggregate_embed, x + ) + + @classmethod + def from_config(cls, config: dict) -> "LTX23GemmaFeaturesExtractor": + cfg = config.get("transformer", config) + return cls( + caption_channels=cfg.get("caption_channels", 3840), + video_dim=cfg.get("cross_attention_dim", 4096), + audio_dim=cfg.get("audio_cross_attention_dim", 2048), + ) + + +def _build_connector( + cfg: dict, + num_attention_heads: int, + attention_head_dim: int, + *, + is_audio: bool = False, +) -> Embeddings1DConnector: + def value(key: str, default: object) -> object: + shared = cfg.get(f"connector_{key}", default) + return cfg.get(f"audio_connector_{key}", shared) if is_audio else shared + + rope_type = value("rope_type", cfg.get("rope_type", "split")) + if isinstance(rope_type, str): + rope_type = rope_type.lower() + rope_type = LTXRopeType(rope_type) + double_precision_rope = ( + value("frequencies_precision", cfg.get("frequencies_precision", False)) == "float64" + ) + return Embeddings1DConnector( + attention_head_dim=attention_head_dim, + num_attention_heads=num_attention_heads, + num_layers=value("num_layers", 8), + positional_embedding_max_pos=cfg.get("connector_positional_embedding_max_pos", [4096]), + num_learnable_registers=value("num_learnable_registers", 128), + rope_type=rope_type, + double_precision_rope=double_precision_rope, + apply_gated_attention=value("apply_gated_attention", True), + ) + + +class LTX23VideoConnectorConfigurator: + """Video connector: 32 x 128 = 4096, 8 layers, gated, 128 registers.""" + + @classmethod + def from_config(cls, config: dict) -> Embeddings1DConnector: + cfg = config.get("transformer", config) + return _build_connector( + cfg, + num_attention_heads=cfg.get("connector_num_attention_heads", 32), + attention_head_dim=cfg.get("connector_attention_head_dim", 128), + ) + + +class LTX23AudioConnectorConfigurator: + """Audio connector: 32 x 64 = 2048, 8 layers, gated, 128 registers.""" + + @classmethod + def from_config(cls, config: dict) -> Embeddings1DConnector: + cfg = config.get("transformer", config) + return _build_connector( + cfg, + num_attention_heads=cfg.get("audio_connector_num_attention_heads", 32), + attention_head_dim=cfg.get("audio_connector_attention_head_dim", 64), + is_audio=True, + ) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/modality.py b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/modality.py new file mode 100644 index 000000000000..2ba23bdac892 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/modality.py @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 Lightricks Ltd. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: LicenseRef-LTX-2 +"""LTX-2.3 modality bundle. + +This is the LTX-2 Modality plus an explicit sigma field, kept distinct from +timesteps because they drive different conditioning: timesteps may vary per +latent token and drives the per-token AdaLN modulation, while sigma is the +global denoising value that derives the prompt_timestep behind the text-context +K/V modulation. +""" + +from dataclasses import dataclass, field + +import torch + +from ...ltx2.ltx2_core.modality import Modality + + +@dataclass(frozen=True) +class LTX23Modality(Modality): + """LTX-2 Modality plus the global denoising sigma.""" + + sigma: torch.Tensor = field(kw_only=True) # (B,) drives prompt K/V + # Optional global sigma for the audio-video cross-attention gate. + cross_modality_sigma: torch.Tensor | None = field(default=None, kw_only=True) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/video_vae.py b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/video_vae.py new file mode 100644 index 000000000000..bed3aaf0228f --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/ltx23/ltx23_core/video_vae.py @@ -0,0 +1,295 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 Lightricks Ltd. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: LicenseRef-LTX-2 +"""LTX-2.3 channel-reducing decoder blocks and overlap-aware tiled encoding.""" + +import itertools +import logging +from typing import Any + +import torch +import torch.nn as nn + +from ...ltx2.ltx2_core.types import ( + VIDEO_SCALE_FACTORS, + SpatioTemporalScaleFactors, + VideoLatentShape, +) +from ...ltx2.ltx2_core.video_vae.enums import NormLayerType, PaddingModeType +from ...ltx2.ltx2_core.video_vae.ops import PerChannelStatistics +from ...ltx2.ltx2_core.video_vae.sampling import DepthToSpaceUpsample +from ...ltx2.ltx2_core.video_vae.tiling import ( + DEFAULT_SPLIT_OPERATION, + Tile, + TilingConfig, + compute_trapezoidal_mask_1d, +) +from ...ltx2.ltx2_core.video_vae.video_vae import VideoDecoder as LTX2VideoDecoder +from ...ltx2.ltx2_core.video_vae.video_vae import VideoEncoder as LTX2VideoEncoder +from ...ltx2.ltx2_core.video_vae.video_vae import ( + _make_decoder_block, + split_temporal_latents, + split_with_symmetric_overlaps, +) + +logger = logging.getLogger(__name__) + +_MIN_SPATIAL_OVERLAP_PIXELS = 64 +_MIN_TEMPORAL_OVERLAP_FRAMES = 16 + +_COMPRESS_STRIDES = { + "compress_time": (2, 1, 1), + "compress_space": (1, 2, 2), + "compress_all": (2, 2, 2), +} + + +def _make_ltx23_decoder_block( + block_name: str, + block_config: dict, + in_channels: int, + convolution_dimensions: int, + norm_layer: NormLayerType, + norm_num_groups: int, + spatial_padding_mode: PaddingModeType, + timestep_conditioning: bool = False, +) -> tuple[nn.Module, int]: + if block_name in ("compress_time", "compress_space"): + multiplier = block_config.get("multiplier", 1) + out_channels = in_channels // multiplier + block = DepthToSpaceUpsample( + dims=convolution_dimensions, + in_channels=in_channels, + stride=_COMPRESS_STRIDES[block_name], + residual=block_config.get("residual", False), + out_channels_reduction_factor=multiplier, + spatial_padding_mode=spatial_padding_mode, + ) + return block, out_channels + return _make_decoder_block( + block_name=block_name, + block_config=block_config, + in_channels=in_channels, + convolution_dimensions=convolution_dimensions, + norm_layer=norm_layer, + timestep_conditioning=timestep_conditioning, + norm_num_groups=norm_num_groups, + spatial_padding_mode=spatial_padding_mode, + ) + + +class LTX23VideoDecoder(LTX2VideoDecoder): + """LTX-2.3 video decoder: LTX-2 primitives, LTX-2.3 channel recipe.""" + + _channel_reducing_blocks = tuple(_COMPRESS_STRIDES) + _make_decoder_block = staticmethod(_make_ltx23_decoder_block) + + def __init__( + self, + convolution_dimensions: int = 3, + in_channels: int = 128, + out_channels: int = 3, + decoder_blocks: list[tuple[str, int | dict[str, Any]]] | None = None, + patch_size: int = 4, + norm_layer: NormLayerType = NormLayerType.PIXEL_NORM, + causal: bool = False, + spatial_padding_mode: PaddingModeType = PaddingModeType.REFLECT, + ) -> None: + super().__init__( + convolution_dimensions=convolution_dimensions, + in_channels=in_channels, + out_channels=out_channels, + decoder_blocks=decoder_blocks or [], + patch_size=patch_size, + norm_layer=norm_layer, + causal=causal, + timestep_conditioning=False, + decoder_spatial_padding_mode=spatial_padding_mode, + ) + self.video_downscale_factors = VIDEO_SCALE_FACTORS + + +class LTX23VideoDecoderConfigurator: + """Create an LTX23VideoDecoder from the native LTX-2.3 config dict.""" + + @classmethod + def from_config(cls, config: dict) -> LTX23VideoDecoder: + vae = config.get("vae", {}) + padding_mode = vae.get( + "spatial_padding_mode", vae.get("decoder_spatial_padding_mode", "reflect") + ) + return LTX23VideoDecoder( + convolution_dimensions=vae.get("dims", 3), + in_channels=vae.get("latent_channels", 128), + out_channels=vae.get("out_channels", 3), + decoder_blocks=vae.get("decoder_blocks", []), + patch_size=vae.get("patch_size", 4), + norm_layer=NormLayerType(vae.get("norm_layer", "pixel_norm")), + causal=vae.get("causal_decoder", False), + spatial_padding_mode=PaddingModeType(padding_mode), + ) + + +class _LTX23PerChannelStatistics(PerChannelStatistics): + """Statistics buffers present in the LTX-2.3 encoder checkpoint.""" + + def __init__(self, latent_channels: int = 128) -> None: + nn.Module.__init__(self) + self.register_buffer("std-of-means", torch.empty(latent_channels)) + self.register_buffer("mean-of-means", torch.empty(latent_channels)) + + +def _prepare_encode_tiles( + video: torch.Tensor, + tiling_config: TilingConfig, + scales: SpatioTemporalScaleFactors, +) -> list[Tile]: + """Split on the latent grid and map causal temporal tiles back to pixels.""" + latent_shape = ( + (video.shape[2] - 1) // scales.time + 1, + video.shape[3] // scales.height, + video.shape[4] // scales.width, + ) + splitters = [DEFAULT_SPLIT_OPERATION] * 3 + if tiling_config.spatial_config is not None: + config = tiling_config.spatial_config + overlap = max(config.tile_overlap_in_pixels, _MIN_SPATIAL_OVERLAP_PIXELS) + for axis, scale in ((1, scales.height), (2, scales.width)): + latent_overlap = overlap // scale + size = max(2, latent_overlap + 1, config.tile_size_in_pixels // scale) + splitters[axis] = split_with_symmetric_overlaps(size, latent_overlap) + if tiling_config.temporal_config is not None: + config = tiling_config.temporal_config + overlap = max(config.tile_overlap_in_frames, _MIN_TEMPORAL_OVERLAP_FRAMES) + latent_overlap = overlap // scales.time + size = max(2, latent_overlap + 1, config.tile_size_in_frames // scales.time) + splitters[0] = split_temporal_latents(size, latent_overlap) + intervals = [ + list(zip(value.starts, value.ends, value.left_ramps, value.right_ramps, strict=True)) + for value in (split(length) for split, length in zip(splitters, latent_shape, strict=True)) + ] + tiles = [] + for axes in itertools.product(*intervals): + inputs = [slice(None), slice(None)] + outputs = [slice(None), slice(None)] + masks = [torch.ones(1), torch.ones(1)] + for axis, ((start, stop, left, right), scale) in enumerate(zip(axes, scales, strict=True)): + pixel_stop = 1 + (stop - 1) * scale if axis == 0 else stop * scale + inputs.append(slice(start * scale, pixel_stop)) + outputs.append(slice(start, stop)) + masks.append(compute_trapezoidal_mask_1d(stop - start, left, right, axis == 0)) + tiles.append(Tile(tuple(inputs), tuple(outputs), tuple(masks))) + return tiles + + +def _encode_masks_are_complementary(tiles: list[Tile], shape: torch.Size) -> bool: + for axis in range(2, len(shape)): + weights = torch.zeros(shape[axis], dtype=torch.float32, device="cpu") + seen = set() + for tile in tiles: + region = tile.out_coords[axis] + key = (region.start, region.stop) + if key not in seen: + weights[region] += tile.masks_1d[axis].float().cpu() + seen.add(key) + if not torch.allclose(weights, torch.ones_like(weights), atol=1e-5, rtol=0.0): + return False + return True + + +def _apply_encode_masks(latent: torch.Tensor, tile: Tile) -> torch.Tensor: + result = latent.float() + for axis, mask in enumerate(tile.masks_1d): + view = [1] * latent.ndim + view[axis] = -1 + result = result * mask.to(device=latent.device, dtype=torch.float32).view(view) + return result + + +class LTX23VideoEncoder(LTX2VideoEncoder): + """LTX-2.3 video encoder with overlap-aware tiled encoding.""" + + def __init__( + self, + convolution_dimensions: int = 3, + in_channels: int = 3, + out_channels: int = 128, + encoder_blocks: list[tuple[str, int | dict[str, Any]]] | None = None, + patch_size: int = 4, + norm_layer: NormLayerType = NormLayerType.PIXEL_NORM, + causal: bool = True, + encoder_spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, + ) -> None: + super().__init__( + convolution_dimensions=convolution_dimensions, + in_channels=in_channels, + out_channels=out_channels, + encoder_blocks=encoder_blocks or [], + patch_size=patch_size, + norm_layer=norm_layer, + causal=causal, + timestep_conditioning=False, + encoder_spatial_padding_mode=encoder_spatial_padding_mode, + ) + self.per_channel_statistics = _LTX23PerChannelStatistics(out_channels) + + def tiled_encode( + self, + video: torch.Tensor, + tiling_config: TilingConfig | None = None, + ) -> torch.Tensor: + if tiling_config is None: + return self(video) + + device = next(self.parameters()).device + dtype = next(self.parameters()).dtype + scales = VIDEO_SCALE_FACTORS + batch, _, frames, height, width = video.shape + remainder = (frames - 1) % scales.time + if remainder: + logger.warning("Cropping %d video frame(s) for causal VAE encode", remainder) + video = video[:, :, :-remainder] + frames = video.shape[2] + + latent_shape = VideoLatentShape( + batch=batch, + channels=self.out_channels, + frames=(frames - 1) // scales.time + 1, + height=height // scales.height, + width=width // scales.width, + ) + latents = torch.zeros(latent_shape.to_torch_shape(), device=device, dtype=dtype) + tiles = _prepare_encode_tiles(video, tiling_config, scales) + weights = ( + None + if _encode_masks_are_complementary(tiles, latents.shape) + else torch.zeros_like(latents) + ) + for tile in tiles: + latent_tile = self(video[tile.in_coords].to(device=device, dtype=dtype)) + latents[tile.out_coords] += _apply_encode_masks(latent_tile, tile) + if weights is not None: + strength = torch.ones_like(latent_tile, dtype=torch.float32) + weights[tile.out_coords] += _apply_encode_masks(strength, tile) + return latents if weights is None else latents / weights.clamp(min=1e-8) + + +class LTX23VideoEncoderConfigurator: + """Create an LTX-2.3 video encoder from the native config.""" + + @classmethod + def from_config(cls, config: dict) -> LTX23VideoEncoder: + vae = config.get("vae", {}) + padding_mode = vae.get( + "spatial_padding_mode", vae.get("encoder_spatial_padding_mode", "zeros") + ) + return LTX23VideoEncoder( + convolution_dimensions=vae.get("dims", 3), + in_channels=vae.get("in_channels", vae.get("out_channels", 3)), + out_channels=vae.get("latent_channels", 128), + encoder_blocks=vae.get("encoder_blocks", []), + patch_size=vae.get("patch_size", 4), + norm_layer=NormLayerType(vae.get("norm_layer", "pixel_norm")), + causal=vae.get("causal_encoder", True), + encoder_spatial_padding_mode=PaddingModeType(padding_mode), + ) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx23/media_io.py b/tensorrt_llm/_torch/visual_gen/models/ltx23/media_io.py new file mode 100644 index 000000000000..aaea9c559c1d --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/ltx23/media_io.py @@ -0,0 +1,211 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 Lightricks Ltd. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: LicenseRef-LTX-2 +"""Source video and audio decoding for the native LTX-2.3 retake pipeline. + +Frames are display-oriented, full-range RGB uint8. Audio retains its original +rate and channel layout on the video's presentation timeline. Model-specific +resizing, resampling and normalization remain in the retake pipeline. +""" + +import math +from collections.abc import Iterator +from fractions import Fraction +from typing import TYPE_CHECKING + +import numpy as np +import torch + +from tensorrt_llm.inputs.multimodal_data import AudioData + +from ..ltx2.ltx2_core.types import VideoPixelShape + +if TYPE_CHECKING: + import av + + +def _rotation(frame: "av.VideoFrame") -> int: + rotation = float(frame.rotation) + if not math.isfinite(rotation): + raise ValueError(f"Video display rotation must be finite: {rotation}.") + turns = round(rotation / 90) + if not math.isclose(rotation, turns * 90, abs_tol=1e-3): + raise ValueError(f"Only quarter-turn video display rotations are supported: {rotation}.") + return turns % 4 + + +def _validate_transfer(stream: "av.video.stream.VideoStream") -> None: + # AVFrame transfer tags are unavailable in PyAV 16; the decoder context + # exposes them. Converting HDR into RGB8 requires explicit tone mapping. + if stream.codec_context.color_trc in (16, 18): + raise ValueError("PQ/HLG video requires HDR tone mapping before RGB8 decoding.") + + +def _rgb_frame( + frame: "av.VideoFrame", + reformatter: "av.video.reformatter.VideoReformatter", + *, + fallback_colorspace: int = 2, + fallback_color_range: int = 0, +) -> np.ndarray: + if frame.format.is_rgb or frame.format.name.startswith("gray"): + pixels = frame.to_ndarray(format="rgb24") + else: + # FFmpeg AVCOL_SPC values differ from SWS_CS values for SMPTE170M. + colorspace = int(frame.colorspace) + if colorspace in (0, 2): + colorspace = fallback_colorspace + matrix = {1: 1, 4: 4, 5: 5, 6: 5, 7: 7, 9: 9}.get(colorspace) + if matrix is None: + if colorspace not in (0, 2): + raise ValueError(f"Unsupported YUV color matrix: {colorspace}.") + matrix = 5 # Explicit BT.601 fallback for untagged SDR video. + color_range = int(frame.color_range) or fallback_color_range + if color_range == 0: + color_range = 2 if frame.format.name.startswith("yuvj") else 1 + if color_range not in (1, 2): + raise ValueError(f"Unsupported YUV color range: {color_range}.") + # RGB has no YUV destination matrix. Choosing a different coefficient + # table forces PyAV 16 to apply the source matrix even for full-range YUV. + destination_matrix = 1 if matrix == 5 else 5 + pixels = reformatter.reformat( + frame, + format="rgb24", + src_colorspace=matrix, + dst_colorspace=destination_matrix, + src_color_range=color_range, + dst_color_range=2, + ).to_ndarray() + turns = _rotation(frame) + if turns: + pixels = np.rot90(pixels, k=turns, axes=(0, 1)) + return np.ascontiguousarray(pixels) + + +def get_videostream_metadata(path: str) -> VideoPixelShape: + """Read display geometry and count frames when the container omits a count.""" + import av + + with av.open(path) as container: + if not container.streams.video: + raise ValueError("Media has no video stream.") + stream = container.streams.video[0] + frames = container.decode(stream) + first = next(frames, None) + if first is None: + raise ValueError("Video stream contains no decoded frames.") + _validate_transfer(stream) + width, height = first.width, first.height + if _rotation(first) % 2: + width, height = height, width + count = int(stream.frames or 0) + if not count: + count = 1 + sum(1 for _ in frames) + rate = stream.average_rate or stream.guessed_rate + if rate is None or rate <= 0 or not math.isfinite(float(rate)): + raise ValueError("Video stream has no positive frame rate.") + return VideoPixelShape(batch=1, frames=count, height=height, width=width, fps=float(rate)) + + +def decode_video_by_frame(path: str) -> Iterator[torch.Tensor]: + """Yield display-oriented RGB uint8 tensors shaped ``(1, H, W, C)``. + + Decode sequentially without seeking or changing the frame rate. + """ + import av + + with av.open(path) as container: + if not container.streams.video: + raise ValueError("Media has no video stream.") + stream = container.streams.video[0] + reformatters: dict[tuple[int, int], av.video.reformatter.VideoReformatter] = {} + for frame in container.decode(stream): + _validate_transfer(stream) + colorspace = int(frame.colorspace) + if colorspace in (0, 2): + colorspace = stream.codec_context.colorspace + color_range = int(frame.color_range) or stream.codec_context.color_range + key = (colorspace, color_range) + if key not in reformatters: + reformatters[key] = av.video.reformatter.VideoReformatter() + pixels = _rgb_frame( + frame, + reformatters[key], + fallback_colorspace=stream.codec_context.colorspace, + fallback_color_range=stream.codec_context.color_range, + ) + yield torch.from_numpy(pixels).unsqueeze(0) + + +def decode_audio_from_file(path: str) -> AudioData | None: + """Decode original-rate audio as ``(1, channels, samples)`` on the video timeline. + + Trim negative preroll and declared trailing padding; retain delayed starts + and interior gaps as silence. Return None if there are no valid audio samples. + """ + import av + + with av.open(path) as container: + if not container.streams.audio: + return None + stream = container.streams.audio[0] + sample_rate = int(stream.codec_context.sample_rate) + if sample_rate <= 0: + raise ValueError("Audio stream has no positive sample rate.") + origin_stream = container.streams.video[0] if container.streams.video else stream + if origin_stream.start_time is not None and origin_stream.time_base is not None: + origin = Fraction(origin_stream.start_time) * origin_stream.time_base + elif container.start_time is not None: + origin = Fraction(container.start_time, 1_000_000) + else: + origin = Fraction(0) + valid_end = None + if stream.duration is not None and stream.time_base is not None: + stream_start = Fraction(stream.start_time or 0) * stream.time_base + valid_end = stream_start + stream.duration * stream.time_base + end_sample = round((valid_end - origin) * sample_rate) if valid_end is not None else None + resampler = av.AudioResampler(format="fltp", layout=stream.layout, rate=sample_rate) + timestamp_tolerance = ( + max(0, math.ceil(stream.time_base * sample_rate) - 1) + if stream.time_base is not None + else 0 + ) + chunks: list[tuple[int, np.ndarray]] = [] + previous_end = None + + def append(frame: "av.AudioFrame") -> None: + nonlocal previous_end + if frame.pts is None or frame.time_base is None: + raise ValueError("Audio frames must carry presentation timestamps.") + offset = round((frame.pts * frame.time_base - origin) * sample_rate) + if previous_end is not None and abs(offset - previous_end) <= timestamp_tolerance: + offset = previous_end + if previous_end is not None and offset < previous_end: + raise ValueError("Audio frame presentation timestamps overlap or run backwards.") + previous_end = offset + frame.samples + left = max(0, -offset) + right = frame.samples if end_sample is None else min(frame.samples, end_sample - offset) + if right > left: + samples = frame.to_ndarray()[:, left:right] + if not np.isfinite(samples).all(): + raise ValueError("Decoded audio contains non-finite samples.") + chunks.append((offset + left, samples)) + + # Decode sequentially: seeking to zero can drop a valid first AAC frame + # depending on the demuxer version and edit-list handling. + for frame in container.decode(stream): + if frame.pts is None or frame.time_base is None: + raise ValueError("Audio frames must carry presentation timestamps.") + if frame.sample_rate != sample_rate or frame.layout.name != stream.layout.name: + raise ValueError("Audio rate or channel layout changes within the stream.") + for converted in resampler.resample(frame): + append(converted) + for converted in resampler.resample(None): + append(converted) + if not chunks: + return None + length = chunks[-1][0] + chunks[-1][1].shape[-1] + samples = np.zeros((chunks[0][1].shape[0], length), dtype=np.float32) + for offset, values in chunks: + samples[:, offset : offset + values.shape[-1]] = values + return AudioData(samples=torch.from_numpy(samples).unsqueeze(0), sample_rate=sample_rate) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx23/pipeline_ltx23.py b/tensorrt_llm/_torch/visual_gen/models/ltx23/pipeline_ltx23.py new file mode 100644 index 000000000000..9e5511bcee6f --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/ltx23/pipeline_ltx23.py @@ -0,0 +1,666 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 Lightricks Ltd. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: LicenseRef-LTX-2 + +import copy +import hashlib +import math +import os +import weakref +from dataclasses import fields, replace +from typing import Any, Dict, List, Optional, Union + +import torch + +from tensorrt_llm._torch.autotuner import autotune +from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput +from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline, ExtraParamSchema +from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline +from tensorrt_llm._torch.visual_gen.utils import postprocess_video_tensor +from tensorrt_llm.logger import logger + +from ..ltx2.ltx2_core.audio_vae import AudioDecoderConfigurator, decode_audio +from ..ltx2.ltx2_core.patchifier import VideoLatentPatchifier, get_pixel_coords +from ..ltx2.ltx2_core.types import ( + VIDEO_SCALE_FACTORS, + AudioLatentShape, + VideoLatentShape, + VideoPixelShape, +) +from ..ltx2.ltx2_core.video_vae import TilingConfig +from ..ltx2.parallel_vae import tile_parallel_decode +from ..ltx2.pipeline_ltx2 import LTX2Pipeline, _load_component_weights, _LTX2CUDAGraphRunner +from ..ltx2.text_cache import TextCache +from .ltx23_core.audio_vae import LTX23VocoderConfigurator +from .ltx23_core.connector import ( + LTX23AudioConnectorConfigurator, + LTX23GemmaFeaturesExtractor, + LTX23VideoConnectorConfigurator, +) +from .ltx23_core.modality import LTX23Modality +from .ltx23_core.video_vae import LTX23VideoDecoderConfigurator +from .transformer_ltx23 import LTX23Model + + +class _LTX23CUDAGraphRunner(_LTX2CUDAGraphRunner): + """Shared generation/retake graph runner for LTX23Modality and TextCache. + + Tensor metadata keys dynamic latents, sigma, timesteps and prompt tensors. + Padding-mask contents also key the immutable FlashInfer plans. Equal masks + and differently valued prompts share a graph; replay copies all tensors. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._mask_fingerprints = {} + + def _key_parts_for(self, prefix, value): + if isinstance(value, (LTX23Modality, TextCache)): + for field in fields(value): + tensor = getattr(value, field.name) + label = f"{prefix}.{field.name}" + yield from self._key_parts_for(label, tensor) + if field.name.endswith("mask") and isinstance(tensor, torch.Tensor): + cached = self._mask_fingerprints.get(label) + # Inference tensors have no version counter, so their small + # masks are re-read outside capture to avoid stale plans. + version = None if torch.is_inference(tensor) else tensor._version + if ( + version is None + or cached is None + or cached[0]() is not tensor + or cached[1] != version + ): + data = ( + tensor.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes() + ) + cached = (weakref.ref(tensor), version, hashlib.sha256(data).hexdigest()) + self._mask_fingerprints[label] = cached + yield (f"{label}.contents", cached[2]) + return + if isinstance(value, torch.Tensor): + yield (prefix, (tuple(value.shape), str(value.dtype), str(value.device))) + return + if isinstance(value, (tuple, list)): + yield (f"{prefix}.length", len(value)) + for index, item in enumerate(value): + yield from self._key_parts_for(f"{prefix}.{index}", item) + return + if isinstance(value, bool): + yield (prefix, value) + return + yield from super()._key_parts_for(prefix, value) + + def capture(self, key, fn, args, kwargs): + outputs = [] + + def remember_outputs(*args, **kwargs): + output = fn(*args, **kwargs) + outputs[:] = [output] + return output + + super().capture(key, remember_outputs, args, kwargs) + # Outputs of distinct mask/shape graphs must not alias freed storage. + self._graph_output_refs[key] = outputs[0] + + @staticmethod + def _clone_value(value): + cloned = _LTX2CUDAGraphRunner._clone_value(value) + if isinstance(value, LTX23Modality): + return replace( + cloned, + sigma=value.sigma.clone(), + cross_modality_sigma=( + value.cross_modality_sigma.clone() + if value.cross_modality_sigma is not None + else None + ), + ) + return cloned + + @staticmethod + def _copy_value(dst, src): + copied = _LTX2CUDAGraphRunner._copy_value(dst, src) + if isinstance(src, LTX23Modality) and isinstance(dst, LTX23Modality): + dst.sigma.copy_(src.sigma) + if dst.cross_modality_sigma is not None and src.cross_modality_sigma is not None: + dst.cross_modality_sigma.copy_(src.cross_modality_sigma) + return copied + + +def detect_native_ltx_pipeline(config: dict) -> str: + """Return LTX2Pipeline or LTX23Pipeline from native safetensors config.""" + ltx23_text_config = { + "caption_proj_before_connector": True, + "caption_projection_first_linear": False, + "caption_projection_second_linear": False, + "caption_proj_input_norm": False, + } + transformer = config.get("transformer", {}) + expected = set(ltx23_text_config) + present = expected.intersection(transformer) + + if not present: + if transformer.get("cross_attention_adaln", False): + raise ValueError( + "Ambiguous LTX checkpoint: cross_attention_adaln=True without " + "the LTX-2.3 text-projection configuration." + ) + return "LTX2Pipeline" + + missing = expected - set(transformer) + mismatched = { + k: transformer[k] + for k, v in ltx23_text_config.items() + if k in transformer and transformer[k] != v + } + if missing or mismatched: + raise ValueError( + "Unsupported/partial LTX text-projection config: " + f"missing={sorted(missing)}, mismatched={mismatched}" + ) + if transformer.get("cross_attention_adaln") is not True: + raise ValueError("LTX-2.3 text projection found without cross_attention_adaln=True.") + return "LTX23Pipeline" + + +@register_pipeline( + "LTX23Pipeline", + hf_ids=["Lightricks/LTX-2.3"], + defaults={ + "text_encoder_path": "google/gemma-3-12b-it-qat-q4_0-unquantized", + "workflow": "generation", + }, + doc="Lightricks LTX-2.3 text-to-video generation and retake with audio.", +) +class LTX23Pipeline(BasePipeline): + """Text-to-video pipeline for the LTX-2.3 checkpoint (native single-file safetensors).""" + + _supports_quantized_attention = True + + # Connectors share the transformer prefix but load as separate components. + _TRANSFORMER_PREFIX = LTX2Pipeline._TRANSFORMER_PREFIX + _TRANSFORMER_EXCLUDE_PREFIXES = LTX2Pipeline._TRANSFORMER_EXCLUDE_PREFIXES + + @classmethod + def resolve_variant(cls, config): + workflow = config.extra_attrs.get("workflow", "generation") + if workflow == "retake": + from .pipeline_ltx23_retake import LTX23RetakePipeline + + return LTX23RetakePipeline + if workflow != "generation": + raise ValueError(f"Unsupported LTX-2.3 workflow: {workflow!r}") + return cls + + default_warmup_resolutions = LTX2Pipeline.default_warmup_resolutions + default_warmup_num_frames = LTX2Pipeline.default_warmup_num_frames + default_generation_params = LTX2Pipeline.default_generation_params + load_transformer_weights = LTX2Pipeline.load_transformer_weights + load_weights = LTX2Pipeline.load_weights + load_standard_components = LTX2Pipeline.load_standard_components + setup_parallel_vae = LTX2Pipeline.setup_parallel_vae + _encode_prompt = LTX2Pipeline._encode_prompt + _extra_transformer_kwargs = LTX2Pipeline._extra_transformer_kwargs + _setup_cuda_graphs = LTX2Pipeline._setup_cuda_graphs + _transformer_class = LTX23Model + _cuda_graph_runner_class = _LTX23CUDAGraphRunner + + def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> None: + self.forward( + prompt="warmup", + negative_prompt="", + height=height, + width=width, + num_frames=num_frames, + num_inference_steps=steps, + guidance_scale=4.0, + seed=42, + ) + + def warmup(self) -> None: + """Tune with capture off, then capture from the warm tactic cache. + + A single pass would record autotuner torch.rand kernels into the graph. + """ + if not ( + self.pipeline_config.torch_compile.enable_autotune + and self.pipeline_config.cuda_graph.enable + and self.world_size == 1 + ): + super().warmup() + return + + shapes, steps = self.resolve_warmup_plan() + if not shapes: + super().warmup() + return + + shape_list = ", ".join(f"{h}x{w}x{f}" for h, w, f in shapes) + logger.info( + f"Running warmup for {self.__class__.__name__}: {len(shapes)} shape(s) " + f"[{shape_list}], {steps} steps, tuning pass then capture pass..." + ) + self._is_warmup = True + try: + with ( + self.disallow_cuda_graph_capture(), + autotune( + cache_path=os.environ.get("TLLM_AUTOTUNER_CACHE_PATH"), + skip_dynamic_tuning_buckets=True, + ), + ): + self._run_warmup_pass(shapes, steps) + self._run_warmup_pass(shapes, steps) + finally: + self._is_warmup = False + + self._warmed_up_shapes = set( + self.warmup_cache_key(h, w, num_frames=f) for h, w, f in shapes + ) + + def _init_transformer(self) -> None: + model_config = self.pipeline_config.model_configs["transformer"] + attention = model_config.attention + quant_attention = attention.quant_attention_config + if ( + attention.backend == "FLASHINFER" + and quant_attention is not None + and quant_attention.qk_dtype in ("mxfp8", "nvfp4") + and quant_attention.v_dtype == "fp8" + ): + if model_config.attention_metadata_state is None: + model_config.attention_metadata_state = {} + model_config.attention_metadata_state["_ltx23_static_global_scale"] = 1.0 + LTX2Pipeline._init_transformer(self) + self.transformer_in_channels = self.transformer._transformer_config.get("in_channels", 128) + + def _load_native_components( + self, + config: Dict[str, Any], + sft_paths: List[str], + device: torch.device, + dtype: torch.dtype, + skip_components: Optional[list] = None, + ) -> None: + skip_components = skip_components or [] + + def _load(module, prefixes, load_dtype): + _load_component_weights(sft_paths, module, prefixes) + return module.to(device=device, dtype=load_dtype) + + if PipelineComponent.VAE not in skip_components: + logger.info("Loading native video decoder...") + self.video_decoder = _load( + LTX23VideoDecoderConfigurator.from_config(config), + ["vae.decoder.", "vae."], + dtype, + ) + + if "connectors" not in skip_components: + logger.info("Loading native text feature extractor + connectors...") + self.feature_extractor = _load( + LTX23GemmaFeaturesExtractor.from_config(config), + "text_embedding_projection.", + dtype, + ) + self.video_connector = _load( + LTX23VideoConnectorConfigurator.from_config(config), + "model.diffusion_model.video_embeddings_connector.", + dtype, + ) + self.audio_connector = _load( + LTX23AudioConnectorConfigurator.from_config(config), + "model.diffusion_model.audio_embeddings_connector.", + dtype, + ) + + if PipelineComponent.VAE not in skip_components and "audio_decoder" not in skip_components: + logger.info("Loading native audio decoder...") + self.audio_decoder = _load( + AudioDecoderConfigurator.from_config(config), + ["audio_vae.decoder.", "audio_vae."], + dtype, + ) + if PipelineComponent.VAE not in skip_components and "vocoder" not in skip_components: + logger.info("Loading native vocoder (BigVGAN-v2 + BWE)...") + self.vocoder = _load( + LTX23VocoderConfigurator.from_config(config), "vocoder.", torch.float32 + ) + + patch_size = self.transformer._transformer_config.get("patch_size", 1) + self.video_patchifier = VideoLatentPatchifier(patch_size=patch_size) + + if PipelineComponent.VAE not in skip_components and hasattr(self, "audio_decoder"): + self.audio_patchifier = self.audio_decoder.patchifier + self.audio_sampling_rate = self.audio_decoder.sample_rate + self.audio_hop_length = self.audio_decoder.mel_hop_length + self.audio_mel_bins = self.audio_decoder.mel_bins + + @staticmethod + def _pack_text_embeds( + text_hidden_states: torch.Tensor, + sequence_lengths: torch.Tensor, + device: Union[str, torch.device], + padding_side: str = "left", + scale_factor: int = 8, + eps: float = 1e-6, + ) -> torch.Tensor: + """LTX-2.3 packs Gemma states with per-token RMS normalization.""" + variance = text_hidden_states.pow(2).mean(dim=2, keepdim=True) + normalized = (text_hidden_states * torch.rsqrt(variance + eps)).flatten(2, 3) + positions = torch.arange(text_hidden_states.shape[1], device=text_hidden_states.device) + if padding_side == "left": + mask = positions[None, :] >= text_hidden_states.shape[1] - sequence_lengths[:, None] + elif padding_side == "right": + mask = positions[None, :] < sequence_lengths[:, None] + else: + raise ValueError(f"Unsupported text padding side: {padding_side!r}") + return torch.where(mask[..., None], normalized, torch.zeros_like(normalized)) + + def _process_connectors(self, prompt_embeds: torch.Tensor, attention_mask: torch.Tensor): + """Feature extractor, then video and audio connectors. Mask is shared.""" + additive_mask = (1 - attention_mask.to(prompt_embeds.dtype)) * -1000000.0 + additive_mask = additive_mask.unsqueeze(1).unsqueeze(1) # [B, 1, 1, S] + + video_proj, audio_proj = self.feature_extractor(prompt_embeds) + video_embeds, video_mask = self.video_connector(video_proj, additive_mask) + audio_embeds, _ = self.audio_connector(audio_proj, additive_mask) + return video_embeds, audio_embeds, video_mask + + @property + def extra_param_specs(self): + return { + "output_type": ExtraParamSchema( + type="str", + default="pt", + description="Output type: 'pt' for tensors, 'latent' for raw latents.", + ), + "guidance_rescale": ExtraParamSchema( + type="float", + default=0.0, + description="Guidance rescale factor to prevent overexposure.", + ), + } + + def infer(self, req): + extra = req.params.extra_params or {} + return self.forward( + prompt=req.prompt, + negative_prompt=req.params.negative_prompt, + height=req.params.height, + width=req.params.width, + num_frames=req.params.num_frames, + frame_rate=req.params.frame_rate, + num_inference_steps=req.params.num_inference_steps, + guidance_scale=req.params.guidance_scale, + seed=req.params.seed, + output_type=extra.get("output_type", "pt"), + guidance_rescale=extra.get("guidance_rescale", 0.0), + max_sequence_length=req.params.max_sequence_length, + ) + + def _decode_video( + self, + latents: torch.Tensor, + video_shape: VideoLatentShape, + tiling_config: TilingConfig, + generator: torch.Generator, + output_type: str = "pt", + ) -> torch.Tensor: + """Unpatchify and decode video latents with the supplied tiling geometry.""" + latents = self.video_patchifier.unpatchify(latents, video_shape) + if output_type == "latent": + return latents + if self._parallel_vae_enabled: + video = tile_parallel_decode( + self.video_decoder, + latents.to(self.dtype), + tiling_config, + pg=self.pipeline_config.visual_gen_mapping.vae_group, + generator=generator, + ) + else: + chunks = self.video_decoder.tiled_decode( + latents.to(self.dtype), tiling_config, generator=generator + ) + video = torch.cat(list(chunks), dim=2) + return postprocess_video_tensor(video) + + @torch.inference_mode() + def forward( + self, + prompt: Union[str, List[str]], + seed: int, + negative_prompt: Optional[Union[str, List[str]]] = None, + height: int = 512, + width: int = 768, + num_frames: int = 121, + frame_rate: float = 24.0, + num_inference_steps: int = 40, + guidance_scale: float = 4.0, + guidance_rescale: float = 0.0, + output_type: str = "pt", + max_sequence_length: int = 1024, + ): + timer = CudaPhaseTimer() + timer.mark_pre_start() + generator = torch.Generator(device=self.device).manual_seed(seed) + + do_cfg = guidance_scale > 1.0 + + logger.info("Encoding prompts...") + prompt_embeds, prompt_attention_mask = self._encode_prompt( + prompt, num_videos_per_prompt=1, max_sequence_length=max_sequence_length + ) + neg_prompt_embeds = neg_prompt_attention_mask = None + if do_cfg: + negative_prompt = negative_prompt or "" + neg_prompt_embeds, neg_prompt_attention_mask = self._encode_prompt( + negative_prompt, num_videos_per_prompt=1, max_sequence_length=max_sequence_length + ) + + if do_cfg: + combined_embeds = torch.cat([neg_prompt_embeds, prompt_embeds], dim=0) + combined_mask = torch.cat([neg_prompt_attention_mask, prompt_attention_mask], dim=0) + video_embeds_combined, audio_embeds_combined, connector_mask_combined = ( + self._process_connectors(combined_embeds, combined_mask) + ) + neg_video_embeds, video_embeds = video_embeds_combined.chunk(2, dim=0) + neg_audio_embeds, audio_embeds = audio_embeds_combined.chunk(2, dim=0) + neg_connector_mask, connector_mask = connector_mask_combined.chunk(2, dim=0) + else: + video_embeds, audio_embeds, connector_mask = self._process_connectors( + prompt_embeds, prompt_attention_mask + ) + neg_video_embeds = neg_audio_embeds = neg_connector_mask = None + + logger.info("Preparing latents...") + pixel_shape = VideoPixelShape( + batch=1, frames=num_frames, height=height, width=width, fps=frame_rate + ) + video_shape = VideoLatentShape.from_pixel_shape( + pixel_shape, latent_channels=self.transformer_in_channels + ) + # Mel bins are pre-patch, so divide by 4 to align with the video timeline. + audio_shape = AudioLatentShape.from_video_pixel_shape( + pixel_shape, + channels=getattr(self.audio_decoder, "z_channels", 8), + mel_bins=self.audio_mel_bins // 4, + sample_rate=self.audio_sampling_rate, + hop_length=self.audio_hop_length, + ) + self.transformer.configure_audio_ulysses(audio_shape.frames) + + latents = torch.randn( + ( + video_shape.batch, + self.video_patchifier.get_token_count(video_shape), + video_shape.channels * math.prod(self.video_patchifier.patch_size), + ), + generator=generator, + device=self.device, + dtype=self.dtype, + ) + + audio_latents = torch.randn( + (audio_shape.batch, audio_shape.frames, audio_shape.channels * audio_shape.mel_bins), + generator=generator, + device=self.device, + dtype=self.dtype, + ) + + video_positions = self.video_patchifier.get_patch_grid_bounds( + video_shape, device=self.device + ) + video_positions = get_pixel_coords( + video_positions.float(), VIDEO_SCALE_FACTORS, causal_fix=True + ) + video_positions[:, 0, ...] = video_positions[:, 0, ...] / frame_rate + audio_positions = self.audio_patchifier.get_patch_grid_bounds( + audio_shape, device=self.device + ) + + # Shared sigma schedule; deepcopy so each stream has its own step index. + self.scheduler.set_timesteps(num_inference_steps) + self.scheduler.sigmas = self.scheduler.sigmas.to(self.device) + audio_scheduler = copy.deepcopy(self.scheduler) + timesteps = self.scheduler.timesteps + + if do_cfg: + v_ctx = torch.cat([neg_video_embeds, video_embeds]) + a_ctx = torch.cat([neg_audio_embeds, audio_embeds]) + v_mask = ( + torch.cat([neg_connector_mask, connector_mask]) + if connector_mask is not None + else None + ) + else: + v_ctx, a_ctx, v_mask = video_embeds, audio_embeds, connector_mask + + text_cache = self.transformer.prepare_text_cache( + video_context=v_ctx, + video_context_mask=v_mask, + video_positions=video_positions, + audio_context=a_ctx, + audio_context_mask=v_mask, + audio_positions=audio_positions, + dtype=self.dtype, + ) + + def _run_transformer(v_latents, a_latents, timestep_val, v_context, a_context, mask): + v_latents_f32 = v_latents.float() + v_latents_bf = v_latents.to(self.dtype) + a_latents_f32 = a_latents.float() + a_latents_bf = a_latents.to(self.dtype) + + video_mod = LTX23Modality( + latent=v_latents_bf, + timesteps=timestep_val, + sigma=timestep_val, + positions=video_positions, + context=v_context, + context_mask=mask, + ) + audio_mod = LTX23Modality( + latent=a_latents_bf, + timesteps=timestep_val, + sigma=timestep_val, + positions=audio_positions, + context=a_context, + context_mask=mask, + ) + vel_v, vel_a = self.transformer( + video=video_mod, + audio=audio_mod, + text_cache=text_cache, + timestep=timestep_val.new_tensor(0.0), + ) + + # x0 prediction (rectified flow): x0 = sample - v * sigma. + def to_x0(latent, velocity): + sigma = timestep_val.float() + while sigma.dim() < velocity.dim(): + sigma = sigma.unsqueeze(-1) + return (latent - velocity.float() * sigma).to(self.dtype) + + return to_x0(v_latents_f32, vel_v), to_x0(a_latents_f32, vel_a) + + def forward_fn( + video_latents, + extra_stream_latents, + step_index, + timestep, + encoder_hidden_states, + extra_tensors, + ): + dn_v, dn_a = _run_transformer( + video_latents, + extra_stream_latents["audio"], + timestep, + encoder_hidden_states, + extra_tensors.get("audio_embeds", audio_embeds), + extra_tensors.get("attention_mask", connector_mask), + ) + return dn_v, {"audio": dn_a} + + timer.mark_denoise_start() + result = self.denoise( + latents=latents, + scheduler=self.scheduler, + prompt_embeds=video_embeds, + neg_prompt_embeds=neg_video_embeds, + guidance_scale=guidance_scale, + forward_fn=forward_fn, + timesteps=timesteps, + guidance_rescale=guidance_rescale, + extra_cfg_tensors=( + { + "audio_embeds": (audio_embeds, neg_audio_embeds), + "attention_mask": (connector_mask, neg_connector_mask), + } + if do_cfg + else None + ), + extra_streams={"audio": (audio_latents, audio_scheduler)}, + ) + latents, extra_stream_latents = result + audio_latents = extra_stream_latents["audio"] + + timer.mark_post_start() + + logger.info("Decoding video and audio...") + + def decode_video_fn(vid_latents): + return self._decode_video( + vid_latents, video_shape, TilingConfig.default(), generator, output_type + ) + + def decode_audio_fn(aud_latents): + aud_latents = self.audio_patchifier.unpatchify(aud_latents, audio_shape) + if output_type == "latent": + return aud_latents + aud_latents = aud_latents.to(self.dtype) + return decode_audio(aud_latents, self.audio_decoder, self.vocoder) + + decoded = self.decode_latents( + latents=latents, + decode_fn=decode_video_fn, + extra_latents={"audio": (audio_latents, decode_audio_fn)}, + ) + if isinstance(decoded, (tuple, list)): + video, audio = decoded[0], decoded[1] + else: + video, audio = decoded, None + + timer.mark_end() + return timer.fill( + PipelineOutput( + video=video, + audio=audio, + frame_rate=float(frame_rate), + audio_sample_rate=( + int(self.vocoder.output_sampling_rate) if audio is not None else None + ), + ) + ) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx23/pipeline_ltx23_retake.py b/tensorrt_llm/_torch/visual_gen/models/ltx23/pipeline_ltx23_retake.py new file mode 100644 index 000000000000..c556c7e97519 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/ltx23/pipeline_ltx23_retake.py @@ -0,0 +1,832 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Native LTX-2.3 retake workflow for VisualGen. + +This pipeline keeps a resident retake model for ``trtllm-serve`` and extends the +native LTX-2.3 pipeline with source-media encoders and masked denoising. +It either loads Gemma for a live prompt or consumes precomputed post-connector +conditioning. Retake generalizes image-to-video masking to condition both sides +of the regenerated window. + +Source video/audio decode, stream metadata, audio/video VAE encode, diffusion, +and VAE decode all use TensorRT-LLM native components. +""" + +from __future__ import annotations + +import copy +import math +from functools import partial +from typing import TYPE_CHECKING, Any + +import safetensors.torch +import torch + +from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput +from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema +from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline +from tensorrt_llm.inputs.multimodal_data import AudioData +from tensorrt_llm.logger import logger + +from ..ltx2.ltx2_core.audio_vae import decode_audio +from ..ltx2.ltx2_core.patchifier import get_pixel_coords +from ..ltx2.ltx2_core.scheduler_adapter import NativeSchedulerAdapter +from ..ltx2.ltx2_core.types import ( + VIDEO_SCALE_FACTORS, + AudioLatentShape, + VideoLatentShape, + VideoPixelShape, +) +from ..ltx2.ltx2_core.video_vae import SpatialTilingConfig, TemporalTilingConfig, TilingConfig +from ..ltx2.pipeline_ltx2 import _load_component_weights +from ..ltx2.text_cache import TextCache +from .ltx23_core.audio_vae import AudioEncoderConfigurator, encode_audio +from .ltx23_core.modality import LTX23Modality +from .ltx23_core.video_vae import LTX23VideoEncoderConfigurator +from .media_io import decode_audio_from_file, decode_video_by_frame, get_videostream_metadata +from .pipeline_ltx23 import LTX23Pipeline + +if TYPE_CHECKING: + from tensorrt_llm._torch.visual_gen.config import DiffusionPipelineConfig + from tensorrt_llm._torch.visual_gen.executor import DiffusionRequest + +# Distilled retake noise schedule (8 Euler steps). +_RETAKE_DISTILLED_SIGMA_VALUES = [ + 1.0, + 0.99375, + 0.9875, + 0.98125, + 0.975, + 0.909375, + 0.725, + 0.421875, + 0.0, +] +_RETAKE_NUM_INFERENCE_STEPS = len(_RETAKE_DISTILLED_SIGMA_VALUES) - 1 + +_DEFAULT_MAX_SEQUENCE_LENGTH = 1024 +_PROMPT_CONDITIONING_KEYS = ("video_embeds", "audio_embeds", "connector_mask") + +# Retake uses the same tiling geometry for source encode and output decode. Keep +# it separate from the generation pipeline's default because tile boundaries are +# part of the retake numerical contract. +_RETAKE_TILING_CONFIG = TilingConfig( + spatial_config=SpatialTilingConfig(tile_size_in_pixels=768, tile_overlap_in_pixels=64), + temporal_config=TemporalTilingConfig(tile_size_in_frames=80, tile_overlap_in_frames=24), +) + + +def _load_prompt_conditioning(path: str) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Load precomputed post-connector prompt tensors from safetensors.""" + tensors = safetensors.torch.load_file(path, device="cpu") + keys = set(tensors) + required = set(_PROMPT_CONDITIONING_KEYS) + if keys != required: + raise ValueError( + "Prompt-conditioning tensor keys do not match the expected schema: " + f"missing={sorted(required - keys)}, unexpected={sorted(keys - required)}" + ) + + video_embeds, audio_embeds, connector_mask = (tensors[key] for key in _PROMPT_CONDITIONING_KEYS) + if video_embeds.ndim != 3 or audio_embeds.ndim != 3: + raise ValueError( + "Prompt embeddings must be rank 3 [batch, sequence, channels]; " + f"got video={tuple(video_embeds.shape)}, audio={tuple(audio_embeds.shape)}" + ) + if connector_mask.ndim != 4: + raise ValueError( + "Connector mask must be rank 4 [batch, 1, 1, sequence]; " + f"got {tuple(connector_mask.shape)}" + ) + if video_embeds.shape[0] != 1: + raise ValueError( + "LTX-2.3 retake prompt conditioning must contain one prompt; " + f"got batch size {video_embeds.shape[0]}." + ) + if video_embeds.shape[:2] != audio_embeds.shape[:2] or ( + connector_mask.shape[0] != video_embeds.shape[0] + or connector_mask.shape[-1] != video_embeds.shape[1] + ): + raise ValueError( + "Prompt-conditioning batch and sequence dimensions do not align: " + f"video={tuple(video_embeds.shape)}, audio={tuple(audio_embeds.shape)}, " + f"mask={tuple(connector_mask.shape)}" + ) + return video_embeds, audio_embeds, connector_mask + + +def _resolve_retake_time_window( + start_time: float | None, + end_time: float | None, + fps: float, + num_frames: int, +) -> tuple[float, float]: + """Resolve an optional retake window against the source video duration.""" + if fps <= 0: + raise ValueError(f"fps must be positive, got {fps}") + if num_frames < 0: + raise ValueError(f"num_frames must be non-negative, got {num_frames}") + + resolved_start = 0.0 if start_time is None else float(start_time) + resolved_end = num_frames / fps if end_time is None else float(end_time) + if not math.isfinite(resolved_start) or not math.isfinite(resolved_end): + raise ValueError("retake_start_time and retake_end_time must be finite.") + if resolved_start >= resolved_end: + raise ValueError( + f"retake_start_time ({resolved_start}) must be less than " + f"retake_end_time ({resolved_end})" + ) + return resolved_start, resolved_end + + +def _init_retake_patchified_latents( + noise_latents: torch.Tensor, + source_latents: torch.Tensor, + denoise_mask: torch.Tensor, +) -> torch.Tensor: + """Initialize retake latents in the denoiser's patchified token layout. + + Sampling in 5D and then patchifying produces a different seeded random field + because the memory order changes, so noise is sampled directly in token + layout. + """ + if noise_latents.shape != source_latents.shape: + raise ValueError( + f"noise/source latent shape mismatch: {tuple(noise_latents.shape)} vs " + f"{tuple(source_latents.shape)}" + ) + if denoise_mask.dim() == 2: + mask = denoise_mask.unsqueeze(-1) + elif denoise_mask.dim() == 3: + mask = denoise_mask + else: + raise ValueError(f"expected denoise mask rank 2 or 3, got {tuple(denoise_mask.shape)}") + if mask.shape[:2] != source_latents.shape[:2] or mask.shape[-1] != 1: + raise ValueError( + f"denoise mask shape {tuple(denoise_mask.shape)} is not compatible with " + f"patchified latents {tuple(source_latents.shape)}" + ) + return torch.lerp(source_latents.float(), noise_latents.float(), mask.float()).to( + source_latents.dtype + ) + + +def _conform_latent_length(latent: torch.Tensor, expected_frames_count: int) -> torch.Tensor: + """Crop or zero-pad audio latent frames along dimension 2.""" + missing_frames = expected_frames_count - latent.shape[2] + if missing_frames < 0: + return latent[:, :, :expected_frames_count] + if missing_frames > 0: + padding = [0] * (2 * (latent.ndim - 2)) + padding[-1] = missing_frames + return torch.nn.functional.pad(latent, padding) + return latent + + +def _retake_temporal_denoise_mask( + positions: torch.Tensor, start_time: float, end_time: float +) -> torch.Tensor: + """Select tokens overlapping a window on a modality's own timeline in seconds.""" + if positions.ndim != 4 or positions.shape[1] < 1 or positions.shape[-1] != 2: + raise ValueError("Positions must have shape (batch, axes, tokens, 2).") + start, end = positions[:, 0].unbind(dim=-1) + return ((end > start_time) & (start < end_time)).to(torch.float32) + + +def _init_retake_audio_state( + clean_audio: torch.Tensor, + generator: torch.Generator, + denoise_mask: torch.Tensor, +) -> torch.Tensor: + """Sample audio after video noise, then apply source conditioning in token layout.""" + noise = torch.randn( + clean_audio.shape, generator=generator, device=clean_audio.device, dtype=clean_audio.dtype + ) + return _init_retake_patchified_latents(noise, clean_audio, denoise_mask) + + +# Retake shares checkpoints with LTX23Pipeline and is selected by ``workflow``. +@register_pipeline("LTX23RetakePipeline", doc="Native LTX-2.3 audio/video retake pipeline.") +class LTX23RetakePipeline(LTX23Pipeline): + """Persistent VisualGen pipeline for native LTX-2.3 retake requests.""" + + def __init__(self, pipeline_config: "DiffusionPipelineConfig") -> None: + if pipeline_config.cache_backend is not None: + raise NotImplementedError("Cache acceleration is not supported for LTX-2.3 retake.") + if pipeline_config.parallel.cfg_size != 1: + raise NotImplementedError("LTX-2.3 retake requires parallel.cfg_size=1.") + if pipeline_config.parallel.tp_size != 1: + raise NotImplementedError("Tensor parallelism is not supported for LTX-2.3 retake.") + if pipeline_config.parallel.seq_parallel_size != pipeline_config.parallel.ulysses_size: + raise NotImplementedError("LTX-2.3 retake supports Ulysses sequence parallelism only.") + super().__init__(pipeline_config) + + @property + def default_generation_params(self) -> dict[str, int]: + return { + "num_inference_steps": _RETAKE_NUM_INFERENCE_STEPS, + "seed": 42, + "max_sequence_length": _DEFAULT_MAX_SEQUENCE_LENGTH, + } + + @property + def extra_param_specs(self) -> dict[str, ExtraParamSchema]: + return { + "retake_video_path": ExtraParamSchema( + type="str", + description="Path to the source video file for retake.", + ), + "retake_prompt_conditioning_path": ExtraParamSchema( + type="str", + default=None, + description="Path to precomputed post-connector prompt conditioning.", + ), + "retake_start_time": ExtraParamSchema( + type="float", + default=0.0, + description=( + "Start time in seconds for the regenerated window. " + "Defaults to the beginning of the video." + ), + ), + "retake_end_time": ExtraParamSchema( + type="float", + default=None, + description=( + "End time in seconds for the regenerated window. " + "Defaults to the end of the video." + ), + ), + "retake_regenerate_audio": ExtraParamSchema( + type="bool", + default=True, + description="Regenerate source audio in the window; missing audio is generated from scratch.", + ), + "retake_regenerate_video": ExtraParamSchema( + type="bool", + default=True, + description="Regenerate video inside the retake window.", + ), + } + + def _load_native_components( + self, + config: dict[str, Any], + safetensors_paths: list[str], + device: torch.device, + dtype: torch.dtype, + skip_components: list | None = None, + ) -> None: + """Extend the LTX-2.3 components with source-media encoders.""" + skip_components = skip_components or [] + super()._load_native_components( + config, + safetensors_paths, + device, + dtype, + skip_components, + ) + + if "video_encoder" not in skip_components: + encoder_blocks = config.get("vae", {}).get("encoder_blocks", []) + if not encoder_blocks: + raise ValueError("LTX-2.3 checkpoint config has no video VAE encoder blocks.") + self.video_encoder = LTX23VideoEncoderConfigurator.from_config(config) + _load_component_weights(safetensors_paths, self.video_encoder, "vae.encoder.") + _load_component_weights( + safetensors_paths, + self.video_encoder.per_channel_statistics, + "vae.per_channel_statistics.", + ) + self.video_encoder = self.video_encoder.to(device=device, dtype=dtype) + else: + self.video_encoder = None + + self._audio_encoder = AudioEncoderConfigurator.from_config(config) + _load_component_weights(safetensors_paths, self._audio_encoder, "audio_vae.encoder.") + _load_component_weights( + safetensors_paths, + self._audio_encoder.per_channel_statistics, + "audio_vae.per_channel_statistics.", + ) + self._audio_encoder = self._audio_encoder.to(device=device, dtype=dtype) + self.audio_patchifier = self._audio_encoder.patchifier + + def _masked_transformer_step( + self, + video_latents: torch.Tensor, + audio_latents: torch.Tensor | None, + step_index: int, + timestep: torch.Tensor, + *, + video_positions: torch.Tensor, + audio_positions: torch.Tensor | None, + denoise_mask: torch.Tensor, + clean_latent: torch.Tensor, + num_steps: int, + text_cache: TextCache, + audio_denoise_mask: torch.Tensor | None = None, + clean_audio_latent: torch.Tensor | None = None, + regenerate_video: bool = True, + regenerate_audio: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Predict masked video x0 and optionally masked audio x0.""" + video_float = video_latents.float() + video_model = video_latents.to(self.dtype) + audio_model = audio_latents.to(self.dtype) if audio_latents is not None else None + # Upstream preserves float32 sigma/mask products through x0 conversion. + # The transformer handles casting inside its native embedding modules. + sigma = timestep.float() + video_sigma = sigma if regenerate_video else torch.zeros_like(sigma) + audio_sigma = sigma if regenerate_audio else torch.zeros_like(sigma) + video_timestep = denoise_mask * video_sigma.unsqueeze(-1) + audio_timestep = ( + audio_denoise_mask * audio_sigma.unsqueeze(-1) + if audio_denoise_mask is not None + else torch.zeros_like(timestep) + ) + video = LTX23Modality( + latent=video_model, + timesteps=video_timestep, + positions=video_positions, + context=text_cache.video_context, + sigma=video_sigma, + cross_modality_sigma=audio_sigma if audio_model is not None else None, + ) + audio = None + if audio_model is not None: + if audio_positions is None: + raise ValueError("Audio positions are required when audio latents are present.") + audio = LTX23Modality( + latent=audio_model, + timesteps=audio_timestep, + positions=audio_positions, + context=text_cache.audio_context, + sigma=audio_sigma, + cross_modality_sigma=video_sigma, + ) + + velocity_video, velocity_audio = self.transformer( + video=video, + audio=audio, + text_cache=text_cache, + timestep=timestep.new_tensor(float(step_index) / num_steps), + step_index=step_index, + ) + if velocity_video is None: + raise RuntimeError("LTX-2.3 retake transformer returned no video prediction.") + denoised_video = (video_float - velocity_video.float() * video_timestep.unsqueeze(-1)).to( + self.dtype + ) + blend = denoise_mask.unsqueeze(-1).float() + denoised_video = (denoised_video.float() * blend + clean_latent.float() * (1.0 - blend)).to( + self.dtype + ) + denoised_audio = None + if audio_denoise_mask is not None: + if audio_latents is None or clean_audio_latent is None or velocity_audio is None: + raise RuntimeError( + "Audio retake requires source latents and an audio velocity prediction." + ) + audio_x0 = ( + audio_latents.float() - velocity_audio.float() * audio_timestep.unsqueeze(-1) + ).to(self.dtype) + audio_blend = audio_denoise_mask.unsqueeze(-1).float() + denoised_audio = ( + audio_x0.float() * audio_blend + clean_audio_latent.float() * (1.0 - audio_blend) + ).to(self.dtype) + return denoised_video, denoised_audio + + def load_standard_components( + self, + checkpoint_dir: str, + device: torch.device, + skip_components: list | None = None, + *, + text_encoder_path: str = "", + **kwargs: Any, + ) -> None: + native_skip_components = list(skip_components or []) + native_skip_components += [PipelineComponent.SCHEDULER] + if not text_encoder_path: + native_skip_components.extend( + [ + PipelineComponent.TOKENIZER, + PipelineComponent.TEXT_ENCODER, + "connectors", + ] + ) + super().load_standard_components( + checkpoint_dir, + device, + text_encoder_path=text_encoder_path, + skip_components=list(dict.fromkeys(native_skip_components)), + ) + + def warmup(self) -> None: + logger.info("LTX-2.3 retake compiles and captures from the first real source request.") + + @torch.inference_mode() + def infer(self, req: "DiffusionRequest") -> PipelineOutput: + extra = req.params.extra_params or {} + video_path = extra.get("retake_video_path") + if video_path is None: + raise ValueError("extra_params['retake_video_path'] is required for LTX-2.3 retake.") + start_time = extra.get("retake_start_time") + end_time = extra.get("retake_end_time") + regenerate_video = extra.get("retake_regenerate_video", True) + regenerate_audio = extra.get("retake_regenerate_audio", True) + for name, value in (("video", regenerate_video), ("audio", regenerate_audio)): + if not isinstance(value, bool): + raise ValueError(f"retake_regenerate_{name} must be a boolean.") + + prompt = self._single_prompt(req.prompt) + timer = CudaPhaseTimer() + timer.mark_pre_start() + video, audio, output_shape = self._run_retake( + req, + video_path, + start_time, + end_time, + prompt, + timer=timer, + regenerate_video=regenerate_video, + regenerate_audio=regenerate_audio, + ) + if video is None: + return timer.fill(PipelineOutput(frame_rate=float(output_shape.fps))) + audio_tensor = audio.samples + sample_rate = audio.sample_rate + if sample_rate <= 0: + raise ValueError(f"sample_rate must be positive, got {sample_rate}") + # The shared encoder uses -shortest; trailing silence preserves all video frames. + required_samples = math.ceil(video.shape[1] * sample_rate / output_shape.fps) + missing_samples = required_samples - audio_tensor.shape[-1] + if missing_samples > 0: + audio_tensor = torch.nn.functional.pad(audio_tensor, (0, missing_samples)) + return timer.fill( + PipelineOutput( + video=video, + audio=audio_tensor, + frame_rate=float(output_shape.fps), + audio_sample_rate=sample_rate, + ) + ) + + def _prepare_prompt_conditioning( + self, + prompt: str, + max_sequence_length: int, + conditioning_path: str | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return supplied prompt conditioning or encode the prompt with Gemma.""" + if conditioning_path is not None: + if prompt: + raise ValueError( + "LTX-2.3 retake accepts either a prompt or precomputed prompt " + "conditioning, not both." + ) + video_embeds, audio_embeds, connector_mask = _load_prompt_conditioning( + conditioning_path + ) + logger.info(f"Using LTX-2.3 retake prompt conditioning from {conditioning_path}.") + return ( + video_embeds.to(device=self.device, dtype=self.dtype), + audio_embeds.to(device=self.device, dtype=self.dtype), + connector_mask.to(device=self.device, dtype=self.dtype), + ) + + text_encoder = getattr(self, "text_encoder", None) + if text_encoder is None: + raise RuntimeError("Gemma is not loaded for LTX-2.3 retake prompt encoding.") + prompt_embeds, prompt_attention_mask = self._encode_prompt( + prompt, + max_sequence_length=max_sequence_length, + ) + return self._process_connectors(prompt_embeds, prompt_attention_mask) + + def _run_retake( + self, + req: "DiffusionRequest", + video_path: str, + start_time: float | None, + end_time: float | None, + prompt: str, + timer: CudaPhaseTimer, + regenerate_video: bool = True, + regenerate_audio: bool = True, + ) -> tuple[torch.Tensor | None, AudioData | None, VideoPixelShape]: + """Encode source media, denoise the masked window, and decode the full clip. + + Returns uint8 ``(1, T, H, W, C)`` video, output audio, and video shape. + Non-decoding ranks return None for both media streams. + """ + if req.params.num_inference_steps != _RETAKE_NUM_INFERENCE_STEPS: + raise ValueError( + "LTX-2.3 native retake uses the fixed distilled " + f"{_RETAKE_NUM_INFERENCE_STEPS}-step schedule; " + f"got num_inference_steps={req.params.num_inference_steps}." + ) + + device = self.device + dtype = self.dtype + seed = req.params.seed + generator = torch.Generator(device=device).manual_seed(seed) + + # ---- 1. Source read + validation -------------------------------- + pixel_shape = get_videostream_metadata(video_path) + num_frames = int(pixel_shape.frames) + height = int(pixel_shape.height) + width = int(pixel_shape.width) + fps = float(pixel_shape.fps) + self._validate_retake_source(num_frames, height, width) + + start_time, end_time = _resolve_retake_time_window(start_time, end_time, fps, num_frames) + + source_norm_5d = self._read_source_video( + video_path, num_frames, height, width, device, dtype + ) + + # ---- 2. Native VAE encode + seed initial latents ---------------- + video_shape = VideoLatentShape.from_pixel_shape( + pixel_shape, latent_channels=self.transformer_in_channels + ) + # Source latents and sampled noise start in the model dtype, while the + # blend/math path widens internally. + if self.video_encoder is None: + raise RuntimeError("LTX-2.3 retake requires native video VAE encoder weights.") + source_window_latents = self.video_encoder.tiled_encode( + source_norm_5d, _RETAKE_TILING_CONFIG + ) + expected_latent_shape = tuple(video_shape.to_torch_shape()) + if tuple(source_window_latents.shape) != expected_latent_shape: + raise ValueError( + "LTX-2.3 native retake: encoded source latent shape " + f"{tuple(source_window_latents.shape)} != expected " + f"{expected_latent_shape}; check source resolution/frame count." + ) + + video_positions = self.video_patchifier.get_patch_grid_bounds(video_shape, device=device) + video_positions = get_pixel_coords( + video_positions.float(), VIDEO_SCALE_FACTORS, causal_fix=True + ) + video_positions[:, 0, ...] = video_positions[:, 0, ...] / fps + # TemporalRegionMask uses overlap in seconds on each modality's own + # causal token bounds, not rounded pixel-frame or shared audio indices. + denoise_mask = ( + _retake_temporal_denoise_mask(video_positions, start_time, end_time) + if regenerate_video + else torch.zeros_like(video_positions[:, 0, :, 0]) + ) + source_patch_latents = self.video_patchifier.patchify(source_window_latents) + # Sample both streams even when frozen to preserve the shared RNG order. + noise_latents = torch.randn( + source_patch_latents.shape, + generator=generator, + device=device, + dtype=dtype, + ) + # Denoising carries the latent in the model dtype and casts it back after + # every Euler step; retaining the float32 blend would change the trajectory. + latents = _init_retake_patchified_latents( + noise_latents, source_patch_latents, denoise_mask + ).to(dtype) + + # Keep the full encoded source in patchified layout. Tokens with + # denoise_mask=1 ignore it; conditioned tokens use it after every + # transformer prediction. + clean_latent = source_patch_latents + + # ---- 4. Native prompt conditioning + text cache ----------------- + max_sequence_length = ( + getattr(req.params, "max_sequence_length", None) or _DEFAULT_MAX_SEQUENCE_LENGTH + ) + conditioning_path = (req.params.extra_params or {}).get("retake_prompt_conditioning_path") + video_embeds, audio_embeds, connector_mask = self._prepare_prompt_conditioning( + prompt, max_sequence_length, conditioning_path + ) + + # RoPE positions remain float32 because position error is amplified by + # the distilled schedule's large final Euler steps. + + # ---- Source audio conditioning ---------------------------------- + source_audio = decode_audio_from_file(video_path) + if getattr(self, "audio_decoder", None) is None or getattr(self, "vocoder", None) is None: + raise RuntimeError( + "Audio retake requires the native audio decoder and vocoder weights." + ) + audio_latent = self._encode_source_audio_latent( + source_audio, + pixel_shape, + ) + scratch_audio = audio_latent is None + regenerate_audio = scratch_audio or regenerate_audio + if scratch_audio: + audio_shape = AudioLatentShape.from_video_pixel_shape(pixel_shape) + audio_latent = torch.zeros(audio_shape.to_torch_shape(), device=device, dtype=dtype) + audio_shape = AudioLatentShape.from_torch_shape(audio_latent.shape) + self.transformer.configure_audio_ulysses(audio_shape.frames) + audio_clean_latents = self.audio_patchifier.patchify(audio_latent.float()) + audio_positions = self.audio_patchifier.get_patch_grid_bounds(audio_shape, device=device) + audio_denoise_mask = ( + _retake_temporal_denoise_mask(audio_positions, start_time, end_time) + if regenerate_audio and not scratch_audio + else torch.ones_like(audio_positions[:, 0, :, 0]) + ) + if not regenerate_audio: + audio_denoise_mask = torch.zeros_like(audio_denoise_mask) + + text_cache = self.transformer.prepare_text_cache( + video_context=video_embeds, + video_context_mask=connector_mask, + video_positions=video_positions, + audio_context=audio_embeds, + audio_context_mask=connector_mask if audio_embeds is not None else None, + audio_positions=audio_positions, + dtype=dtype, + ) + + # ---- 5. Native masked denoise (distilled, non-guided) --------------- + scheduler = NativeSchedulerAdapter() + scheduler.sigmas = torch.tensor( + _RETAKE_DISTILLED_SIGMA_VALUES, dtype=torch.float32, device=device + ) + timesteps = scheduler.timesteps + num_steps = len(timesteps) + + # Frozen audio follows the clean x0 target with conditioning sigma zero; + # regenerated audio instead uses the transformer's masked x0 prediction. + clean_audio = audio_clean_latents.to(dtype) + audio_state = _init_retake_audio_state(clean_audio, generator, audio_denoise_mask) + audio_scheduler = copy.deepcopy(scheduler) + extra_streams = {"audio": (audio_state, audio_scheduler)} + + def retake_forward_fn( + video_latents: torch.Tensor, + extra_stream_latents: dict[str, torch.Tensor] | None, + step_index: int, + timestep: torch.Tensor, + _encoder_hidden_states: torch.Tensor, + _extra_tensors: dict[str, Any], + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + audio_latents = ( + extra_stream_latents.get("audio") if extra_stream_latents is not None else None + ) + denoised_video, denoised_audio = self._masked_transformer_step( + video_latents, + audio_latents, + step_index, + timestep, + video_positions=video_positions, + audio_positions=audio_positions, + denoise_mask=denoise_mask, + clean_latent=clean_latent, + num_steps=num_steps, + text_cache=text_cache, + audio_denoise_mask=audio_denoise_mask, + clean_audio_latent=clean_audio, + regenerate_video=regenerate_video, + regenerate_audio=regenerate_audio, + ) + audio_prediction = {"audio": denoised_audio} + return denoised_video, audio_prediction + + timer.mark_denoise_start() + denoise_result = self.denoise( + latents=latents, + scheduler=scheduler, + prompt_embeds=video_embeds, + guidance_scale=1.0, + forward_fn=retake_forward_fn, + timesteps=timesteps, + extra_streams=extra_streams, + ) + denoised_latents = denoise_result[0] + + # ---- 6. Native decode ------------------------------------------- + timer.mark_post_start() + output_video, output_audio = self.decode_latents( + latents=denoised_latents, + decode_fn=partial( + self._decode_video, + video_shape=video_shape, + tiling_config=_RETAKE_TILING_CONFIG, + generator=generator, + ), + extra_latents={ + "audio": ( + denoise_result[1]["audio"], + partial(self._decode_retake_audio, audio_shape=audio_shape), + ) + }, + ) + if output_video is not None: + output_video = output_video.contiguous() + timer.mark_end() + + return output_video, output_audio, pixel_shape + + def _decode_retake_audio( + self, latents: torch.Tensor, audio_shape: AudioLatentShape + ) -> AudioData: + """Decode the updated audio stream through the existing native audio stack.""" + if getattr(self, "audio_decoder", None) is None or getattr(self, "vocoder", None) is None: + raise RuntimeError( + "Audio retake requires the native audio decoder and vocoder weights." + ) + unpatchified = self.audio_patchifier.unpatchify(latents, audio_shape).to(self.dtype) + samples = decode_audio(unpatchified, self.audio_decoder, self.vocoder) + if samples.ndim != 3 or samples.shape[0] != 1 or samples.shape[-1] == 0: + raise RuntimeError( + f"Audio retake decoder returned invalid waveform shape {tuple(samples.shape)}." + ) + return AudioData(samples=samples[0], sample_rate=int(self.vocoder.output_sampling_rate)) + + @staticmethod + def _validate_retake_source(num_frames: int, height: int, width: int) -> None: + """Fail fast on source video shapes the native VAE cannot round-trip. + + The LTX-2.3 causal video VAE requires ``8k + 1`` pixel frames and spatial + dimensions divisible by the VAE spatial scale factors so encode/decode + preserves shape. + """ + ratio = VIDEO_SCALE_FACTORS.time + if num_frames <= 0: + raise ValueError(f"retake source must have frames; got {num_frames}") + if (num_frames - 1) % ratio != 0: + snapped = ((num_frames - 1) // ratio) * ratio + 1 + raise ValueError( + f"retake source frame count must satisfy {ratio}k+1 (e.g. 97, 193); " + f"got {num_frames}. Use a source with {snapped} frames." + ) + height_scale = VIDEO_SCALE_FACTORS.height + width_scale = VIDEO_SCALE_FACTORS.width + if height % height_scale != 0 or width % width_scale != 0: + raise ValueError( + "retake source resolution must be divisible by the VAE spatial " + f"scale factors {height_scale}x{width_scale}; got {height}x{width}." + ) + + def _read_source_video( + self, + video_path: str, + num_frames: int, + height: int, + width: int, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + """Read and validate the source video as a normalized VAE input. + + Returns ``(1, 3, T, H, W)`` in ``[-1, 1]``. Pixel decoding is + deterministic and uses sequential frame indices. + """ + frames = list(decode_video_by_frame(video_path)) + if not frames: + raise ValueError(f"retake source video decoded no frames: {video_path}") + source_uint8 = torch.cat(frames, dim=0) # (T, H, W, C) + decoded_frames = source_uint8.shape[0] + if decoded_frames != num_frames: + raise ValueError( + f"retake source decoded {decoded_frames} frames but metadata reported {num_frames}." + ) + if source_uint8.shape[1] != height or source_uint8.shape[2] != width: + raise ValueError( + f"retake source frame size {tuple(source_uint8.shape[1:3])} does not " + f"match metadata {(height, width)}." + ) + # uint8 [0, 255] -> [-1, 1], laid out (1, C, T, H, W) for the VAE encoder. + normalized = source_uint8.to(torch.float32) / 127.5 - 1.0 + return normalized.permute(3, 0, 1, 2).unsqueeze(0).to(device=device, dtype=dtype) + + def _encode_source_audio_latent( + self, + audio: AudioData | None, + pixel_shape: VideoPixelShape, + ) -> torch.Tensor | None: + """Encode source audio and align its latent length to the video timeline.""" + if audio is None: + return None + max_samples = round(pixel_shape.frames * audio.sample_rate / pixel_shape.fps) + conditioning_audio = AudioData( + samples=audio.samples[..., :max_samples], + sample_rate=audio.sample_rate, + ) + latents = encode_audio(conditioning_audio, self._audio_encoder).to(self.device, self.dtype) + required_latent_frames = AudioLatentShape.from_video_pixel_shape( + pixel_shape, + sample_rate=self._audio_encoder.sample_rate, + hop_length=self._audio_encoder.mel_hop_length, + audio_latent_downsample_factor=( + self._audio_encoder.patchifier.audio_latent_downsample_factor + ), + ).frames + return _conform_latent_length(latents, required_latent_frames) + + @staticmethod + def _single_prompt(prompt: str | list[str]) -> str: + if isinstance(prompt, str): + return prompt + if len(prompt) == 1: + return prompt[0] + raise ValueError("LTX-2.3 retake workflow supports one prompt per request.") diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx23/transformer_ltx23.py b/tensorrt_llm/_torch/visual_gen/models/ltx23/transformer_ltx23.py new file mode 100644 index 000000000000..a1b56eb19c51 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/ltx23/transformer_ltx23.py @@ -0,0 +1,493 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 Lightricks Ltd. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: LicenseRef-LTX-2 +"""LTX-2.3 transformer using LTX-2 attention, MLP, preprocessing, and fused AdaLN ops. + +Each stream has nine AdaLN slots: self-attention [0:3], FFN [3:6], and text +cross-attention [6:9]. Text K/V is modulated by sigma and projected each step. +Caption projection is handled by the feature extractor before the connectors. +""" + +from __future__ import annotations + +from dataclasses import replace + +import torch +import torch.nn as nn + +from tensorrt_llm._torch.visual_gen.utils import SequenceSharder + +from ..ltx2.ltx2_core.adaln import AdaLayerNormSingle +from ..ltx2.ltx2_core.utils_ltx2 import ( + apply_fused_gate_resid, + apply_fused_gate_resid_rmsnorm_shift_scale, + apply_fused_rmsnorm_shift_scale, + apply_shift_scale, + is_fused_adaln_supported_dim, +) +from ..ltx2.text_cache import TextCache +from ..ltx2.transformer_ltx2 import ( + AudioShardMode, + BasicAVTransformerBlock, + LTX2Attention, + LTXModel, + TransformerConfig, +) +from .ltx23_core.modality import LTX23Modality + +_get_ada_table_ts_pairs = BasicAVTransformerBlock._get_ada_table_ts_pairs +_get_av_ca_ada_table_ts_pairs = BasicAVTransformerBlock._get_av_ca_ada_table_ts_pairs + + +class LTX23TransformerBlock(nn.Module): + """Dual-stream LTX-2.3 block with sigma-dependent text cross-attention.""" + + def __init__( + self, + idx: int, + video: TransformerConfig | None = None, + audio: TransformerConfig | None = None, + rope_type=None, + norm_eps: float = 1e-6, + config=None, + stage2_ulysses_group=None, + stage2_sharder=None, + ): + super().__init__() + if stage2_ulysses_group is not None or stage2_sharder is not None: + raise NotImplementedError("Two-stage topology is not supported for LTX-2.3.") + self.idx = idx + self.norm_eps = norm_eps + vgm = config.visual_gen_mapping if config is not None else None + self._sharder = SequenceSharder.from_vgm(vgm) + self._fuse_adaln = (video is None or is_fused_adaln_supported_dim(video.dim)) and ( + audio is None or is_fused_adaln_supported_dim(audio.dim) + ) + dense_config = config + if config is not None and config.attention.quant_attention_config is not None: + attention = config.attention.model_copy(update={"quant_attention_config": None}) + dense_config = config.model_copy(update={"attention": attention}) + + def attn( + name, query_dim, context_dim, kv, gated, model_config=dense_config, *, parallel=False + ): + return LTX2Attention( + query_dim=query_dim, + context_dim=context_dim, + heads=kv.heads, + dim_head=kv.d_head, + rope_type=rope_type, + norm_eps=norm_eps, + apply_gated_attention=gated, + config=model_config, + layer_idx=idx, + module_name=f"transformer_blocks.{idx}.{name}", + enable_sequence_parallel=parallel, + ) + + if video is not None: + gated_v = video.apply_gated_attention + # Quantized attention applies only to video self-attention; audio + # head dimensions and cross-attention require the dense backend. + self.attn1 = attn("attn1", video.dim, None, video, gated_v, config, parallel=True) + self.attn2 = attn("attn2", video.dim, video.context_dim, video, gated_v) + self.ff = BasicAVTransformerBlock._make_mlp(video, config, idx) + self.scale_shift_table = nn.Parameter(torch.empty(9, video.dim)) + self.prompt_scale_shift_table = nn.Parameter(torch.empty(2, video.dim)) + + if audio is not None: + gated_a = audio.apply_gated_attention + self.audio_attn1 = attn("audio_attn1", audio.dim, None, audio, gated_a) + self.audio_attn2 = attn("audio_attn2", audio.dim, audio.context_dim, audio, gated_a) + self.audio_ff = BasicAVTransformerBlock._make_mlp(audio, config, idx) + self.audio_scale_shift_table = nn.Parameter(torch.empty(9, audio.dim)) + self.audio_prompt_scale_shift_table = nn.Parameter(torch.empty(2, audio.dim)) + + if audio is not None and video is not None: + self.audio_to_video_attn = attn( + "audio_to_video_attn", video.dim, audio.dim, audio, video.apply_gated_attention + ) + self.video_to_audio_attn = attn( + "video_to_audio_attn", + audio.dim, + video.dim, + audio, + audio.apply_gated_attention, + parallel=True, + ) + self.scale_shift_table_a2v_ca_audio = nn.Parameter(torch.empty(5, audio.dim)) + self.scale_shift_table_a2v_ca_video = nn.Parameter(torch.empty(5, video.dim)) + + def _text_cross_attention( + self, + attn_input: torch.Tensor, + context: torch.Tensor, + attn: LTX2Attention, + prompt_scale_shift_table: torch.Tensor, + timestep: torch.Tensor, + prompt_timestep: torch.Tensor, + ) -> torch.Tensor: + """Query is pre-modulated by the caller; K/V are shifted from prompt_timestep.""" + batch_size = attn_input.shape[0] + shift_kv, scale_kv = ( + prompt_scale_shift_table[None, None].to( + device=attn_input.device, dtype=attn_input.dtype + ) + + prompt_timestep.reshape(batch_size, prompt_timestep.shape[1], 2, -1) + ).unbind(dim=2) + enc = apply_shift_scale(context, scale_kv, shift_kv) + k, v = attn.project_kv(enc) + return attn(attn_input, context=enc, pre_projected_kv=(k, v), timestep=timestep) + + def _self_attn_and_text_ca( + self, + x, + stream, + attn1, + attn2, + scale_shift_table, + prompt_table, + prompt_timestep, + ): + shift_msa, scale_msa, gate_msa = _get_ada_table_ts_pairs( + scale_shift_table, x.shape[0], stream.timesteps, slice(0, 3) + ) + shift_ca, scale_ca, gate_ca = _get_ada_table_ts_pairs( + scale_shift_table, x.shape[0], stream.timesteps, slice(6, 9) + ) + norm_x = apply_fused_rmsnorm_shift_scale( + x, + scale_msa[0], + scale_msa[1], + shift_msa[0], + shift_msa[1], + self.norm_eps, + self._fuse_adaln, + ) + msa = attn1( + norm_x, + pe=stream.positional_embeddings, + timestep=stream.timesteps, + key_padding_mask=stream.audio_padding_mask, + ) + x, q_input = apply_fused_gate_resid_rmsnorm_shift_scale( + x, + msa, + gate_msa[0], + gate_msa[1], + scale_ca[0], + scale_ca[1], + shift_ca[0], + shift_ca[1], + self.norm_eps, + self._fuse_adaln, + ) + ca = self._text_cross_attention( + q_input, stream.context, attn2, prompt_table, stream.timesteps, prompt_timestep + ) + return apply_fused_gate_resid(x, ca, gate_ca[0], gate_ca[1], self._fuse_adaln) + + def _ffn(self, x, stream, ff, scale_shift_table): + shift, scale, gate = _get_ada_table_ts_pairs( + scale_shift_table, x.shape[0], stream.timesteps, slice(3, 6) + ) + scaled = apply_fused_rmsnorm_shift_scale( + x, + scale[0], + scale[1], + shift[0], + shift[1], + self.norm_eps, + self._fuse_adaln, + ) + return apply_fused_gate_resid(x, ff(scaled), gate[0], gate[1], self._fuse_adaln) + + def _apply_video_audio_cross_attention( + self, + query_x, + query_pre, + kv_pre, + query_stream, + kv_stream, + query_table, + kv_table, + attn, + is_a2v: bool, + ): + q_scale_a2v, q_shift_a2v, q_scale_v2a, q_shift_v2a, gate = _get_av_ca_ada_table_ts_pairs( + query_table, + query_x.shape[0], + query_stream.cross_scale_shift_timestep, + query_stream.cross_gate_timestep, + ) + kv_scale_a2v, kv_shift_a2v, kv_scale_v2a, kv_shift_v2a, _ = _get_av_ca_ada_table_ts_pairs( + kv_table, + kv_pre.shape[0], + kv_stream.cross_scale_shift_timestep, + kv_stream.cross_gate_timestep, + ) + if is_a2v: + scale_q, shift_q = q_scale_a2v, q_shift_a2v + scale_kv, shift_kv = kv_scale_a2v, kv_shift_a2v + else: + scale_q, shift_q = q_scale_v2a, q_shift_v2a + scale_kv, shift_kv = kv_scale_v2a, kv_shift_v2a + + q_norm = apply_fused_rmsnorm_shift_scale( + query_pre, + scale_q[0], + scale_q[1], + shift_q[0], + shift_q[1], + self.norm_eps, + self._fuse_adaln, + ) + kv_norm = apply_fused_rmsnorm_shift_scale( + kv_pre, + scale_kv[0], + scale_kv[1], + shift_kv[0], + shift_kv[1], + self.norm_eps, + self._fuse_adaln, + ) + k, v = attn.project_kv(kv_norm, pe=kv_stream.cross_positional_embeddings) + query_pe = query_stream.cross_positional_embeddings + query_timestep = query_stream.timesteps + if not is_a2v and self._sharder.is_active: + # Audio stays replicated, as in LTX-2's conditional audio path. + # Only the v2a query is sharded to match the video K/V Ulysses layout. + query_length = q_norm.shape[1] + q_norm = self._sharder.shard(q_norm, dim=1) + query_timestep = self._sharder.shard( + query_timestep, dim=1, expected_seq_len=query_length + ) + if query_pe is not None: + pe_dim = 0 if query_pe[0].ndim == 2 else 1 + query_pe = tuple(self._sharder.shard(pe, dim=pe_dim) for pe in query_pe) + if not attn.is_ulysses: + k = self._sharder.gather(k, dim=1) + v = self._sharder.gather(v, dim=1) + out = attn( + q_norm, + pre_projected_kv=(k, v), + pe=query_pe, + timestep=query_timestep, + key_padding_mask=kv_stream.audio_padding_mask if is_a2v else None, + ) + if not is_a2v and self._sharder.is_active: + out = self._sharder.gather(out, dim=1) + return apply_fused_gate_resid(query_x, out, gate[0], gate[1], self._fuse_adaln) + + def forward( + self, + video, + audio, + video_prompt_timestep: torch.Tensor | None = None, + audio_prompt_timestep: torch.Tensor | None = None, + ): + vx = video.x if video is not None else None + ax = audio.x if audio is not None else None + run_vx = video is not None and video.enabled and vx.numel() > 0 + run_ax = audio is not None and audio.enabled and ax.numel() > 0 + run_a2v = run_vx and audio is not None and ax is not None and ax.numel() > 0 + run_v2a = run_ax and video is not None and vx is not None and vx.numel() > 0 + + if run_vx: + vx = self._self_attn_and_text_ca( + vx, + video, + self.attn1, + self.attn2, + self.scale_shift_table, + self.prompt_scale_shift_table, + video_prompt_timestep, + ) + if run_ax: + ax = self._self_attn_and_text_ca( + ax, + audio, + self.audio_attn1, + self.audio_attn2, + self.audio_scale_shift_table, + self.audio_prompt_scale_shift_table, + audio_prompt_timestep, + ) + + if run_a2v or run_v2a: + vx_pre, ax_pre = vx, ax + if run_a2v: + vx = self._apply_video_audio_cross_attention( + vx, + vx_pre, + ax_pre, + video, + audio, + self.scale_shift_table_a2v_ca_video, + self.scale_shift_table_a2v_ca_audio, + self.audio_to_video_attn, + is_a2v=True, + ) + if run_v2a: + ax = self._apply_video_audio_cross_attention( + ax, + ax_pre, + vx_pre, + audio, + video, + self.scale_shift_table_a2v_ca_audio, + self.scale_shift_table_a2v_ca_video, + self.video_to_audio_attn, + is_a2v=False, + ) + + if run_vx: + vx = self._ffn(vx, video, self.ff, self.scale_shift_table) + if run_ax: + ax = self._ffn(ax, audio, self.audio_ff, self.audio_scale_shift_table) + + return ( + replace(video, x=vx) if video is not None else None, + replace(audio, x=ax) if audio is not None else None, + ) + + +class LTX23Model(LTXModel): + """LTX-2.3 transformer: 9-slot AdaLN, prompt AdaLN, Identity caption projection.""" + + _transformer_block_class = LTX23TransformerBlock + _cache_text_kv = False + + @property + def _audio_shard_mode(self) -> AudioShardMode: + """Replicate the short audio stream and shard its v2a attention queries.""" + return AudioShardMode.CONDITIONAL if self._active_sharder.is_active else AudioShardMode.NONE + + def _init_stream(self, dim, in_channels, out_channels, norm_eps, prefix=""): + setattr(self, f"{prefix}patchify_proj", self._make_linear(in_channels, dim)) + setattr( + self, + f"{prefix}adaln_single", + AdaLayerNormSingle(dim, embedding_coefficient=9, make_linear=self._make_linear), + ) + setattr( + self, + f"{prefix}prompt_adaln_single", + AdaLayerNormSingle(dim, embedding_coefficient=2, make_linear=self._make_linear), + ) + setattr(self, f"{prefix}caption_projection", nn.Identity()) + setattr(self, f"{prefix}scale_shift_table", nn.Parameter(torch.empty(2, dim))) + setattr( + self, + f"{prefix}norm_out", + nn.LayerNorm(dim, elementwise_affine=False, eps=norm_eps), + ) + setattr(self, f"{prefix}proj_out", self._make_linear(dim, out_channels)) + + def _init_video(self, in_channels, out_channels, caption_channels, norm_eps): + self._init_stream(self.inner_dim, in_channels, out_channels, norm_eps) + + def _init_audio(self, in_channels, out_channels, caption_channels, norm_eps): + self._init_stream( + self.audio_inner_dim, in_channels, out_channels, norm_eps, prefix="audio_" + ) + + def _uses_cache_dit(self) -> bool: + # The LTX-2 Cache-DiT wrapper does not accept prompt timestep inputs. + return False + + def _compute_prompt_timestep(self, adaln, sigma, batch_size, dtype): + """prompt_timestep from the global sigma (drives context K/V modulation).""" + sigma_scaled = sigma * self.timestep_scale_multiplier + prompt_ts, _ = adaln(sigma_scaled.flatten(), hidden_dtype=dtype) + return prompt_ts.view(batch_size, -1, prompt_ts.shape[-1]) + + def _prepare_stream(self, modality, text_cache, is_audio: bool): + if modality is None: + return None, None + if is_audio: + preprocessor = self.audio_args_preprocessor + context, mask = text_cache.audio_context, text_cache.audio_mask + pe, cross_pe = text_cache.audio_pe, text_cache.audio_cross_pe + adaln = self.audio_prompt_adaln_single + else: + preprocessor = self.video_args_preprocessor + context, mask = text_cache.video_context, text_cache.video_mask + pe, cross_pe = text_cache.video_pe, text_cache.video_cross_pe + adaln = self.prompt_adaln_single + args = preprocessor.prepare(modality, context, mask, pe, cross_pe) + cross_sigma = modality.cross_modality_sigma + if cross_sigma is not None: + batch_size = args.x.shape[0] + if cross_sigma.ndim > 1 or cross_sigma.numel() != batch_size: + raise ValueError( + "cross modality sigma must be scalar per batch: " + f"got shape {tuple(cross_sigma.shape)} for batch size {batch_size}" + ) + gate, _ = preprocessor.cross_gate_adaln( + (cross_sigma * preprocessor.av_ca_timestep_scale_multiplier).flatten(), + hidden_dtype=modality.latent.dtype, + ) + gate = gate.view(batch_size, -1, gate.shape[-1]) + if gate.shape[1] != args.cross_scale_shift_timestep.shape[1]: + gate = gate.expand(-1, args.cross_scale_shift_timestep.shape[1], -1).contiguous() + args = replace(args, cross_gate_timestep=gate) + prompt_ts = self._compute_prompt_timestep( + adaln, modality.sigma, modality.latent.shape[0], modality.latent.dtype + ) + return args, prompt_ts + + def _stream_output(self, args, is_audio: bool): + if args is None: + return None + if is_audio: + table, norm, proj = ( + self.audio_scale_shift_table, + self.audio_norm_out, + self.audio_proj_out, + ) + else: + table, norm, proj = self.scale_shift_table, self.norm_out, self.proj_out + return self._process_output(table, norm, proj, args.x, args.embedded_timestep) + + def forward( + self, + video: LTX23Modality | None, + audio: LTX23Modality | None, + *, + text_cache: TextCache, + timestep: torch.Tensor | None = None, + step_index=None, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """Run the complete modality batch without splitting CFG branches.""" + audio_length = None + audio_padding_mask = None + if audio is not None and self._audio_pad > 0: + audio_length = audio.latent.shape[1] + audio = self._pad_modality_audio(audio, self._audio_pad) + audio_padding_mask = torch.arange(audio.latent.shape[1], device=audio.latent.device) + audio_padding_mask = (audio_padding_mask < audio_length).expand( + audio.latent.shape[0], -1 + ) + video_args, video_prompt_ts = self._prepare_stream(video, text_cache, is_audio=False) + audio_args, audio_prompt_ts = self._prepare_stream(audio, text_cache, is_audio=True) + if video_args is not None and self._active_sharder.is_active: + video_args = self._shard_transformer_args(video_args) + if audio_args is not None and audio_padding_mask is not None: + audio_args = replace(audio_args, audio_padding_mask=audio_padding_mask) + + for block in self.transformer_blocks: + video_args, audio_args = block( + video_args, + audio_args, + video_prompt_timestep=video_prompt_ts, + audio_prompt_timestep=audio_prompt_ts, + ) + + video_output = self._stream_output(video_args, is_audio=False) + audio_output = self._stream_output(audio_args, is_audio=True) + if video_output is not None and self._active_sharder.is_active: + video_output = self._gather_sequence(video_output) + if audio_output is not None and audio_length is not None: + audio_output = audio_output[:, :audio_length] + return video_output, audio_output diff --git a/tensorrt_llm/_torch/visual_gen/pipeline_registry.py b/tensorrt_llm/_torch/visual_gen/pipeline_registry.py index 7731e1186c16..3743b3c9386e 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline_registry.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline_registry.py @@ -309,10 +309,15 @@ def _detect_from_single_safetensors(checkpoint_dir: str) -> "str | None": return None if "transformer" in config and ("vae" in config or "audio_vae" in config): + # Imported here because pipeline_ltx23 imports register_pipeline from + # this module at import time. + from .models.ltx23.pipeline_ltx23 import detect_native_ltx_pipeline + + detected = detect_native_ltx_pipeline(config) logger.info( - "AutoPipeline: Detected LTX-2 native checkpoint " + f"AutoPipeline: Detected {detected} native checkpoint " f"(safetensors metadata) at {checkpoint_dir}" ) - return "LTX2Pipeline" + return detected return None diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/ltx23_lpips_golden_video.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/ltx23_lpips_golden_video.json new file mode 100644 index 000000000000..6977fa428324 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/ltx23_lpips_golden_video.json @@ -0,0 +1,33 @@ +{ + "video": "ltx23_generation_lpips_golden_video.mp4", + "audio": "ltx23_generation_lpips_golden_audio.wav", + "model": "LTX-2.3 22B dev BF16", + "source": "Official Lightricks/LTX-2 upstream", + "upstream_sha": "be6d6d0a69d11a0a399c462ed92606692e5f8697", + "upstream_pipeline": "ltx_pipelines.ti2vid_one_stage.TI2VidOneStagePipeline", + "checkpoint": "Lightricks/LTX-2.3/ltx-2.3-22b-dev.safetensors", + "text_encoder": "google/gemma-3-12b-it-qat-q4_0-unquantized", + "prompt": "A woman with long brown hair and light skin smiles at the camera while standing in a sunlit park, her hair gently blowing in the breeze as she tilts her head slightly to the side.", + "negative_prompt": "worst quality, inconsistent motion, blurry, jittery, distorted", + "height": 512, + "width": 768, + "num_frames": 49, + "num_inference_steps": 8, + "guidance_scale": 5.0, + "seed": 42, + "frame_rate": 24.0, + "tiling": {"spatial_size": 512, "spatial_overlap": 64, "temporal_size": 64, "temporal_overlap": 24}, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "audio_metric": "multi_resolution_log_stft_l1", + "audio_log_stft_threshold": 0.10, + "torch_version": "2.14.0a0+b2c75dd062.nv26.09", + "cuda_version": "13.4", + "ltx_core_version": "1.4.2", + "ltx_pipelines_version": "1.4.2", + "gpu": "NVIDIA B200", + "video_sha256": "ce1f56c639a9b686587350aaa969d926aa5b07c2feb136775d73fe4a55b98ba6", + "audio_sha256": "582577c83908781d14991664ce2e5f14a89c49c0dcd54dbde405ab3fbdfe0a75" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/ltx23_retake_lpips_golden_video.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/ltx23_retake_lpips_golden_video.json new file mode 100644 index 000000000000..64588a0aee91 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/ltx23_retake_lpips_golden_video.json @@ -0,0 +1,34 @@ +{ + "video": "ltx23_retake_lpips_golden_video.mp4", + "video_sha256": "eeca4dcbe59a6f1ef320b19a416886974b208863c26657dcb8a941edd8474e65", + "audio": "ltx23_retake_lpips_golden_audio.wav", + "audio_sha256": "34354fbfe52fb556f41565c1f926f64e22d579d052b3777e03e88ca3f5d324b3", + "input_video": "ltx23_retake_input.mp4", + "input_video_sha256": "a07d5ac49447c8306377683f3b939ddce0a98859b7a0e5c9a23e8aefd73d58fd", + "source": "https://github.com/Lightricks/LTX-2", + "upstream_sha": "be6d6d0a69d11a0a399c462ed92606692e5f8697", + "upstream_pipeline": "ltx_pipelines.retake.RetakePipeline", + "checkpoint": "Lightricks/LTX-2.3/ltx-2.3-22b-distilled-1.1.safetensors", + "text_encoder": "google/gemma-3-12b-it-qat-q4_0-unquantized", + "tiling": {"spatial_size": 768, "spatial_overlap": 64, "temporal_size": 80, "temporal_overlap": 24}, + "regenerate_video": true, + "regenerate_audio": true, + "prompt": "the duck flaps hard and flies right out of the frame", + "start_time_seconds": 2.0, + "end_time_seconds": 5.0, + "height": 544, + "width": 960, + "num_frames": 121, + "num_inference_steps": 8, + "seed": 42, + "frame_rate": 24.0, + "lpips_frame_start": 48, + "lpips_frame_stop": 120, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "audio_sample_rate": 48000, + "audio_samples": 240480, + "audio_channels": 1, + "audio_metric": "multi_resolution_log_stft_l1", + "audio_log_stft_threshold": 0.10 +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media_ltx23.zip b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media_ltx23.zip new file mode 100644 index 000000000000..af84b1f25e71 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media_ltx23.zip @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:af4782179a3e96ba08808b6e5cdd374cd0c2e9378bab198cdf04dc930a207249 +size 5243408 diff --git a/tests/integration/defs/examples/visual_gen/test_minimax_h3_e2e.py b/tests/integration/defs/examples/visual_gen/test_minimax_h3_e2e.py index 3ac7ee52bf9b..32762e5a637b 100644 --- a/tests/integration/defs/examples/visual_gen/test_minimax_h3_e2e.py +++ b/tests/integration/defs/examples/visual_gen/test_minimax_h3_e2e.py @@ -27,11 +27,11 @@ import lpips as lpips_package import pytest import torch -import torch.nn.functional as F from defs.examples.visual_gen.visual_gen_test_utils import ( _cleanup_cuda, _disable_inductor_compile_worker_quiesce, _llm_models_root, + _multi_resolution_log_stft_distance, ) from diffusers import ModularPipeline from PIL import Image @@ -130,39 +130,6 @@ def _minimax_h3_checkpoint_path() -> str: return path -def _multi_resolution_log_stft_distance( - actual: torch.Tensor, - expected: torch.Tensor, -) -> float: - if actual.shape != expected.shape: - raise AssertionError( - f"Audio shape mismatch: generated {tuple(actual.shape)}, " - f"golden {tuple(expected.shape)}." - ) - actual = actual.float().reshape(-1, actual.shape[-1]) - expected = expected.float().reshape(-1, expected.shape[-1]) - distances: list[torch.Tensor] = [] - for n_fft in (512, 1024, 2048): - hop_length = n_fft // 4 - window = torch.hann_window(n_fft) - actual_stft = torch.stft( - actual, - n_fft=n_fft, - hop_length=hop_length, - window=window, - return_complex=True, - ).abs() - expected_stft = torch.stft( - expected, - n_fft=n_fft, - hop_length=hop_length, - window=window, - return_complex=True, - ).abs() - distances.append(F.l1_loss(torch.log1p(actual_stft), torch.log1p(expected_stft))) - return float(torch.stack(distances).mean()) - - def _generate_diffusers_reference( checkpoint_path: str, keyframes: list[Image.Image], diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_ltx23.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_ltx23.py new file mode 100644 index 000000000000..c608f834f3ec --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_ltx23.py @@ -0,0 +1,458 @@ +# SPDX-FileCopyrightText: Copyright (c) 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. +"""End-to-end smoke and BF16 golden-media tests for LTX-2.3. + +Only BF16 is compared against video LPIPS and audio log-STFT goldens. +Generation and retake FP8/NVFP4 run smoke tests without golden comparisons. Override paths with +LTX23_MODEL_PATH (generation checkpoint), LTX23_RETAKE_CHECKPOINT, +LTX23_TEXT_ENCODER_PATH, and LTX23_GOLDEN_DIR. +""" + +import os +import subprocess +import sys + +import librosa +import numpy as np +import pytest +import soundfile as sf +import torch +from defs.common import venv_check_call +from defs.examples.visual_gen.visual_gen_test_utils import ( + REPO_ROOT, + _assert_lpips_below_threshold, + _cleanup_cuda, + _golden_media_path, + _lpips_deterministic_algorithms, + _lpips_model_path, + _multi_resolution_log_stft_distance, + _preserve_lpips_candidate_on_failure, + _require_exists, + _run_lpips_eval, + _run_reusable_video_lpips_eval, + _save_lpips_video_mp4, + _visual_gen_output_path, +) + +LTX23_MODEL_PATH = os.environ.get( + "LTX23_MODEL_PATH", _lpips_model_path("LTX-2.3", "ltx-2.3-22b-dev.safetensors") +) +if os.path.isdir(LTX23_MODEL_PATH): + LTX23_MODEL_PATH = os.path.join(LTX23_MODEL_PATH, "ltx-2.3-22b-dev.safetensors") +LTX23_TEXT_ENCODER_PATH = os.environ.get( + "LTX23_TEXT_ENCODER_PATH", + _lpips_model_path("gemma-3-12b-it-qat-q4_0-unquantized"), +) + +LTX23_GOLDEN_DIR = os.environ.get("LTX23_GOLDEN_DIR") + +LTX23_LPIPS_PROMPT = ( + "A woman with long brown hair and light skin smiles at the camera while standing in a " + "sunlit park, her hair gently blowing in the breeze as she tilts her head slightly to the side." +) +LTX23_LPIPS_NEGATIVE_PROMPT = "worst quality, inconsistent motion, blurry, jittery, distorted" +LTX23_LPIPS_HEIGHT = 512 +LTX23_LPIPS_WIDTH = 768 +LTX23_LPIPS_NUM_FRAMES = 49 +LTX23_LPIPS_NUM_INFERENCE_STEPS = 8 +LTX23_LPIPS_GUIDANCE_SCALE = 5.0 +LTX23_LPIPS_FRAME_RATE = 24.0 +LTX23_LPIPS_SEED = 42 + +LTX23_LPIPS_THRESHOLD = 0.05 +LTX23_AUDIO_LOG_STFT_THRESHOLD = 0.10 + +LTX23_RETAKE_EXPECTED_FRAMES = 121 +LTX23_RETAKE_EXPECTED_HEIGHT = 544 +LTX23_RETAKE_EXPECTED_WIDTH = 960 +LTX23_RETAKE_EXPECTED_FPS = 24.0 +LTX23_RETAKE_EXPECTED_AUDIO_RATE = 48000 +LTX23_RETAKE_AUDIO_SAMPLES = 240480 +LTX23_RETAKE_CHECKPOINT_ENV = "LTX23_RETAKE_CHECKPOINT" +LTX23_RETAKE_LPIPS_FRAME_START = 48 +LTX23_RETAKE_LPIPS_FRAME_STOP = 120 +LTX23_RETAKE_LPIPS_THRESHOLD = 0.05 +LTX23_RETAKE_PROMPT = "the duck flaps hard and flies right out of the frame" + + +def _ltx23_golden_media_path(tmp_path, media_name: str, label: str) -> str: + if LTX23_GOLDEN_DIR is not None: + path = os.path.join(LTX23_GOLDEN_DIR, media_name) + _require_exists(path, label) + return path + return str(_golden_media_path(tmp_path, media_name, label)) + + +def _save_wav(audio, path, sample_rate): + tensor = torch.as_tensor(audio, dtype=torch.float32).detach().cpu() + mono = librosa.to_mono(tensor.reshape(-1, tensor.shape[-1]).numpy()) + sf.write(str(path), np.clip(mono, -1.0, 1.0), int(sample_rate)) + assert os.path.isfile(path), f"Failed to write WAV {path}" + + +def _audio_log_stft_distance(generated_path: str, reference_path: str) -> float: + """Compare WAV channels with the shared multi-resolution log-STFT metric.""" + generated, gen_rate = sf.read(str(generated_path), dtype="float32", always_2d=True) + reference, ref_rate = sf.read(str(reference_path), dtype="float32", always_2d=True) + if gen_rate != ref_rate: + raise ValueError(f"Sample-rate mismatch: generated {gen_rate} vs golden {ref_rate}") + if generated.shape[0] == 0 or reference.shape[0] == 0: + raise ValueError("Empty audio: a clip decoded to zero samples") + if not np.isfinite(generated).all() or not np.isfinite(reference).all(): + raise ValueError("Non-finite audio samples") + distance = _multi_resolution_log_stft_distance( + torch.from_numpy(generated.T), torch.from_numpy(reference.T) + ) + print(f"[E2E audio log-STFT] distance: {distance:.6f} (samples: {generated.shape[0]})") + return distance + + +def _save_ltx23_retake_audio( + video_path: str | os.PathLike[str], + reference_path: str | os.PathLike[str], + audio_path: str | os.PathLike[str], +) -> None: + """Remove only the fixed Duck case's video/AAC tail padding before WAV comparison.""" + import av + + from tensorrt_llm._torch.visual_gen.models.ltx23.media_io import decode_audio_from_file + + reference = sf.info(str(reference_path)) + assert reference.samplerate == LTX23_RETAKE_EXPECTED_AUDIO_RATE + assert reference.channels == 1 + assert reference.frames == LTX23_RETAKE_AUDIO_SAMPLES + with av.open(str(video_path)) as container: + assert len(container.streams.audio) == 1, "LTX-2.3 retake output must have one audio stream" + stream = container.streams.audio[0] + assert stream.codec_context.name == "aac" and stream.codec_context.profile == "LC", ( + "LTX-2.3 retake tail-padding allowance requires AAC-LC" + ) + assert stream.start_time == 0, "LTX-2.3 retake audio must start at the video time origin" + audio = decode_audio_from_file(str(video_path)) + assert audio is not None, "LTX-2.3 retake output has no decoded audio" + assert audio.sample_rate == reference.samplerate, "LTX-2.3 retake audio sample-rate mismatch" + assert audio.samples.ndim == 3 and audio.samples.shape[:2] == (1, 2), ( + f"LTX-2.3 retake expected (1, 2, samples) audio, got {tuple(audio.samples.shape)}" + ) + assert torch.isfinite(audio.samples).all(), "LTX-2.3 retake audio has non-finite samples" + video_samples = round( + LTX23_RETAKE_EXPECTED_FRAMES * reference.samplerate / LTX23_RETAKE_EXPECTED_FPS + ) + # 1520 samples pad this video's duration; AAC-LC can add at most 1023 tail samples. + # Keep every valid sample from time zero; this is not signal alignment or resampling. + max_samples = video_samples + 1023 + actual_samples = audio.samples.shape[-1] + assert reference.frames <= actual_samples <= max_samples, ( + f"LTX-2.3 retake audio length {actual_samples} outside " + f"[{reference.frames}, {max_samples}] valid samples plus known mux/AAC tail padding" + ) + print( + f"[Retake audio] decoded={actual_samples}, valid={reference.frames}, " + f"removed_tail={actual_samples - reference.frames}; " + f"mux_padding={video_samples - reference.frames}, AAC-LC_tail_max=1023" + ) + _save_wav(audio.samples[..., : reference.frames], audio_path, audio.sample_rate) + + +def _generate_ltx23_av(video_out_path, audio_out_path): + """Render the unquantized BF16 golden-test candidate into an MP4 plus a WAV.""" + from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader + from tensorrt_llm.visual_gen.args import CompilationConfig, CudaGraphConfig, VisualGenArgs + + _require_exists(LTX23_MODEL_PATH, "LTX-2.3 dev checkpoint") + _require_exists( + LTX23_TEXT_ENCODER_PATH, "LTX-2.3 text encoder (Gemma QAT-unquantized)", is_dir=True + ) + + with _lpips_deterministic_algorithms(): + args = VisualGenArgs( + model=LTX23_MODEL_PATH, + pipeline_config={"text_encoder_path": LTX23_TEXT_ENCODER_PATH}, + cuda_graph_config=CudaGraphConfig(enable=True), + compilation_config=CompilationConfig( + resolutions=[(LTX23_LPIPS_HEIGHT, LTX23_LPIPS_WIDTH)], + num_frames=[LTX23_LPIPS_NUM_FRAMES], + ), + ) + pipeline = PipelineLoader(args).load() + try: + with torch.no_grad(): + out = pipeline.forward( + prompt=LTX23_LPIPS_PROMPT, + negative_prompt=LTX23_LPIPS_NEGATIVE_PROMPT, + seed=LTX23_LPIPS_SEED, + height=LTX23_LPIPS_HEIGHT, + width=LTX23_LPIPS_WIDTH, + num_frames=LTX23_LPIPS_NUM_FRAMES, + frame_rate=LTX23_LPIPS_FRAME_RATE, + num_inference_steps=LTX23_LPIPS_NUM_INFERENCE_STEPS, + guidance_scale=LTX23_LPIPS_GUIDANCE_SCALE, + ) + video = out.video.detach().cpu() + audio = out.audio.detach().cpu() if out.audio is not None else None + audio_rate = out.audio_sample_rate + finally: + del pipeline + _cleanup_cuda() + + _save_lpips_video_mp4(video, video_out_path, frame_rate=LTX23_LPIPS_FRAME_RATE) + assert audio is not None, "LTX-2.3 render returned no audio" + _save_wav(audio, audio_out_path, audio_rate) + + +@pytest.fixture(scope="session") +def ltx23_av_candidate(tmp_path_factory): + """Render BF16 once per session; return the candidate (video, audio) paths.""" + out_dir = str(tmp_path_factory.mktemp("ltx23_av")) + video_path = os.path.join(out_dir, "ltx23_candidate_video.mp4") + audio_path = os.path.join(out_dir, "ltx23_candidate_audio.wav") + _generate_ltx23_av(video_path, audio_path) + return video_path, audio_path + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_ltx23_lpips_against_golden(tmp_path, ltx23_av_candidate): + """Check BF16 video LPIPS and audio log-STFT against the frozen goldens.""" + video_golden_path = _ltx23_golden_media_path( + tmp_path, "ltx23_generation_lpips_golden_video.mp4", "LTX-2.3 LPIPS golden video" + ) + audio_golden_path = _ltx23_golden_media_path( + tmp_path, "ltx23_generation_lpips_golden_audio.wav", "LTX-2.3 golden audio" + ) + video_path, audio_path = ltx23_av_candidate + score = _run_lpips_eval( + tmp_path, "ltx23", "video", LTX23_LPIPS_PROMPT, video_golden_path, video_path + ) + distance = _audio_log_stft_distance(audio_path, audio_golden_path) + _assert_lpips_below_threshold(score, LTX23_LPIPS_THRESHOLD) + assert distance < LTX23_AUDIO_LOG_STFT_THRESHOLD, ( + f"LTX-2.3 audio log-STFT too high: {distance:.6f} " + f"(expected < {LTX23_AUDIO_LOG_STFT_THRESHOLD:.6f})" + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize( + "config_name", + ["ltx23-t2v-fp8-1gpu.yaml", "ltx23-t2v-fp4-1gpu.yaml", "ltx23-t2v-bf16-2gpu.yaml"], + ids=["fp8", "nvfp4", "bf16-2gpu"], +) +def test_ltx23_example(tmp_path, config_name): + """Run the configured generation example and verify its video and audio.""" + if config_name.endswith("2gpu.yaml") and torch.cuda.device_count() < 2: + pytest.skip("Two GPUs are required for this parallel example") + _require_exists(LTX23_MODEL_PATH, "LTX-2.3 dev checkpoint") + _require_exists( + LTX23_TEXT_ENCODER_PATH, "LTX-2.3 text encoder (Gemma QAT-unquantized)", is_dir=True + ) + + examples_root = os.path.join(REPO_ROOT, "examples", "visual_gen") + output_path = os.path.join(str(tmp_path), "ltx23_output.mp4") + result = subprocess.run( + [ + sys.executable, + os.path.join(examples_root, "models", "ltx23.py"), + "--model", + LTX23_MODEL_PATH, + "--visual_gen_args", + os.path.join(examples_root, "configs", config_name), + "--text_encoder_path", + LTX23_TEXT_ENCODER_PATH, + "--output_path", + output_path, + ], + check=False, + timeout=1800, + ) + assert result.returncode == 0, "LTX-2.3 example script exited non-zero" + assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" + assert os.path.getsize(output_path) > 0, f"Example produced an empty output at {output_path}" + import av + + with av.open(output_path) as container: + video = container.streams.video[0] + assert video.width == LTX23_LPIPS_WIDTH + assert video.height == LTX23_LPIPS_HEIGHT + assert float(video.average_rate) == pytest.approx(LTX23_LPIPS_FRAME_RATE) + assert sum(1 for _ in container.decode(video=0)) == LTX23_LPIPS_NUM_FRAMES + assert len(container.streams.audio) == 1 + assert container.streams.audio[0].rate == 48000 + + +def _generate_ltx23_retake_video( + llm_root, + llm_venv, + tmp_path, + output_path, + config_filename="ltx23-retake-bf16-1gpu.yaml", + *, + regenerate_audio: bool = True, +): + """Run the native retake example without runtime LoRA and return its output.""" + source = _golden_media_path( + tmp_path, + "ltx23_retake_input.mp4", + "LTX-2.3 official Duck retake input", + ) + _require_exists( + LTX23_TEXT_ENCODER_PATH, "LTX-2.3 text encoder (Gemma QAT-unquantized)", is_dir=True + ) + checkpoint = os.environ.get(LTX23_RETAKE_CHECKPOINT_ENV) or os.path.join( + os.path.dirname(LTX23_MODEL_PATH), "ltx-2.3-22b-distilled-1.1.safetensors" + ) + _require_exists(checkpoint, "LTX-2.3 retake checkpoint") + example = os.path.join(llm_root, "examples", "visual_gen", "models", "ltx23_retake.py") + config = os.path.join(llm_root, "examples", "visual_gen", "configs", config_filename) + venv_check_call( + llm_venv, + [ + example, + "--model", + checkpoint, + "--visual_gen_args", + str(config), + "--source", + str(source), + "--start", + "2.0", + "--end", + "5.0", + "--prompt", + LTX23_RETAKE_PROMPT, + "--regenerate_audio" if regenerate_audio else "--no-regenerate_audio", + "--text_encoder_path", + LTX23_TEXT_ENCODER_PATH, + "--output_path", + str(output_path), + ], + ) + assert os.path.isfile(output_path), f"retake pipeline did not produce {output_path}" + assert os.path.getsize(output_path) > 0, f"retake pipeline produced an empty {output_path}" + return output_path + + +@pytest.fixture(scope="session") +def _ltx23_retake_deps(_visual_gen_deps, llm_venv): + llm_venv.run_cmd( + [ + "-m", + "pip", + "install", + "--no-deps", + "--index-url", + "https://download.pytorch.org/whl/cpu", + "torchaudio==2.11.0+cpu", + ] + ) + + +@pytest.fixture(scope="session") +def ltx23_retake_bf16_video_path(_ltx23_retake_deps, llm_root, llm_venv, tmp_path_factory): + output_path = _visual_gen_output_path(llm_venv, "ltx23_retake_duck_joint_bf16") + if os.path.isfile(output_path): + return output_path + work_dir = tmp_path_factory.mktemp("ltx23_retake") + return _generate_ltx23_retake_video(llm_root, llm_venv, work_dir, output_path) + + +def _assert_ltx23_retake_output(video_path): + """Check the retake output container, decoded frame count, and audio stream.""" + import av + + from tensorrt_llm._torch.visual_gen.models.ltx23.media_io import ( + decode_video_by_frame, + get_videostream_metadata, + ) + + metadata = get_videostream_metadata(str(video_path)) + assert metadata.frames == LTX23_RETAKE_EXPECTED_FRAMES + assert metadata.height == LTX23_RETAKE_EXPECTED_HEIGHT + assert metadata.width == LTX23_RETAKE_EXPECTED_WIDTH + assert metadata.fps == pytest.approx(LTX23_RETAKE_EXPECTED_FPS) + assert sum(1 for _ in decode_video_by_frame(str(video_path))) == LTX23_RETAKE_EXPECTED_FRAMES + + container = av.open(str(video_path)) + try: + audio_streams = list(container.streams.audio) + assert len(audio_streams) == 1 + assert audio_streams[0].rate == LTX23_RETAKE_EXPECTED_AUDIO_RATE + assert audio_streams[0].codec_context.channels == 2 + finally: + container.close() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize( + "case_id, config_filename", + [ + ("fp8", "ltx23-retake-fp8-1gpu.yaml"), + ("nvfp4", "ltx23-retake-fp4-1gpu.yaml"), + ("bf16-2gpu", "ltx23-retake-bf16-2gpu.yaml"), + ], + ids=["fp8", "nvfp4", "bf16-2gpu"], +) +def test_ltx23_retake_example( + tmp_path, _ltx23_retake_deps, llm_root, llm_venv, case_id, config_filename +): + """Run the configured retake example and validate media without a golden comparison.""" + if config_filename.endswith("2gpu.yaml") and torch.cuda.device_count() < 2: + pytest.skip("Two GPUs are required for this parallel example") + output_path = _visual_gen_output_path(llm_venv, f"ltx23_retake_{case_id}") + _generate_ltx23_retake_video( + llm_root, llm_venv, tmp_path, output_path, config_filename=config_filename + ) + _assert_ltx23_retake_output(output_path) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_ltx23_retake_lpips_against_golden( + request, + tmp_path, + ltx23_retake_bf16_video_path, + _visual_gen_lpips_scorer, +): + """Compare BF16 retake video LPIPS and audio log-STFT against the upstream goldens.""" + _assert_ltx23_retake_output(ltx23_retake_bf16_video_path) + golden_path = _golden_media_path( + tmp_path, + "ltx23_retake_lpips_golden_video.mp4", + "LTX-2.3 retake LPIPS golden video", + ) + score = _run_reusable_video_lpips_eval( + "ltx23_retake", + golden_path, + ltx23_retake_bf16_video_path, + _visual_gen_lpips_scorer, + frame_start=LTX23_RETAKE_LPIPS_FRAME_START, + frame_stop=LTX23_RETAKE_LPIPS_FRAME_STOP, + ) + audio_golden_path = _ltx23_golden_media_path( + tmp_path, "ltx23_retake_lpips_golden_audio.wav", "LTX-2.3 retake golden audio" + ) + audio_path = tmp_path / "ltx23_retake_candidate_audio.wav" + _save_ltx23_retake_audio(ltx23_retake_bf16_video_path, audio_golden_path, audio_path) + distance = _audio_log_stft_distance(audio_path, audio_golden_path) + _preserve_lpips_candidate_on_failure( + request, + score, + LTX23_RETAKE_LPIPS_THRESHOLD, + ltx23_retake_bf16_video_path, + "ltx23_retake_generated.mp4", + ) + _assert_lpips_below_threshold(score, LTX23_RETAKE_LPIPS_THRESHOLD) + assert distance < LTX23_AUDIO_LOG_STFT_THRESHOLD, ( + f"LTX-2.3 retake audio log-STFT too high: {distance:.6f} " + f"(expected < {LTX23_AUDIO_LOG_STFT_THRESHOLD:.6f})" + ) diff --git a/tests/integration/defs/examples/visual_gen/visual_gen_test_utils.py b/tests/integration/defs/examples/visual_gen/visual_gen_test_utils.py index 807eb5356d4f..844c050f08a6 100644 --- a/tests/integration/defs/examples/visual_gen/visual_gen_test_utils.py +++ b/tests/integration/defs/examples/visual_gen/visual_gen_test_utils.py @@ -32,6 +32,7 @@ import pytest import torch import torch._inductor.config as inductor_config +import torch.nn.functional as F from torch._inductor.async_compile import shutdown_compile_workers # ============================================================================= @@ -383,7 +384,14 @@ def _decode_video_to_lpips_batch(self, video_path): batch = batch.permute(0, 3, 1, 2).to(device=self.device, dtype=torch.float32) return batch / 127.5 - 1.0 - def score_video(self, reference_path, generated_path): + def score_video( + self, + reference_path, + generated_path, + *, + frame_start=0, + frame_stop=None, + ): reference = self._decode_video_to_lpips_batch(reference_path) generated = self._decode_video_to_lpips_batch(generated_path) if generated.shape[1:] != reference.shape[1:]: @@ -392,10 +400,16 @@ def score_video(self, reference_path, generated_path): f"{tuple(generated.shape[1:])} vs {tuple(reference.shape[1:])}." ) paired_frame_count = min(generated.shape[0], reference.shape[0]) + frame_stop = paired_frame_count if frame_stop is None else frame_stop + if frame_start < 0 or frame_start >= frame_stop or frame_stop > paired_frame_count: + raise ValueError( + f"LPIPS frame window [{frame_start}, {frame_stop}) is outside " + f"the {paired_frame_count} paired frames." + ) with torch.no_grad(): scores = self._get_model()( - generated[:paired_frame_count], - reference[:paired_frame_count], + generated[frame_start:frame_stop], + reference[frame_start:frame_stop], ).flatten() return float(scores.mean().item()) @@ -408,6 +422,40 @@ def close(self): torch.cuda.empty_cache() +def _multi_resolution_log_stft_distance( + actual: torch.Tensor, + expected: torch.Tensor, +) -> float: + """Mean log1p-magnitude STFT L1 over three FFT sizes; the result is not in dB.""" + if actual.shape != expected.shape: + raise AssertionError( + f"Audio shape mismatch: generated {tuple(actual.shape)}, " + f"golden {tuple(expected.shape)}." + ) + actual = actual.float().reshape(-1, actual.shape[-1]) + expected = expected.float().reshape(-1, expected.shape[-1]) + distances: list[torch.Tensor] = [] + for n_fft in (512, 1024, 2048): + hop_length = n_fft // 4 + window = torch.hann_window(n_fft) + actual_stft = torch.stft( + actual, + n_fft=n_fft, + hop_length=hop_length, + window=window, + return_complex=True, + ).abs() + expected_stft = torch.stft( + expected, + n_fft=n_fft, + hop_length=hop_length, + window=window, + return_complex=True, + ).abs() + distances.append(F.l1_loss(torch.log1p(actual_stft), torch.log1p(expected_stft))) + return float(torch.stack(distances).mean()) + + def _llm_models_root(): # Keep imports that load TRT-LLM bindings out of module initialization. # torch.multiprocessing.spawn workers import this utility before the @@ -445,6 +493,8 @@ def _golden_media_family(media_name): return "glm_image" if media_name.startswith("hunyuan_"): return "hunyuan" + if media_name.startswith("ltx23_"): + return "ltx23" if media_name.startswith("ltx2_"): return "ltx2" if media_name.startswith(("qwenimage_", "qwen_image_layered_")): @@ -723,10 +773,23 @@ def _run_reusable_image_lpips_eval(sample_id, reference_path, generated_path, sc return score -def _run_reusable_video_lpips_eval(sample_id, reference_path, generated_path, scorer): +def _run_reusable_video_lpips_eval( + sample_id, + reference_path, + generated_path, + scorer, + *, + frame_start=0, + frame_stop=None, +): try: with _lpips_deterministic_algorithms(): - score = scorer.score_video(reference_path, generated_path) + score = scorer.score_video( + reference_path, + generated_path, + frame_start=frame_start, + frame_stop=frame_stop, + ) except Exception as error: pytest.fail(f"LPIPS video eval failed for {sample_id}: {error}") print(f"\n[E2E {sample_id} LPIPS] score: {score:.6f}") diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index f6f854d4296c..ca20885dee65 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -412,6 +412,10 @@ l0_b200: - examples/visual_gen/test_visual_gen_flux.py::test_flux2_example - examples/visual_gen/test_visual_gen_flux.py::test_flux2_reference_image_example - examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_example + - examples/visual_gen/test_visual_gen_ltx23.py::test_ltx23_example[fp8] + - examples/visual_gen/test_visual_gen_ltx23.py::test_ltx23_example[nvfp4] + - examples/visual_gen/test_visual_gen_ltx23.py::test_ltx23_retake_example[fp8] + - examples/visual_gen/test_visual_gen_ltx23.py::test_ltx23_retake_example[nvfp4] - examples/visual_gen/test_visual_gen_wan.py::test_wan_i2v_example - examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_example[nano-1gpu] - examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_example[super-1gpu] @@ -555,6 +559,8 @@ l0_b200: - examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden[wan22-nvfp4] - examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden[wan22-cuda-graph] - examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_lpips_against_golden + - examples/visual_gen/test_visual_gen_ltx23.py::test_ltx23_lpips_against_golden + - examples/visual_gen/test_visual_gen_ltx23.py::test_ltx23_retake_lpips_against_golden - examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_cuda_graph_lpips_matches_eager - examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_cuda_graph_trtllm_backend - examples/visual_gen/test_visual_gen_wan.py::test_wan21_t2v_lpips_against_golden diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 1a7cee7dd224..6c3bc38df724 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -135,6 +135,9 @@ l0_cpu: - unittest/_torch/visual_gen/test_cosmos3_guardrails.py - unittest/evaluate/visual_gen/test_generation_evaluation_pipeline.py - unittest/_torch/visual_gen/test_ltx2_pipeline.py + - unittest/_torch/visual_gen/test_ltx23_registry.py + - unittest/_torch/visual_gen/test_ltx23_pipeline_config.py + - unittest/_torch/visual_gen/test_ltx23_retake_initialization.py - unittest/_torch/visual_gen/test_minimax_h3_pipeline.py - unittest/_torch/visual_gen/test_minimax_h3_tiled_vae.py - unittest/_torch/visual_gen/test_minimax_h3_scheduler_packing.py diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 8451f018ee21..55a2b01ff18d 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -33,6 +33,8 @@ l0_dgx_b200: - kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2DSv3Lite::test_mtp_chunked_draft_tokens - kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2DSv3Lite::test_mtp_eviction # ------------- VisualGen multi-GPU tests --------------- + - examples/visual_gen/test_visual_gen_ltx23.py::test_ltx23_example[bf16-2gpu] TIMEOUT (30) + - examples/visual_gen/test_visual_gen_ltx23.py::test_ltx23_retake_example[bf16-2gpu] TIMEOUT (30) - unittest/_torch/visual_gen/multi_gpu/test_ltx2_parallel_vae.py - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_ulysses_and_parallel_vae - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_ulysses_parallel_vae_lpips[2-2-VANILLA] TIMEOUT (30) diff --git a/tests/unittest/_torch/visual_gen/test_ltx23_pipeline_config.py b/tests/unittest/_torch/visual_gen/test_ltx23_pipeline_config.py new file mode 100644 index 000000000000..1284ef16325b --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_ltx23_pipeline_config.py @@ -0,0 +1,150 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU pipeline configuration tests for LTX-2.3 generation and retake.""" + +import json +from pathlib import Path + +import pytest +from safetensors.torch import save_file + +# Importing models applies the pipeline registration decorators. +from tensorrt_llm._torch.visual_gen import models # noqa: F401 +from tensorrt_llm._torch.visual_gen.config import DiffusionPipelineConfig +from tensorrt_llm._torch.visual_gen.models.ltx23.pipeline_ltx23 import LTX23Pipeline +from tensorrt_llm._torch.visual_gen.models.ltx23.pipeline_ltx23_retake import LTX23RetakePipeline +from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader +from tensorrt_llm._torch.visual_gen.pipeline_registry import AutoPipeline +from tensorrt_llm.visual_gen.args import VisualGenArgs + +pytestmark = pytest.mark.cpu_only + + +@pytest.fixture +def ltx23_checkpoint(tmp_path: Path) -> str: + """Write only the native metadata needed for detection and configuration.""" + checkpoint = tmp_path / "ltx23.safetensors" + native_config = { + "transformer": { + "cross_attention_adaln": True, + "caption_proj_before_connector": True, + "caption_projection_first_linear": False, + "caption_projection_second_linear": False, + "caption_proj_input_norm": False, + }, + "vae": {}, + } + save_file({}, str(checkpoint), metadata={"config": json.dumps(native_config)}) + return str(checkpoint) + + +def test_ltx23_pipeline_config_defaults(ltx23_checkpoint: str) -> None: + args = VisualGenArgs(model=ltx23_checkpoint) + resolved = PipelineLoader(args, device="cpu")._resolve_pipeline_config(ltx23_checkpoint) + + assert AutoPipeline._detect_from_checkpoint(ltx23_checkpoint) == "LTX23Pipeline" + assert resolved == { + "workflow": "generation", + "text_encoder_path": "google/gemma-3-12b-it-qat-q4_0-unquantized", + } + assert args.pipeline_config == {} + + +def test_ltx23_pipeline_config_user_overrides(ltx23_checkpoint: str) -> None: + overrides = {"workflow": "retake", "text_encoder_path": "/local/gemma"} + args = VisualGenArgs(model=ltx23_checkpoint, pipeline_config=overrides) + + assert ( + PipelineLoader(args, device="cpu")._resolve_pipeline_config(ltx23_checkpoint) == overrides + ) + + +def test_ltx23_pipeline_config_rejects_unknown_keys(ltx23_checkpoint: str) -> None: + args = VisualGenArgs(model=ltx23_checkpoint, pipeline_config={"unknown_option": True}) + with pytest.raises(ValueError, match="Unknown pipeline_config keys for LTX23Pipeline"): + PipelineLoader(args, device="cpu")._resolve_pipeline_config(ltx23_checkpoint) + + +@pytest.mark.parametrize( + "workflow, expected", [("generation", LTX23Pipeline), ("retake", LTX23RetakePipeline)] +) +def test_ltx23_config_preserves_resolved_options( + ltx23_checkpoint: str, workflow: str, expected: type +) -> None: + args = VisualGenArgs( + model=ltx23_checkpoint, + pipeline_config={"workflow": workflow}, + attention_config={"backend": "VANILLA"}, + ) + resolved = PipelineLoader(args, device="cpu")._resolve_pipeline_config(ltx23_checkpoint) + config = DiffusionPipelineConfig.from_pretrained( + ltx23_checkpoint, args=args, pipeline_config=resolved, device="cpu" + ) + + assert config.extra_attrs["workflow"] == workflow + assert config.extra_attrs["text_encoder_path"] == resolved["text_encoder_path"] + assert ( + config.extra_attrs["monolithic_safetensors_config"]["transformer"]["cross_attention_adaln"] + is True + ) + assert LTX23Pipeline.resolve_variant(config) is expected + + +def test_ltx23_config_rejects_unknown_workflow(ltx23_checkpoint: str) -> None: + args = VisualGenArgs(model=ltx23_checkpoint, pipeline_config={"workflow": "unknown"}) + with pytest.raises(ValueError, match="LTX-2.3 workflow must be"): + DiffusionPipelineConfig.from_pretrained(ltx23_checkpoint, args=args, device="cpu") + + +@pytest.mark.parametrize( + "recipe, workflow, quant_algo", + [ + ("ltx23-t2v-bf16-1gpu.yaml", "generation", None), + ("ltx23-t2v-fp8-1gpu.yaml", "generation", "FP8_BLOCK_SCALES"), + ("ltx23-t2v-fp4-1gpu.yaml", "generation", "NVFP4"), + ("ltx23-retake-bf16-1gpu.yaml", "retake", None), + ("ltx23-t2v-bf16-2gpu.yaml", "generation", None), + ("ltx23-retake-bf16-2gpu.yaml", "retake", None), + ("ltx23-retake-fp8-1gpu.yaml", "retake", "FP8_BLOCK_SCALES"), + ("ltx23-retake-fp4-1gpu.yaml", "retake", "NVFP4"), + ], +) +def test_ltx23_example_recipe_config( + ltx23_checkpoint: str, recipe: str, workflow: str, quant_algo: str | None +) -> None: + config_dir = Path(__file__).resolve().parents[4] / "examples" / "visual_gen" / "configs" + args = VisualGenArgs.from_yaml(config_dir / recipe, model=ltx23_checkpoint) + resolved = PipelineLoader(args, device="cpu")._resolve_pipeline_config(ltx23_checkpoint) + + assert resolved["workflow"] == workflow + if isinstance(args.quant_config, dict): + assert args.quant_config["quant_algo"] == quant_algo + assert args.quant_config["dynamic"] is True + else: + assert args.quant_config.quant_algo == quant_algo + assert args.torch_compile_config.enable is True + assert args.cuda_graph_config.enable is True + if recipe.endswith("2gpu.yaml"): + assert args.parallel_config.cfg_size == 1 + assert args.parallel_config.ulysses_size == 2 + assert args.parallel_config.parallel_vae_size == 2 + assert args.attention_config.backend == "VANILLA" + if quant_algo is not None: + assert args.attention_config.backend == "FLASHINFER" + expected_qk = "nvfp4" if quant_algo == "NVFP4" else "mxfp8" + assert args.attention_config.quant_attention_config.qk_dtype == expected_qk + assert args.attention_config.quant_attention_config.v_dtype == "fp8" + if workflow == "retake": + assert args.pipeline_config == {"workflow": "retake"} + config = DiffusionPipelineConfig.from_pretrained( + ltx23_checkpoint, args=args, pipeline_config=resolved, device="cpu" + ) + assert config.quant_config.quant_algo == quant_algo + assert config.dynamic_weight_quant is (quant_algo is not None) + assert config.extra_attrs["workflow"] == workflow + if workflow == "retake": + assert LTX23Pipeline.resolve_variant(config) is LTX23RetakePipeline + assert ( + config.attention.quant_attention_config == args.attention_config.quant_attention_config + ) diff --git a/tests/unittest/_torch/visual_gen/test_ltx23_registry.py b/tests/unittest/_torch/visual_gen/test_ltx23_registry.py new file mode 100644 index 000000000000..2c94432eb6f9 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_ltx23_registry.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU registration and workflow routing tests for LTX-2.3.""" + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +# Importing models applies the pipeline registration decorators. +from tensorrt_llm._torch.visual_gen import models # noqa: F401 +from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2 import LTX2Pipeline +from tensorrt_llm._torch.visual_gen.models.ltx23.pipeline_ltx23 import LTX23Pipeline +from tensorrt_llm._torch.visual_gen.models.ltx23.pipeline_ltx23_retake import LTX23RetakePipeline +from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline +from tensorrt_llm._torch.visual_gen.pipeline_registry import PIPELINE_REGISTRY, AutoPipeline + +pytestmark = pytest.mark.cpu_only + + +def test_ltx23_pipeline_is_registered() -> None: + entry = PIPELINE_REGISTRY["LTX23Pipeline"] + assert entry.pipeline_cls is LTX23Pipeline + assert entry.hf_ids == ["Lightricks/LTX-2.3"] + + +def test_ltx23_retake_is_registered_as_a_workflow_variant() -> None: + entry = PIPELINE_REGISTRY["LTX23RetakePipeline"] + assert entry.pipeline_cls is LTX23RetakePipeline + assert entry.hf_ids == [] + assert entry.defaults == {} + assert issubclass(LTX23RetakePipeline, LTX23Pipeline) + + +def test_ltx23_has_a_separate_pipeline_boundary() -> None: + assert issubclass(LTX23Pipeline, BasePipeline) + assert not issubclass(LTX23Pipeline, LTX2Pipeline) + + +@pytest.mark.parametrize("class_name", ["LTX23Pipeline", "LTX23RetakePipeline"]) +def test_auto_pipeline_detects_registered_ltx23_class_name(tmp_path: Path, class_name: str) -> None: + (tmp_path / "model_index.json").write_text( + json.dumps({"_class_name": class_name}), encoding="utf-8" + ) + assert AutoPipeline._detect_from_checkpoint(str(tmp_path)) == class_name + + +@pytest.mark.parametrize( + "extra_attrs, expected", + [ + ({}, LTX23Pipeline), + ({"workflow": "generation"}, LTX23Pipeline), + ({"workflow": "retake"}, LTX23RetakePipeline), + ], +) +def test_ltx23_routes_workflow_variants(extra_attrs: dict, expected: type) -> None: + assert LTX23Pipeline.resolve_variant(SimpleNamespace(extra_attrs=extra_attrs)) is expected + + +def test_ltx23_rejects_unknown_workflow() -> None: + with pytest.raises(ValueError, match="Unsupported LTX-2.3 workflow"): + LTX23Pipeline.resolve_variant(SimpleNamespace(extra_attrs={"workflow": "unknown"})) diff --git a/tests/unittest/_torch/visual_gen/test_ltx23_retake_initialization.py b/tests/unittest/_torch/visual_gen/test_ltx23_retake_initialization.py new file mode 100644 index 000000000000..ebbaf92a2f70 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_ltx23_retake_initialization.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU regression for frozen-video retake initialization.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from tensorrt_llm._torch.visual_gen.models.ltx2.ltx2_core.patchifier import ( + AudioPatchifier, + VideoLatentPatchifier, +) +from tensorrt_llm._torch.visual_gen.models.ltx2.ltx2_core.types import ( + AudioLatentShape, + VideoLatentShape, + VideoPixelShape, +) +from tensorrt_llm._torch.visual_gen.models.ltx23 import pipeline_ltx23_retake as retake_module +from tensorrt_llm._torch.visual_gen.models.ltx23.pipeline_ltx23_retake import LTX23RetakePipeline + +pytestmark = pytest.mark.cpu_only + + +def test_ltx23_retake_frozen_video_initialization(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep video clean before denoising without changing the audio RNG sequence.""" + dtype = torch.bfloat16 + pixel_shape = VideoPixelShape(batch=1, frames=17, height=32, width=32, fps=8.0) + video_shape = VideoLatentShape.from_pixel_shape(pixel_shape, latent_channels=2) + audio_shape = AudioLatentShape.from_video_pixel_shape(pixel_shape) + source_video = torch.full(video_shape.to_torch_shape(), 2.0, dtype=dtype) + source_audio = torch.full(audio_shape.to_torch_shape(), 3.0, dtype=dtype) + pipeline = LTX23RetakePipeline.__new__(LTX23RetakePipeline) + torch.nn.Module.__init__(pipeline) + pipeline._device = torch.device("cpu") + pipeline.pipeline_config = SimpleNamespace(torch_dtype=dtype) + pipeline.transformer_in_channels = video_shape.channels + pipeline.video_patchifier = VideoLatentPatchifier(patch_size=1) + pipeline.audio_patchifier = AudioPatchifier(patch_size=1) + pipeline.video_encoder = Mock() + pipeline.video_encoder.tiled_encode.return_value = source_video + pipeline._read_source_video = Mock() + pipeline._encode_source_audio_latent = Mock(return_value=source_audio) + pipeline.audio_decoder = Mock() + pipeline.vocoder = Mock() + pipeline.transformer = Mock() + context = torch.zeros(1, 1, 2, dtype=dtype) + pipeline._prepare_prompt_conditioning = Mock(return_value=(context, context, None)) + monkeypatch.setattr(retake_module, "get_videostream_metadata", lambda _: pixel_shape) + monkeypatch.setattr(retake_module, "decode_audio_from_file", lambda _: object()) + + class InitializationComplete(Exception): + """Stop at the denoising boundary without running or decoding the model.""" + + pipeline.denoise = Mock(side_effect=InitializationComplete) + req = SimpleNamespace(params=SimpleNamespace(num_inference_steps=8, seed=42, extra_params={})) + with pytest.raises(InitializationComplete): + pipeline._run_retake( + req, + "source.mp4", + 0.5, + 1.5, + "prompt", + Mock(), + regenerate_video=False, + regenerate_audio=True, + ) + + inputs = pipeline.denoise.call_args.kwargs + clean_video = pipeline.video_patchifier.patchify(source_video) + torch.testing.assert_close(inputs["latents"], clean_video, rtol=0, atol=0) + + audio_state = inputs["extra_streams"]["audio"][0] + pipeline._masked_transformer_step = Mock(return_value=(None, None)) + inputs["forward_fn"]( + inputs["latents"], + {"audio": audio_state}, + 0, + inputs["timesteps"][:1], + inputs["prompt_embeds"], + {}, + ) + conditioning = pipeline._masked_transformer_step.call_args.kwargs + assert conditioning["regenerate_video"] is False + assert conditioning["regenerate_audio"] is True + assert torch.count_nonzero(conditioning["denoise_mask"]) == 0 + audio_mask = conditioning["audio_denoise_mask"] + assert audio_mask.any() and not audio_mask.all() + + clean_audio = pipeline.audio_patchifier.patchify(source_audio) + generator = torch.Generator().manual_seed(req.params.seed) + torch.randn(clean_video.shape, generator=generator, dtype=dtype) + audio_noise = torch.randn(clean_audio.shape, generator=generator, dtype=dtype) + expected_audio = torch.lerp( + clean_audio.float(), audio_noise.float(), audio_mask.unsqueeze(-1) + ).to(dtype) + torch.testing.assert_close(audio_state, expected_audio, rtol=0, atol=0) + pipeline.transformer.assert_not_called()