192 lines
6.9 KiB
Python
192 lines
6.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Optional TAEHV decode helpers for Apple Silicon FastWan experiments.
|
|
|
|
The TAEHV module itself is vendored at ``fastvideo/third_party/taehv`` (MIT,
|
|
madebyollin/taehv), so no source code is downloaded or executed at runtime.
|
|
Only the ``taew2_1.pth`` checkpoint is fetched on demand, and its sha256 is
|
|
verified before use.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import importlib.util
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
TAEW2_1_CHECKPOINT_URL = "https://raw.githubusercontent.com/madebyollin/taehv/main/taew2_1.pth"
|
|
# sha256 of the upstream taew2_1.pth this module was validated against
|
|
# (fetched 2026-07-02). If upstream publishes a new checkpoint, revalidate the
|
|
# decode path and update this pin.
|
|
TAEW2_1_CHECKPOINT_SHA256 = "d26151e76cdc2c9424bef988de874b33d9a53f30ef3060cd556c429c469c797e"
|
|
# Wan2.2 5B (z_dim=48) — see madebyollin/taehv taew2_2.pth; prefer
|
|
# ``fastvideo.mlx_runtime.wan_vae.ensure_taehv_checkpoint(z_dim=48)`` for new code.
|
|
TAEW2_2_CHECKPOINT_URL = ("https://raw.githubusercontent.com/madebyollin/taehv/"
|
|
"563f40bdc820ed86bcad72ea515ee48f06bd22ec/taew2_2.pth")
|
|
|
|
|
|
def _default_cache_dir() -> Path:
|
|
"""Return the default directory used to cache TAEHV checkpoints.
|
|
|
|
Returns:
|
|
Path: The TAEHV checkpoint cache directory under the user's home directory.
|
|
"""
|
|
return Path.home() / ".cache" / "fastvideo" / "taehv"
|
|
|
|
|
|
def _sha256(path: Path) -> str:
|
|
"""
|
|
Compute the SHA-256 digest of a file.
|
|
|
|
Parameters:
|
|
path (Path): The file whose contents are hashed.
|
|
|
|
Returns:
|
|
str: The file's SHA-256 digest in hexadecimal form.
|
|
"""
|
|
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) -> None:
|
|
"""Verify that a TAEW2.1 checkpoint matches the expected SHA-256 digest.
|
|
|
|
Parameters:
|
|
path (Path): Path to the checkpoint file.
|
|
|
|
Raises:
|
|
RuntimeError: If the checkpoint digest does not match the expected value.
|
|
"""
|
|
actual = _sha256(path)
|
|
if actual != TAEW2_1_CHECKPOINT_SHA256:
|
|
raise RuntimeError(f"TAEHV checkpoint at {path} failed sha256 verification "
|
|
f"(expected {TAEW2_1_CHECKPOINT_SHA256}, got {actual}). "
|
|
"Delete the file to re-download it, or pass --taehv-checkpoint-path "
|
|
"pointing at a checkpoint you trust.")
|
|
|
|
|
|
def ensure_taew2_1_checkpoint(checkpoint_path: Path | None = None) -> Path:
|
|
"""
|
|
Ensure the TAEW2.1 checkpoint is available locally.
|
|
|
|
A caller-provided path is treated as trusted and is only checked for existence.
|
|
The module-managed cached checkpoint is verified against the pinned SHA-256 digest
|
|
after downloading or before reuse.
|
|
|
|
Parameters:
|
|
checkpoint_path (Path | None): Optional path to a caller-provided checkpoint.
|
|
|
|
Returns:
|
|
Path: The available checkpoint path.
|
|
|
|
Raises:
|
|
FileNotFoundError: If a caller-provided checkpoint does not exist.
|
|
RuntimeError: If a module-managed checkpoint fails verification.
|
|
"""
|
|
if checkpoint_path is not None:
|
|
if not checkpoint_path.exists():
|
|
raise FileNotFoundError(f"TAEHV checkpoint not found: {checkpoint_path}")
|
|
return checkpoint_path
|
|
|
|
checkpoint_path = _default_cache_dir() / "taew2_1.pth"
|
|
if not checkpoint_path.exists():
|
|
checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
|
|
print(f"Downloading {TAEW2_1_CHECKPOINT_URL} -> {checkpoint_path}")
|
|
import socket
|
|
import tempfile
|
|
# Download to a temporary file, verify, then atomically rename.
|
|
with tempfile.NamedTemporaryFile(
|
|
mode="wb",
|
|
dir=checkpoint_path.parent,
|
|
prefix=".tmp_taew2_1_",
|
|
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(
|
|
TAEW2_1_CHECKPOINT_URL,
|
|
tmp_path, # noqa: S310 - pinned public artifact, hash-verified below.
|
|
)
|
|
finally:
|
|
socket.setdefaulttimeout(old_timeout)
|
|
_verify_checkpoint(tmp_path)
|
|
tmp_path.replace(checkpoint_path)
|
|
except Exception:
|
|
tmp_path.unlink(missing_ok=True)
|
|
raise
|
|
else:
|
|
_verify_checkpoint(checkpoint_path)
|
|
return checkpoint_path
|
|
|
|
|
|
def _load_taehv_class(source_path: Path | None):
|
|
"""Load the TAEHV class from the vendored implementation or a local source override.
|
|
|
|
Parameters:
|
|
source_path (Path | None): Path to a local Python file defining `TAEHV`; `None` selects the vendored implementation.
|
|
|
|
Returns:
|
|
The loaded `TAEHV` class.
|
|
|
|
Raises:
|
|
RuntimeError: If the specified source cannot be loaded.
|
|
"""
|
|
if source_path is None:
|
|
from fastvideo.third_party.taehv import TAEHV
|
|
|
|
return TAEHV
|
|
# Explicit local override for experimenting with a modified TAEHV; this is
|
|
# a user-supplied file on disk, never something this module downloads.
|
|
spec = importlib.util.spec_from_file_location("fastvideo_external_taehv", source_path)
|
|
if spec is None or spec.loader is None:
|
|
raise RuntimeError(f"Could not load TAEHV source from {source_path}")
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module.TAEHV
|
|
|
|
|
|
def decode_latents_to_video_taehv(
|
|
*,
|
|
latents_np: np.ndarray,
|
|
output_path: Path,
|
|
fps: int,
|
|
device,
|
|
dtype,
|
|
parallel: bool,
|
|
source_path: Path | None = None,
|
|
checkpoint_path: Path | None = None,
|
|
) -> None:
|
|
"""Decode Wan/FastWan diffusion latents with TAEW2.1 and export MP4.
|
|
|
|
TAEHV's Wan wrapper expects the diffusion latents directly, without applying
|
|
the standard Wan VAE's `latents_mean` / `latents_std` shift.
|
|
"""
|
|
import torch
|
|
from diffusers.utils import export_to_video
|
|
|
|
checkpoint_path = ensure_taew2_1_checkpoint(checkpoint_path)
|
|
TAEHV = _load_taehv_class(source_path)
|
|
taehv = TAEHV(str(checkpoint_path)).to(device=device, dtype=dtype)
|
|
taehv.eval()
|
|
|
|
latents = torch.from_numpy(latents_np).to(device=device, dtype=dtype)
|
|
with torch.no_grad():
|
|
video_ntchw = taehv.decode_video(
|
|
latents.transpose(1, 2),
|
|
parallel=parallel,
|
|
show_progress_bar=False,
|
|
)
|
|
video = video_ntchw.transpose(1, 2)
|
|
video_np = video[0].permute(1, 2, 3, 0).float().cpu().numpy()
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
export_to_video(video_np, str(output_path), fps=fps)
|