File size: 4,193 Bytes
b4f5f26 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | ---
license: apache-2.0
---
# Image Autoencoder (8x, 32 channels, flow-matching decoder)
Inference-only package for an image autoencoder intended as a latent space for diffusion models.
- **Compression:** 8x spatially, 32 latent channels — an `H x W` RGB image becomes a `32 x H/8 x W/8` latent.
- **Encoder:** deterministic transformer with local window attention (2D RoPE), works at any resolution
whose sides are multiples of 8.
- **Decoder:** conditional flow-matching model (x0-prediction). It starts from Gaussian noise and is
integrated in `N` steps (default **4**) conditioned on the latent.
- **Normalization built in:** `encode()` returns latents already normalized per channel
(`(z - mean) / std`, roughly zero mean / unit variance); `decode()` undoes it. Use
`encode_raw()` / `decode_raw()` for unnormalized latents.
## Install
```bash
pip install torch safetensors huggingface_hub numpy pillow
# optional: triton (CUDA) for the fused window-attention kernel; otherwise PyTorch SDPA is used
```
The model code lives in the `ae/` folder of this repo:
```bash
huggingface-cli download Muinez/f8c32-ae --local-dir f8c32-ae
export PYTHONPATH=$PWD/f8c32-ae:$PYTHONPATH
```
## Usage
```python
import torch
from ae import AutoEncoder
model = AutoEncoder.from_pretrained("Muinez/f8c32-ae", device="cuda") # bf16 compute on CUDA
x = ... # (B, 3, H, W) float in [-1, 1], H and W multiples of 8
z = model.encode(x) # (B, 32, H/8, W/8), normalized latents
y = model.decode(z, steps=4) # (B, 3, H, W), approximately in [-1, 1]
y = y.clamp(-1, 1)
# reproducible decoding
g = torch.Generator("cuda").manual_seed(0)
y = model.decode(z, steps=4, generator=g)
# raw (unnormalized) latents
z_raw = model.encode_raw(x)
y = model.decode_raw(z_raw, steps=4)
```
### Batch of different resolutions / aspect ratios
`encode` / `decode` (and the `_raw` variants) also take a **list** of tensors of different
sizes and return a list. The whole list goes through the model in one forward pass: each
image keeps its own size (no resizing or cropping), is cut into its own grid of 8x8 patch
tokens, and the grids are packed into one sequence batch padded only up to the largest image
area. Attention never crosses into other images or padding, so results match encoding each
image on its own. Only the decoder's small convolutional upsampling head runs per group of
equal-shape images.
```python
from PIL import Image
import numpy as np
def load(path, max_side=1024):
im = Image.open(path).convert("RGB")
s = min(1.0, max_side / max(im.size))
w, h = (int(im.width * s) // 8 * 8, int(im.height * s) // 8 * 8) # sides: multiples of 8
im = im.resize((w, h), Image.LANCZOS)
return torch.from_numpy(np.asarray(im)).permute(2, 0, 1) # uint8 (3, H, W)
images = [load("portrait.jpg"), load("landscape.png"), load("square.webp")] # any sizes
latents = model.encode(images) # list of (32, H_i/8, W_i/8), normalized
recons = model.decode(latents, steps=4) # list of (3, H_i, W_i)
```
Images may be uint8 `[0, 255]` or float `[-1, 1]`.
See `example.py` for a complete round trip (`python example.py input.png recon.png`).
### Notes
- `steps` trades speed for detail: 1 step is fastest and slightly softer, 4 is the default.
Because the decoder samples from noise, different seeds give slightly different fine detail.
- `dtype` in `from_pretrained` sets the compute dtype (default: bfloat16 on CUDA, float32 on CPU).
Encoder weights are cast to it; the decoder keeps float32 weights and runs under autocast.
- On CUDA with Triton installed, a fused window-attention kernel is used for any image sizes;
otherwise (CPU, no Triton) an equivalent masked SDPA path is used. Force SDPA with
`import ae.window_attention as wa; wa.BACKEND = "sdpa"`.
## Files
- `config.json` — architecture and latent normalization statistics
- `model.safetensors` — encoder and decoder weights (float32)
- `ae/` — model code (`AutoEncoder`, `AEConfig`)
- `example.py` — round-trip example
- `LICENSE` — Apache License 2.0
## License
Apache License 2.0 — see `LICENSE`.
|