From 3cf0f5f11bf765b804f39fef1761b833a81ec7d8 Mon Sep 17 00:00:00 2001
From: Mert <101130780+mertalev@users.noreply.github.com>
Date: Tue, 5 Sep 2023 21:48:40 -0400
Subject: [PATCH] fix(ml): model downloading improvements (#3988)

* handle ort `NoSuchFile` error, stricter file check

* keep exception order
---
 machine-learning/app/models/base.py | 4 ++--
 machine-learning/app/models/clip.py | 4 ++++
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/machine-learning/app/models/base.py b/machine-learning/app/models/base.py
index 1207342582..a059568a80 100644
--- a/machine-learning/app/models/base.py
+++ b/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}'."
diff --git a/machine-learning/app/models/clip.py b/machine-learning/app/models/clip.py
index 89d44a573d..fbdded9c72 100644
--- a/machine-learning/app/models/clip.py
+++ b/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: