vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan ¶

logger `module-attribute` ¶

logger = init_logger(__name__)

DistributedAutoencoderKLWan ¶

Bases: OmniAutoencoderKLWan, DistributedVaeMixin

encode_tile_exec ¶

encode_tile_exec(task: TileTask) -> Tensor

Encode a single sample tile into latent space.

encode_tile_merge ¶

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

Merge encoded tiles into a full latent tensor.

encode_tile_split ¶

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

from_pretrained `classmethod` ¶

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

tile_exec ¶

tile_exec(task: TileTask) -> Tensor

Decode a single latent tile into RGB space.

tile_merge ¶

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

Merge decoded tiles into a full image.

tile_split ¶

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

tiled_decode ¶

tiled_decode(z: Tensor, return_dict: bool = True)

tiled_encode ¶

tiled_encode(x: Tensor) -> Tensor

Encode using distributed VAE executor.

Note: x is already patchified by parent's _encode() before calling this method.

OmniAutoencoderKLWan ¶

Bases: AutoencoderKLWan

chunk_value_range `class-attribute` ¶

chunk_value_range: tuple[float, float] = (-1.0, 1.0)

decode ¶

decode(z: Tensor, return_dict: bool = True)

Decode a Wan latent using the Diffusers-compatible full-tensor API.

decode_with_chunks ¶

decode_with_chunks(
    z: Tensor, *, on_chunk: DecodedChunkConsumer
) -> None

Decode z while synchronously delivering temporal chunks.

The ordinary :meth:decode contract remains the Diffusers contract. Chunk callbacks deliberately reject a sliced batch: Diffusers slices each sample before decoding, while a streaming consumer needs a stable batch ownership contract. Callers can disable use_slicing or invoke this capability once per sample.

Every rank participating in distributed VAE execution must invoke this method. Spatial-sharded execution invokes on_chunk only on the output-owning rank. Chunks have shape (B, C, T, H, W) and values clamped to [-1, 1]. After a callback failure, remaining chunks are decoded and discarded before the exception is re-raised.

encode ¶

encode(x: Tensor, return_dict: bool = True)