|
| 1 | +# SPDX-License-Identifier: Apache-2.0 |
| 2 | +"""Optional, approximate MiniMax H3 tiny decoder for CUDA/CPU PyTorch. |
| 3 | +
|
| 4 | +Architecture and temporal mapping adapted from madebyollin/taehv at |
| 5 | +62f7591f59dfbb4c3c02b7a621d180a9eeaba26c (MIT, Ollin Boer Bohan). |
| 6 | +This decoder consumes normalized diffusion latents; it does not use the full |
| 7 | +H3 VAE's latent mean/std or pixel denormalization. |
| 8 | +""" |
| 9 | + |
| 10 | +from __future__ import annotations |
| 11 | + |
| 12 | +import hashlib |
| 13 | +import tempfile |
| 14 | +import urllib.request |
| 15 | +from pathlib import Path |
| 16 | +from typing import Any |
| 17 | + |
| 18 | +import torch |
| 19 | +import torch.nn.functional as F |
| 20 | + |
| 21 | +from fastvideo.logger import init_logger |
| 22 | + |
| 23 | +logger = init_logger(__name__) |
| 24 | + |
| 25 | +TAEH3_REVISION = "62f7591f59dfbb4c3c02b7a621d180a9eeaba26c" |
| 26 | +TAEH3_URL = f"https://raw.githubusercontent.com/madebyollin/taehv/{TAEH3_REVISION}/safetensors/taeh3.safetensors" |
| 27 | +TAEH3_SHA256 = "4fd022bfcab08772fe0536b17ea1a3bbb5625be11e397868d1c5d891863d4c13" |
| 28 | + |
| 29 | +_EXPECTED_SHAPES: dict[str, tuple[int, ...]] = { |
| 30 | + "decoder.1.weight": (256, 24, 3, 3), |
| 31 | + "decoder.1.bias": (256, ), |
| 32 | + "decoder.7.conv.weight": (256, 256, 1, 1), |
| 33 | + "decoder.8.weight": (128, 256, 3, 3), |
| 34 | + "decoder.13.conv.weight": (256, 128, 1, 1), |
| 35 | + "decoder.14.weight": (64, 128, 3, 3), |
| 36 | + "decoder.19.conv.weight": (128, 64, 1, 1), |
| 37 | + "decoder.20.weight": (64, 64, 3, 3), |
| 38 | + "decoder.22.weight": (12, 64, 3, 3), |
| 39 | + "decoder.22.bias": (12, ), |
| 40 | +} |
| 41 | +for _index, _channels in ((3, 256), (4, 256), (5, 256), (9, 128), (10, 128), (11, 128), (15, 64), (16, 64), (17, 64)): |
| 42 | + for _layer in (0, 2, 4): |
| 43 | + _prefix = f"decoder.{_index}.conv.{_layer}" |
| 44 | + _EXPECTED_SHAPES[f"{_prefix}.weight"] = (_channels, _channels * (2 if _layer == 0 else 1), 3, 3) |
| 45 | + _EXPECTED_SHAPES[f"{_prefix}.bias"] = (_channels, ) |
| 46 | + |
| 47 | + |
| 48 | +def ensure_taeh3_checkpoint(checkpoint_path: str | Path | None = None) -> Path: |
| 49 | + """Fetch only pinned weights, atomically; never download executable code.""" |
| 50 | + if checkpoint_path is not None: |
| 51 | + path = Path(checkpoint_path).expanduser() |
| 52 | + if not path.is_file(): |
| 53 | + raise FileNotFoundError(f"TAEH3 checkpoint not found: {path}") |
| 54 | + if path.suffix != ".safetensors": |
| 55 | + raise ValueError("The TAEH3 decoder requires a .safetensors checkpoint.") |
| 56 | + return path |
| 57 | + path = Path.home() / ".cache/fastvideo/taehv/taeh3.safetensors" |
| 58 | + |
| 59 | + def verify(candidate: Path) -> None: |
| 60 | + hasher = hashlib.sha256() |
| 61 | + with candidate.open("rb") as handle: |
| 62 | + for chunk in iter(lambda: handle.read(1 << 20), b""): |
| 63 | + hasher.update(chunk) |
| 64 | + digest = hasher.hexdigest() |
| 65 | + if digest != TAEH3_SHA256: |
| 66 | + raise RuntimeError(f"TAEH3 checkpoint failed SHA-256 verification: {candidate}") |
| 67 | + |
| 68 | + if path.exists(): |
| 69 | + verify(path) |
| 70 | + return path |
| 71 | + path.parent.mkdir(parents=True, exist_ok=True) |
| 72 | + with tempfile.NamedTemporaryFile(dir=path.parent, suffix=".safetensors", delete=False) as temporary_file: |
| 73 | + temporary = Path(temporary_file.name) |
| 74 | + try: |
| 75 | + with urllib.request.urlopen(TAEH3_URL, timeout=60) as response, temporary.open("wb") as handle: |
| 76 | + while chunk := response.read(1 << 20): |
| 77 | + handle.write(chunk) |
| 78 | + verify(temporary) |
| 79 | + temporary.replace(path) |
| 80 | + finally: |
| 81 | + temporary.unlink(missing_ok=True) |
| 82 | + logger.info("Cached TAEH3 checkpoint at %s", path) |
| 83 | + return path |
| 84 | + |
| 85 | + |
| 86 | +class TorchTAEH3Decoder: |
| 87 | + """Decode H3 NTCHW latents with bounded temporal feature memory.""" |
| 88 | + |
| 89 | + def __init__(self, checkpoint_path: str | Path, *, dtype: torch.dtype = torch.float32) -> None: |
| 90 | + from safetensors.torch import load_file |
| 91 | + |
| 92 | + raw = load_file(str(checkpoint_path)) |
| 93 | + actual = {key for key in raw if key.startswith("decoder.")} |
| 94 | + if actual != set(_EXPECTED_SHAPES): |
| 95 | + raise ValueError(f"TAEH3 decoder keys mismatch: missing={set(_EXPECTED_SHAPES) - actual}, " |
| 96 | + f"unexpected={actual - set(_EXPECTED_SHAPES)}") |
| 97 | + self.dtype = dtype |
| 98 | + self.weights: dict[str, torch.Tensor] = {} |
| 99 | + for key, shape in _EXPECTED_SHAPES.items(): |
| 100 | + value = raw[key] |
| 101 | + if tuple(value.shape) != shape: |
| 102 | + raise ValueError(f"TAEH3 weight {key} has shape {tuple(value.shape)}, expected {shape}") |
| 103 | + self.weights[key] = value.detach().to(dtype=dtype).contiguous() |
| 104 | + |
| 105 | + def to(self, device: torch.device) -> "TorchTAEH3Decoder": |
| 106 | + self.weights = {key: value.to(device=device, non_blocking=True) for key, value in self.weights.items()} |
| 107 | + return self |
| 108 | + |
| 109 | + def _conv(self, x: torch.Tensor, name: str) -> torch.Tensor: |
| 110 | + weight = self.weights[f"{name}.weight"] |
| 111 | + bias = self.weights.get(f"{name}.bias") |
| 112 | + padding = weight.shape[-1] // 2 |
| 113 | + return F.conv2d(x, weight, bias, padding=padding) |
| 114 | + |
| 115 | + def _chunk(self, x: torch.Tensor, memory: dict[int, torch.Tensor]) -> torch.Tensor: |
| 116 | + n, t, c, h, w = x.shape |
| 117 | + x = F.relu(self._conv(torch.tanh(x.reshape(n * t, c, h, w) / 3.0) * 3.0, "decoder.1")) |
| 118 | + for indices, grow, projection, stride in (((3, 4, 5), 7, 8, 1), ((9, 10, 11), 13, 14, 2), ((15, 16, 17), 19, 20, |
| 119 | + 2)): |
| 120 | + for index in indices: |
| 121 | + nt, c, h, w = x.shape |
| 122 | + sequence = x.reshape(n, -1, c, h, w) |
| 123 | + previous = memory.get(index) |
| 124 | + if previous is None: |
| 125 | + previous = torch.zeros_like(sequence[:, :1]) |
| 126 | + past = torch.cat([previous, sequence[:, :-1]], dim=1).reshape_as(x) |
| 127 | + memory[index] = sequence[:, -1:].contiguous() |
| 128 | + y = torch.cat([x, past], dim=1) |
| 129 | + for layer in (0, 2, 4): |
| 130 | + y = self._conv(y, f"decoder.{index}.conv.{layer}") |
| 131 | + if layer != 4: |
| 132 | + y = F.relu(y) |
| 133 | + x = F.relu(x + y) |
| 134 | + nt, c_pre, h, w = x.shape |
| 135 | + x = F.interpolate(x, scale_factor=2, mode="nearest") |
| 136 | + x = self._conv(x, f"decoder.{grow}.conv") |
| 137 | + x = x.reshape(nt, stride, c_pre, h * 2, w * 2).reshape(nt * stride, c_pre, h * 2, w * 2) |
| 138 | + x = self._conv(x, f"decoder.{projection}") |
| 139 | + x = self._conv(F.relu(x), "decoder.22") |
| 140 | + x = F.pixel_shuffle(x, 2).clamp(0, 1) |
| 141 | + nt, c, h, w = x.shape |
| 142 | + return x.reshape(n, -1, c, h, w) |
| 143 | + |
| 144 | + def decode_ntchw(self, latents: torch.Tensor, *, chunk_size: int = 5) -> torch.Tensor: |
| 145 | + """Return NTCHW RGB in [0, 1] for H3's valid 5*k-3 latent lengths.""" |
| 146 | + if latents.ndim != 5 or latents.shape[2] != 24 or min(latents.shape) <= 0: |
| 147 | + raise ValueError(f"Expected nonempty NTCHW H3 latents with 24 channels, got {tuple(latents.shape)}") |
| 148 | + if latents.shape[1] % 5 != 2: |
| 149 | + raise ValueError("H3 latent time must be 5*k-3, for example 2, 7, or 37.") |
| 150 | + if chunk_size < 1: |
| 151 | + raise ValueError("TAEH3 chunk_size must be positive.") |
| 152 | + x = latents.to(dtype=self.dtype) |
| 153 | + memory: dict[int, torch.Tensor] = {} |
| 154 | + frames: list[torch.Tensor] = [] |
| 155 | + for start in range(0, x.shape[1], chunk_size): |
| 156 | + decoded = self._chunk(x[:, start:start + chunk_size], memory) |
| 157 | + keep = [i for i in range(decoded.shape[1]) if (start * 4 + i) % 20 >= 3] |
| 158 | + frames.append(decoded[:, keep]) |
| 159 | + return torch.cat(frames, dim=1) |
| 160 | + |
| 161 | + |
| 162 | +_DECODER_CACHE: dict[tuple[str, str, str], TorchTAEH3Decoder] = {} |
| 163 | + |
| 164 | + |
| 165 | +def decode_ncthw_latents_taeh3( |
| 166 | + latents: torch.Tensor, |
| 167 | + *, |
| 168 | + device: torch.device, |
| 169 | + checkpoint_path: str | Path | None = None, |
| 170 | + chunk_size: int = 5, |
| 171 | + dtype: torch.dtype = torch.float32, |
| 172 | +) -> torch.Tensor: |
| 173 | + """Decode normalized NCTHW diffusion latents into NCTHW RGB in [0, 1].""" |
| 174 | + if latents.ndim != 5: |
| 175 | + raise ValueError(f"Expected NCTHW latents, got {tuple(latents.shape)}") |
| 176 | + checkpoint = ensure_taeh3_checkpoint(checkpoint_path) |
| 177 | + cache_key = (str(checkpoint), str(device), str(dtype)) |
| 178 | + decoder = _DECODER_CACHE.get(cache_key) |
| 179 | + if decoder is None: |
| 180 | + decoder = TorchTAEH3Decoder(checkpoint, dtype=dtype).to(device) |
| 181 | + _DECODER_CACHE[cache_key] = decoder |
| 182 | + ntchw = latents.to(device=device, dtype=dtype).permute(0, 2, 1, 3, 4).contiguous() |
| 183 | + rgb = decoder.decode_ntchw(ntchw, chunk_size=chunk_size) |
| 184 | + return rgb.permute(0, 2, 1, 3, 4).contiguous() |
| 185 | + |
| 186 | + |
| 187 | +def taeh3_decoded_pixel_shape(latent_shape: tuple[int, ...] | torch.Size) -> tuple[int, int, int, int, int]: |
| 188 | + """Return NCTHW pixel shape for H3 TAEH3 (16x spatial, drop 3 of every 20 raw frames).""" |
| 189 | + if len(latent_shape) != 5: |
| 190 | + raise ValueError(f"MiniMax-H3 latents must be five-dimensional, got shape {tuple(latent_shape)}.") |
| 191 | + batch, channels, latent_frames, latent_height, latent_width = map(int, latent_shape) |
| 192 | + if channels != 24: |
| 193 | + raise ValueError(f"TAEH3 latents must have 24 channels, got {channels}.") |
| 194 | + if latent_frames % 5 != 2: |
| 195 | + raise ValueError(f"H3 latent time must be 5*k-3, got {latent_frames}.") |
| 196 | + raw_frames = latent_frames * 4 |
| 197 | + kept = sum(1 for index in range(raw_frames) if index % 20 >= 3) |
| 198 | + return (batch, 3, kept, latent_height * 16, latent_width * 16) |
0 commit comments