[Feat] Add vae-decode-parallel for LTX-2.3 (#4277)

Signed-off-by: mglyn <1203789601@qq.com>

docs/user_guide/diffusion_features.md

@@ -148,7 +148,7 @@ The following tables show which models support each feature:
 | **Wan2.2-S2V**               |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |         ❌         |   ❌    |             ✅             |  ✅ (encode/decode)   |       ❌        |        ❌         |
 | **Wan2.1-VACE**              |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LTX-2**                    |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ❌    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **DreamID-Omni**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |

tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py

@@ -97,6 +97,260 @@ class TestPipelineIndependence:
         assert modules.resident_names == ["vocoder"]
         assert len(modules.vaes) == 2
 
+    def test_ltx23_pipeline_has_diffusion_pipeline_profiler_mixin(self):
+        """LTX23Pipeline must support lightweight stage timing."""
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+        from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
+
+        assert issubclass(LTX23Pipeline, DiffusionPipelineProfilerMixin)
+
+
+class TestLTX23VaeDecodeParallel:
+    """Test LTX-2.3 video VAE tiled parallel helpers without loading weights."""
+
+    def test_ltx23_video_vae_is_distributed_tile_only_class(self):
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
+            DistributedAutoencoderKLLTX2Video,
+        )
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeMixin
+
+        assert issubclass(DistributedAutoencoderKLLTX2Video, DistributedVaeMixin)
+        assert not hasattr(DistributedAutoencoderKLLTX2Video, "patch_split")
+
+    def test_ltx23_video_vae_tile_split_uses_native_ltx23_tile_geometry(self):
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
+            DistributedAutoencoderKLLTX2Video,
+        )
+
+        vae = SimpleNamespace(
+            spatial_compression_ratio=32,
+            tile_sample_min_height=512,
+            tile_sample_min_width=512,
+            tile_sample_stride_height=448,
+            tile_sample_stride_width=448,
+            temporal_compression_ratio=8,
+            dtype=torch.float32,
+        )
+
+        z = torch.zeros(1, 2, 5, 16, 24)
+        tasks, grid_spec = DistributedAutoencoderKLLTX2Video.tile_split(vae, z)
+
+        assert grid_spec.grid_shape == (2, 2)
+        assert grid_spec.split_dims == (3, 4)
+        assert grid_spec.tile_spec["sample_height"] == 512
+        assert grid_spec.tile_spec["sample_width"] == 768
+        assert grid_spec.tile_spec["blend_height"] == 64
+        assert grid_spec.tile_spec["blend_width"] == 64
+        assert grid_spec.tile_spec["max_tile_output_shape"] == (1, 3, 33, 512, 512)
+        assert grid_spec.tile_spec["tile_output_shapes"] == {
+            0: (1, 3, 33, 512, 512),
+            1: (1, 3, 33, 512, 320),
+            2: (1, 3, 33, 64, 512),
+            3: (1, 3, 33, 64, 320),
+        }
+        assert [task.grid_coord for task in tasks] == [(0, 0), (0, 1), (1, 0), (1, 1)]
+        assert [tuple(task.tensor.shape) for task in tasks] == [
+            (1, 2, 5, 16, 16),
+            (1, 2, 5, 16, 10),
+            (1, 2, 5, 2, 16),
+            (1, 2, 5, 2, 10),
+        ]
+        assert [task.workload for task in tasks] == [5 * 16 * 16, 5 * 16 * 10, 5 * 2 * 16, 5 * 2 * 10]
+
+    def test_ltx23_video_vae_tile_merge_blends_and_crops_like_tiled_decode(self):
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
+            DistributedAutoencoderKLLTX2Video,
+        )
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import GridSpec
+
+        class FakeVae:
+            def __init__(self):
+                self.blend_calls = []
+
+            def clear_cache(self):
+                pass
+
+            def blend_v(self, _previous, current, blend_height):
+                self.blend_calls.append(("v", blend_height))
+                return current
+
+            def blend_h(self, _previous, current, blend_width):
+                self.blend_calls.append(("h", blend_width))
+                return current
+
+        fake_vae = FakeVae()
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(2, 2),
+            tile_spec={
+                "sample_height": 10,
+                "sample_width": 10,
+                "blend_height": 1,
+                "blend_width": 2,
+                "tile_sample_stride_height": 5,
+                "tile_sample_stride_width": 5,
+            },
+        )
+        tiles = {
+            (0, 0): torch.full((1, 3, 2, 6, 6), 1.0),
+            (0, 1): torch.full((1, 3, 2, 6, 6), 2.0),
+            (1, 0): torch.full((1, 3, 2, 6, 6), 3.0),
+            (1, 1): torch.full((1, 3, 2, 6, 6), 4.0),
+        }
+
+        merged = DistributedAutoencoderKLLTX2Video.tile_merge(fake_vae, tiles, grid_spec)
+
+        assert merged.shape == (1, 3, 2, 10, 10)
+        assert fake_vae.blend_calls == [("h", 2), ("v", 1), ("v", 1), ("h", 2)]
+        torch.testing.assert_close(merged[:, :, :, :5, :5], torch.ones(1, 3, 2, 5, 5))
+        torch.testing.assert_close(merged[:, :, :, :5, 5:], torch.full((1, 3, 2, 5, 5), 2.0))
+        torch.testing.assert_close(merged[:, :, :, 5:, :5], torch.full((1, 3, 2, 5, 5), 3.0))
+        torch.testing.assert_close(merged[:, :, :, 5:, 5:], torch.full((1, 3, 2, 5, 5), 4.0))
+
+    def test_ltx23_video_vae_tiled_decode_dispatches_to_tile_operator(self):
+        from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl_ltx2
+
+        z = torch.zeros(1, 2, 1, 16, 24)
+        expected = torch.ones(1, 3, 1, 512, 768)
+        seen = {}
+
+        class FakeExecutor:
+            def execute(self, tensor, operator, broadcast_result=True):
+                seen["tensor"] = tensor
+                seen["operator"] = operator
+                seen["broadcast_result"] = broadcast_result
+                return expected
+
+        vae = SimpleNamespace(distributed_executor=FakeExecutor(), is_distributed_enabled=lambda: True)
+        vae.tile_split = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_split.__get__(vae)
+        vae.tile_exec = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_exec.__get__(vae)
+        vae.tile_merge = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_merge.__get__(vae)
+
+        output = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tiled_decode(
+            vae,
+            z,
+            temb=torch.tensor(0.5),
+            return_dict=False,
+        )
+
+        assert len(output) == 1
+        assert output[0] is expected
+        assert seen["tensor"] is z
+        assert seen["broadcast_result"] is False
+        assert seen["operator"].split.__name__ == "tile_split"
+        assert seen["operator"].merge.__name__ == "tile_merge"
+
+    def test_ltx23_vae_executor_gathers_known_tile_shapes_and_returns_empty_on_non_rank0(self):
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import LTX2VaeExecutor
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+            DistributedOperator,
+            GridSpec,
+            TileTask,
+        )
+
+        z = torch.zeros(1, 1, 1, 1, 1)
+        tile_output_shapes = {
+            0: (1, 1, 1, 2, 2),
+            1: (1, 1, 1, 2, 1),
+            2: (1, 1, 1, 1, 2),
+            3: (1, 1, 1, 1, 1),
+        }
+        tasks = [
+            TileTask(0, (0, 0), z, workload=4),
+            TileTask(1, (0, 1), z, workload=2),
+            TileTask(2, (1, 0), z, workload=2),
+            TileTask(3, (1, 1), z, workload=1),
+        ]
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(2, 2),
+            tile_spec={
+                "max_tile_output_shape": (1, 1, 1, 2, 2),
+                "tile_output_shapes": tile_output_shapes,
+            },
+            output_dtype=torch.float32,
+        )
+        seen = {}
+
+        def exec_tile(task):
+            return torch.full(tile_output_shapes[task.tile_id], float(task.tile_id + 1))
+
+        def merge_tiles(coord_tensor_map, passed_grid_spec):
+            seen["merged_shapes"] = {coord: tuple(tile.shape) for coord, tile in coord_tensor_map.items()}
+            assert passed_grid_spec is grid_spec
+            return torch.stack(
+                [
+                    coord_tensor_map[(0, 0)].flatten()[0],
+                    coord_tensor_map[(0, 1)].flatten()[0],
+                    coord_tensor_map[(1, 0)].flatten()[0],
+                    coord_tensor_map[(1, 1)].flatten()[0],
+                ]
+            )
+
+        operator = DistributedOperator(split=lambda _z: (tasks, grid_spec), exec=exec_tile, merge=merge_tiles)
+
+        rank0_executor = object.__new__(LTX2VaeExecutor)
+        rank0_executor.parallel_size = 2
+        rank0_executor.world_size = 2
+        rank0_executor.rank = 0
+
+        def gather_rank0(local_tile_tensor):
+            assigned = rank0_executor._balance_tasks(tasks, 2)
+            rank1_results = [(task.tile_id, exec_tile(task)) for task in assigned[1]]
+            rank1_tile_tensor = rank0_executor._pack_local_tiles_without_meta(
+                rank1_results,
+                list(local_tile_tensor.shape),
+                z.device,
+                torch.float32,
+            )
+            seen["rank0_gather_shape"] = tuple(local_tile_tensor.shape)
+            return [local_tile_tensor, rank1_tile_tensor]
+
+        def fail_final_sync(*_args, **_kwargs):
+            raise AssertionError("broadcast_result=False should not sync the final result")
+
+        rank0_executor.gather_tensors = gather_rank0
+        rank0_executor._sync_final_result = fail_final_sync
+
+        rank0_result = rank0_executor.execute(z, operator, broadcast_result=False)
+
+        torch.testing.assert_close(rank0_result, torch.tensor([1.0, 2.0, 3.0, 4.0]))
+        assert seen["rank0_gather_shape"] == (2, 1, 1, 1, 2, 2)
+        assert seen["merged_shapes"] == {
+            (0, 0): (1, 1, 1, 2, 2),
+            (0, 1): (1, 1, 1, 2, 1),
+            (1, 0): (1, 1, 1, 1, 2),
+            (1, 1): (1, 1, 1, 1, 1),
+        }
+
+        non_rank0_executor = object.__new__(LTX2VaeExecutor)
+        non_rank0_executor.parallel_size = 2
+        non_rank0_executor.world_size = 2
+        non_rank0_executor.rank = 1
+
+        def gather_rank1(local_tile_tensor):
+            seen["rank1_gather_shape"] = tuple(local_tile_tensor.shape)
+            return None
+
+        def fail_non_rank0_merge(*_args, **_kwargs):
+            raise AssertionError("non-rank0 should not merge gathered tiles")
+
+        non_rank0_executor.gather_tensors = gather_rank1
+        non_rank0_executor._sync_final_result = fail_final_sync
+
+        empty_result = non_rank0_executor.execute(
+            z,
+            DistributedOperator(
+                split=lambda _z: (tasks, grid_spec),
+                exec=exec_tile,
+                merge=fail_non_rank0_merge,
+            ),
+            broadcast_result=False,
+        )
+
+        assert tuple(empty_result.shape) == (0,)
+        assert seen["rank1_gather_shape"] == (2, 1, 1, 1, 2, 2)
+
 
 class TestCFGParallelHelpers:
     """Test LTX-2.3 CFG helper math without loading model weights."""

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_ltx2.py

@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+import torch
+from diffusers import AutoencoderKLLTX2Video
+from diffusers.models.autoencoders.vae import DecoderOutput
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+    DistributedOperator,
+    DistributedVaeExecutor,
+    DistributedVaeMixin,
+    GridSpec,
+    TileTask,
+)
+
+logger = init_logger(__name__)
+
+
+class LTX2VaeExecutor(DistributedVaeExecutor):
+    def _pack_local_tiles_without_meta(self, local_results, global_padding_shape, device, dtype):
+        tile_tensor = torch.zeros(global_padding_shape, device=device, dtype=dtype)
+        for idx, (_, t_tensor) in enumerate(local_results):
+            slices = tuple(slice(0, s) for s in t_tensor.shape)
+            tile_tensor[idx][slices] = t_tensor
+        return tile_tensor
+
+    def _unpack_tiles_with_known_metadata(self, tile_gather, assigned, tile_output_shapes):
+        coord_tensor_map = {}
+        for r, rank_tasks in enumerate(assigned):
+            tiles_src = tile_gather[r]
+            for idx, task in enumerate(rank_tasks):
+                output_shape = tile_output_shapes[task.tile_id]
+                slices = tuple(slice(0, dim) for dim in output_shape)
+                coord_tensor_map[task.grid_coord] = tiles_src[idx][slices]
+
+        return coord_tensor_map
+
+    def execute(self, z: torch.Tensor, operator: DistributedOperator, broadcast_result: bool = True):
+        pp_size = min(self.parallel_size, self.world_size)
+
+        tiletask_list, grid_spec = operator.split(z)
+        max_tile_output_shape = grid_spec.tile_spec.get("max_tile_output_shape")
+        tile_output_shapes = grid_spec.tile_spec.get("tile_output_shapes")
+        if max_tile_output_shape is None or tile_output_shapes is None:
+            return super().execute(z, operator, broadcast_result=broadcast_result)
+
+        assigned = self._balance_tasks(tiletask_list, pp_size)
+        local_tasks = assigned[self.rank] if self.rank < pp_size else []
+        local_results = [(t.tile_id, operator.exec(t)) for t in local_tasks]
+
+        output_dtype = grid_spec.output_dtype if grid_spec.output_dtype is not None else z.dtype
+        max_local_tile_count = max(len(tasks) for tasks in assigned)
+        global_padding_shape = [max_local_tile_count, *max_tile_output_shape]
+        local_tile_tensor = self._pack_local_tiles_without_meta(
+            local_results,
+            global_padding_shape,
+            z.device,
+            output_dtype,
+        )
+        tile_gather = self.gather_tensors(local_tile_tensor)
+
+        if self.rank != 0:
+            result = torch.empty(0, device=z.device)
+        else:
+            coord_tensor_map = self._unpack_tiles_with_known_metadata(tile_gather, assigned, tile_output_shapes)
+            result = operator.merge(coord_tensor_map, grid_spec)
+
+        if broadcast_result:
+            result = self._sync_final_result(result, z.ndim, z.device, output_dtype)
+        return result
+
+
+class DistributedAutoencoderKLLTX2Video(AutoencoderKLLTX2Video, DistributedVaeMixin):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any):
+        model = super().from_pretrained(*args, **kwargs)
+        model.init_distributed()
+        return model
+
+    def init_distributed(self):
+        self.distributed_executor = LTX2VaeExecutor()
+
+    def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+        _, _, num_frames, height, width = z.shape
+        sample_height = height * self.spatial_compression_ratio
+        sample_width = width * self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+        sample_num_frames = (num_frames - 1) * self.temporal_compression_ratio + 1
+        output_channels = int(getattr(getattr(self, "config", None), "out_channels", 3))
+
+        tiletask_list = []
+        tile_output_shapes = {}
+        for i in range(0, height, tile_latent_stride_height):
+            for j in range(0, width, tile_latent_stride_width):
+                tile = z[:, :, :num_frames, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
+                tile_id = len(tiletask_list)
+                tile_output_shapes[tile_id] = (
+                    z.shape[0],
+                    output_channels,
+                    sample_num_frames,
+                    tile.shape[3] * self.spatial_compression_ratio,
+                    tile.shape[4] * self.spatial_compression_ratio,
+                )
+                tiletask_list.append(
+                    TileTask(
+                        tile_id,
+                        (i // tile_latent_stride_height, j // tile_latent_stride_width),
+                        tile,
+                        workload=tile.shape[2] * tile.shape[3] * tile.shape[4],
+                    )
+                )
+
+        tile_spec = {
+            "sample_height": sample_height,
+            "sample_width": sample_width,
+            "blend_height": self.tile_sample_min_height - self.tile_sample_stride_height,
+            "blend_width": self.tile_sample_min_width - self.tile_sample_stride_width,
+            "tile_sample_stride_height": self.tile_sample_stride_height,
+            "tile_sample_stride_width": self.tile_sample_stride_width,
+            "max_tile_output_shape": (
+                z.shape[0],
+                output_channels,
+                sample_num_frames,
+                self.tile_sample_min_height,
+                self.tile_sample_min_width,
+            ),
+            "tile_output_shapes": tile_output_shapes,
+        }
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+            tile_spec=tile_spec,
+            output_dtype=self.dtype,
+        )
+        return tiletask_list, grid_spec
+
+    def tile_exec(
+        self,
+        task: TileTask,
+        temb: torch.Tensor | None = None,
+        causal: bool | None = None,
+    ) -> torch.Tensor:
+        if hasattr(self, "clear_cache"):
+            self.clear_cache()
+        return self.decoder(task.tensor, temb, causal=causal)
+
+    def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
+        grid_h, grid_w = grid_spec.grid_shape
+        result_rows = []
+
+        if hasattr(self, "clear_cache"):
+            self.clear_cache()
+
+        for i in range(grid_h):
+            result_row = []
+            for j in range(grid_w):
+                tile = coord_tensor_map[(i, j)]
+                if i > 0:
+                    tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+                if j > 0:
+                    tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+                result_row.append(
+                    tile[
+                        :,
+                        :,
+                        :,
+                        : grid_spec.tile_spec["tile_sample_stride_height"],
+                        : grid_spec.tile_spec["tile_sample_stride_width"],
+                    ]
+                )
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        dec = torch.cat(result_rows, dim=3)[
+            :, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"]
+        ]
+        return dec
+
+    def tiled_decode(
+        self,
+        z: torch.Tensor,
+        temb: torch.Tensor | None = None,
+        causal: bool | None = None,
+        return_dict: bool = True,
+    ):
+        if not self.is_distributed_enabled():
+            if causal is None:
+                return super().tiled_decode(z, temb, return_dict=return_dict)
+            return super().tiled_decode(z, temb, causal=causal, return_dict=return_dict)
+
+        logger.debug("LTX2 video VAE decode running with distributed tiled executor")
+        result = self.distributed_executor.execute(
+            z,
+            DistributedOperator(
+                split=self.tile_split,
+                exec=lambda task: self.tile_exec(task, temb=temb, causal=causal),
+                merge=self.tile_merge,
+            ),
+            broadcast_result=False,
+        )
+        if not return_dict:
+            return (result,)
+
+        return DecoderOutput(sample=result)

vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py

@@ -23,7 +23,7 @@ from typing import Any, ClassVar
 
 import numpy as np
 import torch
-from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler
+from diffusers import AutoencoderKLLTX2Audio, FlowMatchEulerDiscreteScheduler
 from diffusers.pipelines.ltx2 import LTX2TextConnectors
 from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
 from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps
@@ -36,6 +36,7 @@ from vllm.logger import init_logger
 from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import DistributedAutoencoderKLLTX2Video
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.parallel_state import (
     get_cfg_group,
@@ -48,6 +49,7 @@ from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_p
 from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
 from .pipeline_ltx2 import (
@@ -113,7 +115,13 @@ def get_ltx2_post_process_func(od_config: OmniDiffusionConfig):
     return post_process_func
 
 
-class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsComponentDiscovery):
+class LTX23Pipeline(
+    nn.Module,
+    CFGParallelMixin,
+    ProgressBarMixin,
+    SupportsComponentDiscovery,
+    DiffusionPipelineProfilerMixin,
+):
     """Fully independent LTX-2.3 pipeline.
 
     Key differences from LTX2Pipeline:
@@ -193,7 +201,7 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsCompo
 
         # --- VAE, Audio VAE ---
         self.vae = from_pretrained_with_prefetch(
-            AutoencoderKLLTX2Video.from_pretrained,
+            DistributedAutoencoderKLLTX2Video.from_pretrained,
             model,
             subfolder="vae",
             prefetch_list=ltx2_subfolders,
@@ -266,6 +274,10 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsCompo
         self._num_timesteps = None
         self._current_timestep = None
 
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
+        )
+
     def _place_aux_components(self) -> None:
         parallel_config = getattr(self.od_config, "parallel_config", None)
         use_managed_placement = bool(
@@ -1253,14 +1265,18 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsCompo
 
             latents = latents.to(self.vae.dtype)
             video = self.vae.decode(latents, timestep_decode, return_dict=False)[0]
-            video = self.video_processor.postprocess_video(video, output_type=output_type)
+            if video.numel() > 0:
+                video = self.video_processor.postprocess_video(video, output_type=output_type)
 
             audio_latents = audio_latents.to(self.audio_vae.dtype)
             generated_mel_spectrograms = self.audio_vae.decode(audio_latents, return_dict=False)[0]
 
             audio = self.vocoder(generated_mel_spectrograms)
 
-        return DiffusionOutput(output=(video, audio))
+        return DiffusionOutput(
+            output=(video, audio),
+            stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+        )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self)