helpy4/immich
1
0
Fork 0
mirror of https://github.com/immich-app/immich synced 2025-11-07 17:27:20 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

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
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
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,