File size: 3,636 Bytes
775b8fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | """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)
|