UniAR-RL BSQ visual tokenizer

Complete BSQ encoder with bit bases, multiscale input/output projections, DeepStack and mergers. Pixel patches -> binary visual tokens -> decoded visual features. Rendering images requires the separate SD3.5 decoder and VAE, which are excluded.

离散视觉 tokenizer：保留 BSQ、四个尺度的投影与特征还原；独立 SD3.5 渲染器与 VAE 不在本包。

Provenance

Extracted from ShareLab-SII/UniAR-RL at 6b02e4eee3d45b34f7f41e6218b6cc3c56332454. Native code revision and attribution are recorded in code_provenance.json and NOTICE. No language model is instantiated by the loader. The default attention backend is eager. This repository is an unofficial component export, not an official release from the original authors.

Tensors: 401; stored parameters/elements: 656,160,305.
Tensor payload: 1,312,321,000 bytes; dtype: ['BF16', 'I64'].
Main checkpoint SHA256: a07a8018053621883d0a988c7ba7d974a1941922d1cb1fe27d76c6420664ad0a.
Full source configuration is preserved separately as source_config.json; the custom root config documents this package only. This is an explicit Python loader, not a registered AutoModel or timm architecture.

Use

Install requirements.txt in a dedicated environment. Source files are vendored at a pinned upstream revision. The loader strictly checks every state key.

import sys, torch
from huggingface_hub import snapshot_download
path = snapshot_download("gatilin/uniar_rl_bsq_tokenizer")
sys.path.insert(0, path)
from component import load_model
model = load_model(path)  # CPU; BF16 for UniAR, FP32 for other components
from vision_encoder.modeling_vision_encoder import _convert_img_to_patch
pixels, grid = _convert_img_to_patch(torch.zeros(1, 3, 448, 448, dtype=torch.bfloat16))
with torch.inference_mode():
    main_bits, multiscale_bits = model(pixels, grid_thw=grid, bsq_only=True, direct_indices=True)
    bits = torch.cat([main_bits] + multiscale_bits, dim=-1)
    features, deepstack = model.bsq_indices_to_features(bits, with_merger=True)
# bits are binary vectors; no learned embedding table is required for BSQ.

Validation

validation.json records all checks: original source tensor bytes match, strict loading passes, 3 native equivalence cases pass, plus the useful-resolution smoke case 448_native_patch_encoding. Tests ran with BF16 on CPU. UniAR defaults to eager attention for exact BF16 bit reproducibility; the separate SDPA diagnostic is recorded in validation.json. Discrete token/bit outputs are checked exactly; floating outputs use tolerance documented in the validation script.

The reference loads original source ranges directly; the exported loader independently reads the assembled Safetensors checkpoint. Tests cover only the component APIs. They do not establish downstream understanding accuracy, text-to-image quality, or GPU/FlashAttention equivalence. Synthetic input images/tensors are used; reconstruction is lossy and no pixel-fidelity benchmark is claimed.

Reproduce extraction with extract_unified_components.py uniar_bsq (default outputs alongside the script); then run python validate_unified_components.py PATH. The byte comparison requires .download source chunks, which extraction creates but this repository intentionally does not upload.

License and attribution

The original weight and code licenses remain applicable. Preserve upstream author notices. See LICENSE, NOTICE, SOURCE_README.md, and source-file headers.

Only the Apache-licensed BSQ component is exported. The upstream SD3 pipeline has separate Stability AI terms and is excluded. BSQ bit bases are signed int64 buffers and remain unchanged; do not coerce them to float or replace them with a learned codebook. Upstream evaluation code requires BF16 in the BSQ projection path; overriding model dtype to FP32 is unsupported.

The exported loader is tested with the pinned Transformers 4.57.6 / tokenizers 0.22.2 / Hugging Face Hub 0.36.0 environment. Use this repository’s requirements in a dedicated environment; newer Qwen3-VL internals remove APIs called by the upstream UniAR implementation.