immich_ml 服务器运行成功

Author: ZHEQIUSHUI

machine-learning/immich_ml/models/base.py

@@ -10,6 +10,7 @@
 import immich_ml.sessions.ann.loader
 import immich_ml.sessions.rknn as rknn
 from immich_ml.sessions.ort import OrtSession
+from immich_ml.sessions.axengine import InferenceSession as AXSession
 
 from ..config import clean_name, log, settings
 from ..schemas import ModelFormat, ModelIdentity, ModelSession, ModelTask, ModelType
@@ -72,13 +73,15 @@ def _download(self) -> None:
             ModelFormat.ARMNN: ["*.rknn"],
             ModelFormat.RKNN: ["*.armnn"],
         }
-
-        snapshot_download(
-            f"immich-app/{clean_name(self.model_name)}",
-            cache_dir=self.cache_dir,
-            local_dir=self.cache_dir,
-            ignore_patterns=ignored_patterns.get(self.model_format, []),
-        )
+        if 'axera' in self.model_name:
+            self.model_format = ModelFormat.AXERA
+        else:
+            snapshot_download(
+                f"immich-app/{clean_name(self.model_name)}",
+                cache_dir=self.cache_dir,
+                local_dir=self.cache_dir,
+                ignore_patterns=ignored_patterns.get(self.model_format, []),
+            )
 
     def _load(self) -> ModelSession:
         return self._make_session(self.model_path)
@@ -116,6 +119,8 @@ def _make_session(self, model_path: Path) -> ModelSession:
                 session = OrtSession(model_path)
             case ".rknn":
                 session = rknn.RknnSession(model_path)
+            case ".axmodel":
+                session = AXSession(str(model_path))
             case _:
                 raise ValueError(f"Unsupported model file type: {model_path.suffix}")
         return session

machine-learning/immich_ml/models/constants.py

@@ -29,6 +29,7 @@
     "ViT-H-14-quickgelu__dfn5b",
     "ViT-H-14__laion2b-s32b-b79k",
     "ViT-L-14-336__openai",
+    "ViT-L-14-336__axera",
     "ViT-L-14-quickgelu__dfn2b",
     "ViT-L-14__laion2b-s32b-b82k",
     "ViT-L-14__laion400m_e31",

machine-learning/immich_ml/schemas.py

@@ -36,6 +36,7 @@ class ModelFormat(StrEnum):
     ARMNN = "armnn"
     ONNX = "onnx"
     RKNN = "rknn"
+    AXERA = "axmodel"
 
 
 class ModelSource(StrEnum):

machine-learning/immich_ml/sessions/axengine/__init__.py

@@ -0,0 +1,22 @@
+# Copyright (c) 2019-2024 Axera Semiconductor Co., Ltd. All Rights Reserved.
+#
+# This source file is the property of Axera Semiconductor Co., Ltd. and
+# may not be copied or distributed in any isomorphic form without the prior
+# written consent of Axera Semiconductor Co., Ltd.
+#
+
+# thanks to community contributors list below:
+#   zylo117: https://github.com/zylo117, first implementation of the axclrt backend
+
+from ._providers import axengine_provider_name, axclrt_provider_name
+from ._providers import get_all_providers, get_available_providers
+
+# check if axclrt is installed, or is a supported chip(e.g. AX650, AX620E etc.)
+_available_providers = get_available_providers()
+if not _available_providers:
+    raise ImportError(
+        f"No providers found. Please make sure you have installed one of the following: {get_all_providers()}")
+print("[INFO] Available providers: ", _available_providers)
+
+from ._node import NodeArg
+from ._session import SessionOptions, InferenceSession