taehv_decode ¶

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)

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