--- title: gen3c_tokenizer_vae ---

GEN3C tokenizer-backed VAE adapter.

This wrapper loads the available tokenizer checkpoint (tokenizer.pth) and adapts it to GEN3C's latent-time contract (T=16 for 121 output frames).

Why this exists: - The converted GEN3C bundle includes tokenizer-style VAE weights, not a standard diffusers Wan VAE contract. - GEN3C diffusion expects 8x temporal compression (121 -> 16), while the available tokenizer checkpoint follows a 4x temporal path.

To bridge this at inference time, we: - keep the inner tokenizer model as-is, - downsample encoded latent time from inner-T to target-T for DiT input, - upsample generated latent time back to inner-T before decoding.

fastvideo.models.vaes.gen3c_tokenizer_vae.init_logger

init_logger(name: str) -> _FastvideoLogger

The main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root fastvideo logger has already been configured.

fastvideo.models.vaes.gen3c_tokenizer_vae.AutoencoderKLGen3CTokenizer

AutoencoderKLGen3CTokenizer(inner: nn.Module, *, target_temporal_compression: int = 8, inner_temporal_compression: int = 4, spatial_compression_factor: int = 8, pixel_chunk_duration: int = 121)

GEN3C VAE wrapper with temporal contract adaptation.

Interface contract: - encode(x) returns normalized latents in the target temporal layout. - decode(z) expects normalized latents in the target temporal layout.

Source
class AutoencoderKLGen3CTokenizer(nn.Module):
    """
    GEN3C VAE wrapper with temporal contract adaptation.

    Interface contract:
    - `encode(x)` returns normalized latents in the *target* temporal layout.
    - `decode(z)` expects normalized latents in the *target* temporal layout.
    """

    handles_latent_norm: bool = True
    handles_latent_denorm: bool = True

    def __init__(
        self,
        inner: nn.Module,
        *,
        target_temporal_compression: int = 8,
        inner_temporal_compression: int = 4,
        spatial_compression_factor: int = 8,
        pixel_chunk_duration: int = 121,
    ) -> None:
        super().__init__()
        self.inner = inner
        self.config = getattr(inner, "config", None)
        self._target_temporal_compression = int(target_temporal_compression)
        self._inner_temporal_compression = int(inner_temporal_compression)
        self._spatial_compression_factor = int(spatial_compression_factor)
        self._pixel_chunk_duration = int(pixel_chunk_duration)

    @staticmethod
    def _extract_latents(encoder_output: Any) -> torch.Tensor:
        if hasattr(encoder_output, "latent_dist"):
            dist = encoder_output.latent_dist
            if hasattr(dist, "mode"):
                return dist.mode()
            if hasattr(dist, "mean"):
                return dist.mean
            return dist.sample()
        if hasattr(encoder_output, "mode"):
            return encoder_output.mode()
        if hasattr(encoder_output, "latents"):
            return encoder_output.latents
        if hasattr(encoder_output, "sample"):
            return encoder_output.sample()
        if isinstance(encoder_output, torch.Tensor):
            return encoder_output
        raise TypeError(f"Unsupported encoder output type: {type(encoder_output)}")

    def _inner_to_target_time(self, z_inner: torch.Tensor) -> torch.Tensor:
        if z_inner.shape[2] <= 1:
            return z_inner

        # Common GEN3C case: inner=4x, target=8x => keep every other latent frame.
        if self._target_temporal_compression == 2 * self._inner_temporal_compression:
            return z_inner[:, :, 0::2, :, :].contiguous()

        # Generic fallback: keep boundary latents and sample uniformly.
        t_inner = z_inner.shape[2]
        t_target = 1 + (t_inner - 1) * self._inner_temporal_compression // self._target_temporal_compression
        idx = torch.linspace(0, t_inner - 1, t_target, device=z_inner.device)
        idx = idx.round().long()
        return z_inner.index_select(2, idx).contiguous()

    def _target_to_inner_time(self, z_target: torch.Tensor) -> torch.Tensor:
        if z_target.shape[2] <= 1:
            return z_target

        # Common GEN3C case: inner=4x, target=8x => insert midpoint frames.
        if self._target_temporal_compression == 2 * self._inner_temporal_compression:
            b, c, t, h, w = z_target.shape
            t_inner = 2 * t - 1
            out = torch.empty(b, c, t_inner, h, w, device=z_target.device, dtype=z_target.dtype)
            out[:, :, 0::2, :, :] = z_target
            out[:, :, 1::2, :, :] = 0.5 * (z_target[:, :, :-1, :, :] + z_target[:, :, 1:, :, :])
            return out.contiguous()

        # Generic fallback: linear index interpolation in time.
        t_target = z_target.shape[2]
        t_inner = 1 + (t_target - 1) * self._target_temporal_compression // self._inner_temporal_compression
        idx = torch.linspace(0, t_target - 1, t_inner, device=z_target.device)
        idx0 = idx.floor().long()
        idx1 = idx.ceil().long().clamp_max(t_target - 1)
        frac = (idx - idx0).view(1, 1, -1, 1, 1)
        z0 = z_target.index_select(2, idx0)
        z1 = z_target.index_select(2, idx1)
        return (z0 * (1.0 - frac) + z1 * frac).contiguous()

    def encode(self, x: torch.Tensor) -> _TensorLatentDist:
        z_inner = self._extract_latents(self.inner.encode(x))
        z_target = self._inner_to_target_time(z_inner)
        return _TensorLatentDist(z_target)

    def decode(self, z: torch.Tensor) -> torch.Tensor:
        z_inner = self._target_to_inner_time(z)
        out = self.inner.decode(z_inner)
        return out.sample if hasattr(out, "sample") else out

    def enable_tiling(self) -> None:
        if hasattr(self.inner, "enable_tiling"):
            self.inner.enable_tiling()

    def disable_tiling(self) -> None:
        if hasattr(self.inner, "disable_tiling"):
            self.inner.disable_tiling()

    def get_latent_num_frames(self, num_pixel_frames: int) -> int:
        num_pixel_frames = int(num_pixel_frames)
        if num_pixel_frames <= 1:
            return 1
        return 1 + (num_pixel_frames - 1) // self._target_temporal_compression

    def get_pixel_num_frames(self, num_latent_frames: int) -> int:
        num_latent_frames = int(num_latent_frames)
        if num_latent_frames <= 1:
            return 1
        return (num_latent_frames - 1) * self._target_temporal_compression + 1

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

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

    @property
    def temporal_compression_ratio(self) -> int:
        return self._target_temporal_compression

    @property
    def pixel_chunk_duration(self) -> int:
        return self._pixel_chunk_duration

    @property
    def latent_chunk_duration(self) -> int:
        return self.get_latent_num_frames(self._pixel_chunk_duration)

    @classmethod
    def from_tokenizer_checkpoint(
        cls,
        checkpoint_path: str,
        *,
        device: torch.device,
        dtype: torch.dtype,
        target_temporal_compression: int = 8,
        pixel_chunk_duration: int = 121,
    ) -> "AutoencoderKLGen3CTokenizer":
        from fastvideo.models.vaes.cosmos25wanvae import Cosmos25WanVAE

        inner = Cosmos25WanVAE(device=device, dtype=dtype)
        loaded = torch.load(checkpoint_path, map_location="cpu")
        if isinstance(loaded, dict):
            for key in ("state_dict", "model", "ema", "model_state_dict"):
                if key in loaded and isinstance(loaded[key], dict):
                    loaded = loaded[key]
                    break
        missing, unexpected = inner.load_state_dict(loaded, strict=False)
        if missing:
            logger.warning(
                "GEN3C tokenizer VAE missing keys (%d). Example: %s",
                len(missing),
                missing[:5],
            )
        if unexpected:
            logger.warning(
                "GEN3C tokenizer VAE unexpected keys (%d). Example: %s",
                len(unexpected),
                unexpected[:5],
            )
        return cls(
            inner,
            target_temporal_compression=target_temporal_compression,
            inner_temporal_compression=4,
            spatial_compression_factor=8,
            pixel_chunk_duration=pixel_chunk_duration,
        )

    @classmethod
    def from_jit_tokenizer(
        cls,
        tokenizer_dir: str,
        *,
        device: torch.device,
        dtype: torch.dtype,
        target_temporal_compression: int = 8,
        pixel_chunk_duration: int = 121,
    ) -> "AutoencoderKLGen3CTokenizer":
        encoder_path = f"{tokenizer_dir}/encoder.jit"
        decoder_path = f"{tokenizer_dir}/decoder.jit"
        mean_std_path = f"{tokenizer_dir}/mean_std.pt"
        inner = _JITGen3CTokenizerInner(
            encoder_path=encoder_path,
            decoder_path=decoder_path,
            mean_std_path=mean_std_path,
            dtype=dtype,
            device=device,
            latent_channels=16,
            latent_chunk_duration=1 + (pixel_chunk_duration - 1) // target_temporal_compression,
        )
        return cls(
            inner,
            target_temporal_compression=target_temporal_compression,
            inner_temporal_compression=target_temporal_compression,
            spatial_compression_factor=8,
            pixel_chunk_duration=pixel_chunk_duration,
        )

fastvideo.models.vaes.gen3c_tokenizer_vae.AutoencoderKLGen3CTokenizer.encode

encode(x: torch.Tensor) -> _TensorLatentDist
Source
def encode(self, x: torch.Tensor) -> _TensorLatentDist:
    z_inner = self._extract_latents(self.inner.encode(x))
    z_target = self._inner_to_target_time(z_inner)
    return _TensorLatentDist(z_target)

fastvideo.models.vaes.gen3c_tokenizer_vae.AutoencoderKLGen3CTokenizer.decode

decode(z: torch.Tensor) -> torch.Tensor
Source
def decode(self, z: torch.Tensor) -> torch.Tensor:
    z_inner = self._target_to_inner_time(z)
    out = self.inner.decode(z_inner)
    return out.sample if hasattr(out, "sample") else out

fastvideo.models.vaes.gen3c_tokenizer_vae.AutoencoderKLGen3CTokenizer.enable_tiling

enable_tiling() -> None
Source
def enable_tiling(self) -> None:
    if hasattr(self.inner, "enable_tiling"):
        self.inner.enable_tiling()

fastvideo.models.vaes.gen3c_tokenizer_vae.AutoencoderKLGen3CTokenizer.disable_tiling

disable_tiling() -> None
Source
def disable_tiling(self) -> None:
    if hasattr(self.inner, "disable_tiling"):
        self.inner.disable_tiling()

fastvideo.models.vaes.gen3c_tokenizer_vae.AutoencoderKLGen3CTokenizer.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:
    num_pixel_frames = int(num_pixel_frames)
    if num_pixel_frames <= 1:
        return 1
    return 1 + (num_pixel_frames - 1) // self._target_temporal_compression

fastvideo.models.vaes.gen3c_tokenizer_vae.AutoencoderKLGen3CTokenizer.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:
    num_latent_frames = int(num_latent_frames)
    if num_latent_frames <= 1:
        return 1
    return (num_latent_frames - 1) * self._target_temporal_compression + 1

fastvideo.models.vaes.gen3c_tokenizer_vae.AutoencoderKLGen3CTokenizer.from_tokenizer_checkpoint

from_tokenizer_checkpoint(checkpoint_path: str, *, device: torch.device, dtype: torch.dtype, target_temporal_compression: int = 8, pixel_chunk_duration: int = 121) -> 'AutoencoderKLGen3CTokenizer'
Source
@classmethod
def from_tokenizer_checkpoint(
    cls,
    checkpoint_path: str,
    *,
    device: torch.device,
    dtype: torch.dtype,
    target_temporal_compression: int = 8,
    pixel_chunk_duration: int = 121,
) -> "AutoencoderKLGen3CTokenizer":
    from fastvideo.models.vaes.cosmos25wanvae import Cosmos25WanVAE

    inner = Cosmos25WanVAE(device=device, dtype=dtype)
    loaded = torch.load(checkpoint_path, map_location="cpu")
    if isinstance(loaded, dict):
        for key in ("state_dict", "model", "ema", "model_state_dict"):
            if key in loaded and isinstance(loaded[key], dict):
                loaded = loaded[key]
                break
    missing, unexpected = inner.load_state_dict(loaded, strict=False)
    if missing:
        logger.warning(
            "GEN3C tokenizer VAE missing keys (%d). Example: %s",
            len(missing),
            missing[:5],
        )
    if unexpected:
        logger.warning(
            "GEN3C tokenizer VAE unexpected keys (%d). Example: %s",
            len(unexpected),
            unexpected[:5],
        )
    return cls(
        inner,
        target_temporal_compression=target_temporal_compression,
        inner_temporal_compression=4,
        spatial_compression_factor=8,
        pixel_chunk_duration=pixel_chunk_duration,
    )

fastvideo.models.vaes.gen3c_tokenizer_vae.AutoencoderKLGen3CTokenizer.from_jit_tokenizer

from_jit_tokenizer(tokenizer_dir: str, *, device: torch.device, dtype: torch.dtype, target_temporal_compression: int = 8, pixel_chunk_duration: int = 121) -> 'AutoencoderKLGen3CTokenizer'
Source
@classmethod
def from_jit_tokenizer(
    cls,
    tokenizer_dir: str,
    *,
    device: torch.device,
    dtype: torch.dtype,
    target_temporal_compression: int = 8,
    pixel_chunk_duration: int = 121,
) -> "AutoencoderKLGen3CTokenizer":
    encoder_path = f"{tokenizer_dir}/encoder.jit"
    decoder_path = f"{tokenizer_dir}/decoder.jit"
    mean_std_path = f"{tokenizer_dir}/mean_std.pt"
    inner = _JITGen3CTokenizerInner(
        encoder_path=encoder_path,
        decoder_path=decoder_path,
        mean_std_path=mean_std_path,
        dtype=dtype,
        device=device,
        latent_channels=16,
        latent_chunk_duration=1 + (pixel_chunk_duration - 1) // target_temporal_compression,
    )
    return cls(
        inner,
        target_temporal_compression=target_temporal_compression,
        inner_temporal_compression=target_temporal_compression,
        spatial_compression_factor=8,
        pixel_chunk_duration=pixel_chunk_duration,
    )