Metadata-Version: 2.4
Name: fastvae
Version: 0.1.0
Summary: A fast video VAE decoder for consumer Blackwell (sm_120).
Author: Vimanyu Taneja
License-Expression: MIT
Project-URL: Homepage, https://github.com/Occipital-Labs/fastvae
Project-URL: Repository, https://github.com/Occipital-Labs/fastvae
Project-URL: Issues, https://github.com/Occipital-Labs/fastvae/issues
Keywords: vae,video,diffusion,triton,cuda,blackwell,wan
Classifier: Development Status :: 4 - Beta
Classifier: Environment :: GPU :: NVIDIA CUDA
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.6
Requires-Dist: triton>=3.0; platform_system == "Linux"
Provides-Extra: bench
Requires-Dist: diffusers>=0.36; extra == "bench"
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Requires-Dist: diffusers>=0.36; extra == "dev"
Dynamic: license-file

# fastvae

A fast video VAE decoder for consumer Blackwell (`sm_120`). Decodes at full resolution
without tiling, 2.4x faster than the reference Wan VAE decoder in eager PyTorch.

## Install

```bash
pip install fastvae
```

Pure Python. The kernels are Triton and compile on first call.

## Use

Patch an existing diffusers pipeline:

```python
from fastvae import accelerate_vae

accelerate_vae(pipe)   # pipe.vae.decode now runs through fastvae
```

Or call the decoder directly:

```python
from fastvae import FastVAEDecoder

dec = FastVAEDecoder.from_vae(vae)
frames = dec.decode(latents)          # (B, 3, T, H, W) in [-1, 1]
```

For ComfyUI, install fastvae into ComfyUI's Python environment, then register the node:

```bash
mkdir -p ComfyUI/custom_nodes/fastvae
echo 'from fastvae.integrations.comfy_node import *' > ComfyUI/custom_nodes/fastvae/__init__.py
```

This adds `FastVAEDecode` under *latent/video*, a drop-in replacement for `VAEDecode`.

## Benchmarks

RTX PRO 6000 Blackwell (`sm_120`), bf16, Wan 2.1 VAE. Median of 3 runs.

| case | eager | fastvae | speedup | peak VRAM |
|---|--:|--:|--:|--:|
| 480p / 81f | 2817 ms | 1171 ms | 2.41x | 6.82 GiB |
| 480p / 161f | 5613 ms | 2340 ms | 2.40x | 7.18 GiB |
| 720p / 81f | 6570 ms | 2782 ms | 2.36x | 15.21 GiB |
| 720p / 161f | 13069 ms | 5419 ms | 2.41x | 16.06 GiB |

Every peak fits a 32 GB RTX 5090 without tiling. `torch.compile` reaches 1.74x on the
shortest case, tiled decode is slower than eager, and the Wan 2.2 VAE speedup is 1.7x.

Reproduce with `python -m fastvae.bench.harness --vae <pipeline-dir>`. Raw output is in
`benchmarks/`.

## How it works

Four changes to the reference decoder:

1. **Fused norm and activation.** `WanRMS_norm` plus SiLU is about seven kernels in eager;
   one Triton kernel does it in a single pass.
2. **Channels-last-3d throughout,** which turns the reference's per-upsample layout
   conversions into free views.
3. **Deferred residual adds and conv biases,** absorbed into the next fused norm.
4. **Temporal chunking instead of spatial tiling.** Causal conv caching is exact for any
   chunk length, so chunk size trades peak VRAM for launch count without changing the
   output.

Ops that only move data are bitwise identical to the reference, asserted in tests. The
channel-wise L2 norm is not, since its fp32 reduction order cannot be matched, which
diffuses to a max absolute error of 4e-2 on a `[-1, 1]` output, against 5.2e-2 for
`torch.compile` and 2.0e-1 for tiled decode.

## Supported models

`AutoencoderKLWan` (Wan 2.1 and Wan 2.2). Support is an adapter that rebuilds a reference
module tree from `fastvae.ops` primitives. See `fastvae/models/wan.py` and
`register_adapter`.

## Development

```bash
pytest tests/                                    # exactness suite
FASTVAE_TEST_MODEL=<pipeline-dir> pytest tests/  # plus real-checkpoint tests
python -m fastvae.bench.harness --vae <pipeline-dir>
python -m fastvae.bench.m0_profile --model <pipeline-dir>
```
