--- title: 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

fastvideo.models.vaes.cosmos25wanvae.Cosmos25WanVAEAdapter

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

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

Source
class Cosmos25WanVAEAdapter(nn.Module):
    """
    Adapter that makes a Wan2.1-style VAE follow Cosmos2.5's latent contract:
    - `encode()` returns **normalized** latents
    - `decode()` expects **normalized** latents
    """

    # Pipeline stage hints (see latent_preparation.py / decoding.py / image_encoding.py)
    handles_latent_norm: bool = True
    handles_latent_denorm: bool = True
    latent_norm_mode: str = "internal"  # informational

    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)

    def _to_latent_stats(self, like: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        mean = self._latents_mean.to(device=like.device, dtype=like.dtype)
        std = self._latents_std.to(device=like.device, dtype=like.dtype)
        return mean, std

    def get_latent_num_frames(self, num_pixel_frames: int) -> int:
        # Keep parity with official interface.
        if hasattr(self.inner, "get_latent_num_frames"):
            return int(self.inner.get_latent_num_frames(num_pixel_frames))
        return 1 + (num_pixel_frames - 1) // 4

    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)

    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.get_latent_num_frames

get_latent_num_frames(num_pixel_frames: int) -> int
Source
def get_latent_num_frames(self, num_pixel_frames: int) -> int:
    # Keep parity with official interface.
    if hasattr(self.inner, "get_latent_num_frames"):
        return int(self.inner.get_latent_num_frames(num_pixel_frames))
    return 1 + (num_pixel_frames - 1) // 4

fastvideo.models.vaes.cosmos25wanvae.Cosmos25WanVAEAdapter.encode

encode(x: torch.Tensor) -> _TensorLatentDist

Returns normalized latents (Cosmos contract).

Source
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)

fastvideo.models.vaes.cosmos25wanvae.Cosmos25WanVAEAdapter.decode

decode(z: torch.Tensor) -> torch.Tensor

Expects normalized latents (Cosmos contract).

Source
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.Cosmos25CausalConv3d

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

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
class Cosmos25CausalConv3d(nn.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.
    """

    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)

    def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
        padding = list(self._padding)
        if cache_x is not None and self._padding[4] > 0:
            cache_x = cache_x.to(x.device)
            x = torch.cat([cache_x, x], dim=2)
            padding[4] -= cache_x.shape[2]
        x = F.pad(x, padding)
        return super().forward(x)

fastvideo.models.vaes.cosmos25wanvae.Cosmos25CausalConv3d.forward

forward(x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor
Source
def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
    padding = list(self._padding)
    if cache_x is not None and self._padding[4] > 0:
        cache_x = cache_x.to(x.device)
        x = torch.cat([cache_x, x], dim=2)
        padding[4] -= cache_x.shape[2]
    x = F.pad(x, padding)
    return super().forward(x)

fastvideo.models.vaes.cosmos25wanvae.Cosmos25RMSNorm

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

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

Source
class Cosmos25RMSNorm(nn.Module):
    """Official-like RMS_norm (uses learnable gamma and optional bias)."""

    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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        dim = 1 if self.channel_first else -1
        return F.normalize(x, dim=dim) * self.scale * self.gamma + self.bias

fastvideo.models.vaes.cosmos25wanvae.Cosmos25RMSNorm.forward

forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
    dim = 1 if self.channel_first else -1
    return F.normalize(x, dim=dim) * self.scale * self.gamma + self.bias

fastvideo.models.vaes.cosmos25wanvae.Cosmos25Upsample

Cosmos25Upsample()

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

Source
class Cosmos25Upsample(nn.Upsample):
    """Official-like Upsample that is safe for bf16 (casts to fp32 internally)."""

    def forward(self, x: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
        return super().forward(x.float()).type_as(x)

fastvideo.models.vaes.cosmos25wanvae.Cosmos25Upsample.forward

forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
    return super().forward(x.float()).type_as(x)

fastvideo.models.vaes.cosmos25wanvae.Cosmos25Resample

Cosmos25Resample(dim: int, mode: str)

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

Source
class Cosmos25Resample(nn.Module):
    """
    Official-like Resample used for both spatial and temporal up/downsampling.
    """

    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()

    def forward(self, x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor:
        b, c, t, h, w = x.size()

        # Temporal upsample uses a time-conv and then interleaves frames.
        if self.mode == "upsample3d" and feat_cache is not None:
            idx = feat_idx[0]
            if feat_cache[idx] is None:
                feat_cache[idx] = "Rep"
                feat_idx[0] += 1
            else:
                cache_x = x[:, :, -CACHE_T:, :, :].clone()
                if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep":
                    cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                                        dim=2)
                if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep":
                    cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2)

                x = self.time_conv(x) if feat_cache[idx] == "Rep" else self.time_conv(x, feat_cache[idx])
                feat_cache[idx] = cache_x
                feat_idx[0] += 1

                x = x.reshape(b, 2, c, t, h, w)
                x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3)
                x = x.reshape(b, c, t * 2, h, w)

        t = x.shape[2]
        x = rearrange(x, "b c t h w -> (b t) c h w")
        x = self.resample(x)
        x = rearrange(x, "(b t) c h w -> b c t h w", t=t)

        # Temporal downsample: time_conv consumes last-frame cache.
        if self.mode == "downsample3d" and feat_cache is not None:
            idx = feat_idx[0]
            if feat_cache[idx] is None:
                feat_cache[idx] = x.clone()
                feat_idx[0] += 1
            else:
                cache_x = x[:, :, -1:, :, :].clone()
                x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
                feat_cache[idx] = cache_x
                feat_idx[0] += 1

        return x

fastvideo.models.vaes.cosmos25wanvae.Cosmos25Resample.forward

forward(x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor
Source
def forward(self, x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor:
    b, c, t, h, w = x.size()

    # Temporal upsample uses a time-conv and then interleaves frames.
    if self.mode == "upsample3d" and feat_cache is not None:
        idx = feat_idx[0]
        if feat_cache[idx] is None:
            feat_cache[idx] = "Rep"
            feat_idx[0] += 1
        else:
            cache_x = x[:, :, -CACHE_T:, :, :].clone()
            if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep":
                cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                                    dim=2)
            if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep":
                cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2)

            x = self.time_conv(x) if feat_cache[idx] == "Rep" else self.time_conv(x, feat_cache[idx])
            feat_cache[idx] = cache_x
            feat_idx[0] += 1

            x = x.reshape(b, 2, c, t, h, w)
            x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3)
            x = x.reshape(b, c, t * 2, h, w)

    t = x.shape[2]
    x = rearrange(x, "b c t h w -> (b t) c h w")
    x = self.resample(x)
    x = rearrange(x, "(b t) c h w -> b c t h w", t=t)

    # Temporal downsample: time_conv consumes last-frame cache.
    if self.mode == "downsample3d" and feat_cache is not None:
        idx = feat_idx[0]
        if feat_cache[idx] is None:
            feat_cache[idx] = x.clone()
            feat_idx[0] += 1
        else:
            cache_x = x[:, :, -1:, :, :].clone()
            x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
            feat_cache[idx] = cache_x
            feat_idx[0] += 1

    return x

fastvideo.models.vaes.cosmos25wanvae.Cosmos25ResidualBlock

Cosmos25ResidualBlock(in_dim: int, out_dim: int, dropout: float = 0.0)
Source
class Cosmos25ResidualBlock(nn.Module):

    def __init__(self, in_dim: int, out_dim: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.residual = nn.Sequential(
            Cosmos25RMSNorm(in_dim, images=False),
            nn.SiLU(),
            Cosmos25CausalConv3d(in_dim, out_dim, 3, padding=1),
            Cosmos25RMSNorm(out_dim, images=False),
            nn.SiLU(),
            nn.Dropout(dropout),
            Cosmos25CausalConv3d(out_dim, out_dim, 3, padding=1),
        )
        self.shortcut = Cosmos25CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity()

    def forward(self, x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor:
        h = self.shortcut(x)
        for layer in self.residual:
            if isinstance(layer, Cosmos25CausalConv3d) and feat_cache is not None:
                idx = feat_idx[0]
                cache_x = x[:, :, -CACHE_T:, :, :].clone()
                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                    cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                                        dim=2)
                x = layer(x, feat_cache[idx])
                feat_cache[idx] = cache_x
                feat_idx[0] += 1
            else:
                x = layer(x)
        return x + h

fastvideo.models.vaes.cosmos25wanvae.Cosmos25ResidualBlock.forward

forward(x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor
Source
def forward(self, x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor:
    h = self.shortcut(x)
    for layer in self.residual:
        if isinstance(layer, Cosmos25CausalConv3d) and feat_cache is not None:
            idx = feat_idx[0]
            cache_x = x[:, :, -CACHE_T:, :, :].clone()
            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                                    dim=2)
            x = layer(x, feat_cache[idx])
            feat_cache[idx] = cache_x
            feat_idx[0] += 1
        else:
            x = layer(x)
    return x + h

fastvideo.models.vaes.cosmos25wanvae.Cosmos25AttentionBlock

Cosmos25AttentionBlock(dim: int)

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

Source
class Cosmos25AttentionBlock(nn.Module):
    """Official-like causal self-attention with a single head."""

    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)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        identity = x
        b, c, t, h, w = x.size()
        x2 = rearrange(x, "b c t h w -> (b t) c h w")
        x2 = self.norm(x2)
        q, k, v = (self.to_qkv(x2).reshape(b * t, 1, c * 3, -1).permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1))
        x2 = F.scaled_dot_product_attention(q, k, v)
        x2 = x2.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
        x2 = self.proj(x2)
        x2 = rearrange(x2, "(b t) c h w-> b c t h w", t=t)
        return x2 + identity

fastvideo.models.vaes.cosmos25wanvae.Cosmos25AttentionBlock.forward

forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
    identity = x
    b, c, t, h, w = x.size()
    x2 = rearrange(x, "b c t h w -> (b t) c h w")
    x2 = self.norm(x2)
    q, k, v = (self.to_qkv(x2).reshape(b * t, 1, c * 3, -1).permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1))
    x2 = F.scaled_dot_product_attention(q, k, v)
    x2 = x2.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
    x2 = self.proj(x2)
    x2 = rearrange(x2, "(b t) c h w-> b c t h w", t=t)
    return x2 + identity

fastvideo.models.vaes.cosmos25wanvae.Cosmos25Encoder3d

Cosmos25Encoder3d(dim: int = 96, z_dim: int = 32, dim_mult: list[int] = [1, 2, 4, 4], num_res_blocks: int = 2, attn_scales: list[float] = [], temperal_downsample: list[bool] = [False, True, True], dropout: float = 0.0)
Source
class Cosmos25Encoder3d(nn.Module):

    def __init__(
        self,
        dim: int = 96,
        z_dim: int = 32,
        dim_mult: list[int] = [1, 2, 4, 4],
        num_res_blocks: int = 2,
        attn_scales: list[float] = [],
        temperal_downsample: list[bool] = [False, True, True],
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        dims = [dim * u for u in [1] + dim_mult]
        scale = 1.0

        self.conv1 = Cosmos25CausalConv3d(3, dims[0], 3, padding=1)

        downsamples: list[nn.Module] = []
        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
            for _ in range(num_res_blocks):
                downsamples.append(Cosmos25ResidualBlock(in_dim, out_dim, dropout))
                if scale in attn_scales:
                    downsamples.append(Cosmos25AttentionBlock(out_dim))
                in_dim = out_dim

            if i != len(dim_mult) - 1:
                mode = "downsample3d" if temperal_downsample[i] else "downsample2d"
                downsamples.append(Cosmos25Resample(out_dim, mode=mode))
                scale /= 2.0
        self.downsamples = nn.Sequential(*downsamples)

        self.middle = nn.Sequential(
            Cosmos25ResidualBlock(out_dim, out_dim, dropout),
            Cosmos25AttentionBlock(out_dim),
            Cosmos25ResidualBlock(out_dim, out_dim, dropout),
        )

        self.head = nn.Sequential(
            Cosmos25RMSNorm(out_dim, images=False),
            nn.SiLU(),
            Cosmos25CausalConv3d(out_dim, z_dim, 3, padding=1),
        )

    def forward(self, x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor:
        if feat_cache is not None:
            idx = feat_idx[0]
            cache_x = x[:, :, -CACHE_T:, :, :].clone()
            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
            x = self.conv1(x, feat_cache[idx])
            feat_cache[idx] = cache_x
            feat_idx[0] += 1
        else:
            x = self.conv1(x)

        for layer in self.downsamples:
            x = layer(x, feat_cache, feat_idx) if feat_cache is not None else layer(x)  # type: ignore[misc]

        for layer in self.middle:
            if isinstance(layer, Cosmos25ResidualBlock) and feat_cache is not None:
                x = layer(x, feat_cache, feat_idx)
            else:
                x = layer(x)  # type: ignore[misc]

        for layer in self.head:
            if isinstance(layer, Cosmos25CausalConv3d) and feat_cache is not None:
                idx = feat_idx[0]
                cache_x = x[:, :, -CACHE_T:, :, :].clone()
                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                    cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                                        dim=2)
                x = layer(x, feat_cache[idx])
                feat_cache[idx] = cache_x
                feat_idx[0] += 1
            else:
                x = layer(x)  # type: ignore[misc]
        return x

fastvideo.models.vaes.cosmos25wanvae.Cosmos25Encoder3d.forward

forward(x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor
Source
def forward(self, x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor:
    if feat_cache is not None:
        idx = feat_idx[0]
        cache_x = x[:, :, -CACHE_T:, :, :].clone()
        if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
            cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
        x = self.conv1(x, feat_cache[idx])
        feat_cache[idx] = cache_x
        feat_idx[0] += 1
    else:
        x = self.conv1(x)

    for layer in self.downsamples:
        x = layer(x, feat_cache, feat_idx) if feat_cache is not None else layer(x)  # type: ignore[misc]

    for layer in self.middle:
        if isinstance(layer, Cosmos25ResidualBlock) and feat_cache is not None:
            x = layer(x, feat_cache, feat_idx)
        else:
            x = layer(x)  # type: ignore[misc]

    for layer in self.head:
        if isinstance(layer, Cosmos25CausalConv3d) and feat_cache is not None:
            idx = feat_idx[0]
            cache_x = x[:, :, -CACHE_T:, :, :].clone()
            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                                    dim=2)
            x = layer(x, feat_cache[idx])
            feat_cache[idx] = cache_x
            feat_idx[0] += 1
        else:
            x = layer(x)  # type: ignore[misc]
    return x

fastvideo.models.vaes.cosmos25wanvae.Cosmos25Decoder3d

Cosmos25Decoder3d(dim: int = 96, z_dim: int = 16, dim_mult: list[int] = [1, 2, 4, 4], num_res_blocks: int = 2, attn_scales: list[float] = [], temperal_upsample: list[bool] = [False, True, True], dropout: float = 0.0)
Source
class Cosmos25Decoder3d(nn.Module):

    def __init__(
        self,
        dim: int = 96,
        z_dim: int = 16,
        dim_mult: list[int] = [1, 2, 4, 4],
        num_res_blocks: int = 2,
        attn_scales: list[float] = [],
        temperal_upsample: list[bool] = [False, True, True],
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
        scale = 1.0 / 2**(len(dim_mult) - 2)

        self.conv1 = Cosmos25CausalConv3d(z_dim, dims[0], 3, padding=1)
        self.middle = nn.Sequential(
            Cosmos25ResidualBlock(dims[0], dims[0], dropout),
            Cosmos25AttentionBlock(dims[0]),
            Cosmos25ResidualBlock(dims[0], dims[0], dropout),
        )

        upsamples: list[nn.Module] = []
        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
            if i in (1, 2, 3):
                in_dim = in_dim // 2
            for _ in range(num_res_blocks + 1):
                upsamples.append(Cosmos25ResidualBlock(in_dim, out_dim, dropout))
                if scale in attn_scales:
                    upsamples.append(Cosmos25AttentionBlock(out_dim))
                in_dim = out_dim
            if i != len(dim_mult) - 1:
                mode = "upsample3d" if temperal_upsample[i] else "upsample2d"
                upsamples.append(Cosmos25Resample(out_dim, mode=mode))
                scale *= 2.0
        self.upsamples = nn.Sequential(*upsamples)

        self.head = nn.Sequential(
            Cosmos25RMSNorm(out_dim, images=False),
            nn.SiLU(),
            Cosmos25CausalConv3d(out_dim, 3, 3, padding=1),
        )

    def forward(self, x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor:
        if feat_cache is not None:
            idx = feat_idx[0]
            cache_x = x[:, :, -CACHE_T:, :, :].clone()
            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
            x = self.conv1(x, feat_cache[idx])
            feat_cache[idx] = cache_x
            feat_idx[0] += 1
        else:
            x = self.conv1(x)

        for layer in self.middle:
            if isinstance(layer, Cosmos25ResidualBlock) and feat_cache is not None:
                x = layer(x, feat_cache, feat_idx)
            else:
                x = layer(x)  # type: ignore[misc]

        for layer in self.upsamples:
            x = layer(x, feat_cache, feat_idx) if feat_cache is not None else layer(x)  # type: ignore[misc]

        for layer in self.head:
            if isinstance(layer, Cosmos25CausalConv3d) and feat_cache is not None:
                idx = feat_idx[0]
                cache_x = x[:, :, -CACHE_T:, :, :].clone()
                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                    cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                                        dim=2)
                x = layer(x, feat_cache[idx])
                feat_cache[idx] = cache_x
                feat_idx[0] += 1
            else:
                x = layer(x)  # type: ignore[misc]
        return x

fastvideo.models.vaes.cosmos25wanvae.Cosmos25Decoder3d.forward

forward(x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor
Source
def forward(self, x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor:
    if feat_cache is not None:
        idx = feat_idx[0]
        cache_x = x[:, :, -CACHE_T:, :, :].clone()
        if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
            cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
        x = self.conv1(x, feat_cache[idx])
        feat_cache[idx] = cache_x
        feat_idx[0] += 1
    else:
        x = self.conv1(x)

    for layer in self.middle:
        if isinstance(layer, Cosmos25ResidualBlock) and feat_cache is not None:
            x = layer(x, feat_cache, feat_idx)
        else:
            x = layer(x)  # type: ignore[misc]

    for layer in self.upsamples:
        x = layer(x, feat_cache, feat_idx) if feat_cache is not None else layer(x)  # type: ignore[misc]

    for layer in self.head:
        if isinstance(layer, Cosmos25CausalConv3d) and feat_cache is not None:
            idx = feat_idx[0]
            cache_x = x[:, :, -CACHE_T:, :, :].clone()
            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                                    dim=2)
            x = layer(x, feat_cache[idx])
            feat_cache[idx] = cache_x
            feat_idx[0] += 1
        else:
            x = layer(x)  # type: ignore[misc]
    return x

fastvideo.models.vaes.cosmos25wanvae.Cosmos25WanVAE

Cosmos25WanVAE(*, 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)

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
class Cosmos25WanVAE(nn.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`.
    """

    handles_latent_norm: bool = True
    handles_latent_denorm: bool = True

    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()

    def clear_cache(self) -> None:
        # Decoder cache
        self._conv_num = _count_cosmos25_conv3d(self.decoder)
        self._conv_idx = [0]
        self._feat_map: list[Any] = [None] * self._conv_num
        # Encoder cache
        self._enc_conv_num = _count_cosmos25_conv3d(self.encoder)
        self._enc_conv_idx = [0]
        self._enc_feat_map: list[Any] = [None] * self._enc_conv_num

    def _scale(self, like: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        mean = self._latents_mean.to(device=like.device, dtype=like.dtype)
        std = self._latents_std.to(device=like.device, dtype=like.dtype)
        return mean, 1.0 / std

    def _i0_encode(self, x: torch.Tensor) -> torch.Tensor:
        return self.encoder(x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)

    def _i0_decode(self, x: torch.Tensor) -> torch.Tensor:
        return self.decoder(x[:, :, 0:1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx)

    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)

    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

    # --- Interface helpers (match official Wan2pt1VAEInterface) ---
    def get_latent_num_frames(self, num_pixel_frames: int) -> int:
        return 1 + (int(num_pixel_frames) - 1) // 4

    def get_pixel_num_frames(self, num_latent_frames: int) -> int:
        return (int(num_latent_frames) - 1) * 4 + 1

    @property
    def spatial_compression_factor(self) -> int:
        return 8

    @property
    def temporal_compression_factor(self) -> int:
        return 4

    @property
    def latent_ch(self) -> int:
        return 16

fastvideo.models.vaes.cosmos25wanvae.Cosmos25WanVAE.clear_cache

clear_cache() -> None
Source
def clear_cache(self) -> None:
    # Decoder cache
    self._conv_num = _count_cosmos25_conv3d(self.decoder)
    self._conv_idx = [0]
    self._feat_map: list[Any] = [None] * self._conv_num
    # Encoder cache
    self._enc_conv_num = _count_cosmos25_conv3d(self.encoder)
    self._enc_conv_idx = [0]
    self._enc_feat_map: list[Any] = [None] * self._enc_conv_num

fastvideo.models.vaes.cosmos25wanvae.Cosmos25WanVAE.encode

encode(x: torch.Tensor) -> _TensorLatentDist

Encode to normalized latents (Cosmos contract).

Source
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.Cosmos25WanVAE.decode

decode(latent: torch.Tensor) -> torch.Tensor

Decode from normalized latents (Cosmos contract).

Source
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.get_latent_num_frames

get_latent_num_frames(num_pixel_frames: int) -> int
Source
def get_latent_num_frames(self, num_pixel_frames: int) -> int:
    return 1 + (int(num_pixel_frames) - 1) // 4

fastvideo.models.vaes.cosmos25wanvae.Cosmos25WanVAE.get_pixel_num_frames

get_pixel_num_frames(num_latent_frames: int) -> int
Source
def get_pixel_num_frames(self, num_latent_frames: int) -> int:
    return (int(num_latent_frames) - 1) * 4 + 1