[Diffusion] Add LTX-2.5 Diffusion VAE decoder support (#6189)

Signed-off-by: mglyn <1203789601@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder_distributed.py

@@ -0,0 +1,429 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright 2026 Lightricks and The HuggingFace Team. All rights reserved.
+#
+# Distributed tiling is copied and modified from Diffusers' serial LTX-2.5
+# tiling at commit d035dcd7cc7c88e0a154609b62887d50bba9fdc2 (Apache-2.0).
+
+"""LTX-2.5-specific distributed execution for the diffusion VAE decoder."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.distributed as dist
+from diffusers.utils import logging
+from diffusers.utils.torch_utils import randn_tensor
+
+from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+    DistributedOperator,
+    DistributedVaeMixin,
+    GridSpec,
+    TileTask,
+)
+
+from .ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def _tile_intervals(length: int, tile_size: int, stride: int, min_size: int) -> list[tuple[int, int]]:
+    """Build overlapping intervals, merging a too-small trailing remnant."""
+    if length <= tile_size:
+        return [(0, length)]
+    starts = list(range(0, length, stride))
+    while len(starts) > 1 and length - starts[-1] < min_size:
+        starts.pop()
+    return [(start, min(start + tile_size, length)) for start in starts[:-1]] + [(starts[-1], length)]
+
+
+@dataclass(frozen=True)
+class LTX2VideoDiffusionTilePlan:
+    """Geometry shared by all ranks for one distributed DiffVAE decode."""
+
+    temporal_tiles: tuple[tuple[int, int], ...]
+    height_tiles: tuple[tuple[int, int], ...]
+    width_tiles: tuple[tuple[int, int], ...]
+    num_frames: int
+    height: int
+    width: int
+    scale_t: int
+    scale_h: int
+    scale_w: int
+    stride_t: int
+    stride_h: int
+    stride_w: int
+    blend_frames: int
+    blend_height: int
+    blend_width: int
+    single_step_x0: bool
+
+
+@dataclass
+class LTX2VideoDiffusionTileTask(TileTask):
+    """One stage-4 + stage-5 tile and its rank-independent noise."""
+
+    drop_leading_frame: bool = False
+    crop_trailing_ghost: bool = False
+    noise_generator: torch.Generator | list[torch.Generator] | None = None
+    x_t: torch.Tensor | None = None
+
+
+class DistributedLTX2VideoDiffusionDecoderModel(LTX2VideoDiffusionDecoderModel, DistributedVaeMixin):
+    """LTX-2.5 diffusion decoder with model-specific tile parallelism.
+
+    Stages 1-3 run over the complete low-resolution feature volume on every
+    participating rank. Stage 4 and the diffusion stage run as independent
+    overlapping tile tasks, and rank 0 performs the reference blend/merge.
+    """
+
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any):
+        model = super().from_pretrained(*args, **kwargs)
+        model.init_distributed()
+        return model
+
+    def set_parallel_size(self, parallel_size: int, mode: str = "tile") -> None:
+        if mode != "tile":
+            raise ValueError(f"LTX-2.5 DiffVAE only supports vae_parallel_mode='tile', got {mode!r}.")
+        super().set_parallel_size(parallel_size, mode=mode)
+
+    def _build_tiled_decode_plan(
+        self,
+        features: torch.Tensor,
+        num_inference_steps: int,
+    ) -> LTX2VideoDiffusionTilePlan:
+        """Build the rank-independent tile schedule around Diffusers' decoder."""
+        decoder = self.decoder
+        patch_size = decoder.patch_size
+        upsample_stride = decoder.upsamples[-1].stride
+        scale_t, scale_h, scale_w = (
+            upsample_stride[0],
+            upsample_stride[1] * patch_size,
+            upsample_stride[2] * patch_size,
+        )
+        tile_t = self.tile_sample_min_num_frames // scale_t
+        tile_h = self.tile_sample_min_height // scale_h
+        tile_w = self.tile_sample_min_width // scale_w
+        stride_t = self.tile_sample_stride_num_frames // scale_t
+        stride_h = self.tile_sample_stride_height // scale_h
+        stride_w = self.tile_sample_stride_width // scale_w
+        min_sizes = [
+            max(kernel_4, -(-kernel_5 // stride))
+            for kernel_4, kernel_5, stride in zip(
+                self.config.decoder_stage_kernels[-1],
+                self.config.decoder_stage5_kernel,
+                upsample_stride,
+            )
+        ]
+
+        ghost_frames = decoder.trailing_pad_latent_frames * math.prod(
+            upsample.stride[0] for upsample in decoder.upsamples[:-1]
+        )
+        num_frames = features.shape[1] - ghost_frames
+        height, width = features.shape[2:4]
+        return LTX2VideoDiffusionTilePlan(
+            temporal_tiles=tuple(_tile_intervals(num_frames, tile_t, stride_t, min_sizes[0])),
+            height_tiles=tuple(_tile_intervals(height, tile_h, stride_h, min_sizes[1])),
+            width_tiles=tuple(_tile_intervals(width, tile_w, stride_w, min_sizes[2])),
+            num_frames=num_frames,
+            height=height,
+            width=width,
+            scale_t=scale_t,
+            scale_h=scale_h,
+            scale_w=scale_w,
+            stride_t=stride_t,
+            stride_h=stride_h,
+            stride_w=stride_w,
+            blend_frames=(tile_t - stride_t) * scale_t,
+            blend_height=(tile_h - stride_h) * scale_h,
+            blend_width=(tile_w - stride_w) * scale_w,
+            single_step_x0=num_inference_steps == 1 and decoder.model_output_type == "x0",
+        )
+
+    @staticmethod
+    def _iter_tiled_decode_coords(plan: LTX2VideoDiffusionTilePlan):
+        for t_idx in range(len(plan.temporal_tiles)):
+            for h_idx in range(len(plan.height_tiles)):
+                for w_idx in range(len(plan.width_tiles)):
+                    yield t_idx, h_idx, w_idx
+
+    @staticmethod
+    def _get_tiled_feature_slice(
+        features: torch.Tensor,
+        coord: tuple[int, int, int],
+        plan: LTX2VideoDiffusionTilePlan,
+    ) -> tuple[torch.Tensor, bool, bool]:
+        t_idx, h_idx, w_idx = coord
+        t0, t1 = plan.temporal_tiles[t_idx]
+        h0, h1 = plan.height_tiles[h_idx]
+        w0, w1 = plan.width_tiles[w_idx]
+        is_origin = t0 == 0
+        is_trailing = t1 == plan.num_frames
+        feature_t1 = features.shape[1] if is_trailing else t1
+        return features[:, t0:feature_t1, h0:h1, w0:w1], is_origin, is_trailing
+
+    def _tiled_pixel_shape_from_features(
+        self,
+        feature_tile: torch.Tensor,
+        *,
+        drop_leading_frame: bool,
+        crop_trailing_ghost: bool,
+    ) -> tuple[int, int, int, int, int]:
+        """Predict one output tile's shape without materializing stage 4."""
+        decoder = self.decoder
+        stride_t, stride_h, stride_w = decoder.upsamples[-1].stride
+        num_frames = feature_tile.shape[1] * stride_t
+        if stride_t == 2 and drop_leading_frame:
+            num_frames -= 1
+        if crop_trailing_ghost and decoder.trailing_pad_latent_frames > 0:
+            content_frames = max(
+                num_frames - decoder.trailing_pad_latent_frames * decoder.temporal_compression_ratio,
+                1,
+            )
+            num_frames = min(num_frames, max(content_frames, decoder.stage5_kernel[0]))
+        return (
+            feature_tile.shape[0],
+            decoder.out_channels,
+            num_frames,
+            feature_tile.shape[2] * stride_h * decoder.patch_size,
+            feature_tile.shape[3] * stride_w * decoder.patch_size,
+        )
+
+    @staticmethod
+    def _slice_tiled_noise(
+        x_t_full: torch.Tensor,
+        coord: tuple[int, int, int],
+        tile_pixel_shape: tuple[int, int, int, int, int],
+        plan: LTX2VideoDiffusionTilePlan,
+    ) -> torch.Tensor:
+        t_idx, h_idx, w_idx = coord
+        t0, _ = plan.temporal_tiles[t_idx]
+        h0, _ = plan.height_tiles[h_idx]
+        w0, _ = plan.width_tiles[w_idx]
+        pixel_t0 = t0 * plan.scale_t - (1 if t0 != 0 and plan.scale_t == 2 else 0)
+        return x_t_full[
+            :,
+            :,
+            pixel_t0 : pixel_t0 + tile_pixel_shape[2],
+            h0 * plan.scale_h : h0 * plan.scale_h + tile_pixel_shape[3],
+            w0 * plan.scale_w : w0 * plan.scale_w + tile_pixel_shape[4],
+        ]
+
+    def _merge_tiled_decode(
+        self,
+        tiles: dict[tuple[int, int, int], torch.Tensor],
+        plan: LTX2VideoDiffusionTilePlan,
+    ) -> torch.Tensor:
+        """Blend distributed RGB tiles in Diffusers' serial order."""
+        frame_groups = []
+        for t_idx in range(len(plan.temporal_tiles)):
+            rows = [
+                [tiles[(t_idx, h_idx, w_idx)] for w_idx in range(len(plan.width_tiles))]
+                for h_idx in range(len(plan.height_tiles))
+            ]
+            result_rows = []
+            for h_idx, row in enumerate(rows):
+                result_row = []
+                for w_idx, tile in enumerate(row):
+                    if h_idx > 0:
+                        tile = self.blend_v(rows[h_idx - 1][w_idx], tile, plan.blend_height)
+                    if w_idx > 0:
+                        tile = self.blend_h(row[w_idx - 1], tile, plan.blend_width)
+                    keep_height = plan.stride_h * plan.scale_h if h_idx < len(rows) - 1 else tile.shape[3]
+                    keep_width = plan.stride_w * plan.scale_w if w_idx < len(row) - 1 else tile.shape[4]
+                    result_row.append(tile[:, :, :, :keep_height, :keep_width])
+                result_rows.append(torch.cat(result_row, dim=4))
+            frame_groups.append(torch.cat(result_rows, dim=3))
+
+        result = []
+        for t_idx, group in enumerate(frame_groups):
+            if t_idx > 0:
+                group = self.blend_t(frame_groups[t_idx - 1], group, plan.blend_frames)
+            if t_idx < len(frame_groups) - 1:
+                keep_frames = plan.stride_t * plan.scale_t - (1 if t_idx == 0 and plan.scale_t == 2 else 0)
+                group = group[:, :, :keep_frames]
+            result.append(group)
+        return torch.cat(result, dim=2)
+
+    def _default_generator(self, device: torch.device) -> torch.Generator:
+        if device.type == "cpu":
+            return torch.default_generator
+        device_module = getattr(torch, device.type, None)
+        default_generators = getattr(device_module, "default_generators", None)
+        if default_generators is None:
+            raise ValueError(
+                f"Distributed LTX-2.5 diffusion decode on {device.type!r} requires an explicit torch.Generator."
+            )
+        device_index = device.index
+        if device_index is None:
+            current_device = getattr(device_module, "current_device", None)
+            device_index = current_device() if current_device is not None else 0
+        return default_generators[device_index]
+
+    def _sync_generators(
+        self,
+        generator: torch.Generator | list[torch.Generator] | None,
+        device: torch.device,
+    ) -> torch.Generator | list[torch.Generator]:
+        """Make every rank start tile-noise generation from rank 0's state."""
+        generators = self._default_generator(device) if generator is None else generator
+        generator_list = generators if isinstance(generators, list) else [generators]
+        for item in generator_list:
+            state_on_device = item.get_state().to(device=device)
+            dist.broadcast(state_on_device, src=0, group=self.distributed_executor.group)
+            item.set_state(state_on_device.cpu())
+        return generators
+
+    @staticmethod
+    def _clone_generator(generator: torch.Generator) -> torch.Generator:
+        cloned = torch.Generator(device=generator.device)
+        cloned.set_state(generator.get_state())
+        return cloned
+
+    def _clone_generators(
+        self,
+        generators: torch.Generator | list[torch.Generator],
+    ) -> torch.Generator | list[torch.Generator]:
+        if isinstance(generators, list):
+            return [self._clone_generator(item) for item in generators]
+        return self._clone_generator(generators)
+
+    def _distributed_tile_split(
+        self,
+        z: torch.Tensor,
+        generator: torch.Generator | list[torch.Generator] | None,
+        num_inference_steps: int,
+    ) -> tuple[list[LTX2VideoDiffusionTileTask], GridSpec]:
+        decoder = self.decoder
+        features = decoder.forward_stages_1_to_3(z)
+        plan = self._build_tiled_decode_plan(features, num_inference_steps)
+        generators = self._sync_generators(generator, z.device)
+
+        x_t_full = None
+        if not plan.single_step_x0:
+            pixel_frames = plan.num_frames * plan.scale_t - (1 if plan.scale_t == 2 else 0)
+            x_t_full = randn_tensor(
+                (
+                    z.shape[0],
+                    decoder.out_channels,
+                    pixel_frames,
+                    plan.height * plan.scale_h,
+                    plan.width * plan.scale_w,
+                ),
+                generator=generators,
+                device=z.device,
+                dtype=z.dtype,
+            )
+
+        tasks = []
+        for tile_id, coord in enumerate(self._iter_tiled_decode_coords(plan)):
+            feature_tile, is_origin, is_trailing = self._get_tiled_feature_slice(features, coord, plan)
+            tile_pixel_shape = self._tiled_pixel_shape_from_features(
+                feature_tile,
+                drop_leading_frame=is_origin,
+                crop_trailing_ghost=is_trailing,
+            )
+            tile_generator = None
+            x_t = None
+            if plan.single_step_x0:
+                # Save the canonical serial state for this tile, then advance
+                # the request generator exactly as the single-rank tiled path
+                # does. Task placement therefore cannot change tile noise.
+                tile_generator = self._clone_generators(generators)
+                discarded_noise = randn_tensor(
+                    tile_pixel_shape,
+                    generator=generators,
+                    device=z.device,
+                    dtype=z.dtype,
+                )
+                del discarded_noise
+            else:
+                x_t = self._slice_tiled_noise(x_t_full, coord, tile_pixel_shape, plan)
+
+            tasks.append(
+                LTX2VideoDiffusionTileTask(
+                    tile_id=tile_id,
+                    grid_coord=coord,
+                    tensor=feature_tile,
+                    workload=math.prod(tile_pixel_shape),
+                    drop_leading_frame=is_origin,
+                    crop_trailing_ghost=is_trailing,
+                    noise_generator=tile_generator,
+                    x_t=x_t,
+                )
+            )
+
+        return tasks, GridSpec(
+            split_dims=(2, 3, 4),
+            grid_shape=(len(plan.temporal_tiles), len(plan.height_tiles), len(plan.width_tiles)),
+            tile_spec={"plan": plan},
+            output_dtype=z.dtype,
+        )
+
+    def _distributed_tile_exec(
+        self,
+        task: LTX2VideoDiffusionTileTask,
+        num_inference_steps: int,
+    ) -> torch.Tensor:
+        context = self.decoder.forward_stage_4(
+            task.tensor,
+            drop_leading_frame=task.drop_leading_frame,
+            crop_trailing_ghost=task.crop_trailing_ghost,
+        )
+        tile_pixel_shape = (
+            context.shape[0],
+            self.decoder.out_channels,
+            context.shape[1],
+            context.shape[2] * self.decoder.patch_size,
+            context.shape[3] * self.decoder.patch_size,
+        )
+        if task.x_t is None:
+            x_t = randn_tensor(
+                tile_pixel_shape,
+                generator=task.noise_generator,
+                device=context.device,
+                dtype=context.dtype,
+            )
+        else:
+            x_t = task.x_t
+        return self.decoder.denoise(context, x_t, num_inference_steps)
+
+    def _distributed_tile_merge(
+        self,
+        tiles: dict[tuple[int, int, int], torch.Tensor],
+        grid_spec: GridSpec,
+    ) -> torch.Tensor:
+        plan = grid_spec.tile_spec["plan"]
+        if not isinstance(plan, LTX2VideoDiffusionTilePlan):
+            raise TypeError(f"Expected an LTX2VideoDiffusionTilePlan, got {type(plan)!r}.")
+        return self._merge_tiled_decode(tiles, plan)
+
+    def tiled_decode(
+        self,
+        z: torch.Tensor,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        num_inference_steps: int | None = None,
+    ) -> torch.Tensor:
+        if not self.is_distributed_enabled():
+            return super().tiled_decode(z, generator=generator, num_inference_steps=num_inference_steps)
+
+        logger.debug("LTX-2.5 diffusion decoder running with distributed stage-4/stage-5 tiles")
+        num_inference_steps = num_inference_steps or self.decoder.default_num_inference_steps
+        result = self.distributed_executor.execute(
+            z,
+            DistributedOperator(
+                split=lambda tensor: self._distributed_tile_split(tensor, generator, num_inference_steps),
+                exec=lambda task: self._distributed_tile_exec(task, num_inference_steps),
+                merge=self._distributed_tile_merge,
+            ),
+            broadcast_result=False,
+        )
+        if result.numel() == 0:
+            # The base decode method crops five dimensions before the LTX
+            # runtime discards non-output-rank results.
+            return torch.empty((0, self.decoder.out_channels, 0, 0, 0), device=z.device, dtype=z.dtype)
+        return result