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

fix(ml): limit load retries (#10494)

Browse Source
This commit is contained in:
Mert
2024-06-20 14:13:18 -04:00
committed by GitHub
parent 79a8ab71ef
commit a42af06889
3 changed files with 26 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -31,6 +31,7 @@ class InferenceModel(ABC):
**model_kwargs: Any,
) -> None:
self.loaded = False
self.load_attempts = 0
self.model_name = clean_name(model_name)
self.cache_dir = Path(cache_dir) if cache_dir is not None else self.cache_dir_default
self.providers = providers if providers is not None else self.providers_default
@@ -48,9 +49,11 @@ class InferenceModel(ABC):
def load(self) -> None:
if self.loaded:
return
self.load_attempts += 1
self.download()
log.info(f"Loading {self.model_type.replace('-', ' ')} model '{self.model_name}' to memory")
attempt = f"Attempt #{self.load_attempts + 1} to load" if self.load_attempts else "Loading"
log.info(f"{attempt} {self.model_type.replace('-', ' ')} model '{self.model_name}' to memory")
self.session = self._load()
self.loaded = True