refactor

docker/hwaccel.ml.yml

@@ -17,10 +17,9 @@ services:
   rknn:
     security_opt:
       - systempaths=unconfined
+      - apparmor=unconfined
     devices:
       - /dev/dri:/dev/dri
-    volumes:
-      - /sys/kernel/debug/:/sys/kernel/debug/:ro
 
   cpu: {}
 

machine-learning/Dockerfile

@@ -106,7 +106,6 @@ RUN echo "hard core 0" >> /etc/security/limits.conf && \
 
 COPY --from=builder /usr/src/app/.venv /usr/src/app/.venv
 COPY ann/ann.py /usr/src/ann/ann.py
-COPY rknn/rknnpool.py /usr/src/rknn/rknnpool.py
 COPY start.sh log_conf.json gunicorn_conf.py ./
 COPY app .
 

machine-learning/app/main.py

@@ -226,9 +226,9 @@ async def load(model: InferenceModel) -> InferenceModel:
             except FileNotFoundError as e:
                 if model.model_format == ModelFormat.ONNX:
                     raise e
-                log.exception(e)
                 log.warning(
-                    f"{model.model_format.upper()} is available, but model '{model.model_name}' does not support it."
+                    f"{model.model_format.upper()} is available, but model '{model.model_name}' does not support it.",
+                    exc_info=e,
                 )
                 model.model_format = ModelFormat.ONNX
                 model.load()

machine-learning/app/models/base.py

@@ -8,9 +8,8 @@ from typing import Any, ClassVar
 from huggingface_hub import snapshot_download
 
 import ann.ann
-import rknn.rknnpool
+import app.sessions.rknn as rknn
 from app.sessions.ort import OrtSession
-from app.sessions.rknn import RknnSession
 
 from ..config import clean_name, log, settings
 from ..schemas import ModelFormat, ModelIdentity, ModelSession, ModelTask, ModelType
@@ -34,6 +33,7 @@ class InferenceModel(ABC):
         self.model_name = clean_name(model_name)
         self.cache_dir = Path(cache_dir) if cache_dir is not None else self._cache_dir_default
         self.model_format = model_format if model_format is not None else self._model_format_default
+        self.model_path_prefix = rknn.model_prefix if self.model_format == ModelFormat.RKNN else None
         if session is not None:
             self.session = session
 
@@ -116,7 +116,7 @@ class InferenceModel(ABC):
             case ".onnx":
                 session = OrtSession(model_path)
             case ".rknn":
-                session = RknnSession(model_path)
+                session = rknn.RknnSession(model_path)
             case _:
                 raise ValueError(f"Unsupported model file type: {model_path.suffix}")
         return session
@@ -127,6 +127,8 @@ class InferenceModel(ABC):
 
     @property
     def model_path(self) -> Path:
+        if self.model_path_prefix:
+            return self.model_dir / self.model_path_prefix / f"model.{self.model_format}"
         return self.model_dir / f"model.{self.model_format}"
 
     @property
@@ -164,7 +166,7 @@ class InferenceModel(ABC):
 
     @property
     def _model_format_default(self) -> ModelFormat:
-        if rknn.rknnpool.is_available and settings.rknn:
+        if rknn.is_available:
             return ModelFormat.RKNN
         elif ann.ann.is_available and settings.ann:
             return ModelFormat.ARMNN

machine-learning/app/sessions/rknn/__init__.py

@@ -6,15 +6,17 @@ from typing import Any, NamedTuple
 import numpy as np
 from numpy.typing import NDArray
 
+from app.config import log, settings
 from app.schemas import SessionNode
-from rknn.rknnpool import RknnPoolExecutor, soc_name
 
-from ..config import log, settings
+from .rknnpool import RknnPoolExecutor, is_available, soc_name
+
+is_available = is_available and settings.rknn
+model_prefix = Path("rknpu") / soc_name if is_available else None
 
 
-def runInference(rknn_lite: Any, input: list[NDArray[np.float32]]) -> list[NDArray[np.float32]]:
+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
 
 
@@ -38,17 +40,13 @@ input_output_mapping: dict[str, dict[str, Any]] = {
 
 
 class RknnSession:
-    def __init__(self, model_path: Path | str):
-        self.model_path = Path(str(model_path).replace("model", soc_name))
-        self.model_type = "detection" if "detection" in self.model_path.as_posix() else "recognition"
+    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 {self.model_path} with {self.tpe} threads.")
-        self.rknnpool = RknnPoolExecutor(rknnModel=self.model_path.as_posix(), tpes=self.tpe, func=runInference)
-        log.info(f"Loaded RKNN model from {self.model_path} with {self.tpe} threads.")
-
-    def __del__(self) -> None:
-        self.rknnpool.release()
+        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.")
 
     def get_inputs(self) -> list[SessionNode]:
         return [RknnNode(name=k, shape=v) for k, v in input_output_mapping[self.model_type]["input"].items()]

machine-learning/app/sessions/rknn/rknnpool.py

@@ -1,49 +1,55 @@
 # This code is from leafqycc/rknn-multi-threaded
 # Following Apache License 2.0
 
-import os
 from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
 from queue import Queue
+
 import numpy as np
-from typing import Callable
 from numpy.typing import NDArray
+
 from app.config import log
 
 supported_socs = ["rk3566", "rk3588"]
-coremask_supported_socs = ["rk3576","rk3588"]
+coremask_supported_socs = ["rk3576", "rk3588"]
 
+
+def get_soc(device_tree_path: Path | str) -> str | None:
+    try:
+        with Path(device_tree_path).open() as f:
+            device_compatible_str = f.read()
+            for soc in supported_socs:
+                if soc in device_compatible_str:
+                    return soc
+            log.warning("Device is not supported for RKNN")
+    except OSError as e:
+        log.warning("Could not read /proc/device-tree/compatible. Reason: %s", e.msg)
+    return None
+
+
+soc_name = None
+is_available = False
 try:
     from rknnlite.api import RKNNLite
 
-    with open("/proc/device-tree/compatible") as f:
-        device_compatible_str = f.read()
-        for soc in supported_socs:
-            if soc in device_compatible_str:
-                is_available = True
-                soc_name = soc
-                break
-            else:
-                is_available = False
-                soc_name = None
-    is_available = is_available and os.path.exists("/sys/kernel/debug/rknpu/load")
-except (FileNotFoundError, ImportError):
+    soc_name = get_soc("/proc/device-tree/compatible")
+    is_available = soc_name is not None
+except ImportError:
     log.debug("RKNN is not available")
-    is_available = False
-    soc_name = None
 
 
-def init_rknn(rknnModel) -> Callable:
+def init_rknn(model_path: str) -> RKNNLite:
     if not is_available:
         raise RuntimeError("rknn is not available!")
     rknn_lite = RKNNLite()
-    ret = rknn_lite.load_rknn(rknnModel)
+    ret = rknn_lite.load_rknn(model_path)
     if ret != 0:
         raise RuntimeError("Load RKNN rknnModel failed")
-    
+
     if soc_name in coremask_supported_socs:
         ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_AUTO)
     else:
-        ret = rknn_lite.init_runtime() # Please do not set this parameter on other platforms.
+        ret = rknn_lite.init_runtime()  # Please do not set this parameter on other platforms.
 
     if ret != 0:
         raise RuntimeError("Init runtime environment failed")
@@ -51,18 +57,11 @@ def init_rknn(rknnModel) -> Callable:
     return rknn_lite
 
 
-def init_rknns(rknnModel, tpes) -> list[Callable]:
-    rknn_list = []
-    for i in range(tpes):
-        rknn_list.append(init_rknn(rknnModel))
-    return rknn_list
-
-
 class RknnPoolExecutor:
-    def __init__(self, rknnModel: str, tpes: int, func):
+    def __init__(self, model_path: str, tpes: int, func):
         self.tpes = tpes
         self.queue = Queue()
-        self.rknn_pool = init_rknns(rknnModel, tpes)
+        self.rknn_pool = [init_rknn(model_path) for _ in range(tpes)]
         self.pool = ThreadPoolExecutor(max_workers=tpes)
         self.func = func
         self.num = 0
@@ -81,3 +80,6 @@ class RknnPoolExecutor:
         self.pool.shutdown()
         for rknn_lite in self.rknn_pool:
             rknn_lite.release()
+
+    def __del__(self) -> None:
+        self.release()

machine-learning/export/rknpu/ViT-B-32__openai/README.md

@@ -0,0 +1,11 @@
+---
+tags:
+ - immich
+ - clip
+---
+# Model Description
+
+This repo contains ONNX exports for the CLIP model [openai/clip-vit-base-patch32](https://huggingface.co/openai/clip-vit-base-patch32). 
+It separates the visual and textual encoders into separate models for the purpose of generating image and text embeddings.
+
+This repo is specifically intended for use with [Immich](https://immich.app/), a self-hosted photo library.

machine-learning/export/rknpu/build_rknn.py

@@ -0,0 +1,69 @@
+import argparse
+from pathlib import Path
+
+from rknn.api import RKNN
+
+parser = argparse.ArgumentParser("ONNX to RKNN model converter")
+parser.add_argument(
+    "model", help="Directory of the model that will be exported to RKNN ex:ViT-B-32__openai.", type=Path
+)
+parser.add_argument("target_platform", help="target platform ex:rk3566", type=str)
+args = parser.parse_args()
+
+
+def ConvertModel(model_dir: Path, target_platform: str, dynamic_input=None):
+    input_path = model_dir / "model.onnx"
+    print(f"Converting model {input_path}")
+    rknn = RKNN(verbose=False)
+
+    rknn.config(
+        target_platform=target_platform,
+        dynamic_input=dynamic_input,
+        enable_flash_attention=True,
+        # remove_reshape=True,
+        # model_pruning=True
+    )
+    ret = rknn.load_onnx(model=input_path.as_posix())
+
+    if ret != 0:
+        print("Load failed!")
+        exit(ret)
+
+    ret = rknn.build(do_quantization=False)
+
+    if ret != 0:
+        print("Build failed!")
+        exit(ret)
+
+    output_path = model_dir / "rknpu" / target_platform / "model.rknn"
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    print(f"Exporting model {model_dir} to {output_path}")
+    ret = rknn.export_rknn(output_path.as_posix())
+    if ret != 0:
+        print("Export rknn model failed!")
+        exit(ret)
+
+
+textual = args.model / "textual"
+visual = args.model / "visual"
+detection = args.model / "detection"
+recognition = args.model / "recognition"
+
+is_dir = [textual.is_dir(), visual.is_dir(), detection.is_dir(), recognition.is_dir()]
+if not any(is_dir):
+    print("Unknown model")
+    exit(1)
+
+is_textual, is_visual, is_detection, is_recognition = is_dir
+
+if is_textual:
+    ConvertModel(textual, target_platform=args.target_platform)
+
+if is_visual:
+    ConvertModel(visual, target_platform=args.target_platform)
+
+if is_detection:
+    ConvertModel(detection, args.target_platform, [[[1, 3, 640, 640]]])
+
+if is_recognition:
+    ConvertModel(recognition, args.target_platform, [[[1, 3, 112, 112]]])

machine-learning/rknn/export/build_rknn.py

@@ -1,53 +0,0 @@
-import argparse
-import os
-
-parser = argparse.ArgumentParser("RKNN model converting")
-parser.add_argument("model", help="Directory of the model that will be exported to RKNN ex:ViT-B-32__openai.", type=str)
-parser.add_argument("target_platform", help="target platform ex:rk3566", type=str)
-args = parser.parse_args()
-
-
-def ConvertModel(model_path='ViT-B-32__openai/textual/model.onnx', target_platform='rk3566', dynamic_input = None):
-    # E build: Repeat call the 'rknn.build' or 'rknn.hybrid_quantization_step1' is not allow!
-    from rknn.api import RKNN
-    rknn = RKNN(verbose=False)
-
-    rknn.config(target_platform=target_platform, dynamic_input=dynamic_input)
-    ret = rknn.load_onnx(model=model_path)
-
-    if ret != 0:
-        print("Load failed!")
-        exit(ret)
-
-    ret = rknn.build(do_quantization=False)
-
-    if ret != 0:
-        print("Build failed!")
-        exit(ret)
-    print(model_path.replace('model.onnx',f'{target_platform}.rknn'))
-    ret = rknn.export_rknn(model_path.replace('model.onnx',f'{target_platform}.rknn'))
-    if ret != 0:
-            print('Export rknn model failed!')
-            exit(ret)
-    print('done')
-    del rknn
-    del RKNN
-
-    if not os.path.isfile(f'{model_path.replace("onnx","rknn")}'):
-        print(f'Dummy model not found at {model_path.replace("onnx","rknn")}, creating one')
-        with open(f'{model_path.replace("onnx","rknn")}', 'w'):
-            pass
-
-
-if os.path.isdir(f'{args.model}/textual') and os.path.isdir(f'{args.model}/visual'): # is a clip model
-    print('Converting Clip model.')
-    ConvertModel(model_path=f'{args.model}/textual/model.onnx', target_platform=args.target_platform)
-    ConvertModel(model_path=f'{args.model}/visual/model.onnx', target_platform=args.target_platform)
-
-elif os.path.isdir(f'{args.model}/detection') and os.path.isdir(f'{args.model}/recognition'): # is a facial model
-    print('Converting facial model.')
-    ConvertModel(f'{args.model}/detection/model.onnx', args.target_platform, [[[1, 3, 640, 640]]])
-    ConvertModel(f'{args.model}/recognition/model.onnx', args.target_platform, [[[1, 3, 112, 112]]])
-
-else:
-    print('Unknown model.')