immich-app/immich
1
0
Fork 0
mirror of https://github.com/immich-app/immich.git synced 2025-12-18 01:11:07 +03:00
Code Issues Packages Projects Releases 16 Wiki Activity

fix(ml): clear model cache on load error (#2951)

Browse Source 
* clear model cache on load error

* updated caught exceptions
This commit is contained in:
Mert
2023-06-27 17:01:24 -04:00
committed by GitHub
parent 39a885a37c
commit 47982641b2
4 changed files with 38 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
machine-learning/app/models/image_classification.py
View File

@@ -1,4 +1,5 @@
from pathlib import Path
from typing import Any
from PIL.Image import Image
from transformers.pipelines import pipeline
@@ -17,10 +18,11 @@ class ImageClassifier(InferenceModel):
min_score: float = settings.min_tag_score,
cache_dir: Path | None = None,
**model_kwargs,
):
super().__init__(model_name, cache_dir)
) -> None:
self.min_score = min_score
super().__init__(model_name, cache_dir, **model_kwargs)
def load(self, **model_kwargs: Any) -> None:
self.model = pipeline(
self.model_type.value,
self.model_name,