Module models 

Registry of deep-learning QSM models and their downloadable ONNX weights.

Query the catalog with models::all_models / models::find_model. A host can supply weights from memory with models::provide_weights (the only source in WASM); native hosts can also use $QSM_MODEL_DIR, the cache, or fetch on use with the download feature. The onnx feature runs them via the pure-Rust tract engine (native and WASM). Deep-learning model registry and weight management.

Several QSM stages have deep-learning implementations (BFRnet, QSMnet, xQSM, χ-sepnet, …). QSM-Core runs them via ONNX rather than bundling a Python/PyTorch/TensorFlow runtime. The trained weights are not vendored in this crate — they are fetched on first use and cached on disk.

This module is the single source of truth for which models exist and where/how to obtain their weights. It is split into three layers:

• Registry (always compiled). [ModelSpec] / [WeightFile] describe each model and its weight files (URL, SHA-256, size, license). [all_models] / [find_model] query the table. This layer has no heavy dependencies and is safe to expose through the WASM bindings so a JavaScript host (e.g. qsmbly) can discover download URLs and fetch weights itself.
• Download & cache (download feature). [download::ensure_model] fetches any missing weight files over HTTP, verifies their SHA-256, and stores them in a local cache. This is the path a native host (e.g. QSMxT) uses to “download on use”. WASM hosts skip this layer.
• Inference (onnx feature). Runs an ONNX graph with the pure-Rust tract engine. The entry points take the model as a byte buffer ([onnx::OnnxModel::load]) rather than a path, so the same inference code runs natively and in WASM. A native host feeds bytes from the download cache; a WASM host (e.g. qsmbly) fetches the weights in JavaScript and supplies them with [provide_weights], after which the same pipeline dispatchers run as natively.

The registry stays runtime-agnostic: tract is the default portable engine, but a host is free to read a model’s URLs from the registry and run it any other way.

Weight resolution order

Every pipeline dispatcher that runs a network (crate::pipeline: inversion, BFRnet, iQFM, iQSM, HD-BET/RS2-Net masking, χ-separation, …) and every host that loads a model through [primary_weight] / [weights] resolves each weight file in this order:

1. Bytes the host supplied with [provide_weights] (checked against the registry’s SHA-256 when they were supplied). This is the only source on wasm32-unknown-unknown, where there is no file system or network: a browser host fetches the files itself and hands them over, and the same dispatchers run.
2. $QSM_MODEL_DIR/<file> — an explicit directory of local weight files (bring-your-own-weights; also how gated models like χ-sepnet are supplied). Native only.
3. The on-disk cache ([cache_dir]). Native only.
4. Download from [WeightFile::url] into the cache (download feature). Native only.

On wasm32-unknown-unknown steps 2–4 are not compiled; a file that was not provided is an error naming the file and [provide_weights].

Memory

Every model in the registry carries [ModelSpec::memory], measured with tract, so a host can estimate a run’s memory before it starts:

• [ModelSpec::resident_bytes]: the parsed model, held for as long as it is loaded.
• [ModelSpec::inference_bytes]: one inference at any input shape, above the resident model (native): a fixed part plus bytes per voxel of the network’s input.
• [ModelSpec::peak_bytes]: loading the model and running it once, for a native process or a wasm32 heap ([Heap]).
• crate::bet::hd_bet_memory_bytes: the whole HD-BET step on a scan, from its field of view. HD-BET resamples to 1 mm, so its network input follows the head’s size in mm, not the acquired matrix (crate::bet::hd_bet_network_input).

The bytes of the weight files themselves ([WeightFile::bytes]) are not included, nor are a wrapper’s own buffers or further tiles in flight (for those, see [ModelSpec::tile_bytes_per_voxel]).

How they were measured. Natively, examples/inference_memory.rs runs each network single-threaded under a counting global allocator, at 1–8 input shapes per network from 32³ to 192×256×128 (to 256³ for MindGrab and 176×256×256 for HD-BET; the fixed-patch networks at their patch). Per shape it takes the larger of a run that plans inside ([onnx::OnnxModel::run]) and a plan built once and then run ([onnx::OnnxModel::plan_for]), since wrappers use both. The figures are the least-squares line through those peaks, with the intercept raised to cover every one. From 64³ up they sit 0–7 % above the measured peaks (0–1.3 % from 96³ up); at 32³ up to 28 % above (at most 25 MB), where the fixed part dominates. MindGrab is fitted from 128³ up: it always runs at 256³, and smaller inputs fall well below the line. Threads (the parallel feature) add 0.1–0.6 % at 8–16 threads. Peak RSS adds the program, the weight file’s bytes and the allocator’s slack to these.

In WASM, the browser_memory::net_* cases in tests/wasm_pipeline.rs run the same inferences under Node (wasm32-unknown-unknown, SIMD128, single-threaded) at one or two shapes per network that fit the 4 GiB heap, and read the size of the linear memory, which only grows. The WASM figures are a line through those points. The WASM heap grows by 0.8–1.6 times the native peak of the same run (median 1.2). The ratio varies with the network and the shape, because the allocator cannot always reuse a freed tensor’s block for the next, larger one. Two runs of the same graph (QSMnet and QSMnet+) differed by 7 %, so budget a margin on top.

src/models/memory_calibration.rs holds every measured peak, and its tests check that each figure covers its measurements and sits within 10 % (or 32 MB) of them. Re-measure after a tract upgrade or a new export. QSMnet+ shares QSMnet’s figures and iQFM shares iQSM’s, since they are the same graphs. SynthSeg 2.0’s weights are not hosted, so its figures are SynthSeg 1.0’s plus its 33rd output channel.

Modules

download

Native weight download and on-disk cache (download feature).

onnx

ONNX inference via the pure-Rust tract engine (onnx feature).

registry 🔒

The static table of known deep-learning QSM models.

Structs

ModelSpec

A deep-learning QSM model and everything needed to obtain and run it.

NetworkMemory

Measured heap of one network in tract (see the module docs for how).

TileCost

What one tile of a tiled network costs, for budgeting how many fit in memory at once.

WeightBytes

The bytes of one weight file, shared rather than copied.

WeightFile

One weight file a model needs at inference time (an exported .onnx, plus any auxiliary file such as normalization statistics).

Enums

Framework

The framework the weights were originally trained in (before ONNX export).

Heap

Which heap a memory figure is for.

ModelStage

The reconstruction stage a model implements.

ProvideWeightsError

Why provide_weights refused a buffer.

WeightStatus

Whether a model’s weights are hosted and ready to run.

Constants

HOST_FS 🔒

Whether this build can reach a file system and network for weights ($QSM_MODEL_DIR, the cache, a download). false on wasm32-unknown-unknown, where provide_weights is the only source.

Statics

PROVIDED 🔒

Weight files supplied from memory, keyed by the registry’s file name.

Functions

all_models

All models known to QSM-Core, in a stable order.

all_weight_bytes

The bytes of every weight file of a model, in registry order.

cache_dir

Root directory for cached weight files (native only).

cache_path

The path a weight file would occupy in the cache (whether or not it exists). Native only.

find_model

Look up a model by its ModelSpec::id (case-sensitive), e.g. "qsmnet".

find_weight_file

Look a weight file up by name across every model in the registry.

forget_all_weights

Drop every buffer supplied with provide_weights.

forget_weights

Drop a buffer supplied with provide_weights, releasing its memory once no running stage still holds it. Returns true if one was held.

host_weight_bytes 🔒

missing_weights_message 🔒

The error for a weight file that is nowhere this build can look.

prefetch_with_progress

Pre-fetch (download + cache) all of a model’s weight files, reporting per-file progress via on_progress(file_name, downloaded_bytes, total_bytes).

primary_weight

Convenience: resolve a model by id and read its primary weight file. Combines find_model + primary_weight_bytes so callers (pipeline runners, external hosts) don’t repeat the lookup+fetch boilerplate. Errors if the id is unknown.

primary_weight_bytes

The bytes of a model’s primary (first) weight file.

provide_weights

Supply a weight file from memory, so every dispatcher that needs it uses these bytes before looking in $QSM_MODEL_DIR, the cache or the network.

provided_map 🔒

provided_weights

The bytes supplied for file_name with provide_weights, if any.

resolve_local

Find an already-present copy of file on disk without downloading, honoring the $QSM_MODEL_DIR override first, then the cache. Returns None if absent. Native only; does not see bytes supplied with provide_weights (see weight_file_present).

sha256_hex

Lowercase hex SHA-256 of a byte slice.

voxels 🔒

weight_file_bytes 🔒

Resolve one weight file to its bytes: supplied bytes, then (native) local override / cache, then download.

weight_file_present

true if file can be resolved without a download: it was supplied with provide_weights, or (native only) it is in $QSM_MODEL_DIR or the cache.

weights

Convenience: resolve a model by id and read all its weight files, in registry order (for multi-file models like NeXtQSM). See all_weight_bytes.