modelport-cli 0.1.0__py3-none-any.whl

modelport/preprocess.py

@@ -0,0 +1,121 @@
+"""Reference implementation of manifest image preprocessing.
+
+The Dart packages implement the same steps, and cross-language tests compare the two.
+See docs/spec.md, "Image preprocessing", for the exact rules.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+from PIL import Image, ImageOps
+
+from .manifest.tensors import DType, ImagePreprocess, InputSpec, ResizeSpec
+
+_PIL_FILTERS = {
+    "bilinear": Image.Resampling.BILINEAR,
+    "bicubic": Image.Resampling.BICUBIC,
+    "nearest": Image.Resampling.NEAREST,
+}
+
+
+def load_image(path: str | Path) -> Image.Image:
+    """Open an image, apply its EXIF orientation, and convert it to 8-bit RGB."""
+    with Image.open(path) as image:
+        return ImageOps.exif_transpose(image).convert("RGB")
+
+
+def resized_size(width: int, height: int, spec: ResizeSpec) -> tuple[int, int]:
+    """Output (width, height) for a resize rule."""
+    if spec.size is not None:
+        out_height, out_width = spec.size
+        return out_width, out_height
+    assert spec.shorter_side is not None
+    short = spec.shorter_side
+    if width <= height:
+        return short, int(short * height / width)
+    return int(short * width / height), short
+
+
+def resize(image: Image.Image, spec: ResizeSpec) -> np.ndarray:
+    """Resize to a float32 array of shape (height, width, 3) with values in 0-255."""
+    width, height = resized_size(image.width, image.height, spec)
+    if spec.antialias or spec.method == "nearest":
+        if (width, height) != image.size:
+            image = image.resize((width, height), resample=_PIL_FILTERS[spec.method])
+        return np.asarray(image, dtype=np.float32)
+    return bilinear_half_pixel(np.asarray(image, dtype=np.float32), height, width)
+
+
+def _axis(in_size: int, out_size: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    scale = in_size / out_size
+    src = np.maximum((np.arange(out_size, dtype=np.float64) + 0.5) * scale - 0.5, 0.0)
+    low = np.minimum(np.floor(src).astype(np.int64), in_size - 1)
+    high = np.minimum(low + 1, in_size - 1)
+    return low, high, src - low
+
+
+def bilinear_half_pixel(array: np.ndarray, out_height: int, out_width: int) -> np.ndarray:
+    """Bilinear resize with half-pixel centers and no antialiasing, in floating point.
+
+    Matches PyTorch `interpolate(mode="bilinear", align_corners=False, antialias=False)`
+    and OpenCV `INTER_LINEAR`.
+    """
+    data = array.astype(np.float64)
+    y_low, y_high, y_frac = _axis(data.shape[0], out_height)
+    x_low, x_high, x_frac = _axis(data.shape[1], out_width)
+    y_frac = y_frac[:, None, None]
+    rows = data[y_low] * (1 - y_frac) + data[y_high] * y_frac
+    x_frac = x_frac[None, :, None]
+    out = rows[:, x_low] * (1 - x_frac) + rows[:, x_high] * x_frac
+    return out.astype(np.float32)
+
+
+def center_crop(array: np.ndarray, crop_height: int, crop_width: int) -> np.ndarray:
+    """Crop the center, padding with zeros first if the image is too small."""
+    height, width = array.shape[:2]
+    if crop_height > height or crop_width > width:
+        dh, dw = max(crop_height - height, 0), max(crop_width - width, 0)
+        array = np.pad(array, ((dh // 2, (dh + 1) // 2), (dw // 2, (dw + 1) // 2), (0, 0)))
+        height, width = array.shape[:2]
+    # Python's round() is round-half-to-even, as in torchvision.
+    top = round((height - crop_height) / 2)
+    left = round((width - crop_width) / 2)
+    return array[top : top + crop_height, left : left + crop_width]
+
+
+def preprocess_array(image: Image.Image, pre: ImagePreprocess) -> np.ndarray:
+    """Run resize, crop, channel order, scale, and normalize. Returns (H, W, 3) float64."""
+    array = resize(image.convert("RGB"), pre.resize)
+    if pre.center_crop is not None:
+        array = center_crop(array, *pre.center_crop)
+    if pre.color == "BGR":
+        array = array[..., ::-1]
+    values = array.astype(np.float64) * pre.scale
+    return (values - np.asarray(pre.mean)) / np.asarray(pre.std)
+
+
+def preprocess_image(image: Image.Image, spec: InputSpec) -> np.ndarray:
+    """Build the input tensor for one image, including the batch dimension."""
+    if spec.preprocess is None:
+        raise ValueError(f"input '{spec.name}' has no image preprocessing")
+    values = preprocess_array(image, spec.preprocess)
+    if spec.dtype is DType.UINT8:
+        tensor = np.clip(np.rint(values), 0, 255).astype(np.uint8)
+    elif spec.dtype is DType.FLOAT16:
+        tensor = values.astype(np.float16)
+    else:
+        tensor = values.astype(np.float32)
+    tensor = tensor.transpose(2, 0, 1)[None] if spec.layout == "NCHW" else tensor[None]
+    for axis, (expected, actual) in enumerate(zip(spec.shape, tensor.shape, strict=True)):
+        if expected != -1 and expected != actual:
+            raise ValueError(
+                f"input '{spec.name}' dimension {axis} is {actual}, manifest says {expected}"
+            )
+    return np.ascontiguousarray(tensor)
+
+
+def to_le_bytes(array: np.ndarray) -> bytes:
+    """Raw little-endian bytes, the format of golden files."""
+    return np.ascontiguousarray(array).astype(array.dtype.newbyteorder("<")).tobytes()

modelport/publish.py

@@ -0,0 +1,191 @@
+"""Upload a bundle to the Hugging Face Hub, where apps can load it with hf://."""
+
+from __future__ import annotations
+
+import shutil
+import subprocess
+import tempfile
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+
+from .bundle import Bundle
+from .errors import MissingDependencyError, ModelPortError
+from .manifest import MANIFEST_FILENAME, Manifest
+
+CARD_NAME = "README.md"
+
+
+@dataclass(frozen=True)
+class PublishResult:
+    repo_id: str
+    url: str
+    location: str
+    """What an app passes to ModelPort.load()."""
+    files: list[str]
+
+
+def model_card(manifest: Manifest, repo_id: str) -> str:
+    """A short Hub model card describing the bundle and how to load it."""
+    rows = "\n".join(
+        f"| `{v.id}` | {v.runtime} | {v.precision} | {v.file.size / 1e6:.1f} MB |"
+        for v in manifest.variants
+    )
+    title = manifest.name or manifest.id
+    description = manifest.description or ""
+    return f"""---
+license: {manifest.license.lower()}
+tags:
+  - modelport
+  - flutter
+  - on-device
+  - {manifest.task}
+---
+
+# {title}
+
+{description}
+
+This is a [ModelPort](https://github.com/ayanparvaiz/modelport) bundle. The
+`{MANIFEST_FILENAME}` file describes inputs, preprocessing, outputs, and checksums, so a
+Flutter app can download and run it with one line.
+
+| Variant | Runtime | Precision | Size |
+|---|---|---|---|
+{rows}
+
+## Use in Flutter
+
+```dart
+final model = await ModelPort.load('hf://{repo_id}');
+```
+
+Source: `{manifest.source or "unknown"}`. License: {manifest.license}.
+"""
+
+
+def publish_bundle(path: str | Path, repo_id: str, *, private: bool = False) -> PublishResult:
+    """Check the bundle, write a model card if missing, and upload the listed files."""
+    if repo_id.count("/") != 1:
+        raise ModelPortError(f"'{repo_id}' is not a repo id like 'org/name'")
+    bundle = Bundle(path)
+    manifest = bundle.read_manifest()
+    problems = bundle.problems(manifest)
+    if problems:
+        raise ModelPortError(
+            "bundle files do not match the manifest, run `modelport pack` first:\n  "
+            + "\n  ".join(problems)
+        )
+    card = bundle.root / CARD_NAME
+    if not card.exists():
+        card.write_text(model_card(manifest, repo_id), encoding="utf-8")
+
+    # Only files inside the bundle are uploaded; https urls stay where they are.
+    files = [MANIFEST_FILENAME, CARD_NAME] + [ref.path for ref in manifest.files() if ref.path]
+    try:
+        from huggingface_hub import HfApi
+    except ImportError as error:
+        raise MissingDependencyError("Publishing to Hugging Face", "hf") from error
+
+    api = HfApi()
+    api.create_repo(repo_id, repo_type="model", private=private, exist_ok=True)
+    api.upload_folder(
+        folder_path=str(bundle.root),
+        repo_id=repo_id,
+        repo_type="model",
+        allow_patterns=files,
+        commit_message=f"modelport: {manifest.id} {manifest.version}",
+    )
+    return PublishResult(
+        repo_id=repo_id,
+        url=f"https://huggingface.co/{repo_id}",
+        location=f"hf://{repo_id}",
+        files=files,
+    )
+
+
+def github_asset_name(model_id: str, path: str) -> str:
+    """Release assets are flat, so 'onnx-fp32/model.onnx' becomes 'id--onnx-fp32--model.onnx'."""
+    return f"{model_id}--{path.replace('/', '--')}"
+
+
+def github_manifest(manifest: Manifest, repo: str, tag: str) -> Manifest:
+    """The manifest with every bundle path replaced by its release download URL."""
+    base = f"https://github.com/{repo}/releases/download/{tag}"
+    data = manifest.model_dump(by_alias=True, exclude_none=True, mode="json")
+
+    def rewrite(node: object) -> object:
+        if isinstance(node, dict):
+            if "path" in node and "sha256" in node and "size" in node:
+                name = github_asset_name(manifest.id, str(node["path"]))
+                return {"url": f"{base}/{name}", "size": node["size"], "sha256": node["sha256"]}
+            return {key: rewrite(value) for key, value in node.items()}
+        if isinstance(node, list):
+            return [rewrite(item) for item in node]
+        return node
+
+    return Manifest.model_validate(rewrite(data))
+
+
+def publish_to_github(
+    path: str | Path,
+    repo: str,
+    tag: str,
+    *,
+    run: Callable[[list[str]], subprocess.CompletedProcess[str]] | None = None,
+) -> PublishResult:
+    """Upload a bundle as assets of a GitHub release, using the `gh` CLI.
+
+    The uploaded manifest, `<id>.json`, points at the other assets by URL, so apps
+    load the model from https://github.com/<repo>/releases/download/<tag>/<id>.json.
+    """
+    if repo.count("/") != 1:
+        raise ModelPortError(f"'{repo}' is not a GitHub repo like 'owner/name'")
+    bundle = Bundle(path)
+    manifest = bundle.read_manifest()
+    problems = bundle.problems(manifest)
+    if problems:
+        raise ModelPortError(
+            "bundle files do not match the manifest, run `modelport pack` first:\n  "
+            + "\n  ".join(problems)
+        )
+    execute = run or _run_gh
+    if shutil.which("gh") is None and run is None:
+        raise ModelPortError("publishing to GitHub needs the gh CLI: https://cli.github.com")
+
+    if execute(["gh", "release", "view", tag, "-R", repo]).returncode != 0:
+        created = execute(
+            [
+                "gh", "release", "create", tag, "-R", repo,
+                "--title", f"Model zoo {tag}",
+                "--notes", "ModelPort model bundles. Load them by their .json URL.",
+                "--latest=false",
+            ]
+        )  # fmt: skip
+        if created.returncode != 0:
+            raise ModelPortError(f"could not create release {tag}: {created.stderr.strip()}")
+
+    remote = github_manifest(manifest, repo, tag)
+    with tempfile.TemporaryDirectory() as tmp:
+        staging = Path(tmp)
+        uploads: list[str] = []
+        for ref in manifest.files():
+            if ref.path is None:
+                continue
+            target = staging / github_asset_name(manifest.id, ref.path)
+            if not target.exists():
+                shutil.copyfile(bundle.root / ref.path, target)
+                uploads.append(str(target))
+        manifest_file = staging / f"{manifest.id}.json"
+        manifest_file.write_text(remote.to_json(), encoding="utf-8")
+        uploads.append(str(manifest_file))
+        result = execute(["gh", "release", "upload", tag, "-R", repo, "--clobber", *uploads])
+        if result.returncode != 0:
+            raise ModelPortError(f"upload to {repo} {tag} failed: {result.stderr.strip()}")
+
+    url = f"https://github.com/{repo}/releases/download/{tag}/{manifest.id}.json"
+    return PublishResult(repo_id=repo, url=url, location=url, files=[Path(u).name for u in uploads])
+
+
+def _run_gh(args: list[str]) -> subprocess.CompletedProcess[str]:
+    return subprocess.run(args, capture_output=True, text=True, check=False)

modelport/quantize.py

@@ -0,0 +1,120 @@
+"""Smaller ONNX variants: fp16 weights, or int8 weights for MatMul and Gemm layers."""
+
+from __future__ import annotations
+
+import math
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+from .bundle import Bundle
+from .errors import MissingDependencyError, ModelPortError
+from .golden import golden_arrays
+from .manifest import ClassificationPostprocess, Manifest, Runtime, Tolerance, Variant
+from .runtimes import run_variant
+from .verify import compare
+
+Kind = Literal["fp16", "int8"]
+
+
+class QuantizeError(ModelPortError):
+    """A smaller variant could not be made or is not accurate enough."""
+
+
+@dataclass(frozen=True)
+class QuantizeResult:
+    variant: Variant
+    size_ratio: float
+    max_abs_diff: float
+    top1_match: bool | None
+
+
+def tolerance_for(max_abs_diff: float, base: Tolerance) -> Tolerance:
+    """Twice the measured error, rounded up to two significant digits, at least base.atol."""
+    limit = max(base.atol, 2 * max_abs_diff)
+    digits = 1 - math.floor(math.log10(limit))
+    return Tolerance(atol=math.ceil(limit * 10**digits) / 10**digits, rtol=base.rtol)
+
+
+def _write_variant(source: Path, target: Path, kind: Kind) -> None:
+    try:
+        import onnx
+        from onnxruntime.quantization import QuantType, quantize_dynamic
+        from onnxruntime.transformers.float16 import convert_float_to_float16
+    except ImportError as error:
+        raise MissingDependencyError("Quantizing ONNX models", "onnx") from error
+
+    model = onnx.load(str(source))
+    # Exporter shape annotations can disagree with ONNX shape inference; drop them.
+    del model.graph.value_info[:]
+    target.parent.mkdir(parents=True, exist_ok=True)
+    if kind == "fp16":
+        # Inputs and outputs stay float32 so apps feed the same tensors to every variant.
+        onnx.save(convert_float_to_float16(model, keep_io_types=True), str(target))
+        return
+    with tempfile.TemporaryDirectory() as tmp:
+        clean = Path(tmp) / "model.onnx"
+        onnx.save(model, str(clean))
+        # Dynamic int8 on convolutions badly hurts CNN accuracy, so only quantize
+        # MatMul and Gemm, where most transformer and classifier-head weights live.
+        quantize_dynamic(
+            str(clean),
+            str(target),
+            weight_type=QuantType.QInt8,
+            op_types_to_quantize=["MatMul", "Gemm"],
+            per_channel=True,
+        )
+
+
+def quantize_bundle(
+    path: str | Path, kinds: list[Kind], *, allow_top1_change: bool = False
+) -> list[QuantizeResult]:
+    """Add one ONNX variant per kind, measure its error, and record a matching tolerance."""
+    bundle = Bundle(path)
+    manifest = bundle.read_manifest()
+    if manifest.golden is None:
+        raise QuantizeError("the bundle needs golden data to measure accuracy")
+    base = next(
+        (v for v in manifest.variants if v.runtime is Runtime.ONNX and v.precision == "fp32"),
+        None,
+    )
+    if base is None or base.file.path is None:
+        raise QuantizeError("the bundle has no onnx fp32 variant to start from")
+
+    inputs, expected = golden_arrays(bundle, manifest)
+    classify = {
+        o.name for o in manifest.outputs if isinstance(o.postprocess, ClassificationPostprocess)
+    }
+    variants = list(manifest.variants)
+    results = []
+    for kind in dict.fromkeys(kinds):
+        variant_id = f"onnx-{kind}"
+        relative = f"{variant_id}/model.onnx"
+        _write_variant(bundle.path(base.file.path), bundle.path(relative), kind)
+        candidate = Variant(
+            id=variant_id, runtime=Runtime.ONNX, precision=kind, file=bundle.add(relative)
+        )
+        got = run_variant(bundle, manifest, candidate, inputs)
+        checks = [
+            compare(name, got[name], want, manifest.golden.tolerance, name in classify)
+            for name, want in expected.items()
+        ]
+        worst = max(check.max_abs_diff for check in checks)
+        top1 = None if not classify else all(c.top1_match is not False for c in checks)
+        if top1 is False and not allow_top1_change:
+            raise QuantizeError(
+                f"{variant_id} changes the top-1 class on the golden input "
+                f"(max |diff| {worst:.3g}). Use --allow-top1-change to keep it anyway."
+            )
+        variant = candidate.model_copy(
+            update={"tolerance": tolerance_for(worst, manifest.golden.tolerance)}
+        )
+        variants = [v for v in variants if v.id != variant_id] + [variant]
+        results.append(QuantizeResult(variant, variant.file.size / base.file.size, worst, top1))
+
+    updated = Manifest.model_validate(
+        {**manifest.model_dump(by_alias=True), "variants": [v.model_dump() for v in variants]}
+    )
+    bundle.write_manifest(updated)
+    return results

modelport/runtimes.py

@@ -0,0 +1,52 @@
+"""Run exported variants in Python, the same engines the Dart adapters wrap."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from .bundle import Bundle
+from .errors import MissingDependencyError, ModelPortError
+from .manifest import Manifest, Runtime, Variant
+
+
+class RuntimeNotSupported(ModelPortError):
+    """This runtime cannot run tensor inputs in Python."""
+
+
+def run_variant(
+    bundle: Bundle, manifest: Manifest, variant: Variant, inputs: dict[str, np.ndarray]
+) -> dict[str, np.ndarray]:
+    """Run one variant on named inputs and return outputs keyed by manifest output name."""
+    if variant.file.path is None:
+        raise ModelPortError(f"{variant.id}: only files inside the bundle can be run")
+    path = bundle.path(variant.file.path)
+    output_names = [o.name for o in manifest.outputs]
+
+    if variant.runtime is Runtime.ONNX:
+        try:
+            import onnxruntime as ort
+        except ImportError as error:
+            raise MissingDependencyError("Running ONNX variants", "onnx") from error
+        session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
+        results = session.run(output_names, inputs)
+        return {name: np.asarray(value) for name, value in zip(output_names, results, strict=True)}
+
+    if variant.runtime is Runtime.EXECUTORCH:
+        try:
+            import torch
+            from executorch.runtime import Runtime as ExecuTorchRuntime
+        except ImportError as error:
+            raise MissingDependencyError("Running ExecuTorch variants", "executorch") from error
+        method = ExecuTorchRuntime.get().load_program(str(path)).load_method("forward")
+        if method is None:
+            raise ModelPortError(f"{variant.id}: the program has no 'forward' method")
+        # ExecuTorch takes inputs by position, in manifest order.
+        args = [torch.from_numpy(np.ascontiguousarray(inputs[s.name])) for s in manifest.inputs]
+        results = method.execute(args)
+        return {
+            name: value.detach().numpy() for name, value in zip(output_names, results, strict=True)
+        }
+
+    raise RuntimeNotSupported(
+        f"{variant.id}: runtime '{variant.runtime}' has no tensor golden test"
+    )

modelport/sources/__init__.py

@@ -0,0 +1,85 @@
+"""Load a PyTorch model and describe it well enough to write a manifest.
+
+Sources:
+    torchvision:<model name>      e.g. torchvision:mobilenet_v3_small
+    hf:<repo id or local folder>  e.g. hf:google/vit-base-patch16-224
+    file:<script.py>[:function]   a function that returns a SourceModel
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+from typing import Any
+
+from ..errors import ModelPortError
+from ..manifest import DetectionPostprocess, InputSpec, Task
+
+_ID_CLEAN = re.compile(r"[^a-z0-9._-]+")
+
+
+class SourceError(ModelPortError):
+    """A model source could not be loaded."""
+
+
+@dataclass
+class SourceModel:
+    """A PyTorch module plus everything needed to describe it in a manifest.
+
+    `example_inputs` must match `inputs` in order and shape. They are made contiguous
+    before export, because ExecuTorch records the memory layout of example inputs.
+    """
+
+    module: Any
+    example_inputs: tuple[Any, ...]
+    id: str
+    task: Task
+    license: str
+    inputs: list[InputSpec]
+    output_names: list[str]
+    labels: list[str] | None = None
+    name: str | None = None
+    description: str | None = None
+    source: str | None = None
+    top_k: int = 5
+    detection: DetectionPostprocess | None = None
+    """For object detection: how to decode the outputs. Labels are attached on export."""
+    extra: dict[str, Any] = field(default_factory=dict)
+
+
+def model_id(text: str) -> str:
+    """Turn a model name or repo id into a manifest id, such as 'google/ViT-B' -> 'vit-b'."""
+    last = text.rstrip("/").split("/")[-1].lower()
+    cleaned = _ID_CLEAN.sub("-", last).strip("-._")
+    if not cleaned:
+        raise SourceError(f"cannot make a model id from '{text}'")
+    return cleaned
+
+
+def load_source(
+    spec: str, *, license: str | None = None, image_size: int | None = None
+) -> SourceModel:
+    """Load a model from a source string such as 'torchvision:mobilenet_v3_small'."""
+    kind, sep, rest = spec.partition(":")
+    if not sep or not rest:
+        raise SourceError(f"'{spec}' is not a source. Use torchvision:, hf:, or file:")
+    if kind == "torchvision":
+        from .torchvision import load_torchvision
+
+        model = load_torchvision(rest)
+    elif kind == "hf":
+        from .hf import load_hf
+
+        model = load_hf(rest, license=license, image_size=image_size)
+    elif kind == "file":
+        from .file import load_file
+
+        model = load_file(rest)
+    else:
+        raise SourceError(f"unknown source type '{kind}'. Use torchvision:, hf:, or file:")
+    if license is not None:
+        model.license = license
+    return model
+
+
+__all__ = ["SourceError", "SourceModel", "load_source", "model_id"]

modelport/sources/file.py

@@ -0,0 +1,46 @@
+"""Models described by a function in the user's own Python file."""
+
+from __future__ import annotations
+
+import importlib.util
+from pathlib import Path
+
+from . import SourceError, SourceModel
+
+DEFAULT_FUNCTION = "build"
+
+
+def split_file_spec(spec: str) -> tuple[Path, str]:
+    """'models/net.py:make' -> (models/net.py, 'make'). The function defaults to 'build'."""
+    path_text, sep, function = spec.rpartition(":")
+    if sep and path_text.endswith(".py") and function:
+        return Path(path_text), function
+    return Path(spec), DEFAULT_FUNCTION
+
+
+def load_file(spec: str) -> SourceModel:
+    """Import a Python file and call a function that returns a SourceModel.
+
+    This runs the file's code, like `python file.py` would. Only use files you trust.
+    """
+    path, function = split_file_spec(spec)
+    if not path.is_file():
+        raise SourceError(f"{path} does not exist")
+    module_spec = importlib.util.spec_from_file_location(f"_modelport_user_{path.stem}", path)
+    if module_spec is None or module_spec.loader is None:
+        raise SourceError(f"cannot import {path}")
+    module = importlib.util.module_from_spec(module_spec)
+    module_spec.loader.exec_module(module)
+
+    factory = getattr(module, function, None)
+    if not callable(factory):
+        raise SourceError(f"{path} has no function named '{function}'")
+    result = factory()
+    if not isinstance(result, SourceModel):
+        raise SourceError(
+            f"{path}:{function} must return modelport.sources.SourceModel, "
+            f"got {type(result).__name__}"
+        )
+    if result.source is None:
+        result.source = f"file:{path.name}:{function}"
+    return result