vllm_omni.diffusion.models.ltx2.vae.distributed ¶

Bases: 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.

from_pretrained(*args: Any, **kwargs: Any)

set_parallel_size(
    parallel_size: int, mode: str = "tile"
) -> None

tiled_decode(
    z: Tensor,
    generator: Generator | list[Generator] | None = None,
    num_inference_steps: int | None = None,
) -> Tensor

Geometry shared by all ranks for one distributed DiffVAE decode.

blend_frames: int

blend_height: int

blend_width: int

height: int

height_tiles: tuple[tuple[int, int], ...]

num_frames: int

scale_h: int

scale_t: int

scale_w: int

single_step_x0: bool

stride_h: int

stride_t: int

stride_w: int

temporal_tiles: tuple[tuple[int, int], ...]

width: int

width_tiles: tuple[tuple[int, int], ...]

Bases: TileTask

One stage-4 + stage-5 tile and its rank-independent noise.

crop_trailing_ghost: bool = False

drop_leading_frame: bool = False

noise_generator: Generator | list[Generator] | None = None

x_t: Tensor | None = None