wan_vae ¶

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_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,
    }

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 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

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