Optional, approximate MiniMax H3 tiny decoder for CUDA/CPU PyTorch.
Architecture and temporal mapping adapted from madebyollin/taehv at 62f7591f59dfbb4c3c02b7a621d180a9eeaba26c (MIT, Ollin Boer Bohan). This decoder consumes normalized diffusion latents; it does not use the full H3 VAE's latent mean/std or pixel denormalization.
Classes
TorchTAEH3Decoder(checkpoint_path: str | Path, *, dtype: dtype = float32)
Decode H3 NTCHW latents with bounded temporal feature memory.
Source code in fastvideo/models/vaes/minimax_h3_taeh3.py
| def __init__(self, checkpoint_path: str | Path, *, dtype: torch.dtype = torch.float32) -> None:
from safetensors.torch import load_file
raw = load_file(str(checkpoint_path))
actual = {key for key in raw if key.startswith("decoder.")}
if actual != set(_EXPECTED_SHAPES):
raise ValueError(f"TAEH3 decoder keys mismatch: missing={set(_EXPECTED_SHAPES) - actual}, "
f"unexpected={actual - set(_EXPECTED_SHAPES)}")
self.dtype = dtype
self.weights: dict[str, torch.Tensor] = {}
for key, shape in _EXPECTED_SHAPES.items():
value = raw[key]
if tuple(value.shape) != shape:
raise ValueError(f"TAEH3 weight {key} has shape {tuple(value.shape)}, expected {shape}")
self.weights[key] = value.detach().to(dtype=dtype).contiguous()
|
Methods:
decode_ntchw(latents: Tensor, *, chunk_size: int = 5) -> Tensor
Return NTCHW RGB in [0, 1] for H3's valid 5*k-3 latent lengths.
Source code in fastvideo/models/vaes/minimax_h3_taeh3.py
| def decode_ntchw(self, latents: torch.Tensor, *, chunk_size: int = 5) -> torch.Tensor:
"""Return NTCHW RGB in [0, 1] for H3's valid 5*k-3 latent lengths."""
if latents.ndim != 5 or latents.shape[2] != 24 or min(latents.shape) <= 0:
raise ValueError(f"Expected nonempty NTCHW H3 latents with 24 channels, got {tuple(latents.shape)}")
if latents.shape[1] % 5 != 2:
raise ValueError("H3 latent time must be 5*k-3, for example 2, 7, or 37.")
if chunk_size < 1:
raise ValueError("TAEH3 chunk_size must be positive.")
x = latents.to(dtype=self.dtype)
memory: dict[int, torch.Tensor] = {}
frames: list[torch.Tensor] = []
for start in range(0, x.shape[1], chunk_size):
decoded = self._chunk(x[:, start:start + chunk_size], memory)
keep = [i for i in range(decoded.shape[1]) if (start * 4 + i) % 20 >= 3]
frames.append(decoded[:, keep])
return torch.cat(frames, dim=1)
|
Functions:
decode_ncthw_latents_taeh3(latents: Tensor, *, device: device, checkpoint_path: str | Path | None = None, chunk_size: int = 5, dtype: dtype = float32) -> Tensor
Decode normalized NCTHW diffusion latents into NCTHW RGB in [0, 1].
Source code in fastvideo/models/vaes/minimax_h3_taeh3.py
| def decode_ncthw_latents_taeh3(
latents: torch.Tensor,
*,
device: torch.device,
checkpoint_path: str | Path | None = None,
chunk_size: int = 5,
dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
"""Decode normalized NCTHW diffusion latents into NCTHW RGB in [0, 1]."""
if latents.ndim != 5:
raise ValueError(f"Expected NCTHW latents, got {tuple(latents.shape)}")
checkpoint = ensure_taeh3_checkpoint(checkpoint_path)
cache_key = (str(checkpoint), str(device), str(dtype))
decoder = _DECODER_CACHE.get(cache_key)
if decoder is None:
decoder = TorchTAEH3Decoder(checkpoint, dtype=dtype).to(device)
_DECODER_CACHE[cache_key] = decoder
ntchw = latents.to(device=device, dtype=dtype).permute(0, 2, 1, 3, 4).contiguous()
rgb = decoder.decode_ntchw(ntchw, chunk_size=chunk_size)
return rgb.permute(0, 2, 1, 3, 4).contiguous()
|
ensure_taeh3_checkpoint(checkpoint_path: str | Path | None = None) -> Path
Fetch only pinned weights, atomically; never download executable code.
Source code in fastvideo/models/vaes/minimax_h3_taeh3.py
| def ensure_taeh3_checkpoint(checkpoint_path: str | Path | None = None) -> Path:
"""Fetch only pinned weights, atomically; never download executable code."""
if checkpoint_path is not None:
path = Path(checkpoint_path).expanduser()
if not path.is_file():
raise FileNotFoundError(f"TAEH3 checkpoint not found: {path}")
if path.suffix != ".safetensors":
raise ValueError("The TAEH3 decoder requires a .safetensors checkpoint.")
return path
path = Path.home() / ".cache/fastvideo/taehv/taeh3.safetensors"
def verify(candidate: Path) -> None:
hasher = hashlib.sha256()
with candidate.open("rb") as handle:
for chunk in iter(lambda: handle.read(1 << 20), b""):
hasher.update(chunk)
digest = hasher.hexdigest()
if digest != TAEH3_SHA256:
raise RuntimeError(f"TAEH3 checkpoint failed SHA-256 verification: {candidate}")
if path.exists():
verify(path)
return path
path.parent.mkdir(parents=True, exist_ok=True)
with tempfile.NamedTemporaryFile(dir=path.parent, suffix=".safetensors", delete=False) as temporary_file:
temporary = Path(temporary_file.name)
try:
with urllib.request.urlopen(TAEH3_URL, timeout=60) as response, temporary.open("wb") as handle:
while chunk := response.read(1 << 20):
handle.write(chunk)
verify(temporary)
temporary.replace(path)
finally:
temporary.unlink(missing_ok=True)
logger.info("Cached TAEH3 checkpoint at %s", path)
return path
|
Return NCTHW pixel shape for H3 TAEH3 (16x spatial, drop 3 of every 20 raw frames).
Source code in fastvideo/models/vaes/minimax_h3_taeh3.py
| def taeh3_decoded_pixel_shape(latent_shape: tuple[int, ...] | torch.Size) -> tuple[int, int, int, int, int]:
"""Return NCTHW pixel shape for H3 TAEH3 (16x spatial, drop 3 of every 20 raw frames)."""
if len(latent_shape) != 5:
raise ValueError(f"MiniMax-H3 latents must be five-dimensional, got shape {tuple(latent_shape)}.")
batch, channels, latent_frames, latent_height, latent_width = map(int, latent_shape)
if channels != 24:
raise ValueError(f"TAEH3 latents must have 24 channels, got {channels}.")
if latent_frames % 5 != 2:
raise ValueError(f"H3 latent time must be 5*k-3, got {latent_frames}.")
raw_frames = latent_frames * 4
kept = sum(1 for index in range(raw_frames) if index % 20 >= 3)
return (batch, 3, kept, latent_height * 16, latent_width * 16)
|