vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 ¶

logger `module-attribute` ¶

logger = init_logger(__name__)

DistributedAutoencoderKLLTX2Video ¶

Bases: AutoencoderKLLTX2Video, DistributedVaeMixin

encode_tile_exec ¶

encode_tile_exec(
    task: TileTask, causal: bool | None = None
) -> Tensor

Encode one sample-space tile into latent moments.

Parameters:

Name	Type	Description	Default
`task`	`TileTask`	Tile task whose tensor is a sample-space slice `(B, C, F, tile_h, tile_w)`.	required
`causal`	`bool \| None`	Optional causal-padding override forwarded to the encoder.	`None`

Returns:

Type	Description
`Tensor`	Encoder output of shape ``(B, 2 * latent_channels,
`Tensor`	(F - 1) // temporal_compression_ratio + 1, tile_h // scr, tile_w // scr)``.

encode_tile_merge ¶

encode_tile_merge(
    coord_tensor_map: dict[tuple[int, ...], Tensor],
    grid_spec: GridSpec,
) -> Tensor

Blend and stitch encoded tiles into the full latent tensor.

Mirrors the sequential blending of diffusers' tiled_encode: each tile is blended against its upper/left neighbors over the latent-space blend_* extents, cropped to tile_latent_stride_*, concatenated, and finally cropped to (latent_height, latent_width). Returns latent moments of shape (B, 2 * latent_channels, F_latent, latent_height, latent_width).

encode_tile_split ¶

encode_tile_split(
    x: Tensor,
) -> tuple[list[TileTask], GridSpec]

Split a sample-space video into overlapping spatial tiles for parallel encoding.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Sample-space input of shape `(B, C, F, H, W)`. Tiles are sliced on the `tile_sample_stride_` grid with `tile_sample_min_` extents, so neighboring tiles overlap by `min - stride` sample pixels.	required

Returns:

Name	Type	Description
	`list[TileTask]`	A `(tiletask_list, grid_spec)` pair. Each task holds one full-frame-range
	`GridSpec`	sample-space tile. `grid_spec.tile_spec` carries the latent-space merge
`metadata`	`tuple[list[TileTask], GridSpec]`	`latent_height`/`latent_width` (final crop), `blend_height`/
	`tuple[list[TileTask], GridSpec]`	`blend_width` (overlap extents), and `tile_latent_stride_*` (per-tile
	`tuple[list[TileTask], GridSpec]`	crop). When every tile's spatial dims divide `spatial_compression_ratio`,
	`tuple[list[TileTask], GridSpec]`	it also carries `max_tile_output_shape` and `tile_output_shapes` — each
	`tuple[list[TileTask], GridSpec]`	predicted as ``(B, 2 * latent_channels,
	`tuple[list[TileTask], GridSpec]`	(F - 1) // temporal_compression_ratio + 1, tile_h // scr, tile_w // scr)`` —
	`tuple[list[TileTask], GridSpec]`	which lets :class:`LTX2VaeExecutor` gather encoder outputs without a shape
	`tuple[list[TileTask], GridSpec]`	all-reduce. For tiles with non-divisible spatial dims the encoder output
	`tuple[list[TileTask], GridSpec]`	shape is not predictable, so both keys are omitted and the executor falls
	`tuple[list[TileTask], GridSpec]`	back to the dynamic metadata-gather path.

from_pretrained `classmethod` ¶

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

init_distributed ¶

init_distributed()

tile_exec ¶

tile_exec(
    task: TileTask,
    temb: Tensor | None = None,
    causal: bool | None = None,
) -> Tensor

tile_merge ¶

tile_merge(
    coord_tensor_map: dict[tuple[int, ...], Tensor],
    grid_spec: GridSpec,
) -> Tensor

tile_split ¶

tile_split(z: Tensor) -> tuple[list[TileTask], GridSpec]

tiled_decode ¶

tiled_decode(
    z: Tensor,
    temb: Tensor | None = None,
    causal: bool | None = None,
    return_dict: bool = True,
)

tiled_encode ¶

tiled_encode(
    x: Tensor, causal: bool | None = None
) -> Tensor

Distributed tiled encode: split across the DiT group, encode, gather, merge.

Drop-in replacement for diffusers' sequential tiled_encode — same signature, same return (latent moments (B, 2 * latent_channels, F_latent, latent_height, latent_width)), numerically identical output. Falls back to the sequential parent implementation when distribution is disabled (vae_patch_parallel_size <= 1, tiling off, or no process group). The merged result is broadcast to all ranks because every rank's denoiser consumes the latents.

LTX2VaeExecutor ¶

Bases: DistributedVaeExecutor

execute ¶

execute(
    z: Tensor,
    operator: DistributedOperator,
    broadcast_result: bool = True,
)