ashipov/immich
1
0
Fork 1
forked from dark_thunder/immich
Code Pull Requests Activity

feat(ml): add face models (#4952)

Browse Source 
added models to config dropdown

fixed downloading

updated tests

use hf for face models

formatting
This commit is contained in:
Mert
2023-11-11 20:04:49 -05:00
committed by GitHub
parent 7fca0d8da5
commit 328a58ac0d
8 changed files with 101 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
machine-learning/app/models/base.py
View File

@ -7,8 +7,9 @@ from shutil import rmtree
from typing import Any
import onnxruntime as ort
from huggingface_hub import snapshot_download
from ..config import get_cache_dir, log, settings
from ..config import get_cache_dir, get_hf_model_name, log, settings
from ..schemas import ModelType
@ -78,9 +79,13 @@ class InferenceModel(ABC):
def configure(self, **model_kwargs: Any) -> None:
pass
@abstractmethod
def _download(self) -> None:
...
snapshot_download(
get_hf_model_name(self.model_name),
cache_dir=self.cache_dir,
local_dir=self.cache_dir,
local_dir_use_symlinks=False,
)
@abstractmethod
def _load(self) -> None: