Bläddra i källkod

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

* handle ort `NoSuchFile` error, stricter file check

* keep exception order
Mert 1 år sedan
förälder
incheckning
3cf0f5f11b
2 ändrade filer med 6 tillägg och 2 borttagningar
  1. 2 2
      machine-learning/app/models/base.py
  2. 4 0
      machine-learning/app/models/clip.py

+ 2 - 2
machine-learning/app/models/base.py

@@ -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}'."

+ 4 - 0
machine-learning/app/models/clip.py

@@ -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: