Skip to content

vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl

logger module-attribute

logger = init_logger(__name__)

DistributedAutoencoderKL

Bases: DistributedAutoencoderKL_base, AutoencoderKL

DistributedAutoencoderKL_base

Bases: DistributedVaeMixin

decode

decode(
    z: Tensor,
    return_dict: bool = True,
    *args: Any,
    **kwargs: Any,
)

from_config classmethod

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

from_pretrained classmethod

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

patch_exec

patch_exec(task: TileTask) -> Tensor

patch_merge

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

patch_split

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

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]