Metadata-Version: 2.4
Name: inkrs
Version: 0.2.1
Classifier: Development Status :: 4 - Beta
Classifier: Environment :: GPU :: NVIDIA CUDA :: 12
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: MacOS
Classifier: Operating System :: Microsoft :: Windows
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Programming Language :: Rust
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Image Processing
Classifier: Typing :: Typed
Requires-Dist: numpy>=1.22
Requires-Dist: nvidia-cudnn-cu12>=9.1,<10 ; (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'cu12') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'cu12')
Requires-Dist: nvidia-cublas-cu12>=12.1,<13 ; (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'cu12') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'cu12')
Requires-Dist: pytest>=7 ; extra == 'test'
Provides-Extra: cu12
Provides-Extra: test
License-File: LICENSE-MIT
License-File: LICENSE-APACHE
License-File: THIRD_PARTY_NOTICES
Summary: GPU-resident ink inference for villa's flat ink readers, from Python
Keywords: ink-detection,vesuvius-challenge,herculaneum,papyrus,cuda,metal,gpu,inference
Author-email: Praveen Perera <me@praveenperera.com>
License-Expression: MIT OR Apache-2.0
Requires-Python: >=3.9
Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
Project-URL: Changelog, https://github.com/praveenperera/inkrs/blob/master/CHANGELOG.md
Project-URL: Documentation, https://github.com/praveenperera/inkrs/blob/master/python/README.md
Project-URL: Repository, https://github.com/praveenperera/inkrs

# inkrs (Python)

inkrs runs villa's flat ink readers on the GPU, from Python: `hybrid_3d2d-seed42` and the readers with its architecture (`reader-v2`, Reader V3, ink_9um seed 43, dense_native, TRACE v3B and Rectina M1, M2 and M4). Give it the same input, layers, direction and stride as villa's `vesuvius.ink_detection.inference.infer` and it returns the uint8 maps that `infer` writes, or writes villa's TIFFs itself. `inkrs.Engine` also runs the v8-in exports, `inkrs.hecate` runs Hecate 9.6 µm and `inkrs.ink3d` runs villa's 3D ink models, each with the CLI's code. The model stays loaded on the GPU between calls, and `Engine.sweep` runs every depth and both directions of a segment in one pass.

## Install

On Linux x86-64 or Windows x86-64:

```
pip install "inkrs[cu12]"
```

On an Apple silicon Mac, `pip install inkrs` (see [macOS](#macos)).

The Linux wheel is for CPython 3.9 and newer, on glibc 2.28 or newer. It doesn't need Rust or a CUDA toolkit, because the kernels ship inside it, compiled for the common NVIDIA GPUs and as PTX. You do need:

- **An NVIDIA GPU of compute capability 7.5 or above** (T4 and RTX 20, RTX 30, 40 and 50, A100, H100) **and a driver of r525 or newer**.
- **cuDNN 9 and cuBLAS 12.** The `cu12` extra installs them from pip (`nvidia-cudnn-cu12`, `nvidia-cublas-cu12`). Without it, inkrs uses the copies already on the system, or the ones a CUDA 12 build of PyTorch installed. A CUDA 13 PyTorch ships cuBLAS 13 under a name inkrs doesn't look for, so there add `pip install nvidia-cublas-cu12`.

The wheel has three PTX tiers: sm_75 for Turing, sm_80 for every GPU from 8.0 up, and sm_120 (r570 or newer) for the RTX 50 series. The engine loads the highest tier the GPU runs, and `INKRS_PTX_TIER=sm80` or `sm75` forces a lower one. Each tier also ships compiled for the GPU families it serves: sm_75; sm_80, sm_86, sm_89, sm_90 and sm_100; and sm_120. On those GPUs the engine loads the compiled kernels and skips the driver's PTX compile, which takes over a second on a machine's first run. A GPU outside those families (compute capability 11.x), a tier forced below the GPU's own, or a driver that refuses the compiled kernels gets the PTX instead, the last with a warning on stderr. `load_report["kernels"]` says which loaded. `inkrs.build_info()` reports each tier and its minimum driver, and the families with compiled kernels. Hecate in bf16 and the 3D ink U-Net need 8.0 or above, so they're refused on Turing. It's tested on an RTX 5090 (every tier), a Tesla T4 (sm75), and an A100, an A10, an L4 and an RTX 4090 (sm80, the readers).

### Which cuDNN and cuBLAS it loads

inkrs takes the system copies first (the linker path, `LD_LIBRARY_PATH`, or a copy already loaded, such as PyTorch's), then the pip packages. `INKRS_CUDA_LIBS=pip` loads the pip copies even when the system has its own, and `INKRS_CUDA_LIBS=cudnn` (or `cublas`) does that for one of them. The engine tries the unversioned name (`libcublas.so`) first, so a system copy under that name still wins; when it would, the forced load raises instead of running on the other copy. If a library is found nowhere, or a pip copy fails to load, `inkrs.Engine` raises `inkrs.CudaLibraryError` naming the package to install. `inkrs.load_cuda_libraries()` and `eng.load_report["cuda_libraries"]` show the file the engine calls for each.

### Windows

The x86-64 Windows wheel is for CPython 3.9 and newer too. It imports only Python's stable ABI DLL and Windows system DLLs, but it hasn't been run on Windows itself yet.

There the libraries are `cudnn64_9.dll` and `cublas64_12.dll`, and the `cu12` extra installs them as well. When the pip packages are installed, inkrs loads them from their `bin` directories before the engine starts, so they win over a CUDA toolkit on `PATH`, and `INKRS_CUDA_LIBS=pip` only makes them required. NVRTC, which cuDNN uses for some of its engines, loads from the `nvidia-cuda-nvrtc-cu12` package's `bin` directory when that's installed. Without the pip packages, the engine looks in the executable's directory, System32, each absolute `PATH` entry, then `%CUDA_PATH%\bin\x64` and `%CUDA_PATH%\bin`.

It never loads a library from the current directory. It raises `inkrs.CudaLibraryError` when a copy there or behind a relative `PATH` entry would load instead, or when a file it tries first, such as a `cudnn.dll` on `PATH`, would load a second copy beside the pip one. The CLI also takes the current directory and `PATH` out of its own DLL search at startup. The Python package can't do that for your interpreter, so a DLL written into the current directory after the engine's checks, under a name it tries first, could still load.

### macOS

```
pip install inkrs
```

The arm64 wheel installs on any Apple silicon Mac, but the engine needs macOS 26 or newer, and on an older macOS it raises `inkrs.UnsupportedError` naming the version. There the engine runs on Metal. It needs no CUDA libraries and no Xcode, because the Metal shaders ship as source and compile when the engine opens. Pass `device=0`, `"mps"`, `"mps:0"` or a torch MPS device; the engine uses the system's default GPU. `algo` does nothing there, and `inkrs.build_info()` reports `"backend": "metal"` and no PTX or compiled kernels. Only the M4 Pro has been tested.

Experimental, and reader-v2 only: the engine can run most batches' 2D network on the Neural Engine through Core ML while the GPU runs every batch's 3D stem. Build the Core ML bodies once with `metal/coreml_body.py`, set `INKRS_COREML_ANE` to the `fp16` `.mlmodelc` it prints, and run at `--batch-size 8` (`batch_size=8`). The engine refuses a body built from another checkpoint, batch or precision. The maps stay within villa's parity limits and are the same on every run, but differ from the GPU-only maps by up to a few 1/255 steps. `src/metal/coreml.rs` lists the tuning variables, including an opt-in CPU lane.

## Build from source

pip builds the source distribution where no wheel fits, which needs Rust but no CUDA toolkit, since the kernels ship compiled and as PTX. From a checkout, `python/build-wheel.sh` builds the Linux x86-64 wheel (manylinux_2_28, through zig) or, on a Mac, the arm64 wheel, into `target/wheels`. It needs Rust and uv. On Windows, run `maturin build --release` in `python/`.

The package's version is the one in `python/Cargo.toml`, which CI keeps equal to the engine's in `Cargo.toml`. The Python code is linted with `uvx ruff@0.17.4 check python`.

## Weights

inkrs reads a safetensors export of villa's checkpoint. Make it once, on the CPU, where villa and torch are installed:

```
python -m inkrs.export_weights step-075000.pth weights/
```

That writes `weights/hybrid_3d2d-seed42.safetensors` and its `.json` sidecar, and `inkrs.Engine` takes the `.safetensors` file directly.

It takes the `.pth` too. inkrs then looks for an export of those exact bytes (by SHA-256) in `weights_dir` (default `$INKRS_WEIGHTS_DIR`, else `./weights`), then in `~/.cache/inkrs/weights/<sha256>/` (or under `$XDG_CACHE_HOME`, or `$INKRS_CACHE_DIR`; `%LOCALAPPDATA%\inkrs\weights\<sha256>\` on Windows). If neither has one and torch and villa are installed, as in a villa environment, it exports the checkpoint into that cache once and reuses it from then on. Without them it raises `inkrs.UnsupportedError` with the command above. A checkpoint the exporter refuses also raises `inkrs.UnsupportedError`: one that isn't a flat model, has non-float32 weights, or whose keys or tensor shapes don't match the model its config builds. The message names the missing and unexpected keys.

## Use

```python
import inkrs

eng = inkrs.Engine("weights/hybrid_3d2d-seed42.safetensors", device=0)

# one villa `infer` call: --layer-start 25 --layer-end 42 --direction both
maps = eng.infer("segment.zarr", layer_start=25, layer_end=42, direction="both")
maps["forward"], maps["reverse"]          # uint8 [H, W]; forward is the blind side, reverse the reading side

# or write villa's files: out.tif and out_reverse.tif
eng.infer("segment.zarr", layer_start=25, layer_end=42, direction="both", output_tiff="out.tif")

# every depth d = -4..4, both directions, in one pass (layers from Z//2 + d - 8)
sweep = eng.sweep("segment.zarr", range(-4, 5))
sweep[(0, "reverse")]

# an in-memory volume: numpy or torch, uint8 [Z, H, W]; a C-contiguous array is read in place
maps = eng.infer(volume, layer_start=25, layer_end=42, occupancy="segment.zarr")
```

`occupancy` controls villa's low-resolution empty-tile skip. For a zarr path the default does what villa does. For an in-memory volume, pass the zarr group it came from (or one pyramid level as an array) to get villa's covered pixels. Pass nothing and only the per-tile skip runs, as villa does for a bare array.

`eng.load_report` and `eng.last_report` hold the load and per-call timings. The GIL is released while the engine loads and runs, and one call runs at a time per engine.

[API.md](API.md) has the full contract, including how it maps onto villa's `infer`.

## Limits

Input, a model or a GPU outside these limits raises `inkrs.UnsupportedError`, a subclass of `inkrs.InkrsError`. So a caller with a PyTorch path can fall back on `UnsupportedError` and treat any other `InkrsError` as a failure.

- In `Engine`, the `hybrid_3d2d-seed42` architecture and contract (its own weights or another reader's of the same architecture) or a v8-in export. Hecate runs through `inkrs.hecate` and 3D ink through `inkrs.ink3d`.
- On macOS, the villa readers, Hecate and the v8-in exports only. `inkrs.ink3d` and `tta=8` raise there.
- uint8 input only.
- The Hann blend only.
- The unpadded tile phase only.
- A volume under 4 GiB, with both spatial dimensions at least 128. That's measured on the array passed to the call, so a 17-layer window can cover up to about 15,800 px square.
- A GPU the embedded PTX serves, with a driver new enough to load it, or on macOS a GPU with Metal 4.

The maps match villa eager's to within 2/255, with the same covered pixels. Two processes can differ from each other by 1/255 on a few pixels when `algo="find"` (the default) times cuDNN's algorithms and picks differently. The timed picks are cached on disk (see [API.md](API.md)), so later processes on the same setup reuse them; `algo="heuristic"` makes the choice repeatable without the cache.

Either the Apache License, Version 2.0 or the MIT licence, at your option. The embedded kernels include NVIDIA libdevice code; see `THIRD_PARTY_NOTICES`.

