| """Fast H3 VAE (non-causal): drop-in replacement for MiniMaxH3VideoVAE.decode_base on whole clips. |
| |
| from fast_h3_vae import FastH3VAEDecoder |
| dec = FastH3VAEDecoder(vae, "fast_h3_vae_noncausal_24l.safetensors") # vae: the released H3 VAE (for its |
| pixels = dec.decode(z) # z: normalized latents (B=1, T, 24, H, W) post_quant_conv / embedder / stats) |
| rgb_u8 = dec.decode_rgb_uint8(z) # (F, 16H, 16W, 3) uint8 on the GPU |
| |
| * any clip length T (t = 0 must be the start of a native 5-latent block, as for the released decoder). The output has |
| the same frames as decode_base: 17 frames per full block, the first latent of every block keeps 1 of its 4 frames |
| (tested T = 3..15: frame count equal to decode_base, fast vs reference 60.7-61.4 dB). For T = 1 and 2 the released |
| decode_base raises an error; this decoder returns 1 and 5 frames by the same rule; |
| * engine="fast" (default): NonCausalFast3 (FlashAttention-4 varlen windows, fused Triton kernels, one CUDA graph per |
| (T, H, W)); engine="ref": the plain PyTorch NonCausalViTDecoder (any GPU with SDPA); |
| * the fast engine is built on first use of a shape and cached (`max_engines` most recent shapes). |
| """ |
| from __future__ import annotations |
|
|
| from collections import OrderedDict |
|
|
| import torch |
| from safetensors.torch import load_file |
|
|
| from noncausal_vit_decoder import KEEP, NonCausalViTDecoder |
|
|
|
|
| class FastH3VAEDecoder: |
| def __init__(self, vae, weights: str, *, layers: int = 24, ffn_keep: int | None = None, time_norm: float = 5, |
| engine: str = "fast", max_engines: int = 4, device="cuda", offload_vae: bool = True): |
| """offload_vae: after building the decoder, move the released VAE (encoder + 36-layer decoder, ~10 GB) to the CPU; |
| decoding only needs its processor (pixel de-normalisation), which works on GPU tensors.""" |
| self.vae = vae |
| m = NonCausalViTDecoder(vae, KEEP[layers], time_norm=time_norm, ffn_keep=ffn_keep) |
| r = m.load_state_dict(load_file(weights, device="cpu"), strict=False) |
| assert not r.missing_keys and not r.unexpected_keys, r |
| self.model = m.to(device).eval().requires_grad_(False) |
| if offload_vae: |
| vae.to("cpu") |
| torch.cuda.empty_cache() |
| self.engine, self.max_engines = engine, max_engines |
| self._engines: OrderedDict = OrderedDict() |
|
|
| def _get(self, T, H, W): |
| key = (T, H, W) |
| if key in self._engines: |
| self._engines.move_to_end(key) |
| return self._engines[key] |
| from noncausal_fast3 import NonCausalFast3 |
| eng = NonCausalFast3(self.model, T, H, W, cuda_graph=True) |
| self._engines[key] = eng |
| while len(self._engines) > self.max_engines: |
| self._engines.popitem(last=False) |
| torch.cuda.empty_cache() |
| return eng |
|
|
| @torch.no_grad() |
| def decode(self, z): |
| """z: normalized latents (1, T, 24, H, W). Returns decode_base-space pixels (1, 3, F, 16H, 16W), fp32.""" |
| assert z.dim() == 5 and z.shape[0] == 1 and z.shape[2] == 24, z.shape |
| if self.engine == "ref": |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| return self.model(z.float()).float() |
| _, T, _, H, W = z.shape |
| return self._get(T, H, W).decode(z.float().contiguous()) |
|
|
| @torch.no_grad() |
| def decode_rgb_uint8(self, z): |
| """Same as decode, then the VAE's revert_tensor, clamp and uint8: (F, 16H, 16W, 3) on the GPU.""" |
| px = self.decode(z) |
| return self.vae.processor.revert_tensor(px).clamp_(0, 1).mul_(255).round_().to(torch.uint8)[0].permute(1, 2, 3, 0) |
|
|
