Metadata-Version: 2.4
Name: flux2-1d-tokenizer
Version: 0.1.1
Summary: Minimal inference library for a 1D image tokenizer over the frozen FLUX.2 VAE latent (1024 tokens x 16 bits, multi-resolution, tail-drop).
Author: fmorgens
Project-URL: Homepage, https://github.com/fmorgens/flux2-1d-tokenizer
Project-URL: Model, https://huggingface.co/fmorgens/flux2-1d-tokenizer-preview
Keywords: image-tokenizer,flux2,vae,discrete-tokens,vq,autoregressive,1d-tokenizer
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.1
Requires-Dist: diffusers>=0.32
Requires-Dist: huggingface_hub>=0.20
Requires-Dist: safetensors>=0.4
Requires-Dist: pillow
Requires-Dist: numpy
Dynamic: license-file

# flux2-1d-tokenizer

Minimal, fast, **inference-only** library for a **1D image tokenizer** that operates in the
frozen **FLUX.2 VAE** latent space. Turns an image into a short sequence of discrete tokens
(and back) — a compact visual vocabulary for downstream autoregressive models.

- **1024 tokens/image**, each one of **65,536 codes (16 bits)** → ~2 KB/image
- **Multi-resolution**: width & height each in `{256, 512, 768}`, any combination (incl.
  rectangular), both multiples of 16. Token identity is resolution-relative (normalized 2D RoPE)
- **Tail-drop**: tokens are emitted coarse-to-fine; keep the first *k* for a lower bit-rate
- Two decoder variants — **`base`** (faithful) and **`sharp`** (adversarially fine-tuned) —
  that emit **identical tokens**
- Pure `encode` / `decode`; no training code, no LPIPS

> ⚠️ **Preview / research release — not state of the art.** Quality sits around the "barely
> usable" line: flat/low-texture content reconstructs well, but faces, small text, and fine
> texture are soft, hard-capped by the 1024×16-bit rate. See the
> [model card](https://huggingface.co/fmorgens/flux2-1d-tokenizer-preview) for honest limits.

## Install

```bash
pip install flux2-1d-tokenizer      # the Python import name is  flux2_tokenizer
```

You also need access to the **FLUX.2 VAE** (`black-forest-labs/FLUX.2-VAE`, Apache-2.0) on
Hugging Face — it is downloaded at runtime (not bundled). `huggingface-cli login` if required.

## Usage

```python
import torch
from flux2_tokenizer import Flux2Tokenizer

tok = Flux2Tokenizer.from_pretrained()                  # sharp decoder (default)
# tok = Flux2Tokenizer.from_pretrained(variant="base")  # faithful/mean decoder
# base & sharp emit IDENTICAL tokens — they differ only in how the decoder renders them.

x = Flux2Tokenizer.load_image("photo.jpg", size=512)    # (1,3,512,512) in [-1,1]
tokens = tok.encode(x)                                   # (1,1024) int64, values 0..65535
recon  = tok.decode(tokens, width=512, height=512)       # (1,3,512,512) in [-1,1]

recon_lowrate = tok.decode(tokens[:, :256], 512, 512)    # tail-drop: keep the first 256 tokens
recon_768     = tok.decode(tokens, 768, 768)             # cross-resolution: decode at a DIFFERENT size
recon_wide    = tok.decode(tokens, 768, 256)             #   ...or a different aspect ratio
```

> **Cross-resolution decode** works because token identity is resolution-relative (normalized 2D
> RoPE): the same tokens re-render at any target `{256,512,768}²` size, not just the source size.
> It's an **emergent side-effect** — the model was never trained to resize — so treat it as a
> curiosity, not a reliable resampler.

### Bulk tokenization (data prep)

`encode` is a pure **batch → ids** function. Bring your own loading (a `DataLoader` with
workers, wherever your data lives) and your own output sink.

```python
import numpy as np
from torch.utils.data import DataLoader

# your dataset yields (3,H,W) float tensors in [-1,1]; same H,W within a batch
loader = DataLoader(my_dataset, batch_size=64, num_workers=8, pin_memory=True)
for i, batch in enumerate(loader):
    ids = tok.encode(batch.cuda())                       # (B,1024) int64 on GPU
    np.save(f"shard_{i}.npy", ids.cpu().numpy().astype(np.uint16))   # store however you like
```

- **Storage dtype `uint16`** — codes are 0..65535 (`int16` overflows). `encode` returns int64.
- **Batch size matters.** The frozen VAE encode dominates each call and batches well — push it
  up toward VRAM for throughput.
- **Encoding is variant-independent** (same encoder + codebook). Encode once; decode with
  `sharp` for the best rendering.

## Variants

| variant | decoder | use it for |
|---|---|---|
| `sharp` (default) | adversarially fine-tuned (latent PatchGAN) | best perceived sharpness; commits to plausible detail at low *k* instead of blurring. Detail is *plausible, not guaranteed faithful* |
| `base` | latent-MSE only | faithful/mean reconstruction; softer |

The encoder and codebook are frozen between the two, so **the token ids are byte-identical** —
you can encode once and decode with either.

## Notes

- **Quality ceiling** is set by the 1024×16-bit rate. More tokens would mean more detail; the
  codebook is already ~fully used, so per-token bits aren't the lever.
- Operates entirely in FLUX.2 VAE latent space (no VAE weights are stored in this package or the
  model checkpoints).

## License & attribution

- This library and the tokenizer weights are **Apache-2.0**.
- Built on the **FLUX.2 VAE** by **Black Forest Labs** (Apache-2.0), downloaded at runtime and
  **not redistributed** here.
- **Unofficial community project. Not affiliated with or endorsed by Black Forest Labs.
  FLUX and FLUX.2 are trademarks of Black Forest Labs.**
