Metadata-Version: 2.5
Name: canvit-nnx
Version: 0.3.1
Summary: CanViT (Canvas Vision Transformer) for JAX and Flax NNX
Project-URL: Repository, https://github.com/m2b3/CanViT
Project-URL: Homepage, https://m2b3.github.io/CanViT/
Project-URL: Paper, https://arxiv.org/abs/2603.22570
Project-URL: Checkpoints, https://huggingface.co/canvit
Author-email: Yohaï-Eliel Berreby <me@yberreby.com>
License-Expression: MIT
License-File: LICENSE
Requires-Python: >=3.12
Requires-Dist: canvit-core<0.3.0,>=0.2.1
Requires-Dist: flax>=0.12.10
Requires-Dist: jax>=0.11.2
Requires-Dist: numpy>=2.2
Requires-Dist: safetensors>=0.8.0
Provides-Extra: train
Requires-Dist: optax>=0.2.8; extra == 'train'
Requires-Dist: pycocotools>=2.0.11; extra == 'train'
Requires-Dist: tfrecord>=1.14.6; extra == 'train'
Requires-Dist: tyro>=1.0.3; extra == 'train'
Description-Content-Type: text/markdown

# canvit-nnx

The JAX and Flax NNX implementation of CanViT (Canvas Vision Transformer). Read the root
[backend documentation](https://github.com/m2b3/CanViT#backend-packages) for
the cross-backend API and tensor conventions.

## Install

```bash
uv add canvit-nnx
```

## Load a Hugging Face checkpoint

Load a pretrained model and an ADE20K probe from Hugging Face. Replace
`path/to/image.jpg` with your image path.

```python
import jax
import jax.numpy as jnp
from PIL import Image

from canvit_nnx import CanViTForSemanticSegmentation, Viewpoint, sample_at_viewpoint
from canvit_nnx.preprocess import preprocess

scene = jnp.asarray(preprocess(512)(Image.open("path/to/image.jpg").convert("RGB")))[None]
model = CanViTForSemanticSegmentation.from_pretrained_with_probe(
    pretrained_repo="canvit/canvitb16-add-vpe-pretrain-g128px-s512px-in21k-dv3b16-2026-02-02-nnx",
    probe_repo="canvit/probe-ade20k-40k-s512-c64-in21k-nnx",
)
model.eval()  # Flax NNX mutates the module and returns None.
viewpoint = Viewpoint.full_scene(batch_size=1)
state = model.init_state(batch_size=1, canvas_grid_size=64)
glimpse = sample_at_viewpoint(spatial=scene, viewpoint=viewpoint, glimpse_size_px=128)
logits, state = model(glimpse=glimpse, state=state, viewpoint=viewpoint)
jax.block_until_ready(logits)
print(logits.shape)  # [1, 64, 64, 150]
```

The model and probe remain separate checkpoints. They use NHWC arrays: scenes
and glimpses are `[B, H, W, 3]`, and the wrapper returns canvas logits as
`[B, G, G, num_classes]`. `Viewpoint` centers use `(row, col)` scene
coordinates in `[-1, 1]`; each scale is the crop's half-side.

The example uses the [ImageNet-21k CanViT-B checkpoint](https://huggingface.co/canvit/canvitb16-add-vpe-pretrain-g128px-s512px-in21k-dv3b16-2026-02-02-nnx)
and the [ADE20K probe](https://huggingface.co/canvit/probe-ade20k-40k-s512-c64-in21k-nnx).
The Hub lists [every NNX checkpoint](https://huggingface.co/models?other=nnx&author=canvit)
by its `nnx` tag.

Constructing a model without a checkpoint requires an explicit parameter RNG,
for example:

```python
from flax import nnx
from canvit_nnx import CanViT, CanViTConfig

model = CanViT(CanViTConfig(), rngs=nnx.Rngs(0))
```

## Fine-tune on ImageNet-1k

`canvit_nnx.specialize.in1k` fine-tunes CanViT-B for ImageNet-1k
classification with the paper's recipe, in bfloat16 with fused Pallas kernels
on TPU. It is new in canvit-nnx 0.3.0 and needs the `train` extra; its
[README](https://github.com/m2b3/CanViT/blob/main/canvit-nnx/canvit_nnx/specialize/in1k/README.md) describes the data and the
command.

## Detect objects on COCO

`canvit_nnx.detect` trains CanViT-B with a CenterNet head for COCO 2017 object
detection and scores average precision after each glimpse, on one host or a
multi-host TPU slice. It is new in canvit-nnx 0.3.0 and needs the `train`
extra; its [README](https://github.com/m2b3/CanViT/blob/main/canvit-nnx/canvit_nnx/detect/README.md) describes the data and the
commands.

## Develop from a checkout

```bash
cd canvit-nnx
uv sync
```

The checkout resolves `canvit-core` from its sibling directory. To convert a
PyTorch checkpoint and run the conversion checks from the repository root:

```bash
uv run --project canvit-nnx python -m tools.convert_checkpoints --help
```
