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

chore(ml): use strict mypy (#5001)

Browse source 
* improved typing

* improved export typing

* strict mypy & check export folder

* formatting

* add formatting checks for export folder

* re-added init call
This commit is contained in:
Mert 2023-11-13 11:18:46 -05:00 • committed by GitHub
parent 9fa9ad05b1
commit 935f471ccb
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 70 additions and 55 deletions
Download patch file Download diff file Expand all files Collapse all files

10
machine-learning/app/models/base.py
View file

@ -8,6 +8,7 @@ from typing import Any
import onnxruntime as ort
from huggingface_hub import snapshot_download
from typing_extensions import Buffer
from ..config import get_cache_dir, get_hf_model_name, log, settings
from ..schemas import ModelType
@ -139,11 +140,12 @@ class InferenceModel(ABC):
# HF deep copies configs, so we need to make session options picklable
class PicklableSessionOptions(ort.SessionOptions):
class PicklableSessionOptions(ort.SessionOptions): # type: ignore[misc]
def __getstate__(self) -> bytes:
return pickle.dumps([(attr, getattr(self, attr)) for attr in dir(self) if not callable(getattr(self, attr))])
def __setstate__(self, state: Any) -> None:
self.__init__() # type: ignore
for attr, val in pickle.loads(state):
def __setstate__(self, state: Buffer) -> None:
self.__init__() # type: ignore[misc]
attrs: list[tuple[str, Any]] = pickle.loads(state)
for attr, val in attrs:
setattr(self, attr, val)