| """Portable eager inference for the released V2 mean-pooling model. |
| |
| The CLI used to ship a pair of TorchScript graphs (``encoder.ts`` / ``classifier.ts``). |
| TorchScript is a *code* container - it stores serialised Python plus pickle that runs on |
| load - so third-party model scanners flag it as potentially executable no matter what is |
| inside (Protect AI Guardian, ``PAIT-TCHST-301``). This module rebuilds the same two halves |
| from the code-free ``model.safetensors`` and ``config.json`` instead: |
| |
| * **encoder** - the DINOv3 + MoE trunk cast to bfloat16, exactly the dtype and weights the |
| former ``encoder.ts`` carried, so the embeddings are bit-identical (verified at |
| ``0.000e+00`` max-abs against the old graphs); |
| * **classifier** - the fp32 Sub-Centre ArcFace and taxonomic heads. |
| |
| The public surface mirrors the old graphs so callers only swap ``torch.jit.load`` for |
| :func:`load_bundle`. |
| |
| Built with DINOv3. |
| """ |
| from __future__ import annotations |
|
|
| import copy |
| import json |
| from pathlib import Path |
|
|
| import torch |
| from torch import nn |
|
|
| from .plant_unified import PlantUnified |
|
|
|
|
| class _Encoder(nn.Module): |
| """The DINOv3 + MoE trunk and feature fusion, without the training-only aux losses. |
| |
| Mirrors the module the removed ``encoder.ts`` wrapped: it runs the stem, one MoE layer |
| and the fusion, and skips the load-balance statistics the training forward computes |
| (they do not affect the embedding). Keeping the op set identical to the old graph also |
| keeps the cost identical rather than paying for eager-only bookkeeping. |
| """ |
|
|
| def __init__(self, model: PlantUnified): |
| super().__init__() |
| self.stem = model.stem |
| self.norm = model.moe_norms[0] |
| self.moe = model.moe_layers[0] |
| self.fusion = model.fusion |
| self.n_prefix = model.n_prefix |
|
|
| def forward(self, image: torch.Tensor) -> torch.Tensor: |
| maps = self.stem.forward_intermediates( |
| image, indices=[7, 11], output_fmt="NLC", norm=True, |
| return_prefix_tokens=True, intermediates_only=True) |
| mid, deep = [torch.cat((prefix, spatial), dim=1) for spatial, prefix in maps] |
| routed = self.moe(self.norm(deep), ctx=deep[:, 0])[0] |
| deep = deep + routed |
| return self.fusion(mid[:, self.n_prefix:], deep[:, self.n_prefix:]).float() |
|
|
|
|
| class V2Bundle: |
| """The encoder/classifier pair the CLI and harness scripts call. |
| |
| ``encoder`` takes the same bfloat16 image batch the graphs took and returns a float32 |
| embedding; ``classifier`` takes a float32 embedding and returns the |
| ``(species, genus, family)`` logits. Keeping the exact signatures means the call sites |
| do not change. |
| """ |
|
|
| def __init__(self, model_fp32: PlantUnified, encoder_fp32: PlantUnified): |
| self.model = model_fp32 |
| |
| |
| self._encoder = _Encoder(encoder_fp32).to(torch.bfloat16).eval() |
|
|
| @torch.inference_mode() |
| def encoder(self, images: torch.Tensor) -> torch.Tensor: |
| return self._encoder(images) |
|
|
| @torch.inference_mode() |
| def classifier(self, embedding: torch.Tensor): |
| logits = self.model.classify_embeddings(embedding) |
| return logits["species"], logits["genus"], logits["family"] |
|
|
|
|
| def build_model(config: dict) -> PlantUnified: |
| """Instantiate the architecture from a released ``config.json`` (no pretrained load).""" |
| model_config = dict(config["model_config"]) |
| model_config.pop("pretrained", None) |
| return PlantUnified(pretrained=False, **model_config) |
|
|
|
|
| def load_bundle(directory, device="cpu") -> V2Bundle: |
| """Build the inference bundle from a release directory containing ``config.json``. |
| |
| ``model.safetensors`` is loaded with ``strict=True``; the shared ``moe_layer`` alias in |
| the released file is resolved by ``safetensors.torch.load_model``. |
| """ |
| from safetensors.torch import load_model |
|
|
| directory = Path(directory) |
| config = json.loads((directory / "config.json").read_text(encoding="utf-8")) |
| model = build_model(config) |
| load_model(model, str(directory / "model.safetensors"), strict=True, device=str(device)) |
| model = model.to(device).eval() |
| return V2Bundle(model, copy.deepcopy(model)) |
|
|
