onnx¶

async def load(self) -> None:
    """Load ONNX model and prepare for inference (async).

    Performs:
        - Configure session options
        - Load model from disk (in thread pool)
        - Extract input/output names
        - Warmup inference (in thread pool)

    Raises:
        FileNotFoundError: If model file does not exist.
        RuntimeError: If ONNX Runtime fails to load model.
    """
    logger.info(f"Loading ONNX model from {self.model_path}")

    # Session options
    sess_options = ort.SessionOptions()
    sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

    # Load in thread pool
    loop = asyncio.get_event_loop()
    self.session = await loop.run_in_executor(
        None,
        lambda _: ort.InferenceSession(
            str(self.model_path),
            sess_options=sess_options,
            providers=self.providers,
        ),
        None,
    )

    # Get input/output names
    self.input_name = self.session.get_inputs()[0].name
    self.output_names = [output.name for output in self.session.get_outputs()]

    logger.info(
        f"Model loaded:  input={self.input_name}, "
        f"outputs={self.output_names}, "
        f"providers={self.session.get_providers()}"
    )

    # Warmup
    await self._warmup()

async def infer(self, input: np.ndarray) -> t.Any:
    """Run inference on a single input (async).

    Automatically handles:
        - Converting to correct dtype

    Args:
        input: Input numpy array.

    Returns:
        Output array or tuple of arrays (if multi-output model).

    Raises:
        RuntimeError: If model is not loaded.

    Example:
        ```python
        async with runtime:
            input = np.random.randn(1, 3, 224, 224).astype(
                np.float32
            )
            output = await runtime.infer(input)
        ```
    """
    if self.session is None:
        raise RuntimeError("Model not loaded. Call load() first.")

    # Prepare input (reuse mixin)
    input = self._prepare_onnx_input(input)

    # Run inference in thread pool
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(None, self._infer_sync, input)

async def infer_batch(self, inputs: list[np.ndarray]) -> list[t.Any]:
    """Run inference on a batch of inputs (async).

    Concatenates inputs into a single batch array for efficient
    processing, then splits the output back into individual results.

    Args:
        inputs: List of input arrays. Each should have shape (1, C, H, W).

    Returns:
        List of outputs, one per input. Each maintains batch dimension.

    Raises:
        RuntimeError: If model is not loaded.

    Example:
        ```python
        async with runtime:
            batch = [
                np.random.randn(1, 3, 224, 224).astype(np.float32),
                np.random.randn(1, 3, 224, 224).astype(np.float32),
            ]
            outputs = await runtime.infer_batch(batch)
        ```
    """
    if self.session is None:
        raise RuntimeError("Model not loaded. Call load() first.")

    # Concatenate inputs
    batch = np.concatenate(inputs, axis=0)

    # Prepare input (reuse mixin)
    batch = self._prepare_onnx_input(batch)

    # Run batch inference
    batch_output = await self.infer(batch)

    # Split outputs (reuse mixin)
    return self._split_onnx_batch_output(batch_output, len(inputs))

async def unload(self) -> None:
    """Unload model and free resources (async).

    Performs:
        - Release session from memory

    Safe to call multiple times.
    """
    logger.info("Unloading model")
    self.session = None
    self.input_name = None
    self.output_names = None
    logger.info("Model unloaded")