cosmos25wanvae ¶

Cosmos 2.5 / Wan2.1 VAE adapter.

Why this exists: - Cosmos2.5 uses a Wan2.1-style VAE, but the diffusion model operates in a normalized latent space: z_norm = (z - mean) / std

Meanwhile, FastVideo's AutoencoderKLWan operates in the VAE's native latent space (denormalized): z = z_norm * std + mean

This adapter provides a single, stable interface for FastVideo pipelines: - encode(x) returns an object with .mean / .sample() / .mode() - decode(z) returns a tensor in pixel space

It also exposes flags used by pipeline stages to avoid double (de)normalization: - handles_latent_norm = True -> stages should NOT normalize encoder latents - handles_latent_denorm = True -> stages should NOT denormalize before decode

Classes¶

fastvideo.models.vaes.cosmos25wanvae.Cosmos25AttentionBlock ¶

Cosmos25AttentionBlock(dim: int)

Bases: Module

Official-like causal self-attention with a single head.

Source code in fastvideo/models/vaes/cosmos25wanvae.py

def __init__(self, dim: int) -> None:
    super().__init__()
    self.dim = dim
    self.norm = Cosmos25RMSNorm(dim)
    self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
    self.proj = nn.Conv2d(dim, dim, 1)
    nn.init.zeros_(self.proj.weight)

fastvideo.models.vaes.cosmos25wanvae.Cosmos25CausalConv3d ¶

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

Bases: Conv3d

Official-like causal 3D convolution.

Matches CausalConv3d in the official tokenizer: uses explicit F.pad and supports a cache_x prefix for causal chunking.

Source code in fastvideo/models/vaes/cosmos25wanvae.py

def __init__(self, *args: Any, **kwargs: Any) -> None:
    super().__init__(*args, **kwargs)
    # padding order for F.pad: (W_left, W_right, H_left, H_right, T_left, T_right)
    self._padding: tuple[int, ...] = (
        self.padding[2],
        self.padding[2],
        self.padding[1],
        self.padding[1],
        2 * self.padding[0],
        0,
    )
    self.padding = (0, 0, 0)

fastvideo.models.vaes.cosmos25wanvae.Cosmos25RMSNorm ¶

Cosmos25RMSNorm(dim: int, channel_first: bool = True, images: bool = True, bias: bool = False)

Bases: Module

Official-like RMS_norm (uses learnable gamma and optional bias).

Source code in fastvideo/models/vaes/cosmos25wanvae.py

def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False) -> None:
    super().__init__()
    broadcastable_dims = (1, 1, 1) if not images else (1, 1)
    shape = (dim, *broadcastable_dims) if channel_first else (dim,)

    self.channel_first = channel_first
    self.scale = dim**0.5
    self.gamma = nn.Parameter(torch.ones(shape))
    self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0

fastvideo.models.vaes.cosmos25wanvae.Cosmos25Resample ¶

Cosmos25Resample(dim: int, mode: str)

Bases: Module

Official-like Resample used for both spatial and temporal up/downsampling.

Source code in fastvideo/models/vaes/cosmos25wanvae.py

def __init__(self, dim: int, mode: str) -> None:
    assert mode in ("none", "upsample2d", "upsample3d", "downsample2d", "downsample3d")
    super().__init__()
    self.dim = dim
    self.mode = mode

    if mode == "upsample2d":
        self.resample = nn.Sequential(
            Cosmos25Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
            nn.Conv2d(dim, dim // 2, 3, padding=1),
        )
    elif mode == "upsample3d":
        self.resample = nn.Sequential(
            Cosmos25Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
            nn.Conv2d(dim, dim // 2, 3, padding=1),
        )
        self.time_conv = Cosmos25CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
    elif mode == "downsample2d":
        self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
    elif mode == "downsample3d":
        self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
        self.time_conv = Cosmos25CausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
    else:
        self.resample = nn.Identity()

fastvideo.models.vaes.cosmos25wanvae.Cosmos25Upsample ¶

Bases: Upsample

Official-like Upsample that is safe for bf16 (casts to fp32 internally).

fastvideo.models.vaes.cosmos25wanvae.Cosmos25WanVAE ¶

Cosmos25WanVAE(*, device: device | str = 'cpu', dtype: dtype = float32, temporal_window: int = 4, latents_mean: Optional[Tensor] = None, latents_std: Optional[Tensor] = None)

Bases: Module

A FastVideo-native copy of the official-like Wan2.1 VAE core.

Key properties: - Module naming matches official tokenizer (encoder, decoder, conv1, conv2) so it can consume tokenizer.pth keys directly. - encode() returns normalized latents and decode() expects normalized latents (Cosmos2.5 contract), matching Wan2pt1VAEInterface.

Source code in fastvideo/models/vaes/cosmos25wanvae.py

def __init__(
    self,
    *,
    device: torch.device | str = "cpu",
    dtype: torch.dtype = torch.float32,
    temporal_window: int = 4,
    latents_mean: Optional[torch.Tensor] = None,
    latents_std: Optional[torch.Tensor] = None,
) -> None:
    super().__init__()

    # Official hyperparams for Cosmos2.5 tokenizer (Wan2.1 VAE).
    cfg = dict(
        dim=96,
        z_dim=16,
        dim_mult=[1, 2, 4, 4],
        num_res_blocks=2,
        attn_scales=[],
        temperal_downsample=[False, True, True],
        dropout=0.0,
        temporal_window=temporal_window,
    )
    self.z_dim = 16
    self.temporal_window = temporal_window

    self.encoder = Cosmos25Encoder3d(
        dim=cfg["dim"],
        z_dim=cfg["z_dim"] * 2,
        dim_mult=cfg["dim_mult"],
        num_res_blocks=cfg["num_res_blocks"],
        attn_scales=cfg["attn_scales"],
        temperal_downsample=cfg["temperal_downsample"],
        dropout=cfg["dropout"],
    )
    self.conv1 = Cosmos25CausalConv3d(self.z_dim * 2, self.z_dim * 2, 1)
    self.conv2 = Cosmos25CausalConv3d(self.z_dim, self.z_dim, 1)
    self.decoder = Cosmos25Decoder3d(
        dim=cfg["dim"],
        z_dim=cfg["z_dim"],
        dim_mult=cfg["dim_mult"],
        num_res_blocks=cfg["num_res_blocks"],
        attn_scales=cfg["attn_scales"],
        temperal_upsample=list(cfg["temperal_downsample"])[::-1],
        dropout=cfg["dropout"],
    )

    # Default Cosmos2.5 latent stats (shared with configs).
    if latents_mean is None:
        latents_mean = torch.tensor(
            [
                -0.7571,
                -0.7089,
                -0.9113,
                0.1075,
                -0.1745,
                0.9653,
                -0.1517,
                1.5508,
                0.4134,
                -0.0715,
                0.5517,
                -0.3632,
                -0.1922,
                -0.9497,
                0.2503,
                -0.2921,
            ],
            dtype=torch.float32,
        ).view(1, 16, 1, 1, 1)
    if latents_std is None:
        latents_std = torch.tensor(
            [
                2.8184,
                1.4541,
                2.3275,
                2.6558,
                1.2196,
                1.7708,
                2.6052,
                2.0743,
                3.2687,
                2.1526,
                2.8652,
                1.5579,
                1.6382,
                1.1253,
                2.8251,
                1.9160,
            ],
            dtype=torch.float32,
        ).view(1, 16, 1, 1, 1)

    self.register_buffer("_latents_mean", latents_mean, persistent=False)
    self.register_buffer("_latents_std", latents_std, persistent=False)

    self.to(device=device, dtype=dtype)
    self.clear_cache()

Methods:¶

fastvideo.models.vaes.cosmos25wanvae.Cosmos25WanVAE.decode ¶

decode(latent: Tensor) -> Tensor

Decode from normalized latents (Cosmos contract).

Source code in fastvideo/models/vaes/cosmos25wanvae.py

def decode(self, latent: torch.Tensor) -> torch.Tensor:
    """
    Decode from *normalized* latents (Cosmos contract).
    """
    self.clear_cache()
    mean, inv_std = self._scale(latent)
    z = latent / inv_std + mean  # z = z_norm * std + mean

    iter_ = z.shape[2]
    x = self.conv2(z)
    for i in range(iter_):
        self._conv_idx = [0]
        if i == 0:
            out = self._i0_decode(x)
        else:
            out_ = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx)
            out = torch.cat([out, out_], 2)

    self.clear_cache()
    return out

fastvideo.models.vaes.cosmos25wanvae.Cosmos25WanVAE.encode ¶

encode(x: Tensor) -> _TensorLatentDist

Encode to normalized latents (Cosmos contract).

Source code in fastvideo/models/vaes/cosmos25wanvae.py

def encode(self, x: torch.Tensor) -> _TensorLatentDist:
    """
    Encode to *normalized* latents (Cosmos contract).
    """
    self.clear_cache()
    t = x.shape[2]
    iters = 1 + (t - 1) // self.temporal_window

    for i in range(iters):
        self._enc_conv_idx = [0]
        if i == 0:
            out = self._i0_encode(x)
        else:
            out_ = self.encoder(
                x[:, :, 1 + self.temporal_window * (i - 1) : 1 + self.temporal_window * i, :, :],
                feat_cache=self._enc_feat_map,
                feat_idx=self._enc_conv_idx,
            )
            out = torch.cat([out, out_], 2)

    if (t - 1) % self.temporal_window:
        self._enc_conv_idx = [0]
        out_ = self.encoder(
            x[:, :, 1 + self.temporal_window * (iters - 1) :, :, :],
            feat_cache=self._enc_feat_map,
            feat_idx=self._enc_conv_idx,
        )
        out = torch.cat([out, out_], 2)

    mu, _log_var = self.conv1(out).chunk(2, dim=1)
    mean, inv_std = self._scale(mu)
    z_norm = (mu - mean) * inv_std
    self.clear_cache()
    return _TensorLatentDist(z_norm)

fastvideo.models.vaes.cosmos25wanvae.Cosmos25WanVAEAdapter ¶

Cosmos25WanVAEAdapter(inner: Any, *, latents_mean: Optional[Tensor] = None, latents_std: Optional[Tensor] = None)

Bases: Module

Adapter that makes a Wan2.1-style VAE follow Cosmos2.5's latent contract: - encode() returns normalized latents - decode() expects normalized latents

Source code in fastvideo/models/vaes/cosmos25wanvae.py

def __init__(
    self,
    inner: Any,
    *,
    latents_mean: Optional[torch.Tensor] = None,
    latents_std: Optional[torch.Tensor] = None,
) -> None:
    super().__init__()
    self.inner = inner

    # Preserve `config` when available; some pipeline utilities expect it.
    self.config = getattr(inner, "config", None)

    # If not provided, try to derive from `config.latents_mean/std`.
    cfg = self.config
    if latents_mean is None and cfg is not None and hasattr(cfg, "latents_mean"):
        latents_mean = torch.tensor(cfg.latents_mean, dtype=torch.float32).view(1, -1, 1, 1, 1)
    if latents_std is None and cfg is not None and hasattr(cfg, "latents_std"):
        latents_std = torch.tensor(cfg.latents_std, dtype=torch.float32).view(1, -1, 1, 1, 1)

    if latents_mean is None or latents_std is None:
        raise RuntimeError(
            "Cosmos25WanVAEAdapter requires latents_mean/latents_std (either passed explicitly or available on inner.config)."
        )

    # Register as buffers so `.to(...)` moves them with the module.
    self.register_buffer("_latents_mean", latents_mean, persistent=False)
    self.register_buffer("_latents_std", latents_std, persistent=False)

Methods:¶

fastvideo.models.vaes.cosmos25wanvae.Cosmos25WanVAEAdapter.decode ¶

decode(z: Tensor) -> Tensor

Expects normalized latents (Cosmos contract).

Source code in fastvideo/models/vaes/cosmos25wanvae.py

def decode(self, z: torch.Tensor) -> torch.Tensor:
    """
    Expects *normalized* latents (Cosmos contract).
    """
    mean, std = self._to_latent_stats(z)
    z_denorm = z * std + mean
    out = self.inner.decode(z_denorm)
    return out.sample if hasattr(out, "sample") else out

fastvideo.models.vaes.cosmos25wanvae.Cosmos25WanVAEAdapter.encode ¶

encode(x: Tensor) -> _TensorLatentDist

Returns normalized latents (Cosmos contract).

Source code in fastvideo/models/vaes/cosmos25wanvae.py

def encode(self, x: torch.Tensor) -> _TensorLatentDist:
    """
    Returns *normalized* latents (Cosmos contract).
    """
    enc_out = self.inner.encode(x)

    # Support common encoder output shapes:
    # - DiagonalGaussianDistribution (FastVideo VAE): has `.mean` / `.sample()` / `.mode()`
    # - diffusers EncoderOutput: has `.latent_dist`
    # - raw tensor
    if hasattr(enc_out, "latent_dist"):
        dist = enc_out.latent_dist
        z_mean = dist.mode() if hasattr(dist, "mode") else dist.mean
    elif hasattr(enc_out, "mode") and hasattr(enc_out, "mean"):
        z_mean = enc_out.mode()
    elif isinstance(enc_out, torch.Tensor):
        z_mean = enc_out
    else:
        attrs = [a for a in dir(enc_out) if not a.startswith("_")]
        raise RuntimeError(
            f"Unsupported VAE encoder output type: {type(enc_out)}. attrs={attrs}"
        )

    mean, std = self._to_latent_stats(z_mean)
    z_norm = (z_mean - mean) / std
    return _TensorLatentDist(z_norm)