fix(ml): model downloading improvements (#3988)

* handle ort `NoSuchFile` error, stricter file check

* keep exception order
This commit is contained in:
Mert 2023-09-05 21:48:40 -04:00 committed by GitHub
parent 454737ca79
commit 3cf0f5f11b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 6 additions and 2 deletions

View file

@ -8,7 +8,7 @@ from typing import Any
from zipfile import BadZipFile
import onnxruntime as ort
from onnxruntime.capi.onnxruntime_pybind11_state import InvalidProtobuf # type: ignore
from onnxruntime.capi.onnxruntime_pybind11_state import InvalidProtobuf, NoSuchFile # type: ignore
from ..config import get_cache_dir, log, settings
from ..schemas import ModelType
@ -57,7 +57,7 @@ class InferenceModel(ABC):
try:
loader(**model_kwargs)
except (OSError, InvalidProtobuf, BadZipFile):
except (OSError, InvalidProtobuf, BadZipFile, NoSuchFile):
log.warn(
(
f"Failed to load {self.model_type.replace('_', ' ')} model '{self.model_name}'."

View file

@ -131,6 +131,10 @@ class CLIPEncoder(InferenceModel):
os.remove(file)
return True
@property
def cached(self) -> bool:
return (self.cache_dir / "textual.onnx").is_file() and (self.cache_dir / "visual.onnx").is_file()
# same as `_transform_blob` without `_blob2image`
def _transform_pil_image(n_px: int) -> Compose: