Cutlery/immich
1
0
Fork 1
mirror of https://github.com/immich-app/immich.git synced 2025-06-22 06:50:54 -04:00
Code Issues Packages Projects Releases Wiki Activity

all infrencing works with 1 max job concurrency

Browse Source
This commit is contained in:
yoni13 2025-01-09 10:38:40 +00:00
parent 082c426e34
commit a94fad543b
1 changed files with 37 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

68
machine-learning/app/sessions/rknn.py
View File

@ -4,8 +4,10 @@ from pathlib import Path
from typing import Any, List
import numpy as np
import onnxruntime as ort
from numpy.typing import NDArray
from rknnlite.api import RKNNLite # Importing RKNN API
from rknnlite.api import RKNNLite
from app.models.constants import SUPPORTED_PROVIDERS
from app.schemas import SessionNode
@ -16,55 +18,55 @@ from ..config import log, settings
class RknnSession:
def __init__(self, model_path: Path | str):
self.model_path = Path(model_path)
self.rknn = RKNNLite() # Initialize RKNN object
# self.rknn.config(target_platform='rk3566')
# Load the RKNN model
log.info(f"Loading RKNN model from {self.model_path}")
self._load_model()
def _load_model(self) -> None:
# ret = self.rknn.load_onnx(self.model_path.as_posix())
self.ort_model_path = str(self.model_path).replace(".rknn", ".onnx")
self.rknn = RKNNLite()
log.info(f"Loading RKNN model from {self.model_path}")
self.ort_session = ort.InferenceSession(
self.ort_model_path,
)
# ret = self.rknn.load_onnx(self.model_path)
print('--> Load RKNN model')
ret = self.rknn.load_rknn(self.model_path.as_posix())
if ret != 0:
raise RuntimeError("Failed to load RKNN model")
# print('--> Building model')
# ret = self.rknn.build(do_quantization=False)
if ret != 0:
print('Build model failed!')
exit(ret)
ret = self.rknn.init_runtime()
if ret != 0:
raise RuntimeError("Failed to initialize RKNN runtime")
def get_inputs(self) -> List[SessionNode]:
input_attrs = self.rknn.query_inputs()
return input_attrs # RKNN does not provide direct SessionNode equivalent
def get_inputs(self) -> list[SessionNode]:
inputs: list[SessionNode] = self.ort_session.get_inputs()
return inputs
def get_outputs(self) -> List[SessionNode]:
output_attrs = self.rknn.query_outputs()
return output_attrs
def get_outputs(self) -> list[SessionNode]:
outputs: list[SessionNode] = self.ort_session.get_outputs()
return outputs
def run(
self,
output_names: list[str] | None,
input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]],
run_options: Any = None,
) -> List[NDArray[np.float32]]:
inputs = [v for v in input_feed.values()]
):
# Run inference
log.debug(f"Running inference on RKNN model")
outputs = self.rknn.inference(inputs=inputs)
input_data = [np.ascontiguousarray(v) for v in input_feed.values()][0]
# log.info(f"Running inference on RKNN model")
outputs = self.rknn.inference(inputs=[input_data])
if not outputs:
outputs = self.rknn.inference(inputs=[input_data], data_format='nchw')
# log.info("inputs:")
# log.info(input_data)
# log.info("outputs:")
# log.info(outputs)
# log.info("RKNN END")
return outputs
def release(self) -> None:
log.info("Releasing RKNN resources")
self.rknn.release()
# Example Usage:
# session = RknnSession(model_path="path/to/model.rknn")
# outputs = session.run(input_feed={"input_name": input_data})
# session.release()