Skip to content

vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath.install

Instance-level installer for the Wan VAE decoder fast path.

REPORT_ATTR module-attribute 

REPORT_ATTR = '_vllm_omni_wan_fastpath_report'

VAE_FAST_PATH_LEVELS module-attribute 

VAE_FAST_PATH_LEVELS: tuple[str, ...] = (
    "off",
    "lossless",
    "channels_last",
)

logger module-attribute 

logger = init_logger(__name__)

WanVaeFastPathReport dataclass

What :func:install_wan_vae_fastpath did to one VAE instance.

channels_last class-attribute instance-attribute 

channels_last: bool = False

fused_silu_dtypes class-attribute instance-attribute 

fused_silu_dtypes: tuple[str, ...] = ()

installed instance-attribute 

installed: bool

level instance-attribute 

level: str

patched class-attribute instance-attribute 

patched: Mapping[str, int] = field(default_factory=dict)

reason class-attribute instance-attribute 

reason: str | None = None

install_wan_vae_fastpath 

install_wan_vae_fastpath(
    vae: Module, *, level: str = "lossless"
) -> WanVaeFastPathReport

Install the Wan decoder fast path on one loaded AutoencoderKLWan. Idempotent.

level: * "off": do nothing. * "lossless": bind the bit-exact replacement forwards (Tier 1). * "channels_last": Tier 1 plus channels-last conv weights (not bit-exact).

The installer refuses (and reports why) when the VAE is not a diffusers Wan VAE, or when spatially-sharded decode is configured or already installed: wan_spatial_shard rebinds the same forwards and replaces the causal convolutions with halo-exchanging variants.

is_installed 

is_installed(vae: Module) -> bool

uninstall_wan_vae_fastpath 

uninstall_wan_vae_fastpath(vae: Module) -> None

Restore original forwards and tensor layouts, retaining current weight values.