feat(ml): optimized model inference (#31750)
* add immich_model * graph per shape * decouple ocr from ort * prepared graphs * use f16 for accelerators * preload all shapes * warm-up pass after preload * keep precision for fp32 graphs * ensure padding for text recognition * optimize postprocessing for fp16 outputs * handle raw ctc logits * use f16 output for rknn * mmap token embedding * read device metadata
This commit is contained in:
Mert
@@ -154,33 +154,33 @@ Redis (Sentinel) URL example JSON before encoding:
|
||||
|
||||
## Machine Learning
|
||||
|
||||
| Variable | Description | Default | Containers |
|
||||
| :---------------------------------------------------------- | :----------------------------------------------------------------------------------------------------------------------------------------------------------- | :--------------------------: | :--------------- |
|
||||
| `MACHINE_LEARNING_MODEL_TTL` | Inactivity time (s) before a model is unloaded (disabled if \<= 0) | `300` | machine learning |
|
||||
| `MACHINE_LEARNING_MODEL_TTL_POLL_S` | Interval (s) between checks for the model TTL (disabled if \<= 0) | `10` | machine learning |
|
||||
| `MACHINE_LEARNING_CACHE_FOLDER` | Directory where models are downloaded | `/cache` | machine learning |
|
||||
| `MACHINE_LEARNING_REQUEST_THREADS`<sup>\*1</sup> | Thread count of the request thread pool (disabled if \<= 0) | number of CPU cores | machine learning |
|
||||
| `MACHINE_LEARNING_MODEL_INTER_OP_THREADS` | Number of parallel model operations | `1` | machine learning |
|
||||
| `MACHINE_LEARNING_MODEL_INTRA_OP_THREADS` | Number of threads for each model operation | `2` | machine learning |
|
||||
| `MACHINE_LEARNING_WORKERS`<sup>\*2</sup> | Number of worker processes to spawn | `1` | machine learning |
|
||||
| `MACHINE_LEARNING_HTTP_KEEPALIVE_TIMEOUT_S`<sup>\*3</sup> | HTTP Keep-alive time in seconds | `2` | machine learning |
|
||||
| `MACHINE_LEARNING_WORKER_TIMEOUT` | Maximum time (s) of unresponsiveness before a worker is killed | `300` (`900` if using ROCm) | machine learning |
|
||||
| `MACHINE_LEARNING_PRELOAD__CLIP__TEXTUAL` | Comma-separated list of (textual) CLIP model(s) to preload and cache | | machine learning |
|
||||
| `MACHINE_LEARNING_PRELOAD__CLIP__VISUAL` | Comma-separated list of (visual) CLIP model(s) to preload and cache | | machine learning |
|
||||
| `MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION__RECOGNITION` | Comma-separated list of (recognition) facial recognition model(s) to preload and cache | | machine learning |
|
||||
| `MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION__DETECTION` | Comma-separated list of (detection) facial recognition model(s) to preload and cache | | machine learning |
|
||||
| `MACHINE_LEARNING_PRELOAD__OCR__RECOGNITION` | Comma-separated list of (recognition) OCR model(s) to preload and cache | | machine learning |
|
||||
| `MACHINE_LEARNING_PRELOAD__OCR__DETECTION` | Comma-separated list of (detection) OCR model(s) to preload and cache | | machine learning |
|
||||
| `MACHINE_LEARNING_ANN` | Enable ARM-NN hardware acceleration if supported | `True` | machine learning |
|
||||
| `MACHINE_LEARNING_ANN_FP16_TURBO` | Execute operations in FP16 precision: increasing speed, reducing precision (applies only to ARM-NN) | `False` | machine learning |
|
||||
| `MACHINE_LEARNING_ANN_TUNING_LEVEL` | ARM-NN GPU tuning level (1: rapid, 2: normal, 3: exhaustive) | `2` | machine learning |
|
||||
| `MACHINE_LEARNING_DEVICE_IDS`<sup>\*4</sup> | Device IDs to use in multi-GPU environments | `0` | machine learning |
|
||||
| `MACHINE_LEARNING_MAX_BATCH_SIZE__FACIAL_RECOGNITION` | Set the maximum number of faces that will be processed at once by the facial recognition model | None (`1` if using OpenVINO) | machine learning |
|
||||
| `MACHINE_LEARNING_MAX_BATCH_SIZE__OCR` | Set the maximum number of boxes that will be processed at once by the OCR model | `6` | machine learning |
|
||||
| `MACHINE_LEARNING_RKNN` | Enable RKNN hardware acceleration if supported | `True` | machine learning |
|
||||
| `MACHINE_LEARNING_RKNN_THREADS` | How many threads of RKNN runtime should be spun up while inferencing. | `1` | machine learning |
|
||||
| `MACHINE_LEARNING_MODEL_ARENA` | Pre-allocates CPU memory to avoid memory fragmentation | true | machine learning |
|
||||
| `MACHINE_LEARNING_OPENVINO_PRECISION` | If set to FP16, uses half-precision floating-point operations for faster inference with reduced accuracy (one of [`FP16`, `FP32`], applies only to OpenVINO) | `FP32` | machine learning |
|
||||
| Variable | Description | Default | Containers |
|
||||
| :---------------------------------------------------------- | :-------------------------------------------------------------------------------------------------- | :-------------------------: | :--------------- |
|
||||
| `MACHINE_LEARNING_MODEL_TTL` | Inactivity time (s) before a model is unloaded (disabled if \<= 0) | `300` | machine learning |
|
||||
| `MACHINE_LEARNING_MODEL_TTL_POLL_S` | Interval (s) between checks for the model TTL (disabled if \<= 0) | `10` | machine learning |
|
||||
| `MACHINE_LEARNING_CACHE_FOLDER` | Directory where models are downloaded | `/cache` | machine learning |
|
||||
| `MACHINE_LEARNING_REQUEST_THREADS`<sup>\*1</sup> | Thread count of the request thread pool (disabled if \<= 0) | number of CPU cores | machine learning |
|
||||
| `MACHINE_LEARNING_MODEL_INTER_OP_THREADS` | Number of parallel model operations | `1` | machine learning |
|
||||
| `MACHINE_LEARNING_MODEL_INTRA_OP_THREADS` | Number of threads for each model operation | `2` | machine learning |
|
||||
| `MACHINE_LEARNING_WORKERS`<sup>\*2</sup> | Number of worker processes to spawn | `1` | machine learning |
|
||||
| `MACHINE_LEARNING_HTTP_KEEPALIVE_TIMEOUT_S`<sup>\*3</sup> | HTTP Keep-alive time in seconds | `2` | machine learning |
|
||||
| `MACHINE_LEARNING_WORKER_TIMEOUT` | Maximum time (s) of unresponsiveness before a worker is killed | `300` (`900` if using ROCm) | machine learning |
|
||||
| `MACHINE_LEARNING_PRELOAD__CLIP__TEXTUAL` | Comma-separated list of (textual) CLIP model(s) to preload and cache | | machine learning |
|
||||
| `MACHINE_LEARNING_PRELOAD__CLIP__VISUAL` | Comma-separated list of (visual) CLIP model(s) to preload and cache | | machine learning |
|
||||
| `MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION__RECOGNITION` | Comma-separated list of (recognition) facial recognition model(s) to preload and cache | | machine learning |
|
||||
| `MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION__DETECTION` | Comma-separated list of (detection) facial recognition model(s) to preload and cache | | machine learning |
|
||||
| `MACHINE_LEARNING_PRELOAD__OCR__RECOGNITION` | Comma-separated list of (recognition) OCR model(s) to preload and cache | | machine learning |
|
||||
| `MACHINE_LEARNING_PRELOAD__OCR__DETECTION` | Comma-separated list of (detection) OCR model(s) to preload and cache | | machine learning |
|
||||
| `MACHINE_LEARNING_PRELOAD__OCR__MAX_RESOLUTION` | Resolution to preload the OCR detection model(s) for (must match the server's OCR resolution) | `736` | machine learning |
|
||||
| `MACHINE_LEARNING_ANN` | Enable ARM-NN hardware acceleration if supported | `True` | machine learning |
|
||||
| `MACHINE_LEARNING_ANN_FP16_TURBO` | Execute operations in FP16 precision: increasing speed, reducing precision (applies only to ARM-NN) | `False` | machine learning |
|
||||
| `MACHINE_LEARNING_ANN_TUNING_LEVEL` | ARM-NN GPU tuning level (1: rapid, 2: normal, 3: exhaustive) | `2` | machine learning |
|
||||
| `MACHINE_LEARNING_DEVICE_IDS`<sup>\*4</sup> | Device IDs to use in multi-GPU environments | `0` | machine learning |
|
||||
| `MACHINE_LEARNING_MAX_BATCH_SIZE__FACIAL_RECOGNITION` | Set the maximum number of faces that will be processed at once by the facial recognition model | `4` | machine learning |
|
||||
| `MACHINE_LEARNING_MAX_BATCH_SIZE__OCR` | Set the maximum number of boxes that will be processed at once by the OCR model | `6` | machine learning |
|
||||
| `MACHINE_LEARNING_RKNN` | Enable RKNN hardware acceleration if supported | `True` | machine learning |
|
||||
| `MACHINE_LEARNING_RKNN_THREADS` | How many threads of RKNN runtime should be spun up while inferencing. | `1` | machine learning |
|
||||
| `MACHINE_LEARNING_MODEL_ARENA` | Pre-allocates CPU memory to avoid memory fragmentation | true | machine learning |
|
||||
|
||||
\*1: It is recommended to begin with this parameter when changing the concurrency levels of the machine learning service and then tune the other ones.
|
||||
|
||||
|
||||
@@ -31,7 +31,7 @@ ENV PYTHONDONTWRITEBYTECODE=1 \
|
||||
PYTHONUNBUFFERED=1 \
|
||||
VIRTUAL_ENV=/opt/venv
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends g++
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends g++ git
|
||||
|
||||
COPY --from=ghcr.io/astral-sh/uv:0.8.15@sha256:a5727064a0de127bdb7c9d3c1383f3a9ac307d9f2d8a391edc7896c54289ced0 /uv /uvx /bin/
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, Callable, Iterator
|
||||
from unittest import mock
|
||||
@@ -10,6 +11,15 @@ from PIL import Image
|
||||
from immich_ml.config import log
|
||||
from immich_ml.main import app
|
||||
from immich_ml.models.base import InferenceModel
|
||||
from immich_ml.schemas import Shape
|
||||
from immich_ml.sessions.ort import Device
|
||||
|
||||
TEST_ASSETS = Path(__file__).parent.parent / "e2e/test-assets"
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def asset() -> Callable[[str], bytes]:
|
||||
return lambda path: (TEST_ASSETS / path).read_bytes()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -20,6 +30,7 @@ def pil_image() -> Image.Image:
|
||||
@pytest.fixture
|
||||
def mock_get_model() -> Iterator[mock.Mock]:
|
||||
with mock.patch("immich_ml.models.cache.get_model_class") as mocked:
|
||||
mocked.return_value.graph_options = ()
|
||||
yield mocked
|
||||
|
||||
|
||||
@@ -103,6 +114,15 @@ def providers(request: pytest.FixtureRequest) -> Iterator[mock.Mock]:
|
||||
yield providers
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def gpus() -> Iterator[None]:
|
||||
with (
|
||||
mock.patch("immich_ml.sessions.ort._intel_gpu", return_value=Device("12.71.4-128eu", "26.22.38646.4")),
|
||||
mock.patch("immich_ml.sessions.ort._amd_gpu", return_value=Device("gfx1100", "7.2.0")),
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def ort_pybind() -> Iterator[mock.Mock]:
|
||||
with mock.patch("immich_ml.sessions.ort.ort.capi._pybind_state") as mocked:
|
||||
@@ -120,15 +140,29 @@ def ov_device_ids(request: pytest.FixtureRequest, ort_pybind: mock.Mock) -> Iter
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def ort_session() -> Iterator[mock.Mock]:
|
||||
with mock.patch("immich_ml.sessions.ort.ort.InferenceSession") as mocked:
|
||||
# a graph is opened as it is: preparing one takes a child process and real files
|
||||
with (
|
||||
mock.patch("immich_ml.sessions.ort.ort.InferenceSession") as mocked,
|
||||
mock.patch("immich_ml.sessions.ort.prepared", side_effect=lambda spec: spec.model_path),
|
||||
):
|
||||
yield mocked
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def stub_session() -> Callable[..., mock.Mock]:
|
||||
def _make(shape: tuple[Any, ...], outputs: Any = None, name: str = "input.1") -> mock.Mock:
|
||||
def _make(
|
||||
shape: tuple[Any, ...],
|
||||
outputs: Any = None,
|
||||
name: str = "input.1",
|
||||
normalizes_input: bool = False,
|
||||
shapes: tuple[Shape, ...] = (Shape(batch=1),),
|
||||
) -> mock.Mock:
|
||||
session = mock.Mock()
|
||||
session.get_inputs.return_value = [SimpleNamespace(name=name, shape=shape)]
|
||||
session.normalizes_input = normalizes_input
|
||||
session.shapes = shapes
|
||||
session.batches = tuple(sorted({shape.batch for shape in shapes}, reverse=True))
|
||||
session.for_shape.return_value = session # a stub is the session and the one graph in it
|
||||
if outputs is not None:
|
||||
session.run.return_value = outputs
|
||||
return session
|
||||
@@ -139,12 +173,16 @@ def stub_session() -> Callable[..., mock.Mock]:
|
||||
@pytest.fixture(scope="function")
|
||||
def ann_session() -> Iterator[mock.Mock]:
|
||||
with mock.patch("immich_ml.sessions.ann.Ann") as mocked:
|
||||
mocked.return_value.input_shapes.__getitem__.return_value = [(1, 3, 112, 112)] # armnn compiles one batch
|
||||
yield mocked
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def rknn_session() -> Iterator[mock.Mock]:
|
||||
with mock.patch("immich_ml.sessions.rknn.RknnPoolExecutor") as mocked:
|
||||
mocked.return_value.custom_string = "" # an unstamped binary, i.e. one that wants normalized input
|
||||
compiled = SimpleNamespace(name="input", shape=(1, 224, 224, 3)) # the compiler reports NHWC either way
|
||||
mocked.return_value.inputs = [compiled, compiled]
|
||||
yield mocked
|
||||
|
||||
|
||||
@@ -188,6 +226,7 @@ def exception() -> Iterator[mock.Mock]:
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def snapshot_download() -> Iterator[mock.Mock]:
|
||||
# a download that leaves the model cached, as one that found the artifact does
|
||||
with (
|
||||
mock.patch("immich_ml.models.base.snapshot_download") as mocked,
|
||||
mock.patch.object(InferenceModel, "cached", new_callable=mock.PropertyMock, side_effect=lambda: mocked.called),
|
||||
|
||||
@@ -13,7 +13,7 @@ from rich.logging import RichHandler
|
||||
from uvicorn import Server
|
||||
from uvicorn.workers import UvicornWorker
|
||||
|
||||
from .schemas import ModelOrganization, ModelPrecision
|
||||
from .schemas import ModelOrganization
|
||||
|
||||
|
||||
class ClipSettings(BaseModel):
|
||||
@@ -29,6 +29,7 @@ class FacialRecognitionSettings(BaseModel):
|
||||
class OcrSettings(BaseModel):
|
||||
recognition: str | None = None
|
||||
detection: str | None = None
|
||||
max_resolution: int = 736 # must match the server's OCR setting, or the preloaded model will not be used
|
||||
|
||||
|
||||
class PreloadModelData(BaseModel):
|
||||
@@ -38,8 +39,8 @@ class PreloadModelData(BaseModel):
|
||||
|
||||
|
||||
class MaxBatchSize(BaseModel):
|
||||
facial_recognition: int | None = None
|
||||
ocr: int | None = None
|
||||
facial_recognition: int = 4
|
||||
ocr: int = 6
|
||||
|
||||
|
||||
def default_worker_timeout() -> int:
|
||||
@@ -71,9 +72,7 @@ class Settings(BaseSettings):
|
||||
rknn: bool = True
|
||||
rknn_threads: int = 1
|
||||
preload: PreloadModelData | None = None
|
||||
max_batch_size: MaxBatchSize | None = None
|
||||
openvino_precision: ModelPrecision = ModelPrecision.FP32
|
||||
rocm_precision: ModelPrecision = ModelPrecision.FP32
|
||||
max_batch_size: MaxBatchSize = MaxBatchSize()
|
||||
model_organization: ModelOrganization = ModelOrganization.APP
|
||||
model_revision: str = "main"
|
||||
|
||||
|
||||
@@ -91,6 +91,7 @@ async def preload_models(preload: PreloadModelData) -> None:
|
||||
model_name = model_name.strip()
|
||||
model = model_cache.get(model_name, model_type, model_task, **options)
|
||||
await load(model)
|
||||
await attempt(model, model.build, MODEL_FILE_ERRORS)
|
||||
|
||||
if preload.clip.textual is not None:
|
||||
await load_models(preload.clip.textual, ModelType.TEXTUAL, ModelTask.SEARCH)
|
||||
@@ -117,6 +118,7 @@ async def preload_models(preload: PreloadModelData) -> None:
|
||||
preload.ocr.detection,
|
||||
ModelType.DETECTION,
|
||||
ModelTask.OCR,
|
||||
maxResolution=preload.ocr.max_resolution,
|
||||
)
|
||||
|
||||
if preload.ocr.recognition is not None:
|
||||
|
||||
@@ -15,6 +15,7 @@ from immich_ml.sessions.ort import OrtSession
|
||||
from ..config import clean_name, log, settings
|
||||
from ..schemas import ModelFormat, ModelIdentity, ModelSession, ModelTask, ModelType
|
||||
from ..sessions.ann import AnnSession
|
||||
from ..sessions.policy import ShapePolicy
|
||||
|
||||
_IGNORED_PATTERNS: dict[ModelFormat, list[str]] = {
|
||||
ModelFormat.ONNX: ["*.armnn", "*.rknn"],
|
||||
@@ -26,6 +27,10 @@ _IGNORED_PATTERNS: dict[ModelFormat, list[str]] = {
|
||||
class InferenceModel(ABC):
|
||||
depends: ClassVar[list[ModelIdentity]]
|
||||
identity: ClassVar[ModelIdentity]
|
||||
# options a graph is built for, so one instance cannot serve two values of them
|
||||
graph_options: ClassVar[tuple[str, ...]] = ()
|
||||
# the shapes this model feeds, which pick the artifact and tell the engine what it may pin
|
||||
shape_policy: ShapePolicy = ShapePolicy()
|
||||
threads: ClassVar[int] = 2 # how many threads one run takes for CPU inference
|
||||
|
||||
def __init__(
|
||||
@@ -78,6 +83,10 @@ class InferenceModel(ABC):
|
||||
del self.session
|
||||
self.loaded = False
|
||||
|
||||
def build(self) -> None:
|
||||
self.load()
|
||||
self.session.warm()
|
||||
|
||||
def predict(self, *inputs: Any, **model_kwargs: Any) -> Any:
|
||||
self.load()
|
||||
return self._predict(*inputs, **model_kwargs)
|
||||
@@ -115,23 +124,19 @@ class InferenceModel(ABC):
|
||||
case ModelFormat.ARMNN:
|
||||
return AnnSession(self.model_path)
|
||||
case ModelFormat.ONNX:
|
||||
return OrtSession(self.model_path, threads=self.threads)
|
||||
return OrtSession(self.model_path, self.shape_policy, threads=self.threads)
|
||||
case ModelFormat.RKNN:
|
||||
return rknn.RknnSession(self.model_path)
|
||||
|
||||
def model_path_for_format(self, model_format: ModelFormat) -> Path:
|
||||
model_path_prefix = rknn.model_prefix if model_format == ModelFormat.RKNN else None
|
||||
if model_path_prefix:
|
||||
return self.model_dir / model_path_prefix / f"model.{model_format}"
|
||||
return self.model_dir / f"model.{model_format}"
|
||||
|
||||
@property
|
||||
def model_dir(self) -> Path:
|
||||
return self.cache_dir / self.model_type.value
|
||||
|
||||
@property
|
||||
def model_path(self) -> Path:
|
||||
return self.model_path_for_format(self.model_format)
|
||||
if self.model_format == ModelFormat.RKNN:
|
||||
return rknn.model_path(self.model_dir, self.shape_policy.label)
|
||||
return self.model_dir / f"model.{self.model_format}"
|
||||
|
||||
@property
|
||||
def model_task(self) -> ModelTask:
|
||||
@@ -151,7 +156,8 @@ class InferenceModel(ABC):
|
||||
|
||||
@property
|
||||
def _cache_dir_default(self) -> Path:
|
||||
return settings.cache_folder / self.model_task.value / self.model_name
|
||||
cache_dir = settings.cache_folder / self.model_task.value / self.model_name
|
||||
return cache_dir if settings.legacy_models else cache_dir / settings.model_revision
|
||||
|
||||
@property
|
||||
def cached(self) -> bool:
|
||||
|
||||
@@ -17,7 +17,7 @@ class ModelCache:
|
||||
self, model_name: str, model_type: ModelType, model_task: ModelTask, ttl: int | None = None, **options: Any
|
||||
) -> InferenceModel:
|
||||
model_cls = get_model_class(model_name, model_type, model_task)
|
||||
key = (model_name, model_type, model_task)
|
||||
key = (model_name, model_type, model_task, *(options.get(option) for option in model_cls.graph_options))
|
||||
if key not in self._models:
|
||||
self._models[key] = model_cls(model_name, **options)
|
||||
elif key not in self._expiries:
|
||||
|
||||
@@ -15,18 +15,33 @@ from immich_ml.models.transforms import clean_text, serialize_np_array
|
||||
from immich_ml.schemas import ModelSession, ModelTask, ModelType
|
||||
|
||||
|
||||
def _mapped(model_path: Path, name: str) -> NDArray[Any]:
|
||||
import onnx
|
||||
|
||||
initializer = next(i for i in onnx.load(model_path, load_external_data=False).graph.initializer if i.name == name)
|
||||
data = {entry.key: entry.value for entry in initializer.external_data}
|
||||
dtype = onnx.helper.tensor_dtype_to_np_dtype(initializer.data_type)
|
||||
offset, shape = int(data.get("offset", 0)), tuple(initializer.dims)
|
||||
return np.memmap(model_path.parent / data["location"], dtype=dtype, mode="r", offset=offset, shape=shape)
|
||||
|
||||
|
||||
class BaseCLIPTextualEncoder(InferenceModel):
|
||||
depends = []
|
||||
identity = (ModelType.TEXTUAL, ModelTask.SEARCH)
|
||||
threads = 4 # lower search latency
|
||||
|
||||
def _predict(self, inputs: str, language: str | None = None) -> str:
|
||||
tokens = self.tokenize(inputs, language=language)
|
||||
res: NDArray[np.float32] = self.session.run(None, tokens)[0][0]
|
||||
tokens: dict[str, NDArray[Any]] = self.tokenize(inputs, language=language)
|
||||
graph = self.session.for_shape(self.shape_policy.dims[0])
|
||||
if self.embedding is not None: # a graph that left its token table to the host
|
||||
tokens["token_embeds"] = self.embedding[next(iter(tokens.values()))]
|
||||
res: NDArray[np.float32] = graph.run(None, {node.name: tokens[node.name] for node in graph.get_inputs()})[0][0]
|
||||
return serialize_np_array(res)
|
||||
|
||||
def _load(self) -> ModelSession:
|
||||
session = super()._load()
|
||||
table = session.for_shape(self.shape_policy.dims[0]).get_metadata().get("embedding")
|
||||
self.embedding = _mapped(self.model_dir / "model.onnx", table) if table is not None else None
|
||||
log.debug(f"Loading tokenizer for CLIP model '{self.model_name}'")
|
||||
self.tokenizer = self._load_tokenizer()
|
||||
tokenizer_kwargs: dict[str, Any] | None = self.text_cfg.get("tokenizer_kwargs")
|
||||
|
||||
@@ -19,7 +19,7 @@ from immich_ml.models.transforms import (
|
||||
serialize_np_array,
|
||||
to_numpy,
|
||||
)
|
||||
from immich_ml.schemas import ModelSession, ModelTask, ModelType
|
||||
from immich_ml.schemas import ModelGraph, ModelInput, ModelSession, ModelTask, ModelType
|
||||
|
||||
|
||||
class BaseCLIPVisualEncoder(InferenceModel):
|
||||
@@ -28,11 +28,12 @@ class BaseCLIPVisualEncoder(InferenceModel):
|
||||
|
||||
def _predict(self, inputs: Image.Image | bytes) -> str:
|
||||
image = decode_pil(inputs)
|
||||
res: NDArray[np.float32] = self.session.run(None, self.transform(image))[0][0]
|
||||
session = self.session.for_shape(self.shape_policy.dims[0])
|
||||
res: NDArray[np.float32] = session.run(None, self.transform(session, image))[0][0]
|
||||
return serialize_np_array(res)
|
||||
|
||||
@abstractmethod
|
||||
def transform(self, image: Image.Image) -> dict[str, NDArray[np.float32]]:
|
||||
def transform(self, session: ModelGraph, image: Image.Image) -> ModelInput:
|
||||
pass
|
||||
|
||||
@property
|
||||
@@ -69,11 +70,12 @@ class OpenClipVisualEncoder(BaseCLIPVisualEncoder):
|
||||
|
||||
return super()._load()
|
||||
|
||||
def transform(self, image: Image.Image) -> dict[str, NDArray[np.float32]]:
|
||||
def transform(self, session: ModelGraph, image: Image.Image) -> ModelInput:
|
||||
image = self._resize(image)
|
||||
image_np = to_numpy(image)
|
||||
image_np = normalize(image_np, self.mean, self.std)
|
||||
return {"image": np.expand_dims(image_np.transpose(2, 0, 1), 0)}
|
||||
if session.normalizes_input:
|
||||
rgb: NDArray[np.uint8] = np.asarray(image if image.mode == "RGB" else image.convert("RGB"))
|
||||
return {"image": rgb[None]}
|
||||
return {"image": normalize(to_numpy(image), self.mean, self.std).transpose(2, 0, 1)[None]}
|
||||
|
||||
def _resize(self, image: Image.Image) -> Image.Image:
|
||||
match self.preprocess_cfg.get("resize_mode", "shortest"): # open_clip's default for older configs
|
||||
|
||||
@@ -8,11 +8,11 @@ from functools import lru_cache
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from immich_model.constants import FACE_DETECTION_SIZE as DET_SIZE
|
||||
from numpy.typing import NDArray
|
||||
|
||||
from immich_ml.models.transforms import ensure_dims
|
||||
|
||||
DET_SIZE = 640
|
||||
ALIGNED_SIZE = 112
|
||||
|
||||
# the FPN levels for which the fused detector emits a (scores, boxes, kps) triple,
|
||||
@@ -69,6 +69,5 @@ def umeyama(src: NDArray[np.float32], dst: NDArray[np.float32]) -> NDArray[np.fl
|
||||
return np.hstack([scale * rotation, translation[:, None]], dtype=np.float32)
|
||||
|
||||
|
||||
def align_face(image: NDArray[np.uint8], kps: NDArray[np.float32]) -> NDArray[np.float32]:
|
||||
matrix = umeyama(kps, ARCFACE_DST)
|
||||
return cv2.warpAffine(image, matrix, (ALIGNED_SIZE, ALIGNED_SIZE)).astype(np.float32)
|
||||
def align_face(image: NDArray[np.uint8], kps: NDArray[np.float32], crop: NDArray[np.uint8]) -> None:
|
||||
cv2.warpAffine(image, umeyama(kps, ARCFACE_DST), (ALIGNED_SIZE, ALIGNED_SIZE), dst=crop)
|
||||
|
||||
@@ -1,23 +1,30 @@
|
||||
import numpy as np
|
||||
from immich_model.constants import FACE_DETECTION_SIZE as DET_SIZE
|
||||
from numpy.typing import NDArray
|
||||
|
||||
from immich_ml.models.base import InferenceModel
|
||||
from immich_ml.models.transforms import decode_pil, letterbox, normalize
|
||||
from immich_ml.schemas import FaceDetectionOutput, ModelTask, ModelType
|
||||
from immich_ml.models.transforms import decode_pil, letterbox, normalize, widen
|
||||
from immich_ml.schemas import FaceDetectionOutput, ModelTask, ModelType, Shape
|
||||
from immich_ml.sessions.policy import ShapePolicy
|
||||
|
||||
from ._ops import DET_SIZE, decode_scrfd, nms
|
||||
from ._ops import decode_scrfd, nms
|
||||
|
||||
|
||||
class FaceDetector(InferenceModel):
|
||||
depends = []
|
||||
identity = (ModelType.DETECTION, ModelTask.FACIAL_RECOGNITION)
|
||||
shape_policy = ShapePolicy(dims=(Shape(batch=1, height=DET_SIZE, width=DET_SIZE),))
|
||||
|
||||
def _predict(self, inputs: NDArray[np.uint8] | bytes, minScore: float) -> FaceDetectionOutput:
|
||||
session = self.session.for_shape(self.shape_policy.dims[0])
|
||||
canvas, scale = letterbox(decode_pil(inputs), DET_SIZE)
|
||||
blob = normalize(canvas.astype(np.float32), mean=127.5, std=128).transpose(2, 0, 1)[None]
|
||||
blob: NDArray[np.float32] | NDArray[np.uint8] = (
|
||||
canvas[None]
|
||||
if session.normalizes_input
|
||||
else normalize(canvas.astype(np.float32), 127.5, 128).transpose(2, 0, 1)[None]
|
||||
)
|
||||
|
||||
input_name = self.session.get_inputs()[0].name
|
||||
heads = self.session.run(None, {input_name: blob})
|
||||
heads = [widen(head) for head in session.run(None, {session.get_inputs()[0].name: blob})]
|
||||
scores, boxes, kps = decode_scrfd(heads, DET_SIZE)
|
||||
|
||||
candidates = scores >= minScore
|
||||
|
||||
@@ -1,26 +1,16 @@
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import onnx
|
||||
import onnxruntime as ort
|
||||
from numpy.typing import NDArray
|
||||
from onnx.tools.update_model_dims import update_inputs_outputs_dims
|
||||
from PIL import Image
|
||||
|
||||
from immich_ml.config import log, settings
|
||||
from immich_ml.config import settings
|
||||
from immich_ml.models.base import InferenceModel
|
||||
from immich_ml.models.transforms import decode_pil, normalize, serialize_np_array
|
||||
from immich_ml.schemas import (
|
||||
FaceDetectionOutput,
|
||||
FacialRecognitionOutput,
|
||||
ModelFormat,
|
||||
ModelSession,
|
||||
ModelTask,
|
||||
ModelType,
|
||||
)
|
||||
from immich_ml.schemas import FaceDetectionOutput, FacialRecognitionOutput, ModelTask, ModelType, Shape
|
||||
from immich_ml.sessions.policy import ShapePolicy, batches, runs
|
||||
|
||||
from ._ops import align_face
|
||||
from ._ops import ALIGNED_SIZE, align_face
|
||||
|
||||
|
||||
class FaceRecognizer(InferenceModel):
|
||||
@@ -29,15 +19,8 @@ class FaceRecognizer(InferenceModel):
|
||||
|
||||
def __init__(self, model_name: str, **model_kwargs: Any) -> None:
|
||||
super().__init__(model_name, **model_kwargs)
|
||||
max_batch_size = settings.max_batch_size and settings.max_batch_size.facial_recognition
|
||||
self.batch_size = max_batch_size if max_batch_size else self._batch_size_default
|
||||
|
||||
def _load(self) -> ModelSession:
|
||||
session = self._make_session()
|
||||
if (not self.batch_size or self.batch_size > 1) and str(session.get_inputs()[0].shape[0]) != "batch":
|
||||
self._add_batch_axis(self.model_path)
|
||||
session = self._make_session()
|
||||
return session
|
||||
sizes = batches(settings.max_batch_size.facial_recognition)
|
||||
self.shape_policy = ShapePolicy(dims=tuple(Shape(batch) for batch in sizes))
|
||||
|
||||
def _predict(
|
||||
self, inputs: NDArray[np.uint8] | bytes | Image.Image, faces: FaceDetectionOutput
|
||||
@@ -45,20 +28,26 @@ class FaceRecognizer(InferenceModel):
|
||||
if faces["boxes"].shape[0] == 0:
|
||||
return []
|
||||
image = np.asarray(decode_pil(inputs), dtype=np.uint8)
|
||||
crops = np.stack([align_face(image, kps).transpose(2, 0, 1) for kps in faces["landmarks"]])
|
||||
embeddings = self._predict_batch(normalize(crops, mean=127.5, std=127.5))
|
||||
landmarks = faces["landmarks"]
|
||||
sizes = runs(len(landmarks), self.session.batches)
|
||||
crops: NDArray[np.float32] | NDArray[np.uint8] = np.empty(
|
||||
(len(landmarks), ALIGNED_SIZE, ALIGNED_SIZE, 3), dtype=np.uint8
|
||||
)
|
||||
for crop, kps in zip(crops, landmarks):
|
||||
align_face(image, kps, crop)
|
||||
if not self.session.for_shape(Shape(sizes[0])).normalizes_input:
|
||||
crops = normalize(crops.transpose(0, 3, 1, 2).astype(np.float32, order="C"), mean=127.5, std=127.5)
|
||||
embeddings = self._predict_batch(crops, sizes)
|
||||
return self.postprocess(faces, embeddings)
|
||||
|
||||
def _predict_batch(self, crops: NDArray[np.float32]) -> NDArray[np.float32]:
|
||||
input_name = self.session.get_inputs()[0].name
|
||||
if not self.batch_size or crops.shape[0] <= self.batch_size:
|
||||
return self.session.run(None, {input_name: crops})[0]
|
||||
|
||||
batches = [
|
||||
self.session.run(None, {input_name: crops[i : i + self.batch_size]})[0]
|
||||
for i in range(0, crops.shape[0], self.batch_size)
|
||||
]
|
||||
return np.concatenate(batches, axis=0)
|
||||
def _predict_batch(self, crops: NDArray[np.float32] | NDArray[np.uint8], sizes: list[int]) -> NDArray[np.float32]:
|
||||
embeddings = []
|
||||
start = 0
|
||||
for size in sizes:
|
||||
session = self.session.for_shape(Shape(size))
|
||||
embeddings.append(session.run(None, {session.get_inputs()[0].name: crops[start : start + size]})[0])
|
||||
start += size
|
||||
return embeddings[0] if len(embeddings) == 1 else np.concatenate(embeddings, axis=0)
|
||||
|
||||
def postprocess(self, faces: FaceDetectionOutput, embeddings: NDArray[np.float32]) -> FacialRecognitionOutput:
|
||||
return [
|
||||
@@ -69,24 +58,3 @@ class FaceRecognizer(InferenceModel):
|
||||
}
|
||||
for (x1, y1, x2, y2), embedding, score in zip(faces["boxes"], embeddings, faces["scores"])
|
||||
]
|
||||
|
||||
def _add_batch_axis(self, model_path: Path) -> None:
|
||||
log.debug(f"Adding batch axis to model {model_path}")
|
||||
proto = onnx.load(model_path)
|
||||
static_input_dims = [shape.dim_value for shape in proto.graph.input[0].type.tensor_type.shape.dim[1:]]
|
||||
static_output_dims = [shape.dim_value for shape in proto.graph.output[0].type.tensor_type.shape.dim[1:]]
|
||||
input_dims = {proto.graph.input[0].name: ["batch"] + static_input_dims}
|
||||
output_dims = {proto.graph.output[0].name: ["batch"] + static_output_dims}
|
||||
updated_proto = update_inputs_outputs_dims(proto, input_dims, output_dims)
|
||||
onnx.save(updated_proto, model_path)
|
||||
|
||||
@property
|
||||
def _batch_size_default(self) -> int | None:
|
||||
providers = ort.get_available_providers()
|
||||
if (
|
||||
self.model_format == ModelFormat.ONNX
|
||||
and "MIGraphXExecutionProvider" not in providers
|
||||
and "OpenVINOExecutionProvider" not in providers
|
||||
):
|
||||
return None
|
||||
return 1
|
||||
|
||||
@@ -1,38 +1,83 @@
|
||||
from collections.abc import Sequence
|
||||
from collections.abc import Callable, Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any, Self
|
||||
|
||||
import numpy as np
|
||||
from numpy.typing import NDArray
|
||||
|
||||
from immich_ml.models.transforms import widen
|
||||
|
||||
def greedy(outputs: Sequence[NDArray[Any]]) -> tuple[NDArray[np.int32], NDArray[np.float32]]:
|
||||
"""Models with an in-graph argmax head emit (indices, probs); the rest emit raw logits."""
|
||||
if len(outputs) == 2:
|
||||
return outputs[0], outputs[1]
|
||||
(probs,) = outputs
|
||||
indices = probs.argmax(axis=2)
|
||||
return indices.astype(np.int32), np.take_along_axis(probs, indices[:, :, None], axis=2)[:, :, 0]
|
||||
|
||||
def picked(indices: NDArray[np.int32], confidence: NDArray[Any]) -> tuple[NDArray[np.int32], NDArray[np.float32]]:
|
||||
return indices, widen(confidence)
|
||||
|
||||
|
||||
def kept(indices: NDArray[np.integer[Any]]) -> NDArray[np.bool_]:
|
||||
"""The steps a greedy CTC decode reads a character from: not blank, and not a repeat of the step before."""
|
||||
keep = np.empty(indices.shape, dtype=bool)
|
||||
keep[:, 0] = True
|
||||
np.not_equal(indices[:, 1:], indices[:, :-1], out=keep[:, 1:])
|
||||
keep &= indices != 0
|
||||
return keep
|
||||
|
||||
|
||||
def logits(raw: NDArray[Any]) -> tuple[NDArray[np.int32], NDArray[np.float32]]:
|
||||
steps = raw.reshape(-1, raw.shape[-1]) # a binary keeps a unit axis between the steps and the classes
|
||||
indices = _argmax(steps)
|
||||
keep = kept(indices.reshape(raw.shape[:2])).ravel()
|
||||
rows = widen(steps[keep]) # a confidence is only read where a character is kept
|
||||
best = rows[np.arange(len(rows)), indices[keep]]
|
||||
# the winner's softmax, summing only the few classes within exp(-16) of it rather than every class
|
||||
near = rows > (best - 16)[:, None]
|
||||
counts = np.count_nonzero(near, axis=1)
|
||||
confidence = np.zeros(len(steps), np.float32)
|
||||
if len(rows):
|
||||
weights = np.exp(rows[near] - np.repeat(best, counts))
|
||||
confidence[keep] = 1 / np.add.reduceat(weights, np.cumsum(counts) - counts)
|
||||
return indices.reshape(raw.shape[:2]).astype(np.int32), confidence.reshape(raw.shape[:2])
|
||||
|
||||
|
||||
def _argmax(steps: NDArray[Any]) -> NDArray[np.intp]:
|
||||
if steps.dtype != np.float16:
|
||||
return steps.argmax(axis=1)
|
||||
bits = steps.view(np.int16) # a half's bits order like its value where it is not negative
|
||||
indices = bits.argmax(axis=1)
|
||||
negative = np.flatnonzero(bits[np.arange(len(bits)), indices] < 0) # a step with no logit above zero
|
||||
if len(negative):
|
||||
flipped = bits[negative] ^ ((bits[negative] >> 15) & np.int16(0x7FFF)) # negatives' order restored
|
||||
indices[negative] = flipped.argmax(axis=1)
|
||||
return indices
|
||||
|
||||
|
||||
def probabilities(probs: NDArray[Any]) -> tuple[NDArray[np.int32], NDArray[np.float32]]:
|
||||
# probabilities are non-negative, so half-precision bits order like their values and need no widening to compare
|
||||
indices = (probs.view(np.uint16) if probs.dtype == np.float16 else probs).argmax(axis=2)
|
||||
return indices.astype(np.int32), widen(np.take_along_axis(probs, indices[:, :, None], axis=2)[:, :, 0])
|
||||
|
||||
|
||||
Greedy = Callable[..., tuple[NDArray[np.int32], NDArray[np.float32]]]
|
||||
GREEDY: dict[str, Greedy] = {"ctc_indices": picked, "ctc_logits": logits, "logits_softmax": probabilities}
|
||||
|
||||
|
||||
class CtcDecoder:
|
||||
def __init__(self, charset: list[str]):
|
||||
def __init__(self, charset: list[str], greedy: Greedy):
|
||||
self.charset = ["", *charset, " "] # PP-OCR: blank at 0, space last
|
||||
self.greedy = greedy
|
||||
|
||||
@classmethod
|
||||
def from_file(cls, charset_path: Path) -> Self:
|
||||
def from_file(cls, charset_path: Path, greedy: Greedy) -> Self:
|
||||
if not charset_path.is_file():
|
||||
raise FileNotFoundError(f"Recognition charset not found: {charset_path}")
|
||||
return cls(charset_path.read_text(encoding="utf-8").splitlines())
|
||||
return cls(charset_path.read_text(encoding="utf-8").splitlines(), greedy)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.charset)
|
||||
|
||||
def __call__(self, outputs: Sequence[NDArray[Any]]) -> tuple[list[str], NDArray[np.float32]]:
|
||||
return self.decode(*self.greedy(*outputs))
|
||||
|
||||
def decode(self, indices: NDArray[np.int32], probs: NDArray[np.float32]) -> tuple[list[str], NDArray[np.float32]]:
|
||||
keep = np.empty(indices.shape, dtype=bool)
|
||||
keep[:, 0] = True
|
||||
np.not_equal(indices[:, 1:], indices[:, :-1], out=keep[:, 1:]) # repeats before blanks
|
||||
keep &= indices != 0
|
||||
keep = kept(indices)
|
||||
|
||||
scores: NDArray[np.float32] = np.where(keep, probs, 0).sum(1)
|
||||
scores /= np.maximum(keep.sum(1), 1) # an all-blank row sums to 0, so it stays 0
|
||||
|
||||
@@ -1,29 +1,31 @@
|
||||
from functools import cached_property
|
||||
from collections.abc import Sequence
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
from immich_model.constants import ocr_canvases
|
||||
from numpy.typing import NDArray
|
||||
from PIL import Image
|
||||
from rapidocr.inference_engine.base import FileInfo, InferSession
|
||||
from rapidocr.utils.download_file import DownloadFile, DownloadFileInput
|
||||
from rapidocr.utils.typings import EngineType, LangDet, OCRVersion, TaskType
|
||||
from rapidocr.utils.typings import ModelType as RapidModelType
|
||||
|
||||
from immich_ml.config import log
|
||||
from immich_ml.models.base import InferenceModel
|
||||
from immich_ml.schemas import ModelFormat, ModelSession, ModelTask, ModelType
|
||||
from immich_ml.sessions.ort import OrtSession
|
||||
from immich_ml.models.transforms import widen
|
||||
from immich_ml.schemas import ModelGraph, ModelTask, ModelType, Shape
|
||||
from immich_ml.sessions.policy import ShapePolicy
|
||||
|
||||
from .legacy import TextModel
|
||||
from .postprocess import DBPostProcess
|
||||
from .schemas import TextDetectionOutput
|
||||
|
||||
|
||||
class TextDetector(InferenceModel):
|
||||
class TextDetector(TextModel):
|
||||
depends = []
|
||||
identity = (ModelType.DETECTION, ModelTask.OCR)
|
||||
graph_options = ("maxResolution",)
|
||||
|
||||
def __init__(self, model_name: str, **model_kwargs: Any) -> None:
|
||||
super().__init__(model_name, **model_kwargs, model_format=ModelFormat.ONNX)
|
||||
super().__init__(model_name, **model_kwargs)
|
||||
short_side = model_kwargs.get("maxResolution", 736)
|
||||
# RKNPU ships a binary per short side, which the label picks
|
||||
canvases = tuple(Shape(batch=1, **canvas) for canvas in ocr_canvases(short_side))
|
||||
self.shape_policy = ShapePolicy(dims=canvases, label=f"res{short_side}")
|
||||
self.scale = np.float32(1.0 / 127.5) # (x/255 - 0.5) / 0.5
|
||||
self._empty: TextDetectionOutput = {
|
||||
"boxes": np.empty(0, dtype=np.float32),
|
||||
@@ -31,28 +33,6 @@ class TextDetector(InferenceModel):
|
||||
}
|
||||
self.postprocess = DBPostProcess(thresh=0.3, max_candidates=1000, unclip_ratio=1.6, use_dilation=True)
|
||||
|
||||
def _download(self) -> None:
|
||||
model_info = InferSession.get_model_url(
|
||||
FileInfo(
|
||||
engine_type=EngineType.ONNXRUNTIME,
|
||||
ocr_version=OCRVersion.PPOCRV5,
|
||||
task_type=TaskType.DET,
|
||||
lang_type=LangDet.CH,
|
||||
model_type=RapidModelType.MOBILE if "mobile" in self.model_name else RapidModelType.SERVER,
|
||||
)
|
||||
)
|
||||
download_params = DownloadFileInput(
|
||||
file_url=model_info["model_dir"],
|
||||
sha256=model_info["SHA256"],
|
||||
save_path=self.model_path,
|
||||
logger=log,
|
||||
)
|
||||
DownloadFile.run(download_params)
|
||||
|
||||
def _load(self) -> ModelSession:
|
||||
# TODO: support other runtime sessions
|
||||
return OrtSession(self.model_path)
|
||||
|
||||
def _predict(
|
||||
self, inputs: Image.Image, maxResolution: int = 736, minScore: float = 0.5, scoreMode: str = "fast"
|
||||
) -> TextDetectionOutput:
|
||||
@@ -60,9 +40,8 @@ class TextDetector(InferenceModel):
|
||||
if width < 32 or height < 32:
|
||||
return self._empty
|
||||
|
||||
image, content = self._transform(inputs, maxResolution)
|
||||
input_name = self.session.get_inputs()[0].name
|
||||
probs = self.session.run(None, {input_name: image})[0][0]
|
||||
session, image, content = self._transform(inputs, maxResolution)
|
||||
probs = widen(session.run(None, {session.get_inputs()[0].name: image})[0][0])
|
||||
if probs.ndim == 3:
|
||||
probs = probs[0]
|
||||
boxes, scores = self.postprocess(probs[: content[0], : content[1]], (height, width), minScore, scoreMode)
|
||||
@@ -71,27 +50,46 @@ class TextDetector(InferenceModel):
|
||||
order = self.reading_order(boxes)
|
||||
return {"boxes": boxes[order], "scores": scores[order]}
|
||||
|
||||
@cached_property
|
||||
def raw_input(self) -> bool:
|
||||
return self.session.get_inputs()[0].shape[-1] == 3 # NHWC models handle normalization and transpose internally
|
||||
|
||||
def _transform(self, img: Image.Image, max_resolution: int) -> tuple[NDArray[Any], tuple[int, int]]:
|
||||
def _transform(self, img: Image.Image, max_resolution: int) -> tuple[ModelGraph, NDArray[Any], tuple[int, int]]:
|
||||
ratio = min(max_resolution / min(img.height, img.width), 1.0)
|
||||
resize_h = max(self._round32(img.height * ratio), 32)
|
||||
resize_w = max(self._round32(img.width * ratio), 32)
|
||||
resized = img.resize((resize_w, resize_h), resample=Image.Resampling.LANCZOS)
|
||||
content_h = max(self._round32(img.height * ratio), 32)
|
||||
content_w = max(self._round32(img.width * ratio), 32)
|
||||
canvas = self._canvas(self.session.shapes, content_h, content_w)
|
||||
session = self.session.for_shape(Shape(batch=1, height=canvas[0], width=canvas[1]))
|
||||
# a compiled canvas may be smaller than the image was sized for, so it shrinks again
|
||||
scale = min(canvas[0] / content_h, canvas[1] / content_w, 1.0)
|
||||
content_h = min(max(self._round32(content_h * scale), 32), canvas[0])
|
||||
content_w = min(max(self._round32(content_w * scale), 32), canvas[1])
|
||||
resized = img.resize((content_w, content_h), resample=Image.Resampling.LANCZOS)
|
||||
if resized.mode != "RGB":
|
||||
resized = resized.convert("RGB")
|
||||
array = np.asarray(resized)
|
||||
if self.raw_input:
|
||||
return array[None], (resize_h, resize_w)
|
||||
content = (content_h, content_w)
|
||||
height, width = canvas
|
||||
if session.normalizes_input:
|
||||
if canvas == content:
|
||||
return session, array[None], content
|
||||
image = np.zeros((1, height, width, 3), dtype=np.uint8) # what is left of the canvas stays black
|
||||
image[0, :content_h, :content_w] = array
|
||||
return session, image, content
|
||||
# reverse plane order gets the BGR swap and the CHW transpose for free
|
||||
out = np.empty((1, 3, resize_h, resize_w), dtype=np.float32)
|
||||
out = np.empty((1, 3, height, width), dtype=np.float32)
|
||||
if canvas != content: # -1.0 is what the graph would have made of a black pixel
|
||||
out[0, :, content_h:] = -1.0
|
||||
out[0, :, :content_h, content_w:] = -1.0
|
||||
for channel in range(3):
|
||||
plane = out[0, 2 - channel]
|
||||
plane = out[0, 2 - channel, :content_h, :content_w]
|
||||
np.multiply(array[:, :, channel], self.scale, out=plane)
|
||||
plane -= 1.0
|
||||
return out, (resize_h, resize_w)
|
||||
return session, out, content
|
||||
|
||||
@staticmethod
|
||||
def _canvas(shapes: Sequence[Shape], height: int, width: int) -> tuple[int, int]:
|
||||
return min(
|
||||
((shape.height, shape.width) for shape in shapes if shape.height and shape.width),
|
||||
key=lambda canvas: (-min(canvas[0] / height, canvas[1] / width, 1.0), canvas[0] * canvas[1]),
|
||||
default=(height, width),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _round32(value: float) -> int:
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
from typing import Any
|
||||
|
||||
from immich_ml.config import log, settings
|
||||
from immich_ml.models.base import InferenceModel
|
||||
from immich_ml.schemas import ModelFormat, ModelType
|
||||
|
||||
|
||||
# TODO: Remove once everything uses the new model graphs
|
||||
class TextModel(InferenceModel):
|
||||
def __init__(self, model_name: str, **model_kwargs: Any) -> None:
|
||||
if settings.legacy_models:
|
||||
model_kwargs["model_format"] = ModelFormat.ONNX # the older exports come in no other format
|
||||
super().__init__(model_name, **model_kwargs)
|
||||
|
||||
def download(self) -> None:
|
||||
if not settings.legacy_models or self.cached:
|
||||
return super().download()
|
||||
# here, as only an install still on the older exports has a use for it
|
||||
from rapidocr.inference_engine.base import FileInfo, InferSession
|
||||
from rapidocr.utils.download_file import DownloadFile, DownloadFileInput
|
||||
from rapidocr.utils.typings import EngineType, LangDet, LangRec, OCRVersion, TaskType
|
||||
from rapidocr.utils.typings import ModelType as RapidModelType
|
||||
|
||||
detects = self.model_type == ModelType.DETECTION
|
||||
language = self.model_name.split("__")[0] if "__" in self.model_name else "CH"
|
||||
model_info = InferSession.get_model_url(
|
||||
FileInfo(
|
||||
engine_type=EngineType.ONNXRUNTIME,
|
||||
ocr_version=OCRVersion.PPOCRV5,
|
||||
task_type=TaskType.DET if detects else TaskType.REC,
|
||||
lang_type=LangDet.CH if detects else LangRec[language],
|
||||
model_type=RapidModelType.MOBILE if "mobile" in self.model_name else RapidModelType.SERVER,
|
||||
)
|
||||
)
|
||||
log.info(
|
||||
f"Downloading {self.model_type} model '{self.model_name}' to {self.model_path}. This may take a while."
|
||||
)
|
||||
DownloadFile.run(
|
||||
DownloadFileInput(
|
||||
file_url=model_info["model_dir"], sha256=model_info["SHA256"], save_path=self.model_path, logger=log
|
||||
)
|
||||
)
|
||||
@@ -1,77 +1,57 @@
|
||||
import math
|
||||
from functools import cached_property
|
||||
from bisect import bisect_left
|
||||
from typing import Any
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from immich_model.constants import OCR_RECOGNITION_WIDTHS
|
||||
from numpy.typing import NDArray
|
||||
from PIL import Image
|
||||
from rapidocr import LangRec
|
||||
from rapidocr.inference_engine.base import FileInfo, InferSession
|
||||
from rapidocr.utils.download_file import DownloadFile, DownloadFileInput
|
||||
from rapidocr.utils.typings import EngineType, OCRVersion, TaskType
|
||||
from rapidocr.utils.typings import ModelType as RapidModelType
|
||||
|
||||
from immich_ml.config import log, settings
|
||||
from immich_ml.models.base import InferenceModel
|
||||
from immich_ml.schemas import ModelFormat, ModelSession, ModelTask, ModelType
|
||||
from immich_ml.sessions.ort import OrtSession
|
||||
from immich_ml.config import settings
|
||||
from immich_ml.schemas import ModelGraph, ModelSession, ModelTask, ModelType, Shape
|
||||
from immich_ml.sessions.policy import ShapePolicy, batches, runs
|
||||
|
||||
from .ctc import CtcDecoder, greedy
|
||||
from .ctc import GREEDY, CtcDecoder, probabilities
|
||||
from .legacy import TextModel
|
||||
from .schemas import TextDetectionOutput, TextRecognitionOutput
|
||||
|
||||
REC_HEIGHT = 48
|
||||
REC_BASE_RATIO = 320 / REC_HEIGHT # PP-OCR's rec_image_shape floor
|
||||
SCALE = np.float32(1.0 / 127.5)
|
||||
|
||||
|
||||
class TextRecognizer(InferenceModel):
|
||||
class TextRecognizer(TextModel):
|
||||
depends = [(ModelType.DETECTION, ModelTask.OCR)]
|
||||
identity = (ModelType.RECOGNITION, ModelTask.OCR)
|
||||
|
||||
def __init__(self, model_name: str, **model_kwargs: Any) -> None:
|
||||
self.language = LangRec[model_name.split("__")[0]] if "__" in model_name else LangRec.CH
|
||||
self._empty: TextRecognitionOutput = {
|
||||
"box": np.empty(0, dtype=np.float32),
|
||||
"boxScore": np.empty(0, dtype=np.float32),
|
||||
"text": [],
|
||||
"textScore": np.empty(0, dtype=np.float32),
|
||||
}
|
||||
super().__init__(model_name, **model_kwargs, model_format=ModelFormat.ONNX)
|
||||
max_batch_size = settings.max_batch_size and settings.max_batch_size.ocr
|
||||
self.batch_size = max_batch_size if max_batch_size else 6
|
||||
|
||||
def _download(self) -> None:
|
||||
model_info = InferSession.get_model_url(
|
||||
FileInfo(
|
||||
engine_type=EngineType.ONNXRUNTIME,
|
||||
ocr_version=OCRVersion.PPOCRV5,
|
||||
task_type=TaskType.REC,
|
||||
lang_type=self.language,
|
||||
model_type=RapidModelType.MOBILE if "mobile" in self.model_name else RapidModelType.SERVER,
|
||||
)
|
||||
super().__init__(model_name, **model_kwargs)
|
||||
self.decoder: CtcDecoder | None = None
|
||||
sizes = batches(settings.max_batch_size.ocr)
|
||||
self.shape_policy = ShapePolicy(
|
||||
dims=tuple(Shape(batch, width=width) for width in OCR_RECOGNITION_WIDTHS for batch in sizes)
|
||||
)
|
||||
download_params = DownloadFileInput(
|
||||
file_url=model_info["model_dir"],
|
||||
sha256=model_info["SHA256"],
|
||||
save_path=self.model_path,
|
||||
logger=log,
|
||||
)
|
||||
DownloadFile.run(download_params)
|
||||
|
||||
def _load(self) -> ModelSession:
|
||||
# TODO: support other runtimes
|
||||
session = OrtSession(self.model_path)
|
||||
self.decoder = (
|
||||
CtcDecoder(character.splitlines())
|
||||
if (character := session.get_metadata().get("character")) is not None
|
||||
else CtcDecoder.from_file(self.model_dir / "charset.txt")
|
||||
)
|
||||
session = super()._load()
|
||||
# the widths it was compiled for, narrowest first: none where it takes any
|
||||
self.widths = tuple(sorted({shape.width for shape in session.shapes if shape.width is not None}))
|
||||
return session
|
||||
|
||||
@cached_property
|
||||
def raw_input(self) -> bool:
|
||||
return self.session.get_inputs()[0].shape[-1] == 3 # NHWC models handle normalization and transpose internally
|
||||
def _decoder(self, session: ModelGraph) -> CtcDecoder:
|
||||
character = session.get_metadata().get("character")
|
||||
greedy = GREEDY.get(session.get_outputs()[0].name, probabilities)
|
||||
return (
|
||||
CtcDecoder(character.splitlines(), greedy)
|
||||
if character is not None
|
||||
else CtcDecoder.from_file(self.model_dir / "charset.txt", greedy)
|
||||
)
|
||||
|
||||
def _predict(self, img: Image.Image, texts: TextDetectionOutput, minScore: float = 0.9) -> TextRecognitionOutput:
|
||||
boxes, box_scores = texts["boxes"], texts["scores"]
|
||||
@@ -85,28 +65,35 @@ class TextRecognizer(InferenceModel):
|
||||
|
||||
text_list: list[str] = [""] * len(order)
|
||||
score_list = np.zeros(len(order), dtype=np.float32)
|
||||
input_name = self.session.get_inputs()[0].name
|
||||
for start in range(0, len(order), self.batch_size):
|
||||
chunk = order[start : start + self.batch_size]
|
||||
width = int(REC_HEIGHT * max(REC_BASE_RATIO, ratios[chunk[-1]]))
|
||||
images: NDArray[Any] = (
|
||||
np.full((len(chunk), REC_HEIGHT, width, 3), 127, dtype=np.uint8)
|
||||
if self.raw_input
|
||||
else np.zeros((len(chunk), 3, REC_HEIGHT, width), dtype=np.float32)
|
||||
)
|
||||
sized = runs(len(order), self.session.batches)
|
||||
bounds = [sum(sized[:index]) for index in range(len(sized))]
|
||||
chunked = [order[start : start + size] for start, size in zip(bounds, sized)]
|
||||
batch_widths = [self._width(int(REC_HEIGHT * ratios[chunk[-1]])) for chunk in chunked]
|
||||
graphs = [self.session.for_shape(Shape(len(chunk), width=width)) for chunk, width in zip(chunked, batch_widths)]
|
||||
raw = graphs[0].normalizes_input
|
||||
# one buffer per request, sized for the largest batch and reshaped for each
|
||||
sizes = [len(chunk) * REC_HEIGHT * width * 3 for chunk, width in zip(chunked, batch_widths)]
|
||||
buffer = np.empty(max(sizes), dtype=np.uint8 if raw else np.float32)
|
||||
for chunk, width, size, session in zip(chunked, batch_widths, sizes, graphs):
|
||||
shape = (len(chunk), REC_HEIGHT, width, 3) if raw else (len(chunk), 3, REC_HEIGHT, width)
|
||||
images: NDArray[Any] = buffer[:size].reshape(shape)
|
||||
for i, index in enumerate(chunk):
|
||||
crop = self._crop(img, coeffs[index], widths[index], heights[index], upright[index])
|
||||
resized_w = max(1, min(width, math.ceil(REC_HEIGHT * ratios[index])))
|
||||
if self.raw_input:
|
||||
if raw:
|
||||
cv2.resize(crop, (resized_w, REC_HEIGHT), dst=images[i, :, :resized_w])
|
||||
images[i, :, resized_w:] = 127 # the pad the crop is read against, over whatever was there
|
||||
continue
|
||||
resized = cv2.resize(crop, (resized_w, REC_HEIGHT))
|
||||
view = images[i, :, :, :resized_w]
|
||||
np.multiply(resized.transpose(2, 0, 1)[::-1], SCALE, out=view) # [::-1] is the RGB -> BGR swap
|
||||
view -= 1.0
|
||||
images[i, :, :, resized_w:] = 0
|
||||
|
||||
out_indices, out_probs = greedy(self.session.run(None, {input_name: images}))
|
||||
chunk_texts, chunk_scores = self.decoder.decode(out_indices, out_probs)
|
||||
outputs = session.run(None, {session.get_inputs()[0].name: images})
|
||||
if self.decoder is None:
|
||||
self.decoder = self._decoder(session)
|
||||
chunk_texts, chunk_scores = self.decoder(outputs)
|
||||
for index, text, score in zip(chunk, chunk_texts, chunk_scores):
|
||||
text_list[index] = text
|
||||
score_list[index] = score
|
||||
@@ -123,6 +110,11 @@ class TextRecognizer(InferenceModel):
|
||||
"textScore": score_list[valid],
|
||||
}
|
||||
|
||||
def _width(self, width: int) -> int:
|
||||
if not self.widths: # room past the text in proportion to it, as the recognizer reads a line worse without
|
||||
return min(OCR_RECOGNITION_WIDTHS[-1], max(OCR_RECOGNITION_WIDTHS[0], math.ceil(width * 1.25)))
|
||||
return self.widths[min(bisect_left(self.widths, width), len(self.widths) - 1)]
|
||||
|
||||
def _crop_geometry(
|
||||
self, boxes: NDArray[np.float32]
|
||||
) -> tuple[NDArray[np.int32], NDArray[np.int32], NDArray[np.float32]]:
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import string
|
||||
from io import BytesIO
|
||||
from typing import Any
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
@@ -43,10 +44,6 @@ def get_pil_resampling(resample: str) -> Image.Resampling:
|
||||
return _PIL_RESAMPLING_METHODS[resample.lower()]
|
||||
|
||||
|
||||
def pil_to_cv2(image: Image.Image) -> NDArray[np.uint8]:
|
||||
return cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) # type: ignore
|
||||
|
||||
|
||||
def decode_pil(image_bytes: bytes | Image.Image | NDArray[np.uint8]) -> Image.Image:
|
||||
image: Image.Image
|
||||
match image_bytes:
|
||||
@@ -75,6 +72,13 @@ def serialize_np_array(arr: NDArray[np.float32]) -> str:
|
||||
return orjson.dumps(arr, option=orjson.OPT_SERIALIZE_NUMPY).decode()
|
||||
|
||||
|
||||
def widen(array: NDArray[Any]) -> NDArray[Any]:
|
||||
"""Single precision from half through OpenCV, as numpy converts one value at a time without native half support."""
|
||||
if array.dtype != np.float16:
|
||||
return array
|
||||
return cv2.convertFp16(array.reshape(-1, array.shape[-1]).view(np.int16)).reshape(array.shape) # int16 carries half
|
||||
|
||||
|
||||
def letterbox(image: NDArray[np.uint8] | Image.Image, size: int) -> tuple[NDArray[np.uint8], float]:
|
||||
if isinstance(image, Image.Image):
|
||||
image = np.asarray(image)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from collections.abc import Sequence
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import asdict, dataclass
|
||||
from enum import Enum
|
||||
from typing import Any, Literal, Protocol, TypeVar
|
||||
from typing import Any, Literal, Protocol, TypeAlias, TypeVar
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
@@ -44,11 +45,6 @@ class ModelSource(StrEnum):
|
||||
PADDLE = "paddle"
|
||||
|
||||
|
||||
class ModelPrecision(StrEnum):
|
||||
FP16 = "FP16"
|
||||
FP32 = "FP32"
|
||||
|
||||
|
||||
class ModelOrganization(StrEnum):
|
||||
APP = "immich-app"
|
||||
TESTING = "immich-testing"
|
||||
@@ -62,16 +58,29 @@ class SessionNode(Protocol):
|
||||
def name(self) -> str: ...
|
||||
|
||||
@property
|
||||
def shape(self) -> tuple[int, ...]: ...
|
||||
def shape(self) -> tuple[int | str, ...]: ... # ORT names a symbolic dim rather than sizing it
|
||||
|
||||
|
||||
class ModelSession(Protocol):
|
||||
ModelTensor: TypeAlias = npt.NDArray[np.float32] | npt.NDArray[np.int32] | npt.NDArray[np.uint8]
|
||||
ModelInput = Mapping[str, ModelTensor]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Shape:
|
||||
batch: int
|
||||
height: int | None = None
|
||||
width: int | None = None
|
||||
|
||||
@property
|
||||
def pins(self) -> dict[str, int]:
|
||||
return {name: size for name, size in asdict(self).items() if size is not None}
|
||||
|
||||
|
||||
class ModelGraph(Protocol):
|
||||
def run(
|
||||
self,
|
||||
output_names: list[str] | None,
|
||||
input_feed: dict[str, npt.NDArray[np.float32]]
|
||||
| dict[str, npt.NDArray[np.int32]]
|
||||
| dict[str, npt.NDArray[np.uint8]],
|
||||
input_feed: ModelInput,
|
||||
run_options: Any = None,
|
||||
) -> list[npt.NDArray[np.float32]]: ...
|
||||
|
||||
@@ -81,6 +90,21 @@ class ModelSession(Protocol):
|
||||
|
||||
def get_metadata(self) -> dict[str, str]: ...
|
||||
|
||||
@property
|
||||
def normalizes_input(self) -> bool: ...
|
||||
|
||||
|
||||
class ModelSession(Protocol):
|
||||
@property
|
||||
def shapes(self) -> tuple[Shape, ...]: ...
|
||||
|
||||
@property
|
||||
def batches(self) -> tuple[int, ...]: ...
|
||||
|
||||
def for_shape(self, shape: Shape) -> ModelGraph: ...
|
||||
|
||||
def warm(self) -> None: ...
|
||||
|
||||
|
||||
class FaceDetectionOutput(TypedDict):
|
||||
boxes: npt.NDArray[np.float32]
|
||||
|
||||
@@ -7,7 +7,7 @@ import numpy as np
|
||||
from numpy.typing import NDArray
|
||||
|
||||
from immich_ml.config import log, settings
|
||||
from immich_ml.schemas import SessionNode
|
||||
from immich_ml.schemas import ModelInput, SessionNode, Shape
|
||||
|
||||
from .loader import Ann
|
||||
|
||||
@@ -29,6 +29,9 @@ class AnnSession:
|
||||
fp16=settings.ann_fp16_turbo,
|
||||
)
|
||||
log.info("Loaded ANN model with ID %d", self.model)
|
||||
batch = self.ann.input_shapes[self.model][0][0]
|
||||
self.shapes = (Shape(batch),) # the artifact is compiled for one shape
|
||||
self.batches = (batch,)
|
||||
|
||||
def __del__(self) -> None:
|
||||
self.ann.unload(self.model)
|
||||
@@ -46,10 +49,20 @@ class AnnSession:
|
||||
def get_metadata(self) -> dict[str, str]:
|
||||
return {}
|
||||
|
||||
@property
|
||||
def normalizes_input(self) -> bool:
|
||||
return False # no ARM-NN artifact is built from a graph that carries its own preprocessing
|
||||
|
||||
def for_shape(self, shape: Shape) -> AnnSession:
|
||||
return self
|
||||
|
||||
def warm(self) -> None:
|
||||
pass # the artifact is loaded whole when the session opens
|
||||
|
||||
def run(
|
||||
self,
|
||||
output_names: list[str] | None,
|
||||
input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]] | dict[str, NDArray[np.uint8]],
|
||||
input_feed: ModelInput,
|
||||
run_options: Any = None,
|
||||
) -> list[NDArray[np.float32]]:
|
||||
inputs: list[NDArray[np.float32]] = [np.ascontiguousarray(v) for v in input_feed.values()]
|
||||
|
||||
@@ -1,77 +1,56 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import ctypes
|
||||
import pickle
|
||||
import platform
|
||||
import subprocess
|
||||
import sys
|
||||
from collections.abc import Sequence
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from functools import cache, cached_property
|
||||
from pathlib import Path
|
||||
from threading import Lock
|
||||
from typing import Any
|
||||
from typing import Any, NamedTuple
|
||||
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
from immich_model.runtime import RewriteContext, RewritePlan, plan_rewrites
|
||||
from numpy.typing import NDArray
|
||||
from onnxruntime.capi.onnxruntime_pybind11_state import InvalidProtobuf
|
||||
from pydantic import BaseModel
|
||||
|
||||
from immich_ml.models.constants import SUPPORTED_PROVIDERS
|
||||
from immich_ml.schemas import ModelPrecision, SessionNode
|
||||
from immich_ml.schemas import ModelInput, SessionNode, Shape
|
||||
|
||||
from ..config import log, settings
|
||||
from .policy import ShapePolicy
|
||||
|
||||
MigraphxInputSignature = tuple[tuple[str, str, tuple[int, ...]], ...]
|
||||
|
||||
_migraphx_registry_lock = Lock()
|
||||
_migraphx_model_locks: dict[str, Lock] = {}
|
||||
_migraphx_compiled_inputs: set[tuple[str, MigraphxInputSignature]] = set()
|
||||
# the one provider that handles a free dim well; the rest miscompile it, recompile inside the run,
|
||||
# or partition around it, so they are handed one graph per shape instead
|
||||
DYNAMIC_PROVIDERS = frozenset({"CPUExecutionProvider"})
|
||||
|
||||
|
||||
def _migraphx_get_model_lock(model_key: str) -> Lock:
|
||||
with _migraphx_registry_lock:
|
||||
lock = _migraphx_model_locks.get(model_key)
|
||||
if lock is None:
|
||||
lock = Lock()
|
||||
_migraphx_model_locks[model_key] = lock
|
||||
return lock
|
||||
def _label(pins: Mapping[str, int]) -> str:
|
||||
return "_".join(f"{name}{size}" for name, size in pins.items())
|
||||
|
||||
|
||||
def _migraphx_has_compiled_input(key: tuple[str, MigraphxInputSignature]) -> bool:
|
||||
with _migraphx_registry_lock:
|
||||
return key in _migraphx_compiled_inputs
|
||||
UNREADABLE = 3 # how the preparing child says its source does not parse
|
||||
DTYPES = {
|
||||
"tensor(float)": np.float32,
|
||||
"tensor(float16)": np.float16,
|
||||
"tensor(uint8)": np.uint8,
|
||||
"tensor(int32)": np.int32,
|
||||
"tensor(int64)": np.int64,
|
||||
}
|
||||
|
||||
|
||||
def _migraphx_mark_compiled_input(key: tuple[str, MigraphxInputSignature]) -> None:
|
||||
with _migraphx_registry_lock:
|
||||
_migraphx_compiled_inputs.add(key)
|
||||
|
||||
|
||||
def _migraphx_input_signature(input_feed: dict[str, NDArray[Any]]) -> MigraphxInputSignature:
|
||||
return tuple((name, str(value.dtype), tuple(value.shape)) for name, value in sorted(input_feed.items()))
|
||||
|
||||
|
||||
class OrtSession:
|
||||
session: ort.InferenceSession
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_path: Path | str,
|
||||
providers: list[str] | None = None,
|
||||
provider_options: list[dict[str, Any]] | None = None,
|
||||
sess_options: ort.SessionOptions | None = None,
|
||||
threads: int = 2,
|
||||
):
|
||||
self.model_path = Path(model_path)
|
||||
self.threads = threads
|
||||
self.providers = providers if providers is not None else self._providers_default
|
||||
self.disabled_optimizers = _disabled_optimizers_default(self.providers)
|
||||
log.debug(f"Setting disabled_optimizers to {self.disabled_optimizers}")
|
||||
self.provider_options = provider_options if provider_options is not None else self._provider_options_default
|
||||
self.sess_options = sess_options if sess_options is not None else self._sess_options_default
|
||||
self.session = ort.InferenceSession(
|
||||
self.model_path.as_posix(),
|
||||
providers=self.providers,
|
||||
provider_options=self.provider_options,
|
||||
sess_options=self.sess_options,
|
||||
disabled_optimizers=self.disabled_optimizers,
|
||||
)
|
||||
class OrtGraph:
|
||||
def __init__(self, spec: GraphSpec) -> None:
|
||||
self.session = spec.session(prepared(spec))
|
||||
image = self.session.get_inputs()[0]
|
||||
# whether the graph scales and shifts the image itself, so the host hands over raw pixels
|
||||
self.normalizes_input = image.type == "tensor(uint8)"
|
||||
# CoreML's rewrite moves the layout change out of the graph, so whoever applied it performs it
|
||||
self.channels_first = self.normalizes_input and len(image.shape) == 4 and image.shape[1] == 3
|
||||
|
||||
def get_inputs(self) -> Sequence[SessionNode]:
|
||||
inputs: Sequence[SessionNode] = self.session.get_inputs()
|
||||
@@ -88,52 +67,91 @@ class OrtSession:
|
||||
def run(
|
||||
self,
|
||||
output_names: list[str] | None,
|
||||
input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]] | dict[str, NDArray[np.uint8]],
|
||||
input_feed: ModelInput,
|
||||
run_options: Any = None,
|
||||
) -> list[NDArray[np.float32]]:
|
||||
if "MIGraphXExecutionProvider" in self.providers:
|
||||
model_key = self.model_path.resolve().as_posix()
|
||||
input_key = (model_key, _migraphx_input_signature(input_feed))
|
||||
if not _migraphx_has_compiled_input(input_key):
|
||||
model_lock = _migraphx_get_model_lock(model_key)
|
||||
with model_lock:
|
||||
if not _migraphx_has_compiled_input(input_key):
|
||||
outputs: list[NDArray[np.float32]] = self.session.run(output_names, input_feed, run_options)
|
||||
_migraphx_mark_compiled_input(input_key)
|
||||
return outputs
|
||||
|
||||
outputs = self.session.run(output_names, input_feed, run_options)
|
||||
return outputs
|
||||
|
||||
outputs = self.session.run(output_names, input_feed, run_options)
|
||||
if self.channels_first:
|
||||
input_feed = {name: frames.transpose(0, 3, 1, 2) for name, frames in input_feed.items()}
|
||||
outputs: list[NDArray[Any]] = self.session.run(output_names, input_feed, run_options)
|
||||
return outputs
|
||||
|
||||
@property
|
||||
def providers(self) -> list[str]:
|
||||
return self._providers
|
||||
|
||||
@providers.setter
|
||||
def providers(self, providers: list[str]) -> None:
|
||||
log.info(f"Setting execution providers to {providers}, in descending order of preference")
|
||||
self._providers = providers
|
||||
class Device(NamedTuple):
|
||||
kind: str # what the provider compiles for, which names the directory
|
||||
version: str # what the facts go stale under
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GraphSpec:
|
||||
model_path: Path
|
||||
pins: Mapping[str, int]
|
||||
overrides: Sequence[tuple[str, int]]
|
||||
providers: list[str]
|
||||
disabled_optimizers: list[str]
|
||||
threads: int = 2 # for the CPU, where they are not configured
|
||||
|
||||
@property
|
||||
def _providers_default(self) -> list[str]:
|
||||
available_providers = set(ort.get_available_providers())
|
||||
log.debug(f"Available ORT providers: {available_providers}")
|
||||
return [provider for provider in SUPPORTED_PROVIDERS if provider in available_providers]
|
||||
def provider(self) -> str:
|
||||
return self.providers[0] # the rest take only what it cannot run
|
||||
|
||||
@cached_property
|
||||
def directory(self) -> Path:
|
||||
provider = self.provider.removesuffix("ExecutionProvider").lower()
|
||||
kind = self.device.kind if self.device else ""
|
||||
return self.model_path.parent / provider / kind / _label(self.pins)
|
||||
|
||||
@cached_property
|
||||
def device(self) -> Device | None:
|
||||
match self.provider:
|
||||
case "CPUExecutionProvider":
|
||||
return Device(platform.machine(), _cpu()) # a prepack built for other instruction sets fails at runtime
|
||||
case "OpenVINOExecutionProvider" if self.openvino_device.startswith("GPU"):
|
||||
return _intel_gpu(self.openvino_device)
|
||||
case "MIGraphXExecutionProvider":
|
||||
return _amd_gpu(int(settings.device_id))
|
||||
return None
|
||||
|
||||
@cached_property
|
||||
def openvino_device(self) -> str:
|
||||
device_ids: list[str] = ort.capi._pybind_state.get_available_openvino_device_ids()
|
||||
return f"GPU.{settings.device_id}" if any(d.startswith("GPU") for d in device_ids) else "CPU"
|
||||
|
||||
@cached_property
|
||||
def manifest(self) -> Path:
|
||||
return self.directory / "manifest.json"
|
||||
|
||||
@property
|
||||
def plan(self) -> RewritePlan:
|
||||
return _plan(self.provider)
|
||||
|
||||
@cached_property
|
||||
def half(self) -> bool:
|
||||
return not (settings.legacy_models or self.provider == "CPUExecutionProvider")
|
||||
|
||||
@cached_property
|
||||
def facts(self) -> Facts:
|
||||
return Facts(
|
||||
half=self.half,
|
||||
onnxruntime=ort.__version__,
|
||||
# workers on devices of one kind share what is prepared
|
||||
options={key: value for key, value in self.provider_options[0].items() if not key.startswith("device")},
|
||||
overrides=list(self.overrides),
|
||||
disabled_optimizers=self.disabled_optimizers,
|
||||
rewrite=None if settings.legacy_models else self.plan.digest, # which were never tested with the rewriter
|
||||
device=self.device.version if self.device else None,
|
||||
)
|
||||
|
||||
def session(self, graph: Path, sess_options: ort.SessionOptions | None = None) -> ort.InferenceSession:
|
||||
return ort.InferenceSession(
|
||||
graph.as_posix(),
|
||||
providers=self.providers,
|
||||
provider_options=self.provider_options,
|
||||
sess_options=sess_options or self.sess_options(),
|
||||
disabled_optimizers=self.disabled_optimizers,
|
||||
)
|
||||
|
||||
@cached_property
|
||||
def provider_options(self) -> list[dict[str, Any]]:
|
||||
return self._provider_options
|
||||
|
||||
@provider_options.setter
|
||||
def provider_options(self, provider_options: list[dict[str, Any]]) -> None:
|
||||
log.debug(f"Setting execution provider options to {provider_options}")
|
||||
self._provider_options = provider_options
|
||||
|
||||
@property
|
||||
def _provider_options_default(self) -> list[dict[str, Any]]:
|
||||
provider_options = []
|
||||
for provider in self.providers:
|
||||
match provider:
|
||||
@@ -142,80 +160,245 @@ class OrtSession:
|
||||
case "CUDAExecutionProvider":
|
||||
options = {"arena_extend_strategy": "kSameAsRequested", "device_id": settings.device_id}
|
||||
case "MIGraphXExecutionProvider":
|
||||
migraphx_dir = self.model_path.parent / "migraphx"
|
||||
# MIGraphX does not create the underlying folder and will crash if it does not exist
|
||||
migraphx_dir.mkdir(parents=True, exist_ok=True)
|
||||
options = {
|
||||
"device_id": settings.device_id,
|
||||
"migraphx_model_cache_dir": migraphx_dir.as_posix(),
|
||||
"migraphx_fp16_enable": "1" if settings.rocm_precision == ModelPrecision.FP16 else "0",
|
||||
}
|
||||
options = {"device_id": settings.device_id, "migraphx_model_cache_dir": self.directory.as_posix()}
|
||||
case "OpenVINOExecutionProvider":
|
||||
device_ids: list[str] = ort.capi._pybind_state.get_available_openvino_device_ids()
|
||||
# Check for available devices, preferring GPU over CPU
|
||||
gpu_devices = [d for d in device_ids if d.startswith("GPU")]
|
||||
if gpu_devices:
|
||||
device_type = f"GPU.{settings.device_id}"
|
||||
log.debug(f"OpenVINO: Using GPU device {device_type}")
|
||||
else:
|
||||
device_type = "CPU"
|
||||
log.debug("OpenVINO: No GPU found, using CPU")
|
||||
options = {
|
||||
"device_type": device_type,
|
||||
"precision": settings.openvino_precision.value,
|
||||
"cache_dir": (self.model_path.parent / "openvino").as_posix(),
|
||||
}
|
||||
options = {"device_type": self.openvino_device, "cache_dir": self.directory.as_posix()}
|
||||
if not self.half: # the GPU would otherwise run an fp32 graph at fp16
|
||||
options["precision"] = "FP32"
|
||||
case "CoreMLExecutionProvider":
|
||||
options = {
|
||||
"ModelFormat": "MLProgram",
|
||||
"MLComputeUnits": "ALL",
|
||||
"SpecializationStrategy": "FastPrediction",
|
||||
"AllowLowPrecisionAccumulationOnGPU": "1",
|
||||
"ModelCacheDirectory": (self.model_path.parent / "coreml").as_posix(),
|
||||
"ModelCacheDirectory": self.directory.as_posix(),
|
||||
}
|
||||
case _:
|
||||
options = {}
|
||||
provider_options.append(options)
|
||||
log.debug(f"Setting execution provider options to {provider_options}")
|
||||
return provider_options
|
||||
|
||||
@property
|
||||
def sess_options(self) -> ort.SessionOptions:
|
||||
return self._sess_options
|
||||
|
||||
@sess_options.setter
|
||||
def sess_options(self, sess_options: ort.SessionOptions) -> None:
|
||||
log.debug(f"Setting execution_mode to {sess_options.execution_mode.name}")
|
||||
log.debug(f"Setting inter_op_num_threads to {sess_options.inter_op_num_threads}")
|
||||
log.debug(f"Setting intra_op_num_threads to {sess_options.intra_op_num_threads}")
|
||||
self._sess_options = sess_options
|
||||
|
||||
@property
|
||||
def _sess_options_default(self) -> ort.SessionOptions:
|
||||
sess_options = ort.SessionOptions()
|
||||
sess_options.enable_cpu_mem_arena = settings.model_arena
|
||||
# some CPUs slow down many times over on subnormal operands, and ORT clears the flush in its threads otherwise
|
||||
sess_options.add_session_config_entry("session.set_denormal_as_zero", "1")
|
||||
sess_options.enable_cpu_mem_arena = settings.model_arena
|
||||
|
||||
# avoid thread contention between models
|
||||
# Set inter_op threads
|
||||
if settings.model_inter_op_threads > 0:
|
||||
sess_options.inter_op_num_threads = settings.model_inter_op_threads
|
||||
# these defaults work well for CPU, but bottleneck GPU
|
||||
elif settings.model_inter_op_threads == 0 and self.providers == ["CPUExecutionProvider"]:
|
||||
elif settings.model_inter_op_threads == 0 and self.provider == "CPUExecutionProvider":
|
||||
sess_options.inter_op_num_threads = 1
|
||||
|
||||
# Set intra_op threads
|
||||
if settings.model_intra_op_threads > 0:
|
||||
sess_options.intra_op_num_threads = settings.model_intra_op_threads
|
||||
elif settings.model_intra_op_threads == 0 and self.providers == ["CPUExecutionProvider"]:
|
||||
elif settings.model_intra_op_threads == 0 and self.provider == "CPUExecutionProvider":
|
||||
sess_options.intra_op_num_threads = self.threads
|
||||
|
||||
if sess_options.inter_op_num_threads > 1:
|
||||
sess_options.execution_mode = ort.ExecutionMode.ORT_PARALLEL
|
||||
|
||||
for dimension, size in self.overrides:
|
||||
sess_options.add_free_dimension_override_by_name(dimension, size)
|
||||
|
||||
log.debug(f"Setting execution_mode to {sess_options.execution_mode.name}")
|
||||
log.debug(f"Setting inter_op_num_threads to {sess_options.inter_op_num_threads}")
|
||||
log.debug(f"Setting intra_op_num_threads to {sess_options.intra_op_num_threads}")
|
||||
return sess_options
|
||||
|
||||
|
||||
class Facts(BaseModel):
|
||||
onnxruntime: str
|
||||
options: dict[str, Any]
|
||||
overrides: list[tuple[str, int]]
|
||||
disabled_optimizers: list[str]
|
||||
rewrite: str | None = None
|
||||
device: str | None = None
|
||||
half: bool = False
|
||||
|
||||
|
||||
class Manifest(BaseModel):
|
||||
facts: Facts
|
||||
graph: str # relative to the artifact's directory
|
||||
|
||||
|
||||
def fresh(spec: GraphSpec) -> Path | None:
|
||||
try:
|
||||
manifest = Manifest.model_validate_json(spec.manifest.read_bytes())
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
return spec.model_path.parent / manifest.graph if manifest.facts == spec.facts else None
|
||||
|
||||
|
||||
def prepared(spec: GraphSpec) -> Path:
|
||||
if (graph := fresh(spec)) is None:
|
||||
log.info(f"Preparing {spec.model_path} for {spec.provider} {dict(spec.pins)}")
|
||||
# child process prepares the graph to avoid memory overhead in the main process
|
||||
child = subprocess.run([sys.executable, "-m", "immich_ml.sessions.prepare"], input=pickle.dumps(spec))
|
||||
if (graph := fresh(spec)) is None:
|
||||
if child.returncode == UNREADABLE:
|
||||
raise InvalidProtobuf(f"{spec.model_path} could not be read")
|
||||
raise RuntimeError(f"Preparing {spec.model_path} failed; the output above says why")
|
||||
return graph
|
||||
|
||||
|
||||
@cache
|
||||
def _plan(provider: str) -> RewritePlan:
|
||||
version = tuple(int(piece) for piece in ort.__version__.split(".")[:3])
|
||||
return plan_rewrites(RewriteContext(target=provider, ort_version=version))
|
||||
|
||||
|
||||
@cache
|
||||
def _cpu() -> str:
|
||||
try:
|
||||
cpuinfo = Path("/proc/cpuinfo").read_text()
|
||||
except OSError:
|
||||
return platform.machine()
|
||||
return next(line for line in cpuinfo.splitlines() if line.startswith(("flags", "Features")))
|
||||
|
||||
|
||||
class OrtSession:
|
||||
def __init__(
|
||||
self, model_path: Path | str, policy: ShapePolicy, providers: list[str] | None = None, threads: int = 2
|
||||
) -> None:
|
||||
self.model_path = Path(model_path)
|
||||
self.policy = policy
|
||||
self.providers = providers if providers is not None else _providers_default()
|
||||
log.info(f"Setting execution providers to {self.providers}, in descending order of preference")
|
||||
self.disabled_optimizers = _disabled_optimizers_default(self.providers)
|
||||
log.debug(f"Setting disabled_optimizers to {self.disabled_optimizers}")
|
||||
self.threads = threads
|
||||
self.dynamic = self.providers[0] in DYNAMIC_PROVIDERS
|
||||
# the shapes to snap onto, where a dim a shape leaves out takes any size
|
||||
self.shapes = (policy.pinned,) if self.dynamic else policy.dims
|
||||
self.batches = tuple(sorted({shape.batch for shape in self.shapes}, reverse=True))
|
||||
self.graphs: dict[Shape, OrtGraph] = {}
|
||||
self.locks: dict[Shape, Lock] = {}
|
||||
if self.dynamic or len(policy.dims) == 1:
|
||||
self.for_shape(policy.dims[0]) # the one graph there is, so nothing waits on a shape
|
||||
|
||||
def for_shape(self, shape: Shape) -> OrtGraph:
|
||||
if self.dynamic: # a batch gains nothing here, and a graph that knows it has one row optimizes further
|
||||
shape = self.shapes[0]
|
||||
if (graph := self.graphs.get(shape)) is not None:
|
||||
return graph
|
||||
if not self.dynamic and shape not in self.policy.dims: # every shape admitted is another graph to compile
|
||||
raise ValueError(f"Asked for {shape}, which is none of {list(self.policy.dims)}")
|
||||
with self.locks.setdefault(shape, Lock()): # requests arrive together for a graph none of them has
|
||||
if (graph := self.graphs.get(shape)) is None:
|
||||
log.debug(f"Building a graph for {shape}")
|
||||
spec = GraphSpec(
|
||||
self.model_path,
|
||||
shape.pins,
|
||||
_overrides(self.policy, shape.pins),
|
||||
self.providers,
|
||||
self.disabled_optimizers,
|
||||
self.threads,
|
||||
)
|
||||
graph = self.graphs[shape] = OrtGraph(spec)
|
||||
return graph
|
||||
|
||||
def warm(self) -> None:
|
||||
"""Opens a graph for every shape and runs it once before opening the next: OpenVINO defers work to the first
|
||||
run of a graph it imports from its cache, which importing another graph first can corrupt."""
|
||||
for shape in self.policy.dims[:1] if self.dynamic else self.policy.dims:
|
||||
session = self.for_shape(shape).session
|
||||
sizes = dict(_overrides(self.policy, shape.pins)) # for a dim the graph leaves free
|
||||
feed = {
|
||||
node.name: np.zeros(
|
||||
[dim if isinstance(dim, int) else sizes.get(dim, 1) for dim in node.shape], DTYPES[node.type]
|
||||
)
|
||||
for node in session.get_inputs()
|
||||
}
|
||||
session.run(None, feed)
|
||||
|
||||
|
||||
def _overrides(policy: ShapePolicy, pins: Mapping[str, int]) -> list[tuple[str, int]]:
|
||||
order = tuple(policy.dims[0].pins)
|
||||
named = []
|
||||
for dimension, size in pins.items():
|
||||
position = order.index(dimension) # a graph that names several dims takes the first pin for all
|
||||
named += [(dimension, size), (f"DynamicDimension.{position}", size)]
|
||||
named.append(("None", size) if position == 0 else ("?", size))
|
||||
return named
|
||||
|
||||
|
||||
@cache
|
||||
def _intel_gpu(device: str) -> Device:
|
||||
"""The IP version and execution units OpenVINO keys its blobs by, and the driver it checks them against."""
|
||||
ov = ctypes.CDLL(str(Path(ort.__file__).parent / "capi" / "libopenvino_c.so"))
|
||||
core, wanted = ctypes.c_void_p(), ctypes.c_char_p()
|
||||
ov.ov_core_create(ctypes.byref(core))
|
||||
ov.ov_core_get_property(core, device.encode(), b"DEVICE_UUID", ctypes.byref(wanted))
|
||||
uuid = wanted.value
|
||||
ov.ov_free(wanted)
|
||||
ov.ov_core_free(core)
|
||||
cl, gpu = ctypes.CDLL("libOpenCL.so.1"), ctypes.c_uint64(1 << 2) # CL_DEVICE_TYPE_GPU
|
||||
platforms, devices, count = (ctypes.c_void_p * 8)(), (ctypes.c_void_p * 8)(), ctypes.c_uint32()
|
||||
cl.clGetPlatformIDs(8, platforms, ctypes.byref(count))
|
||||
for vendor in platforms[: count.value]:
|
||||
found = ctypes.c_uint32()
|
||||
cl.clGetDeviceIDs(ctypes.c_void_p(vendor), gpu, 8, devices, ctypes.byref(found))
|
||||
for handle in devices[: found.value]:
|
||||
own, driver = (ctypes.c_ubyte * 16)(), ctypes.create_string_buffer(64)
|
||||
ip, units = ctypes.c_uint32(), ctypes.c_uint32()
|
||||
# CL_DEVICE_UUID_KHR, CL_DRIVER_VERSION, CL_DEVICE_IP_VERSION_INTEL, CL_DEVICE_MAX_COMPUTE_UNITS
|
||||
for key, value in ((0x106A, own), (0x102D, driver), (0x4250, ip), (0x1002, units)):
|
||||
cl.clGetDeviceInfo(ctypes.c_void_p(handle), key, ctypes.sizeof(value), ctypes.byref(value), None)
|
||||
if bytes(own).hex().encode() == uuid:
|
||||
version = f"{ip.value >> 22}.{ip.value >> 14 & 0xFF}.{ip.value & 0x3FFF}"
|
||||
return Device(f"{version}-{units.value}eu", driver.value.decode())
|
||||
raise LookupError(f"OpenCL has no GPU with the UUID of OpenVINO's {device}")
|
||||
|
||||
|
||||
@cache
|
||||
def _amd_gpu(index: int) -> Device:
|
||||
"""The arch MIGraphX keys its programs by, and the MIGraphX and HIP it loads, which its keys leave out."""
|
||||
# the provider finds ROCm through its runpath
|
||||
ctypes.CDLL(str(Path(ort.__file__).parent / "capi" / "libonnxruntime_providers_migraphx.so"))
|
||||
with open("/proc/self/maps") as maps:
|
||||
loaded = {Path(path).name.split(".so")[0]: Path(path) for *_, path in map(str.split, maps) if ".so" in path}
|
||||
hip, hsa = ctypes.CDLL(str(loaded["libamdhip64"])), ctypes.CDLL(str(loaded["libhsa-runtime64"]))
|
||||
wanted, arches = ctypes.create_string_buffer(64), {}
|
||||
hip.hipDeviceGetPCIBusId(wanted, 64, index) # honors HIP_VISIBLE_DEVICES, unlike HSA
|
||||
|
||||
def visit(agent: int, _: int | None) -> int:
|
||||
name, bdf, domain = ctypes.create_string_buffer(64), ctypes.c_uint32(), ctypes.c_uint32()
|
||||
# HSA_AGENT_INFO_NAME, HSA_AMD_AGENT_INFO_BDFID, HSA_AMD_AGENT_INFO_DOMAIN
|
||||
for key, value in ((0, name), (0xA006, bdf), (0xA00F, domain)):
|
||||
hsa.hsa_agent_get_info(ctypes.c_uint64(agent), key, ctypes.byref(value))
|
||||
bus, device, function = bdf.value >> 8 & 0xFF, bdf.value >> 3 & 0x1F, bdf.value & 0x7
|
||||
arches[f"{domain.value:04x}:{bus:02x}:{device:02x}.{function:01x}"] = name.value.decode()
|
||||
return 0
|
||||
|
||||
hsa.hsa_init()
|
||||
hsa.hsa_iterate_agents(ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_uint64, ctypes.c_void_p)(visit), None)
|
||||
return Device(arches[wanted.value.decode()], f"{loaded['libmigraphx'].name} {loaded['libamdhip64'].name}")
|
||||
|
||||
|
||||
def flush_denormals() -> None:
|
||||
"""Reads subnormal floats as zero on the calling thread, which ORT does only on its own threads."""
|
||||
if sys.platform != "linux" or platform.machine() != "x86_64":
|
||||
return
|
||||
libm = ctypes.CDLL("libm.so.6")
|
||||
env = (ctypes.c_uint32 * 8)() # glibc's fenv_t here: the x87 environment, then MXCSR
|
||||
libm.fegetenv(env)
|
||||
env[7] |= 0x8040 # denormals are zero, flush to zero
|
||||
libm.fesetenv(env)
|
||||
|
||||
|
||||
def _providers_default() -> list[str]:
|
||||
from immich_ml.models.constants import SUPPORTED_PROVIDERS # here, as the preparing child has no use for the models
|
||||
|
||||
available_providers = set(ort.get_available_providers())
|
||||
log.debug(f"Available ORT providers: {available_providers}")
|
||||
return [provider for provider in SUPPORTED_PROVIDERS if provider in available_providers]
|
||||
|
||||
|
||||
def _disabled_optimizers_default(providers: list[str]) -> list[str]:
|
||||
disabled_optimizers: list[str] = []
|
||||
if platform.machine() in ("arm64", "aarch64"): # as macOS and Linux name the same architecture
|
||||
@@ -226,14 +409,3 @@ def _disabled_optimizers_default(providers: list[str]) -> list[str]:
|
||||
disabled_optimizers.append("MatMulAddFusion")
|
||||
|
||||
return disabled_optimizers
|
||||
|
||||
|
||||
def flush_denormals() -> None:
|
||||
"""Reads subnormal floats as zero on the calling thread, which ORT does only on its own threads."""
|
||||
if sys.platform != "linux" or platform.machine() != "x86_64":
|
||||
return
|
||||
libm = ctypes.CDLL("libm.so.6")
|
||||
env = (ctypes.c_uint32 * 8)() # glibc's fenv_t here: the x87 environment, then MXCSR
|
||||
libm.fegetenv(env)
|
||||
env[7] |= 0x8040 # denormals are zero, flush to zero
|
||||
libm.fesetenv(env)
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
|
||||
from immich_ml.schemas import Shape
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ShapePolicy:
|
||||
dims: tuple[Shape, ...] = (Shape(batch=1),)
|
||||
label: str = "" # the artifact subdirectory, for a submodel that ships more than one set
|
||||
|
||||
@property
|
||||
def pinned(self) -> Shape:
|
||||
first = self.dims[0]
|
||||
return Shape(
|
||||
batch=1,
|
||||
height=first.height if all(shape.height == first.height for shape in self.dims) else None,
|
||||
width=first.width if all(shape.width == first.width for shape in self.dims) else None,
|
||||
)
|
||||
|
||||
|
||||
def batches(maximum: int) -> tuple[int, ...]:
|
||||
return (1, maximum) if maximum > 1 else (1,)
|
||||
|
||||
|
||||
def runs(total: int, batches: Sequence[int]) -> list[int]:
|
||||
runs = []
|
||||
while total > 0:
|
||||
runs.append(max(size for size in batches if size <= total))
|
||||
total -= runs[-1]
|
||||
return runs
|
||||
@@ -0,0 +1,107 @@
|
||||
import fcntl
|
||||
import pickle
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from shutil import rmtree
|
||||
from typing import Any
|
||||
|
||||
import onnx
|
||||
from google.protobuf.message import DecodeError
|
||||
from immich_model.onnx.f16 import derive
|
||||
from immich_model.runtime import apply_rewrites
|
||||
from onnxruntime.capi.onnxruntime_pybind11_state import InvalidProtobuf, NoSuchFile
|
||||
|
||||
from ..config import settings
|
||||
from .ort import UNREADABLE, GraphSpec, Manifest, fresh
|
||||
|
||||
|
||||
def main() -> None:
|
||||
spec: GraphSpec = pickle.loads(sys.stdin.buffer.read())
|
||||
with spec.model_path.with_suffix(".lock").open("w") as lock:
|
||||
fcntl.flock(lock, fcntl.LOCK_EX) # several workers may want the same graph
|
||||
if fresh(spec) is None:
|
||||
try:
|
||||
prepare(spec)
|
||||
except (InvalidProtobuf, NoSuchFile, DecodeError):
|
||||
sys.exit(UNREADABLE)
|
||||
|
||||
|
||||
def prepare(spec: GraphSpec) -> None:
|
||||
rmtree(spec.directory, ignore_errors=True) # whatever is here was made under other facts, or not finished
|
||||
spec.directory.mkdir(parents=True)
|
||||
graph = rewritten(spec, narrowed(spec))
|
||||
if spec.provider == "CPUExecutionProvider":
|
||||
graph = prepacked(spec, graph) # allows mmap'ing the file instead of prepacking to dirty heap memory
|
||||
elif spec.directory.as_posix() in spec.provider_options[0].values():
|
||||
spec.session(graph) # the provider keeps what it compiles, in the directory it was given
|
||||
manifest = Manifest(facts=spec.facts, graph=graph.relative_to(spec.model_path.parent).as_posix())
|
||||
spec.manifest.write_text(manifest.model_dump_json()) # last, so that it means complete
|
||||
|
||||
|
||||
def narrowed(spec: GraphSpec) -> Path:
|
||||
if not spec.half:
|
||||
return spec.model_path
|
||||
destination = spec.directory / "model_fp16.onnx"
|
||||
derive(spec.model_path, destination)
|
||||
return destination
|
||||
|
||||
|
||||
def rewritten(spec: GraphSpec, graph: Path) -> Path:
|
||||
written: Path = graph if settings.legacy_models else apply_rewrites(graph, spec.plan)
|
||||
return written
|
||||
|
||||
|
||||
def prepacked(spec: GraphSpec, graph: Path) -> Path:
|
||||
destination = spec.directory / "model.onnx"
|
||||
options = spec.sess_options()
|
||||
options.optimized_model_filepath = destination.as_posix()
|
||||
options.add_session_config_entry("session.optimized_model_external_initializers_file_name", "model.data")
|
||||
options.add_session_config_entry("session.save_external_prepacked_constant_initializers", "1")
|
||||
options.add_session_config_entry("session.optimized_model_external_initializers_min_size_in_bytes", "1024")
|
||||
spec.session(graph, options)
|
||||
align(destination)
|
||||
return destination
|
||||
|
||||
|
||||
def align(destination: Path) -> None:
|
||||
"""Re-lays the data file so every region starts aligned. ORT pads only large regions, so one odd length skews
|
||||
all that follow, and MLAS reads a packed weight with aligned loads: it faults, and reads the rest slower."""
|
||||
model = onnx.load(destination.as_posix(), load_external_data=False)
|
||||
length_at: dict[int, int] = {}
|
||||
offsets: list[Any] = []
|
||||
packs: list[tuple[Any, str, list[list[str]]]] = []
|
||||
for tensor in model.graph.initializer:
|
||||
fields = {entry.key: entry for entry in tensor.external_data}
|
||||
for key, entry in fields.items():
|
||||
if key == "offset":
|
||||
length_at[int(entry.value)] = int(fields["length"].value)
|
||||
offsets.append(entry)
|
||||
elif key.startswith("prepacked_"):
|
||||
blob_key, *rest = entry.value.split("|")
|
||||
triples = [item.split(";") for item in rest]
|
||||
length_at.update((int(offset), int(length)) for offset, length, _ in triples)
|
||||
packs.append((entry, blob_key, triples))
|
||||
|
||||
moved, cursor = {}, 0
|
||||
for offset in sorted(length_at):
|
||||
cursor += -cursor % (4096 if length_at[offset] > 1 << 20 else 64)
|
||||
moved[offset] = cursor
|
||||
cursor += length_at[offset]
|
||||
|
||||
data = destination.with_suffix(".data")
|
||||
with data.open("rb") as source, data.with_suffix(".aligned").open("wb") as target:
|
||||
for offset in sorted(length_at):
|
||||
source.seek(offset)
|
||||
target.seek(moved[offset]) # what it skips reads as zeros
|
||||
target.write(source.read(length_at[offset]))
|
||||
data.with_suffix(".aligned").replace(data)
|
||||
|
||||
for entry in offsets:
|
||||
entry.value = str(moved[int(entry.value)])
|
||||
for entry, blob_key, triples in packs:
|
||||
entry.value = "|".join([blob_key, *(f"{moved[int(o)]};{n};{c}" for o, n, c in triples)])
|
||||
onnx.save(model, destination.as_posix())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,78 +1,84 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any, NamedTuple
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import orjson
|
||||
from numpy.typing import NDArray
|
||||
|
||||
from immich_ml.config import log, settings
|
||||
from immich_ml.schemas import SessionNode
|
||||
from immich_ml.schemas import ModelInput, ModelTensor, SessionNode, Shape
|
||||
|
||||
from .rknnpool import RknnPoolExecutor, is_available, soc_name
|
||||
from .rknnpool import RknnNode, RknnPoolExecutor, is_available, native_outputs, soc_name
|
||||
|
||||
is_available = is_available and settings.rknn
|
||||
model_prefix = Path("rknpu") / soc_name if is_available and soc_name is not None else None
|
||||
|
||||
|
||||
def run_inference(rknn_lite: Any, input: list[NDArray[np.float32]]) -> list[NDArray[np.float32]]:
|
||||
outputs: list[NDArray[np.float32]] = rknn_lite.inference(inputs=input, data_format="nchw")
|
||||
return outputs
|
||||
def model_path(model_dir: Path, variant: str = "") -> Path:
|
||||
return (model_dir / model_prefix if model_prefix else model_dir) / variant / "model.rknn"
|
||||
|
||||
|
||||
input_output_mapping: dict[str, dict[str, Any]] = {
|
||||
"detection": {
|
||||
"input": {"norm_tensor:0": (1, 3, 640, 640)},
|
||||
"output": {
|
||||
"norm_tensor:1": (12800, 1),
|
||||
"norm_tensor:2": (3200, 1),
|
||||
"norm_tensor:3": (800, 1),
|
||||
"norm_tensor:4": (12800, 4),
|
||||
"norm_tensor:5": (3200, 4),
|
||||
"norm_tensor:6": (800, 4),
|
||||
"norm_tensor:7": (12800, 10),
|
||||
"norm_tensor:8": (3200, 10),
|
||||
"norm_tensor:9": (800, 10),
|
||||
},
|
||||
},
|
||||
"recognition": {"input": {"norm_tensor:0": (1, 3, 112, 112)}, "output": {"norm_tensor:1": (1, 512)}},
|
||||
}
|
||||
def run_inference(rknn_lite: Any, inputs: list[ModelTensor], data_format: str | None) -> list[NDArray[np.float32]]:
|
||||
rknn_lite.rknn_runtime.set_inputs(inputs, None, data_format)
|
||||
rknn_lite.rknn_runtime.run(False)
|
||||
return native_outputs(rknn_lite)
|
||||
|
||||
|
||||
def input_layout(compiled: tuple[int, ...], array: ModelTensor) -> str | None:
|
||||
if array.ndim != 4:
|
||||
return None # token ids carry no layout
|
||||
if array.shape[3] == compiled[3]:
|
||||
return "nhwc"
|
||||
if array.shape[1] == compiled[3]:
|
||||
return "nchw"
|
||||
raise ValueError(f"binary takes {compiled[3]} channels, fed {list(array.shape)}")
|
||||
|
||||
|
||||
class RknnSession:
|
||||
def __init__(self, model_path: Path) -> None:
|
||||
self.model_type = "detection" if "detection" in model_path.parts else "recognition"
|
||||
self.tpe = settings.rknn_threads
|
||||
|
||||
log.info(f"Loading RKNN model from {model_path} with {self.tpe} threads.")
|
||||
self.rknnpool = RknnPoolExecutor(model_path=model_path.as_posix(), tpes=self.tpe, func=run_inference)
|
||||
log.info(f"Loaded RKNN model from {model_path} with {self.tpe} threads.")
|
||||
# the shapes the binary was compiled for, among which the runtime routes itself
|
||||
batch = self.rknnpool.inputs[0].shape[0]
|
||||
declared = orjson.loads(self.rknnpool.custom_string) if self.rknnpool.custom_string else {}
|
||||
self.shapes = tuple(Shape(batch, **dims) for dims in declared.get("dims", [{}]))
|
||||
self.metadata = {key: value for key, value in declared.items() if key != "dims"}
|
||||
self.batches = (batch,)
|
||||
|
||||
def get_inputs(self) -> list[SessionNode]:
|
||||
return [RknnNode(name=k, shape=v) for k, v in input_output_mapping[self.model_type]["input"].items()]
|
||||
def for_shape(self, shape: Shape) -> RknnSession:
|
||||
return self
|
||||
|
||||
def get_outputs(self) -> list[SessionNode]:
|
||||
return [RknnNode(name=k, shape=v) for k, v in input_output_mapping[self.model_type]["output"].items()]
|
||||
def warm(self) -> None:
|
||||
pass # the binary is loaded whole when the session opens
|
||||
|
||||
def get_inputs(self) -> Sequence[SessionNode]:
|
||||
return self.rknnpool.inputs
|
||||
|
||||
def get_outputs(self) -> Sequence[SessionNode]:
|
||||
return self.rknnpool.outputs
|
||||
|
||||
def get_metadata(self) -> dict[str, str]:
|
||||
return {}
|
||||
return self.metadata
|
||||
|
||||
@property
|
||||
def normalizes_input(self) -> bool:
|
||||
# the compiler transposes and casts every binary, so only what the exporter declared can say
|
||||
return bool(self.rknnpool.custom_string)
|
||||
|
||||
def run(
|
||||
self,
|
||||
output_names: list[str] | None,
|
||||
input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]] | dict[str, NDArray[np.uint8]],
|
||||
input_feed: ModelInput,
|
||||
run_options: Any = None,
|
||||
) -> list[NDArray[np.float32]]:
|
||||
input_data: list[NDArray[np.float32]] = [np.ascontiguousarray(v) for v in input_feed.values()]
|
||||
res = self.rknnpool.run(input_data)
|
||||
if res is None:
|
||||
raise RuntimeError("RKNN inference failed!")
|
||||
return res
|
||||
inputs = [array if array.flags.c_contiguous else array.copy() for array in input_feed.values()]
|
||||
return self.rknnpool.run(inputs, input_layout(self.rknnpool.inputs[0].shape, inputs[0]))
|
||||
|
||||
|
||||
class RknnNode(NamedTuple):
|
||||
name: str
|
||||
shape: tuple[int, ...]
|
||||
|
||||
|
||||
__all__ = ["RknnSession", "RknnNode", "is_available", "soc_name", "model_prefix"]
|
||||
__all__ = ["RknnSession", "RknnNode", "is_available", "soc_name", "model_path"]
|
||||
|
||||
@@ -1,17 +1,19 @@
|
||||
# This code is from leafqycc/rknn-multi-threaded
|
||||
# Following Apache License 2.0
|
||||
|
||||
import ctypes
|
||||
import logging
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
from typing import Any, Callable, NamedTuple
|
||||
|
||||
import numpy as np
|
||||
from numpy.typing import NDArray
|
||||
|
||||
from immich_ml.config import log
|
||||
from immich_ml.models.constants import RKNN_COREMASK_SUPPORTED_SOCS, RKNN_SUPPORTED_SOCS
|
||||
from immich_ml.schemas import ModelTensor
|
||||
|
||||
|
||||
def get_soc(device_tree_path: Path | str) -> str | None:
|
||||
@@ -30,7 +32,7 @@ def get_soc(device_tree_path: Path | str) -> str | None:
|
||||
soc_name = None
|
||||
is_available = False
|
||||
try:
|
||||
from rknnlite.api import RKNNLite
|
||||
from rknnlite.api import RKNNLite, rknn_runtime
|
||||
|
||||
soc_name = get_soc("/proc/device-tree/compatible")
|
||||
is_available = soc_name is not None
|
||||
@@ -38,6 +40,82 @@ except ImportError:
|
||||
log.debug("RKNN is not available")
|
||||
|
||||
|
||||
class RknnNode(NamedTuple):
|
||||
name: str
|
||||
shape: tuple[int, ...]
|
||||
|
||||
|
||||
_CUSTOM_STRING_QUERY = 7 # RKNN_QUERY_CUSTOM_STRING
|
||||
# RKNN_QUERY_OUTPUT_ATTR, and RKNN_QUERY_CURRENT_OUTPUT_ATTR for the shape a dynamic binary last produced
|
||||
_OUTPUT_QUERY, _CURRENT_OUTPUT_QUERY = 2, 15
|
||||
_DTYPES = {0: np.float32, 1: np.float16} # rknn_tensor_type
|
||||
|
||||
|
||||
class _CustomString(ctypes.Structure):
|
||||
_fields_ = [("string", ctypes.c_char * 1024)]
|
||||
|
||||
|
||||
class _TensorAttr(ctypes.Structure): # rknn_tensor_attr
|
||||
_fields_ = [
|
||||
("index", ctypes.c_uint32),
|
||||
("n_dims", ctypes.c_uint32),
|
||||
("dims", ctypes.c_uint32 * 16),
|
||||
("name", ctypes.c_char * 256),
|
||||
("n_elems", ctypes.c_uint32),
|
||||
("size", ctypes.c_uint32),
|
||||
("fmt", ctypes.c_int),
|
||||
("type", ctypes.c_int),
|
||||
("qnt_type", ctypes.c_int),
|
||||
("fl", ctypes.c_int8),
|
||||
("zp", ctypes.c_int32),
|
||||
("scale", ctypes.c_float),
|
||||
("w_stride", ctypes.c_uint32),
|
||||
("size_with_stride", ctypes.c_uint32),
|
||||
("pass_through", ctypes.c_uint8),
|
||||
("h_stride", ctypes.c_uint32),
|
||||
]
|
||||
|
||||
|
||||
def native_outputs(rknn_lite: "RKNNLite") -> list[NDArray[Any]]:
|
||||
"""The outputs in the precision the NPU computed them in, which RKNNLite would widen to fp32 on the CPU."""
|
||||
runtime: Any = rknn_lite.rknn_runtime
|
||||
count = runtime.get_in_out_num()[1]
|
||||
buffers = (rknn_runtime.RKNNOutput * count)()
|
||||
for index, buffer in enumerate(buffers):
|
||||
buffer.index = index
|
||||
if runtime.lib.rknn_outputs_get(runtime.context, count, buffers, None) != 0:
|
||||
raise RuntimeError("RKNN inference failed!")
|
||||
query = _CURRENT_OUTPUT_QUERY if runtime.is_dynamic_shape else _OUTPUT_QUERY
|
||||
outputs = []
|
||||
for buffer in buffers:
|
||||
attr = _TensorAttr(index=buffer.index)
|
||||
runtime.lib.rknn_query(runtime.context, query, ctypes.byref(attr), ctypes.sizeof(attr))
|
||||
data = ctypes.cast(buffer.buf, ctypes.POINTER(ctypes.c_char * buffer.size)).contents
|
||||
outputs.append(np.frombuffer(data, _DTYPES[attr.type]).reshape(attr.dims[: attr.n_dims]).copy())
|
||||
runtime.lib.rknn_outputs_release(runtime.context, count, buffers)
|
||||
return outputs
|
||||
|
||||
|
||||
def tensor_nodes(rknn_lite: "RKNNLite") -> tuple[list[RknnNode], list[RknnNode]]:
|
||||
runtime = rknn_lite.rknn_runtime
|
||||
n_inputs, n_outputs = runtime.get_in_out_num()
|
||||
|
||||
def node(index: int, is_output: bool) -> RknnNode:
|
||||
attr = runtime.get_tensor_attr(index, is_output=is_output)
|
||||
return RknnNode(attr.name.decode(), tuple(attr.dims[: attr.n_dims]))
|
||||
|
||||
return [node(i, False) for i in range(n_inputs)], [node(i, True) for i in range(n_outputs)]
|
||||
|
||||
|
||||
def custom_string(rknn_lite: "RKNNLite") -> str:
|
||||
runtime: Any = rknn_lite.rknn_runtime
|
||||
buffer = _CustomString()
|
||||
# rknnlite declares argtypes on the shared handle, so pass plain values and let ctypes convert
|
||||
runtime.lib.rknn_query(runtime.context, _CUSTOM_STRING_QUERY, ctypes.byref(buffer), ctypes.sizeof(buffer))
|
||||
declared: str = buffer.string.decode()
|
||||
return declared
|
||||
|
||||
|
||||
def init_rknn(model_path: str) -> "RKNNLite":
|
||||
if not is_available:
|
||||
raise RuntimeError("rknn is not available!")
|
||||
@@ -63,21 +141,23 @@ class RknnPoolExecutor:
|
||||
self,
|
||||
model_path: str,
|
||||
tpes: int,
|
||||
func: Callable[["RKNNLite", list[NDArray[np.float32]]], list[NDArray[np.float32]]],
|
||||
func: Callable[["RKNNLite", list[ModelTensor], str | None], list[NDArray[np.float32]]],
|
||||
) -> None:
|
||||
self.tpes = tpes
|
||||
self.rknn_pool = [init_rknn(model_path) for _ in range(tpes)]
|
||||
self.inputs, self.outputs = tensor_nodes(self.rknn_pool[0])
|
||||
self.custom_string = custom_string(self.rknn_pool[0])
|
||||
self.pool = ThreadPoolExecutor(max_workers=tpes)
|
||||
self.func = func
|
||||
self.num = 0
|
||||
self.lock = threading.Lock()
|
||||
|
||||
def run(self, inputs: list[NDArray[np.float32]]) -> list[NDArray[np.float32]]:
|
||||
def run(self, inputs: list[ModelTensor], data_format: str | None) -> list[NDArray[np.float32]]:
|
||||
with self.lock:
|
||||
idx = self.num % self.tpes
|
||||
self.num += 1
|
||||
|
||||
fut = self.pool.submit(self.func, self.rknn_pool[idx], inputs)
|
||||
fut = self.pool.submit(self.func, self.rknn_pool[idx], inputs, data_format)
|
||||
return fut.result()
|
||||
|
||||
def release(self) -> None:
|
||||
|
||||
@@ -9,6 +9,7 @@ dependencies = [
|
||||
"fastapi>=0.95.2,<1.0",
|
||||
"gunicorn>=21.1.0",
|
||||
"huggingface-hub>=1.0,<2.0",
|
||||
"immich-model",
|
||||
"numpy>=2.4.0,<3.0",
|
||||
"onnx>=1.22.0",
|
||||
"opencv-python-headless>=4.7.0.72,<5.0",
|
||||
@@ -58,6 +59,7 @@ compile-bytecode = true
|
||||
override-dependencies = ["opencv-python ; sys_platform == 'never'"] # only install the headless version
|
||||
|
||||
[tool.uv.sources]
|
||||
immich-model = { git = "https://github.com/immich-app/ml-models.git", rev = "88d2b64ff95ce7117760660681d95e0e1308819b" }
|
||||
onnxruntime-gpu = { index = "onnxruntime-cuda-12" }
|
||||
|
||||
[[tool.uv.index]]
|
||||
|
||||
File diff suppressed because one or more lines are too long
+1073
-504
File diff suppressed because it is too large
Load Diff
Generated
+72
-3
@@ -682,6 +682,7 @@ dependencies = [
|
||||
{ name = "fastapi" },
|
||||
{ name = "gunicorn" },
|
||||
{ name = "huggingface-hub" },
|
||||
{ name = "immich-model" },
|
||||
{ name = "numpy" },
|
||||
{ name = "onnx" },
|
||||
{ name = "opencv-python-headless" },
|
||||
@@ -762,6 +763,7 @@ requires-dist = [
|
||||
{ name = "fastapi", specifier = ">=0.95.2,<1.0" },
|
||||
{ name = "gunicorn", specifier = ">=21.1.0" },
|
||||
{ name = "huggingface-hub", specifier = ">=1.0,<2.0" },
|
||||
{ name = "immich-model", git = "https://github.com/immich-app/ml-models.git?rev=88d2b64ff95ce7117760660681d95e0e1308819b" },
|
||||
{ name = "numpy", specifier = ">=2.4.0,<3.0" },
|
||||
{ name = "onnx", specifier = ">=1.22.0" },
|
||||
{ name = "onnxruntime", marker = "extra == 'armnn'", specifier = ">=1.23.2,<2" },
|
||||
@@ -824,6 +826,18 @@ types = [
|
||||
{ name = "types-ujson", specifier = ">=5.10.0.20240515" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "immich-model"
|
||||
version = "0.2.0"
|
||||
source = { git = "https://github.com/immich-app/ml-models.git?rev=88d2b64ff95ce7117760660681d95e0e1308819b#88d2b64ff95ce7117760660681d95e0e1308819b" }
|
||||
dependencies = [
|
||||
{ name = "numpy" },
|
||||
{ name = "onnx" },
|
||||
{ name = "onnx-ir" },
|
||||
{ name = "onnxscript" },
|
||||
{ name = "safetensors" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "iniconfig"
|
||||
version = "2.0.0"
|
||||
@@ -1284,6 +1298,22 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/00/50/257a880384a1dd502d543b0067945074d63cd17d0840e958355bc8197da8/onnx-1.22.0-cp314-cp314t-win_arm64.whl", hash = "sha256:8e268cdc0547e3949799ffd4a44451dc2b9080b57d0824a2db680b6ec65506f0", size = 17231391, upload-time = "2026-06-15T12:50:03.047Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "onnx-ir"
|
||||
version = "1.0.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "ml-dtypes" },
|
||||
{ name = "numpy" },
|
||||
{ name = "onnx" },
|
||||
{ name = "sympy" },
|
||||
{ name = "typing-extensions" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/d6/c2/61194cec0dbc5622273c0ebd592d37cc1dca0d7f1a744f02edd45ac905a3/onnx_ir-1.0.0.tar.gz", hash = "sha256:9e261f25fde8da9612ae5cb43b3b374d5ff469c04af0363cad588b2bb000b812", size = 163121, upload-time = "2026-08-11T14:49:46.895Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/91/cd/6d1637172eb59c7b18ac90ed089d1f599a11fe0e63b4db2d017f3bb38a32/onnx_ir-1.0.0-py3-none-any.whl", hash = "sha256:e578f0d608d3062866b48223616eb2d10a6d6d01f8b8faac596129034f483cc7", size = 185849, upload-time = "2026-08-11T14:49:45.524Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "onnxruntime"
|
||||
version = "1.30.0"
|
||||
@@ -1371,6 +1401,23 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/3e/92/46ae2cd565961a89189900f385bb2f13a9fa731ea4674001d23720fbb1e0/onnxruntime_openvino-1.24.1-cp313-cp313-win_amd64.whl", hash = "sha256:434bf49aa71393c577a456c9d76c98e6d6958a833fa0876793e3d5437b5a511a", size = 13658485, upload-time = "2026-02-26T13:44:53.889Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "onnxscript"
|
||||
version = "0.7.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "ml-dtypes" },
|
||||
{ name = "numpy" },
|
||||
{ name = "onnx" },
|
||||
{ name = "onnx-ir" },
|
||||
{ name = "packaging" },
|
||||
{ name = "typing-extensions" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/0a/01/3e3fab8d643ca097ea4aa9e51246643699dfaaa0650589744fe44bc46651/onnxscript-0.7.2.tar.gz", hash = "sha256:2c664f6383d10f332a4d47b2876dcab16dba84909fe703656b19abc281fda165", size = 646719, upload-time = "2026-09-09T17:06:44.567Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/b9/3b/06260997cdc41138e58718588a6c87d0eb342bbe0dda8a6aae91d163c384/onnxscript-0.7.2-py3-none-any.whl", hash = "sha256:d0e7121c6a1eefd608058928e111cbdb76709f70d269ff0d07aee493bd1d13c9", size = 754215, upload-time = "2026-09-09T17:06:46.442Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "opencv-python"
|
||||
version = "4.11.0.86"
|
||||
@@ -2034,6 +2081,28 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/15/19/016553f86f207450aebebc2b2b5088d086b901cc8186c02ac4284db3bd88/ruff-0.15.16-py3-none-win_arm64.whl", hash = "sha256:8cd61783afb39638a7133ef0d2dfb1e91277593962f81b5a8423eb0b888a6121", size = 11134555, upload-time = "2026-06-04T16:33:00.136Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "safetensors"
|
||||
version = "0.7.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/29/9c/6e74567782559a63bd040a236edca26fd71bc7ba88de2ef35d75df3bca5e/safetensors-0.7.0.tar.gz", hash = "sha256:07663963b67e8bd9f0b8ad15bb9163606cd27cc5a1b96235a50d8369803b96b0", size = 200878, upload-time = "2025-11-19T15:18:43.199Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/fa/47/aef6c06649039accf914afef490268e1067ed82be62bcfa5b7e886ad15e8/safetensors-0.7.0-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c82f4d474cf725255d9e6acf17252991c3c8aac038d6ef363a4bf8be2f6db517", size = 467781, upload-time = "2025-11-19T15:18:35.84Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/e8/00/374c0c068e30cd31f1e1b46b4b5738168ec79e7689ca82ee93ddfea05109/safetensors-0.7.0-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:94fd4858284736bb67a897a41608b5b0c2496c9bdb3bf2af1fa3409127f20d57", size = 447058, upload-time = "2025-11-19T15:18:34.416Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/f1/06/578ffed52c2296f93d7fd2d844cabfa92be51a587c38c8afbb8ae449ca89/safetensors-0.7.0-cp38-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e07d91d0c92a31200f25351f4acb2bc6aff7f48094e13ebb1d0fb995b54b6542", size = 491748, upload-time = "2025-11-19T15:18:09.79Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ae/33/1debbbb70e4791dde185edb9413d1fe01619255abb64b300157d7f15dddd/safetensors-0.7.0-cp38-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:8469155f4cb518bafb4acf4865e8bb9d6804110d2d9bdcaa78564b9fd841e104", size = 503881, upload-time = "2025-11-19T15:18:16.145Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/8e/1c/40c2ca924d60792c3be509833df711b553c60effbd91da6f5284a83f7122/safetensors-0.7.0-cp38-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:54bef08bf00a2bff599982f6b08e8770e09cc012d7bba00783fc7ea38f1fb37d", size = 623463, upload-time = "2025-11-19T15:18:21.11Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/9b/3a/13784a9364bd43b0d61eef4bea2845039bc2030458b16594a1bd787ae26e/safetensors-0.7.0-cp38-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:42cb091236206bb2016d245c377ed383aa7f78691748f3bb6ee1bfa51ae2ce6a", size = 532855, upload-time = "2025-11-19T15:18:25.719Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/a0/60/429e9b1cb3fc651937727befe258ea24122d9663e4d5709a48c9cbfceecb/safetensors-0.7.0-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dac7252938f0696ddea46f5e855dd3138444e82236e3be475f54929f0c510d48", size = 507152, upload-time = "2025-11-19T15:18:33.023Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/3c/a8/4b45e4e059270d17af60359713ffd83f97900d45a6afa73aaa0d737d48b6/safetensors-0.7.0-cp38-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:1d060c70284127fa805085d8f10fbd0962792aed71879d00864acda69dbab981", size = 541856, upload-time = "2025-11-19T15:18:31.075Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/06/87/d26d8407c44175d8ae164a95b5a62707fcc445f3c0c56108e37d98070a3d/safetensors-0.7.0-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:cdab83a366799fa730f90a4ebb563e494f28e9e92c4819e556152ad55e43591b", size = 674060, upload-time = "2025-11-19T15:18:37.211Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/11/f5/57644a2ff08dc6325816ba7217e5095f17269dada2554b658442c66aed51/safetensors-0.7.0-cp38-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:672132907fcad9f2aedcb705b2d7b3b93354a2aec1b2f706c4db852abe338f85", size = 771715, upload-time = "2025-11-19T15:18:38.689Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/86/31/17883e13a814bd278ae6e266b13282a01049b0c81341da7fd0e3e71a80a3/safetensors-0.7.0-cp38-abi3-musllinux_1_2_i686.whl", hash = "sha256:5d72abdb8a4d56d4020713724ba81dac065fedb7f3667151c4a637f1d3fb26c0", size = 714377, upload-time = "2025-11-19T15:18:40.162Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/4a/d8/0c8a7dc9b41dcac53c4cbf9df2b9c83e0e0097203de8b37a712b345c0be5/safetensors-0.7.0-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b0f6d66c1c538d5a94a73aa9ddca8ccc4227e6c9ff555322ea40bdd142391dd4", size = 677368, upload-time = "2025-11-19T15:18:41.627Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/05/e5/cb4b713c8a93469e3c5be7c3f8d77d307e65fe89673e731f5c2bfd0a9237/safetensors-0.7.0-cp38-abi3-win32.whl", hash = "sha256:c74af94bf3ac15ac4d0f2a7c7b4663a15f8c2ab15ed0fc7531ca61d0835eccba", size = 326423, upload-time = "2025-11-19T15:18:45.74Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/5d/e6/ec8471c8072382cb91233ba7267fd931219753bb43814cbc71757bfd4dab/safetensors-0.7.0-cp38-abi3-win_amd64.whl", hash = "sha256:d1239932053f56f3456f32eb9625590cc7582e905021f94636202a864d470755", size = 341380, upload-time = "2025-11-19T15:18:44.427Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "setuptools"
|
||||
version = "80.9.0"
|
||||
@@ -2123,14 +2192,14 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "sympy"
|
||||
version = "1.12"
|
||||
version = "1.14.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "mpmath" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/e5/57/3485a1a3dff51bfd691962768b14310dae452431754bfc091250be50dd29/sympy-1.12.tar.gz", hash = "sha256:ebf595c8dac3e0fdc4152c51878b498396ec7f30e7a914d6071e674d49420fb8", size = 6722203, upload-time = "2023-05-10T18:23:00.378Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/d2/05/e6600db80270777c4a64238a98d442f0fd07cc8915be2a1c16da7f2b9e74/sympy-1.12-py3-none-any.whl", hash = "sha256:c3588cd4295d0c0f603d0f2ae780587e64e2efeedb3521e46b9bb1d08d184fa5", size = 5742435, upload-time = "2023-05-10T18:22:14.76Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
Reference in New Issue
Block a user