flux2-1d-tokenizer 0.1.0

Minimal inference library for a 1D image tokenizer over the frozen FLUX.2 VAE latent (1024 tokens x 16 bits, multi-resolution, tail-drop).

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 for honest limits.

Install

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

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.

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.