Skip to content

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

decode

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

encode

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