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) -> _FastvideoLoggerThe 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,
)</code></pre></details>
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,
)</code></pre></details>
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,
)