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