558 lines
21 KiB
Python
558 lines
21 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
"""Wan VAE decode helpers for Apple Silicon MLX inference.
|
||
|
||
Two decode backends:
|
||
|
||
1. **TAEHV (primary / fast)** — Tiny AutoEncoder (madebyollin/taehv). Fully
|
||
MLX-native Conv2d path. ``taew2_1.pth`` for Wan2.1 (z_dim=16),
|
||
``taew2_2.pth`` for Wan2.2 5B (z_dim=48, patch_size=2). Expected decode
|
||
wall-clock ~seconds vs ~minutes for the full 3D VAE on MPS.
|
||
|
||
2. **Full AutoencoderKLWan (reference / quality)** — denormalize with
|
||
``latents_mean`` / ``latents_std`` then torch decode (MPS preferred). Used
|
||
for parity gates and when TAEHV is unavailable. A pure-MLX 3D-conv port of
|
||
the residual Wan2.2 decoder is left as follow-up (causal feat-cache +
|
||
residual up blocks are large); TAEHV covers the product latency path.
|
||
|
||
Diffusion latents from the DiT are **not** mean/std-normalized for TAEHV
|
||
(matching ``taehv_decode.py``); full VAE decode **does** denormalize first
|
||
(matching ``mlx_wan_prompt_to_video.decode_latents_to_video``).
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import json
|
||
import urllib.request
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Any, Literal
|
||
|
||
import numpy as np
|
||
|
||
from fastvideo.mlx_runtime.memory import cleanup_mlx, cleanup_torch_mps
|
||
|
||
GIB = 1024**3
|
||
|
||
TAEW2_1_URL = "https://raw.githubusercontent.com/madebyollin/taehv/main/taew2_1.pth"
|
||
TAEW2_2_URL = ("https://raw.githubusercontent.com/madebyollin/taehv/"
|
||
"563f40bdc820ed86bcad72ea515ee48f06bd22ec/taew2_2.pth")
|
||
# Validated 2026-07-02 / 2026-07-09 against upstream madebyollin/taehv.
|
||
TAEW2_1_SHA256 = "d26151e76cdc2c9424bef988de874b33d9a53f30ef3060cd556c429c469c797e"
|
||
TAEW2_2_SHA256 = "d053e216ca50e2bb837bbcd79b85f0366bea00e5938025572382a773b74c559a"
|
||
|
||
DecodeBackend = Literal["taehv", "taehv-torch", "wan-vae"]
|
||
|
||
|
||
def _cache_dir() -> Path:
|
||
"""Return the local directory used to cache TAEHV files."""
|
||
return Path.home() / ".cache" / "fastvideo" / "taehv"
|
||
|
||
|
||
def _sha256(path: Path) -> str:
|
||
"""Compute the SHA-256 digest of a file.
|
||
|
||
Parameters:
|
||
path (Path): Path to the file to hash.
|
||
|
||
Returns:
|
||
str: Lowercase hexadecimal SHA-256 digest.
|
||
"""
|
||
digest = hashlib.sha256()
|
||
with path.open("rb") as handle:
|
||
for chunk in iter(lambda: handle.read(1 << 20), b""):
|
||
digest.update(chunk)
|
||
return digest.hexdigest()
|
||
|
||
|
||
def _verify_checkpoint(path: Path, expected_digest: str) -> None:
|
||
"""
|
||
Verify a checkpoint's SHA-256 digest against a required expected digest.
|
||
|
||
Parameters:
|
||
path (Path): Path to the checkpoint file.
|
||
expected_digest (str): Expected lowercase SHA-256 digest.
|
||
|
||
Raises:
|
||
RuntimeError: If verification is enabled and the checkpoint digest does not match.
|
||
"""
|
||
if len(expected_digest) != 64 or any(char not in "0123456789abcdef" for char in expected_digest):
|
||
raise ValueError("A valid lowercase SHA-256 digest is required for bundled TAEHV checkpoints")
|
||
actual = _sha256(path)
|
||
if actual != expected_digest:
|
||
raise RuntimeError(f"TAEHV checkpoint at {path} failed sha256 verification "
|
||
f"(expected {expected_digest}, got {actual}). "
|
||
"Delete the file to re-download it.")
|
||
|
||
|
||
def ensure_taehv_checkpoint(*, z_dim: int, checkpoint_path: Path | None = None) -> Path:
|
||
"""
|
||
Return a validated TAEHV checkpoint for the specified latent channel count.
|
||
|
||
Parameters:
|
||
z_dim (int): Number of latent channels, supported values are 16 and 48.
|
||
checkpoint_path (Path | None): Optional existing checkpoint path to validate and use.
|
||
|
||
Returns:
|
||
Path: Path to the validated TAEHV checkpoint.
|
||
|
||
Raises:
|
||
FileNotFoundError: If the supplied checkpoint path does not exist.
|
||
ValueError: If no checkpoint is mapped to the specified latent channel count.
|
||
"""
|
||
if checkpoint_path is not None:
|
||
if not checkpoint_path.exists():
|
||
raise FileNotFoundError(f"TAEHV checkpoint not found: {checkpoint_path}")
|
||
return checkpoint_path
|
||
if z_dim == 16:
|
||
name, url, expect = "taew2_1.pth", TAEW2_1_URL, TAEW2_1_SHA256
|
||
elif z_dim == 48:
|
||
name, url, expect = "taew2_2.pth", TAEW2_2_URL, TAEW2_2_SHA256
|
||
else:
|
||
raise ValueError(f"No TAEHV checkpoint mapped for z_dim={z_dim} (supported: 16, 48)")
|
||
path = _cache_dir() / name
|
||
if not path.exists():
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
print(f"Downloading {url} -> {path}")
|
||
import socket
|
||
import tempfile
|
||
# Download to a temporary file, verify, then atomically rename.
|
||
with tempfile.NamedTemporaryFile(
|
||
mode="wb",
|
||
dir=path.parent,
|
||
prefix=f".tmp_{name}_",
|
||
suffix=".pth",
|
||
delete=False,
|
||
) as tmp_file:
|
||
tmp_path = Path(tmp_file.name)
|
||
try:
|
||
old_timeout = socket.getdefaulttimeout()
|
||
socket.setdefaulttimeout(300)
|
||
try:
|
||
urllib.request.urlretrieve(url, tmp_path) # noqa: S310 - public pinned artifact.
|
||
finally:
|
||
socket.setdefaulttimeout(old_timeout)
|
||
_verify_checkpoint(tmp_path, expect)
|
||
tmp_path.replace(path)
|
||
except Exception:
|
||
tmp_path.unlink(missing_ok=True)
|
||
raise
|
||
else:
|
||
_verify_checkpoint(path, expect)
|
||
return path
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class WanVAEConfigView:
|
||
"""Minimal config fields needed for denormalize + spatial scale."""
|
||
|
||
z_dim: int
|
||
latents_mean: tuple[float, ...]
|
||
latents_std: tuple[float, ...]
|
||
scale_factor_spatial: int = 8
|
||
scale_factor_temporal: int = 4
|
||
patch_size: int | None = None
|
||
vae_dir: Path | None = None
|
||
|
||
@classmethod
|
||
def from_vae_dir(cls, vae_dir: Path) -> WanVAEConfigView:
|
||
"""
|
||
Load Wan VAE configuration values from a directory.
|
||
|
||
Parameters:
|
||
vae_dir (Path): Directory containing the VAE ``config.json`` file.
|
||
|
||
Returns:
|
||
WanVAEConfigView: Configuration loaded from the VAE directory.
|
||
"""
|
||
cfg = json.loads((vae_dir / "config.json").read_text())
|
||
return cls(
|
||
z_dim=int(cfg["z_dim"]),
|
||
latents_mean=tuple(float(x) for x in cfg["latents_mean"]),
|
||
latents_std=tuple(float(x) for x in cfg["latents_std"]),
|
||
scale_factor_spatial=int(cfg.get("scale_factor_spatial", 8)),
|
||
scale_factor_temporal=int(cfg.get("scale_factor_temporal", 4)),
|
||
patch_size=cfg.get("patch_size"),
|
||
vae_dir=vae_dir,
|
||
)
|
||
|
||
|
||
def denormalize_latents_np(latents: np.ndarray, config: WanVAEConfigView) -> np.ndarray:
|
||
"""
|
||
Denormalize Wan VAE latent values using the configured means and standard deviations.
|
||
|
||
Parameters:
|
||
latents (np.ndarray): Latent values in normalized form.
|
||
config (WanVAEConfigView): Wan VAE latent statistics.
|
||
|
||
Returns:
|
||
np.ndarray: Denormalized latent values as float32.
|
||
"""
|
||
mean = np.asarray(config.latents_mean, dtype=np.float32).reshape(1, -1, 1, 1, 1)
|
||
std = np.asarray(config.latents_std, dtype=np.float32).reshape(1, -1, 1, 1, 1)
|
||
return latents.astype(np.float32) * std + mean
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# MLX TAEHV decoder (Conv2d stack — primary fully-MLX product path)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _mlx_conv2d(x: Any, weight: Any, bias: Any, *, stride: int = 1) -> Any:
|
||
import mlx.core as mx
|
||
|
||
# x: NCHW, weight: OIHW
|
||
y = mx.conv2d(x.transpose(0, 2, 3, 1), weight.transpose(0, 2, 3, 1), stride=stride, padding=1)
|
||
y = y.transpose(0, 3, 1, 2)
|
||
if bias is not None:
|
||
y = y + bias.reshape(1, -1, 1, 1)
|
||
return y
|
||
|
||
|
||
def _mlx_conv2d_1x1(x: Any, weight: Any, bias: Any = None) -> Any:
|
||
"""
|
||
Applies a 1×1 convolution to an MLX tensor in channel-first layout.
|
||
|
||
Parameters:
|
||
x (Any): Input tensor with shape [batch, channels, height, width].
|
||
weight (Any): Convolution weights.
|
||
bias (Any, optional): Optional output-channel bias.
|
||
|
||
Returns:
|
||
Any: The convolved tensor with shape [batch, output_channels, height, width].
|
||
"""
|
||
import mlx.core as mx
|
||
|
||
y = mx.conv2d(x.transpose(0, 2, 3, 1), weight.transpose(0, 2, 3, 1), stride=1, padding=0)
|
||
y = y.transpose(0, 3, 1, 2)
|
||
if bias is not None:
|
||
y = y + bias.reshape(1, -1, 1, 1)
|
||
return y
|
||
|
||
|
||
def _load_torch_state(path: Path) -> dict[str, np.ndarray]:
|
||
"""Load a PyTorch state dictionary as NumPy arrays.
|
||
|
||
Parameters:
|
||
path (Path): Path to the PyTorch checkpoint.
|
||
|
||
Returns:
|
||
dict[str, np.ndarray]: State dictionary with tensors converted to NumPy arrays.
|
||
"""
|
||
import torch
|
||
|
||
sd = torch.load(path, map_location="cpu", weights_only=True)
|
||
return {k: v.detach().float().cpu().numpy() for k, v in sd.items()}
|
||
|
||
|
||
class MLXTAEHVDecoder:
|
||
"""Minimal MLX port of TAEHV ``decoder`` (parallel-over-time MemBlocks)."""
|
||
|
||
def __init__(self, checkpoint_path: Path, *, z_dim: int) -> None:
|
||
"""Initialize the TAEHV decoder from a checkpoint for the specified latent dimensionality.
|
||
|
||
Parameters:
|
||
checkpoint_path (Path): Path to the TAEHV checkpoint.
|
||
z_dim (int): Number of latent channels, determining the decoder patch size.
|
||
"""
|
||
import mlx.core as mx
|
||
|
||
self.checkpoint_path = Path(checkpoint_path)
|
||
self.latent_channels = z_dim
|
||
# Derive patch_size from z_dim: 48 channels → patch_size=2, 16 → patch_size=1
|
||
self.patch_size = 2 if z_dim == 48 else 1
|
||
self.image_channels = 3
|
||
self.frames_to_trim = 3 # TGrow strides (1,2,2) → 2**2 - 1 for w2.1/w2.2 defaults
|
||
|
||
sd = _load_torch_state(self.checkpoint_path)
|
||
# Patch TGrow kernels like upstream TAEHV.patch_tgrow_layers.
|
||
self.weights = {k: mx.array(v) for k, v in sd.items()}
|
||
self._n_f = [256, 128, 64, 64]
|
||
|
||
def decode_ntchw(self, latents_ntchw: Any) -> Any:
|
||
"""
|
||
Decode latent video batches into clipped RGB frames.
|
||
|
||
Parameters:
|
||
latents_ntchw (Any): Latents with shape ``[N, T, C, H, W]`` and the
|
||
decoder's configured latent channel count.
|
||
|
||
Returns:
|
||
Any: Decoded frames with shape ``[N, T_out, 3, H_out, W_out]`` and values
|
||
clipped to the range ``[0, 1]``.
|
||
"""
|
||
import mlx.core as mx
|
||
|
||
x = latents_ntchw
|
||
n, t, c, h, w = x.shape
|
||
if c != self.latent_channels:
|
||
raise ValueError(f"expected C={self.latent_channels}, got {c}")
|
||
x = x.reshape(n * t, c, h, w)
|
||
x = self._run_decoder_parallel(x, n=n)
|
||
# Pixel-shuffle if patch_size > 1: (NT, 3*p*p, H, W) -> (NT, 3, H*p, W*p)
|
||
if self.patch_size > 1:
|
||
p = self.patch_size
|
||
nt, c_out, hh, ww = x.shape
|
||
x = x.reshape(nt, self.image_channels, p, p, hh, ww)
|
||
x = x.transpose(0, 1, 4, 2, 5, 3).reshape(nt, self.image_channels, hh * p, ww * p)
|
||
_, c_out, hh, ww = x.shape
|
||
t_out = x.shape[0] // n
|
||
x = x.reshape(n, t_out, c_out, hh, ww)
|
||
if self.frames_to_trim > 0 and t_out > self.frames_to_trim:
|
||
x = x[:, self.frames_to_trim:]
|
||
return mx.clip(x, 0.0, 1.0)
|
||
|
||
def _run_decoder_parallel(self, x: Any, *, n: int) -> Any:
|
||
"""
|
||
Apply the TAEHV decoder stack to flattened batch and temporal frames while preserving temporal memory.
|
||
"""
|
||
import mlx.core as mx
|
||
|
||
w = self.weights
|
||
|
||
def memblock(base: int, xx: Any, past: Any) -> Any:
|
||
"""
|
||
Apply a temporal memory block to the current and past feature tensors.
|
||
|
||
Parameters:
|
||
base (int): Decoder block index used to select the block weights.
|
||
xx (Any): Current feature tensor.
|
||
past (Any): Past feature tensor concatenated with the current features.
|
||
|
||
Returns:
|
||
Any: Activated feature tensor produced by the memory block.
|
||
"""
|
||
cat = mx.concatenate([xx, past], axis=1)
|
||
h = _mlx_conv2d(cat, w[f"decoder.{base}.conv.0.weight"], w.get(f"decoder.{base}.conv.0.bias"))
|
||
h = mx.maximum(h, 0.0)
|
||
h = _mlx_conv2d(h, w[f"decoder.{base}.conv.2.weight"], w.get(f"decoder.{base}.conv.2.bias"))
|
||
h = mx.maximum(h, 0.0)
|
||
h = _mlx_conv2d(h, w[f"decoder.{base}.conv.4.weight"], w.get(f"decoder.{base}.conv.4.bias"))
|
||
skip_key = f"decoder.{base}.skip.weight"
|
||
skip = _mlx_conv2d_1x1(xx, w[skip_key], None) if skip_key in w else xx
|
||
return mx.maximum(h + skip, 0.0)
|
||
|
||
def upsample2(xx: Any) -> Any:
|
||
"""
|
||
Upsample a four-dimensional tensor by a factor of two along its spatial dimensions.
|
||
|
||
Parameters:
|
||
xx (Any): Tensor with shape `(N, C, H, W)`.
|
||
|
||
Returns:
|
||
Any: Tensor with shape `(N, C, 2H, 2W)` containing replicated spatial values.
|
||
"""
|
||
nt, c, h, ww = xx.shape
|
||
xx = xx.reshape(nt, c, h, 1, ww, 1)
|
||
xx = mx.broadcast_to(xx, (nt, c, h, 2, ww, 2))
|
||
return xx.reshape(nt, c, h * 2, ww * 2)
|
||
|
||
def tgrow(base: int, xx: Any, stride: int) -> Any:
|
||
wt = w[f"decoder.{base}.conv.weight"]
|
||
out_ch = int(xx.shape[1]) * stride
|
||
if int(wt.shape[0]) > out_ch:
|
||
wt = wt[-out_ch:]
|
||
y = _mlx_conv2d_1x1(xx, wt, None)
|
||
if stride == 1:
|
||
return y
|
||
# TGrow.forward: (NT, C*stride, H, W) -> (NT*stride, C, H, W)
|
||
nt, c, h, ww = y.shape
|
||
c_in = c // stride
|
||
y = y.reshape(nt, stride, c_in, h, ww).transpose(0, 1, 2, 3, 4)
|
||
return y.reshape(nt * stride, c_in, h, ww)
|
||
|
||
def mem_past(xx: Any) -> Any:
|
||
"""
|
||
Build a temporal memory tensor containing a zero frame followed by the preceding frame at each time step.
|
||
|
||
Parameters:
|
||
xx (Any): Flattened batch and temporal tensor with shape ``(batch * time, channels, height, width)``.
|
||
|
||
Returns:
|
||
Any: Tensor with the same shape as ``xx`` containing the preceding frame for each temporal position.
|
||
"""
|
||
nt, c, h, ww = xx.shape
|
||
t_cur = nt // n
|
||
x_ = xx.reshape(n, t_cur, c, h, ww)
|
||
# pad one zero frame at t=0, align past[t] = x[t-1]
|
||
past = mx.concatenate([mx.zeros_like(x_[:, :1]), x_[:, :-1]], axis=1)
|
||
return past.reshape(nt, c, h, ww)
|
||
|
||
# 0 Clamp, 1 conv, 2 ReLU
|
||
x = mx.tanh(x / 3.0) * 3.0
|
||
x = _mlx_conv2d(x, w["decoder.1.weight"], w.get("decoder.1.bias"))
|
||
x = mx.maximum(x, 0.0)
|
||
for mem_idx in (3, 4, 5):
|
||
x = memblock(mem_idx, x, mem_past(x))
|
||
x = upsample2(x)
|
||
x = tgrow(7, x, 1)
|
||
x = _mlx_conv2d(x, w["decoder.8.weight"], w.get("decoder.8.bias"))
|
||
for mem_idx in (9, 10, 11):
|
||
x = memblock(mem_idx, x, mem_past(x))
|
||
x = upsample2(x)
|
||
x = tgrow(13, x, 2)
|
||
x = _mlx_conv2d(x, w["decoder.14.weight"], w.get("decoder.14.bias"))
|
||
for mem_idx in (15, 16, 17):
|
||
x = memblock(mem_idx, x, mem_past(x))
|
||
x = upsample2(x)
|
||
x = tgrow(19, x, 2)
|
||
x = _mlx_conv2d(x, w["decoder.20.weight"], w.get("decoder.20.bias"))
|
||
x = mx.maximum(x, 0.0)
|
||
x = _mlx_conv2d(x, w["decoder.22.weight"], w.get("decoder.22.bias"))
|
||
return x
|
||
|
||
|
||
def decode_latents_taehv_mlx(
|
||
latents_np: np.ndarray,
|
||
*,
|
||
z_dim: int | None = None,
|
||
checkpoint_path: Path | None = None,
|
||
) -> np.ndarray:
|
||
"""
|
||
Decode latent representations with the MLX TAEHV decoder.
|
||
|
||
Parameters:
|
||
latents_np (np.ndarray): Latents arranged as [B, C, T, H, W].
|
||
z_dim (int | None): Latent channel dimension used to select the decoder checkpoint.
|
||
checkpoint_path (Path | None): Optional path to a TAEHV checkpoint.
|
||
|
||
Returns:
|
||
np.ndarray: Decoded pixels arranged as [B, T, H, W, 3] with values in [0, 1].
|
||
|
||
Raises:
|
||
ValueError: If `latents_np` does not have five dimensions.
|
||
"""
|
||
import mlx.core as mx
|
||
|
||
if latents_np.ndim != 5:
|
||
raise ValueError(f"expected [B,C,T,H,W], got {latents_np.shape}")
|
||
c = latents_np.shape[1]
|
||
z = z_dim if z_dim is not None else c
|
||
ckpt = ensure_taehv_checkpoint(z_dim=z, checkpoint_path=checkpoint_path)
|
||
dec = MLXTAEHVDecoder(ckpt, z_dim=z)
|
||
# NTCHW
|
||
x = mx.array(latents_np.transpose(0, 2, 1, 3, 4).astype(np.float32))
|
||
out = dec.decode_ntchw(x) # N T C H W
|
||
mx.eval(out)
|
||
arr = np.array(out)
|
||
# B T H W C
|
||
return arr.transpose(0, 1, 3, 4, 2)
|
||
|
||
|
||
def decode_latents_wan_vae_torch(
|
||
latents_np: np.ndarray,
|
||
*,
|
||
vae_dir: Path,
|
||
device: str = "auto",
|
||
dtype_name: str = "fp16",
|
||
) -> np.ndarray:
|
||
"""Full AutoencoderKLWan decode on torch (MPS/CPU) with mean/std denormalize.
|
||
|
||
Returns pixels ``[B, T, H, W, 3]`` float in [0, 1].
|
||
"""
|
||
import torch
|
||
from diffusers import AutoencoderKLWan
|
||
from diffusers.video_processor import VideoProcessor
|
||
|
||
if device == "auto":
|
||
device = "mps" if torch.backends.mps.is_available() else "cpu"
|
||
dtype = torch.float16 if dtype_name == "fp16" and device == "mps" else torch.float32
|
||
config = WanVAEConfigView.from_vae_dir(vae_dir)
|
||
vae = AutoencoderKLWan.from_pretrained(vae_dir, torch_dtype=dtype, local_files_only=True).to(device)
|
||
vae.eval()
|
||
latents = torch.from_numpy(latents_np.astype(np.float32)).to(device=device, dtype=dtype)
|
||
mean = torch.tensor(config.latents_mean, device=device, dtype=dtype).view(1, -1, 1, 1, 1)
|
||
inv_std = (1.0 / torch.tensor(config.latents_std, device=device, dtype=dtype)).view(1, -1, 1, 1, 1)
|
||
latents = latents / inv_std + mean # matches prompt_to_video path
|
||
with torch.no_grad():
|
||
video = vae.decode(latents, return_dict=False)[0]
|
||
video = VideoProcessor(vae_scale_factor=config.scale_factor_spatial).postprocess_video(video, output_type="np")
|
||
return video # [B, T, H, W, 3]
|
||
|
||
|
||
def decode_latents_to_video(
|
||
latents_np: np.ndarray,
|
||
output_path: Path,
|
||
*,
|
||
fps: int = 16,
|
||
backend: DecodeBackend = "taehv",
|
||
vae_dir: Path | None = None,
|
||
z_dim: int | None = None,
|
||
taehv_checkpoint: Path | None = None,
|
||
torch_device: str = "auto",
|
||
) -> dict[str, Any]:
|
||
"""Decode latent video frames and export them as an MP4 file.
|
||
|
||
Parameters:
|
||
latents_np (np.ndarray): Latent video representation to decode.
|
||
output_path (Path): Destination path for the MP4 file.
|
||
fps (int): Output video frame rate.
|
||
backend (DecodeBackend): Decoder backend to use.
|
||
vae_dir (Path | None): Directory containing the full Wan VAE when using
|
||
the ``wan-vae`` backend.
|
||
z_dim (int | None): Latent channel count for TAEHV decoding.
|
||
taehv_checkpoint (Path | None): Optional TAEHV checkpoint path.
|
||
torch_device (str): PyTorch device selection for PyTorch-based decoding.
|
||
|
||
Returns:
|
||
dict[str, Any]: Decode time in seconds, backend name, output path, frame
|
||
count, and video resolution.
|
||
|
||
Raises:
|
||
ValueError: If the full VAE backend lacks ``vae_dir`` or the backend is
|
||
unknown.
|
||
"""
|
||
import time
|
||
|
||
from diffusers.utils import export_to_video
|
||
|
||
t0 = time.perf_counter()
|
||
if backend in ("taehv", "taehv-torch"):
|
||
c = latents_np.shape[1] if z_dim is None else z_dim
|
||
if backend == "taehv":
|
||
video = decode_latents_taehv_mlx(latents_np, z_dim=c, checkpoint_path=taehv_checkpoint)
|
||
else:
|
||
# torch TAEHV (regression / parity reference)
|
||
import torch
|
||
from fastvideo.third_party.taehv import TAEHV
|
||
|
||
ckpt = ensure_taehv_checkpoint(z_dim=c, checkpoint_path=taehv_checkpoint)
|
||
if torch_device == "auto":
|
||
torch_device = "mps" if torch.backends.mps.is_available() else "cpu"
|
||
dtype = torch.float16 if torch_device == "mps" else torch.float32
|
||
model = TAEHV(str(ckpt)).to(device=torch_device, dtype=dtype).eval()
|
||
lat = torch.from_numpy(latents_np).to(device=torch_device, dtype=dtype)
|
||
with torch.no_grad():
|
||
out = model.decode_video(lat.transpose(1, 2), parallel=True, show_progress_bar=False)
|
||
video = out[0].permute(0, 2, 3, 1).float().cpu().numpy()[None, ...]
|
||
# out is NTCHW -> need BTHWC; decode_video returns NTCHW for batch
|
||
if video.ndim == 4:
|
||
video = video[None]
|
||
elif backend == "wan-vae":
|
||
if vae_dir is None:
|
||
raise ValueError("vae_dir required for wan-vae backend")
|
||
video = decode_latents_wan_vae_torch(latents_np, vae_dir=vae_dir, device=torch_device)
|
||
else:
|
||
raise ValueError(f"unknown backend {backend}")
|
||
|
||
decode_s = time.perf_counter() - t0
|
||
output_path = Path(output_path)
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
# export_to_video expects list/array of frames HxWxC
|
||
frames = video[0]
|
||
frames = np.clip(frames, 0.0, 1.0)
|
||
export_to_video(frames, str(output_path), fps=fps)
|
||
if backend == "taehv":
|
||
cleanup_mlx()
|
||
else:
|
||
if backend == "taehv-torch":
|
||
del model, lat, out
|
||
cleanup_torch_mps()
|
||
return {
|
||
"decode_s": decode_s,
|
||
"backend": backend,
|
||
"output_path": str(output_path),
|
||
"num_frames": int(frames.shape[0]),
|
||
"resolution": f"{frames.shape[2]}x{frames.shape[1]}" if frames.ndim == 4 else None,
|
||
}
|