diff --git a/machine-learning/immich_ml/sessions/rknn/immich_session.py b/machine-learning/immich_ml/sessions/rknn/immich_session.py index 827a3ec1d2..c5c4457278 100644 --- a/machine-learning/immich_ml/sessions/rknn/immich_session.py +++ b/machine-learning/immich_ml/sessions/rknn/immich_session.py @@ -24,9 +24,15 @@ _native: ModuleType | None = _native_mod if TYPE_CHECKING: class NativeRKNNExecutor: - def __init__(self, model_path: str, num_workers: int = 1) -> None: ... + def __init__( + self, + model_path: str, + num_workers: int = 1, + want_float: bool = True, + uint8_input: bool = False, + ) -> None: ... - def infer(self, inputs: list[NDArray[np.float32]]) -> list[NDArray[np.float32]]: ... + def infer(self, inputs: list[NDArray[np.float32]]) -> list[NDArray[Any]]: ... def get_io_info(self) -> dict[str, Any]: ... else: @@ -74,7 +80,7 @@ class RKNNInferenceResult(NamedTuple): outputs: list[NDArray[np.float32]] -Array32 = NDArray[np.float32] | NDArray[np.int32] +Array32 = NDArray[np.float32] | NDArray[np.int32] | NDArray[np.uint8] class InferenceExecutor(Protocol): @@ -86,19 +92,23 @@ def run_inference(executor: InferenceExecutor, inputs: list[NDArray[np.float32]] class RknnPoolExecutor: - def __init__(self, model_path: str | Path, tpes: int) -> None: + def __init__( + self, model_path: str | Path, tpes: int, *, want_float: bool = True, uint8_input: bool = False + ) -> None: if NativeRKNNExecutor is None: raise RuntimeError("RKNN native extension is not available") if tpes < 1: raise ValueError("tpes must be >= 1") model_path_str = Path(model_path).as_posix() - self._native = NativeRKNNExecutor(model_path_str, num_workers=tpes) + self._native = NativeRKNNExecutor( + model_path_str, num_workers=tpes, want_float=want_float, uint8_input=uint8_input + ) self._executor = ThreadPoolExecutor(max_workers=tpes, thread_name_prefix="rknn-worker") self._closed = False def _run_inference(self, inputs: list[Array32], tag: Any) -> RKNNInferenceResult: start = time.perf_counter() - outputs = self._native.infer(cast(list[NDArray[np.float32]], inputs)) + outputs = cast(list[NDArray[np.float32]], self._native.infer(cast(list[NDArray[np.float32]], inputs))) end = time.perf_counter() return RKNNInferenceResult( tag=tag, @@ -144,6 +154,8 @@ class RknnSession: model_path: Path | str, *, num_workers: Optional[int] = None, + want_float: bool = True, + uint8_input: bool = False, logger: Any = None, ) -> None: if not is_available: @@ -153,12 +165,18 @@ class RknnSession: self.tpe = num_workers or settings.rknn_threads if self.tpe < 1: raise ValueError("num_workers must be >= 1") + self.want_float = want_float + self.uint8_input = uint8_input self.log.info( - "Loading RKNN model from %s with %s worker(s).", + "Loading RKNN model from %s with %s worker(s) (want_float=%s, uint8_input=%s).", self.model_path, self.tpe, + want_float, + uint8_input, + ) + self.rknnpool = RknnPoolExecutor( + self.model_path, self.tpe, want_float=want_float, uint8_input=uint8_input ) - self.rknnpool = RknnPoolExecutor(self.model_path, self.tpe) self._io_info = self._normalize_io_info(self.rknnpool.executor.get_io_info()) self._input_nodes: list[SessionNode] = self._build_nodes("inputs") self._output_nodes: list[SessionNode] = self._build_nodes("outputs") @@ -177,7 +195,7 @@ class RknnSession: def run( self, _output_names: Sequence[str] | None, - input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]], + input_feed: dict[str, Array32], _run_options: Any = None, ) -> list[NDArray[np.float32]]: return self.run_async(_output_names, input_feed, _run_options).result().outputs @@ -185,7 +203,7 @@ class RknnSession: def run_async( self, _output_names: Sequence[str] | None, - input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]], + input_feed: dict[str, Array32], _run_options: Any = None, ) -> Future[RKNNInferenceResult]: inputs_list: list[Array32] = list(input_feed.values()) diff --git a/machine-learning/immich_ml/sessions/rknn/native/rknn_pool.cpp b/machine-learning/immich_ml/sessions/rknn/native/rknn_pool.cpp index e217ba3c6b..aa80858eda 100644 --- a/machine-learning/immich_ml/sessions/rknn/native/rknn_pool.cpp +++ b/machine-learning/immich_ml/sessions/rknn/native/rknn_pool.cpp @@ -4,7 +4,6 @@ #include #include -#include #include #include #include @@ -74,7 +73,10 @@ static py::list dims_to_list(const uint32_t* dims, uint32_t n_dims) { static py::dict make_tensor_info(uint32_t index, const rknn_tensor_attr& attr) { return py::dict("index"_a=index, "name"_a=py::str(attr.name), "fmt"_a=static_cast(attr.fmt), "type"_a=static_cast(attr.type), "n_dims"_a=attr.n_dims, - "dims"_a=dims_to_list(attr.dims, attr.n_dims)); + "dims"_a=dims_to_list(attr.dims, attr.n_dims), + "n_elems"_a=attr.n_elems, "size"_a=attr.size, + "qnt_type"_a=static_cast(attr.qnt_type), + "zp"_a=attr.zp, "scale"_a=attr.scale); } static py::dict make_dynamic_dict(const rknn_input_range& rng) { @@ -113,17 +115,39 @@ static py::array align_layout(py::array arr, const rknn_tensor_attr& attr) { return arr; } -static PreparedInput prepare_input_tensor(py::handle handle, const rknn_tensor_attr& attr, bool capture_shape) { +// Opt-in UINT8 feed for image tensors when ATTR reports FP16/FP32 but mean_values +// are baked at compile (PP-OCR). Non-image inputs keep the ATTR dtype path. +static bool is_uint8_image(const py::array& arr, const rknn_tensor_attr& attr) { + if (!arr.dtype().is(py::dtype::of())) return false; + py::buffer_info bi = arr.request(); + if (bi.ndim != 4) return false; + auto channels_ok = [](ssize_t c) { return c == 1 || c == 3 || c == 4; }; + if (attr.fmt == RKNN_TENSOR_NCHW) return channels_ok(bi.shape[1]); + // NHWC (OCR) or unknown: channels on last dim + return channels_ok(bi.shape[3]); +} + +static PreparedInput prepare_input_tensor(py::handle handle, const rknn_tensor_attr& attr, bool capture_shape, + bool uint8_input) { py::array arr = handle.cast(); py::array contiguous = py::array::ensure(arr, py::array::c_style); contiguous = align_layout(contiguous, attr); - contiguous = ensure_dtype(contiguous, attr.type); - auto bi = contiguous.request(); - + rknn_input tensor{}; tensor.index = attr.index; - tensor.type = attr.type; - tensor.fmt = attr.fmt; + tensor.pass_through = 0; + + if (uint8_input && is_uint8_image(contiguous, attr)) { + // Keep host UINT8; driver applies baked mean_values (pass_through=0). + tensor.type = RKNN_TENSOR_UINT8; + tensor.fmt = (attr.fmt == RKNN_TENSOR_NCHW) ? RKNN_TENSOR_NCHW : RKNN_TENSOR_NHWC; + } else { + contiguous = ensure_dtype(contiguous, attr.type); + tensor.type = attr.type; + tensor.fmt = attr.fmt; + } + + auto bi = contiguous.request(); tensor.size = static_cast(contiguous.nbytes()); tensor.buf = const_cast(bi.ptr); @@ -148,16 +172,34 @@ static int find_matching_shape(const rknn_input_range& rng, const std::vector shape(attr.n_dims == 0 ? 1 : attr.n_dims); +static std::vector output_shape(const rknn_tensor_attr& attr) { if (attr.n_dims == 0) { - shape[0] = static_cast(out.size / sizeof(float)); - } else { - std::copy(attr.dims, attr.dims + attr.n_dims, shape.begin()); + return {static_cast(attr.n_elems ? attr.n_elems : 0)}; } - py::array arr(py::dtype::of(), shape); - std::memcpy(arr.mutable_data(), out.buf, out.size); - return arr; + return {attr.dims, attr.dims + attr.n_dims}; +} + +static py::dtype dtype_for_tensor_type(rknn_tensor_type t) { + switch (t) { + case RKNN_TENSOR_FLOAT16: return py::dtype("float16"); + case RKNN_TENSOR_FLOAT32: return py::dtype::of(); + case RKNN_TENSOR_UINT8: return py::dtype::of(); + case RKNN_TENSOR_INT8: return py::dtype::of(); + case RKNN_TENSOR_UINT16: return py::dtype::of(); + case RKNN_TENSOR_INT16: return py::dtype::of(); + case RKNN_TENSOR_UINT32: return py::dtype::of(); + case RKNN_TENSOR_INT32: return py::dtype::of(); + case RKNN_TENSOR_INT64: return py::dtype::of(); + default: + throw std::runtime_error("Unsupported RKNN tensor type for native output"); + } +} + +static py::array make_empty_output(const rknn_tensor_attr& attr, bool want_float) { + if (want_float) { + return py::array(py::dtype::of(), output_shape(attr)); + } + return py::array(dtype_for_tensor_type(attr.type), output_shape(attr)); } @@ -223,15 +265,21 @@ static const rknn_tensor_attr& resolve_output_attr(bool is_dynamic, RknnCtx& ctx class NativeRKNNExecutor { public: - explicit NativeRKNNExecutor(const std::string& model_path, int num_workers) + explicit NativeRKNNExecutor(const std::string& model_path, int num_workers, bool want_float = true, + bool uint8_input = false) : rr_index_(0), - is_dynamic_model_(false) { + is_dynamic_model_(false), + want_float_(want_float), + uint8_input_(uint8_input) { if (num_workers < 1) throw std::invalid_argument("num_workers must be >= 1"); if (num_workers > 3) throw std::invalid_argument("num_workers must be <= 3"); const bool debug_ctor = (std::getenv("RKNN_EXEC_DEBUG") != nullptr); RknnCtx master; - if (rknn_init(&master.ctx, const_cast(model_path.c_str()), 0, 0, nullptr) != RKNN_SUCC) + // Prefer GPU over CPU when an op is not supported by the NPU. If GPU + // also cannot run it, the runtime still falls back to CPU. + constexpr uint32_t init_flags = RKNN_FLAG_EXECUTE_FALLBACK_PRIOR_DEVICE_GPU; + if (rknn_init(&master.ctx, const_cast(model_path.c_str()), 0, init_flags, nullptr) != RKNN_SUCC) throw std::runtime_error("rknn_init failed"); master.query_io(); if (debug_ctor) debug_print_io_info(master); @@ -268,6 +316,8 @@ public: py::dict info; const RknnCtx& master = contexts_.front(); info["is_dynamic"] = is_dynamic_model_; + info["want_float"] = want_float_; + info["uint8_input"] = uint8_input_; py::list inputs(master.io_num.n_input); for (uint32_t i = 0; i < master.io_num.n_input; ++i) { py::dict desc = make_tensor_info(i, master.input_attrs[i]); @@ -301,37 +351,46 @@ public: if (is_dynamic_model_) input_shapes.reserve(c.io_num.n_input); for (uint32_t i = 0; i < c.io_num.n_input; ++i) { - PreparedInput prepared = prepare_input_tensor(inputs[i], c.input_attrs[i], is_dynamic_model_); + PreparedInput prepared = + prepare_input_tensor(inputs[i], c.input_attrs[i], is_dynamic_model_, uint8_input_); in[i] = prepared.tensor; if (is_dynamic_model_) input_shapes.push_back(std::move(prepared.shape)); keep_alive.push_back(std::move(prepared.buffer)); } if (is_dynamic_model_) set_dynamic_shapes(c, input_shapes); - std::vector out(c.io_num.n_output); { py::gil_scoped_release nogil; if (rknn_inputs_set(c.ctx, c.io_num.n_input, in.data()) != RKNN_SUCC) throw std::runtime_error("rknn_inputs_set failed"); if (rknn_run(c.ctx, nullptr) != RKNN_SUCC) throw std::runtime_error("rknn_run failed"); - for (uint32_t i = 0; i < c.io_num.n_output; ++i) { - out[i] = {}; - out[i].want_float = 1; - out[i].index = i; - } - if (rknn_outputs_get(c.ctx, c.io_num.n_output, out.data(), nullptr) != RKNN_SUCC) - throw std::runtime_error("rknn_outputs_get failed"); } + // want_float=true: runtime dequants into float32 (CLIP/faces). want_float=false: + // native INT8/FP16 buffers — needed for OCR CTC logits (~4× smaller than float32). + std::vector out(c.io_num.n_output); py::list result(c.io_num.n_output); rknn_tensor_attr scratch{}; + std::vector keep_alive_out; + keep_alive_out.reserve(c.io_num.n_output); for (uint32_t i = 0; i < c.io_num.n_output; ++i) { const auto& attr = resolve_output_attr(is_dynamic_model_, c, i, scratch); - result[i] = make_output_array(attr, out[i]); + py::array arr = make_empty_output(attr, want_float_); + auto bi = arr.request(); + out[i] = {}; + out[i].want_float = want_float_ ? 1 : 0; + out[i].is_prealloc = 1; + out[i].index = i; + out[i].buf = bi.ptr; + out[i].size = static_cast(bi.size * bi.itemsize); + result[i] = arr; + keep_alive_out.push_back(std::move(arr)); } { py::gil_scoped_release nogil; + if (rknn_outputs_get(c.ctx, c.io_num.n_output, out.data(), nullptr) != RKNN_SUCC) + throw std::runtime_error("rknn_outputs_get failed"); rknn_outputs_release(c.ctx, c.io_num.n_output, out.data()); } return result; @@ -383,6 +442,8 @@ private: std::vector contexts_; std::vector ctx_busy_; bool is_dynamic_model_; + bool want_float_; + bool uint8_input_; std::vector input_ranges_; }; @@ -407,9 +468,11 @@ void NativeRKNNExecutor::set_dynamic_shapes(RknnCtx& ctx, const std::vector(m, "NativeRKNNExecutor") - .def(py::init(), + .def(py::init(), py::arg("model_path"), - py::arg("num_workers") = 1) + py::arg("num_workers") = 1, + py::arg("want_float") = true, + py::arg("uint8_input") = false) .def("infer", &NativeRKNNExecutor::infer, py::arg("inputs"), "Run inference with a list of numpy arrays, returns list of numpy arrays.") .def("get_io_info", &NativeRKNNExecutor::get_io_info,