Cutlery/immich
1
0
Fork 1
mirror of https://github.com/immich-app/immich.git synced 2025-07-07 18:24:10 -04:00
Code Issues Packages Projects Releases Wiki Activity

Set running threads from env

Browse Source
This commit is contained in:
yoni13 2025-01-12 01:02:16 +08:00
parent 665718b09e
commit efaf70eb9d
2 changed files with 11 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
machine-learning/app/config.py
View File

@ -45,6 +45,9 @@ class Settings(BaseSettings):
ann_fp16_turbo: bool = False ann_fp16_turbo: bool = False
ann_tuning_level: int = 2 ann_tuning_level: int = 2
rknn: bool = True rknn: bool = True
rknn_textual_threads: int = 1
rknn_visual_threads: int = 2
rknn_facial_detection_threads: int = 2
preload: PreloadModelData | None = None preload: PreloadModelData | None = None
max_batch_size: MaxBatchSize | None = None max_batch_size: MaxBatchSize | None = None

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

@ -10,7 +10,7 @@ from numpy.typing import NDArray
from app.schemas import SessionNode from app.schemas import SessionNode
from rknn.rknnpool import rknnPoolExecutor from rknn.rknnpool import rknnPoolExecutor
from ..config import log from ..config import log, settings
def runInfrence(rknn_lite, input): def runInfrence(rknn_lite, input):
@ -23,7 +23,13 @@ class RknnSession:
def __init__(self, model_path: Path | str): def __init__(self, model_path: Path | str):
self.model_path = Path(model_path) self.model_path = Path(model_path)
self.ort_model_path = str(self.model_path).replace(".rknn", ".onnx") self.ort_model_path = str(self.model_path).replace(".rknn", ".onnx")
self.tpe = 1 if "textual" in str(self.model_path) else 2
if 'textual' in self.model_path.name:
self.tpe = settings.rknn_textual_threads
elif 'visual' in self.model_path.name:
self.tpe = settings.rknn_visual_threads
else:
self.tpe = settings.rknn_facial_detection_threads
log.info(f"Loading RKNN model from {self.model_path} with {self.tpe} 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=runInfrence) self.rknnpool = rknnPoolExecutor(rknnModel=self.model_path.as_posix(), TPEs=self.tpe, func=runInfrence)