From bf03413c3651bc01dc1f690f951cf98eee61dd45 Mon Sep 17 00:00:00 2001 From: Mu GuanLin <39453182+mglyn@users.noreply.github.com> Date: Tue, 9 Jun 2026 21:14:16 +0800 Subject: [PATCH] [Feat] Add vae-decode-parallel for LTX-2.3 (#4277) Signed-off-by: mglyn <1203789601@qq.com> --- docs/user_guide/diffusion_features.md | 2 +- .../models/ltx2/test_ltx2_3_pipeline.py | 254 ++++++++++++++++++ .../autoencoders/autoencoder_kl_ltx2.py | 210 +++++++++++++++ .../diffusion/models/ltx2/pipeline_ltx2_3.py | 26 +- 4 files changed, 486 insertions(+), 6 deletions(-) create mode 100644 vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_ltx2.py diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index d4b0c9177..234446c01 100644 --- a/docs/user_guide/diffusion_features.md +++ b/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** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py index 9943d092b..bf6a75636 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py +++ b/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.""" diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_ltx2.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_ltx2.py new file mode 100644 index 000000000..0e7410808 --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py index caa40c93f..f333db8ce 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py +++ b/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)