Download model from HuggingFace to tmp folder for running on Lambda

README.md

@@ -26,5 +26,5 @@ Before running this project, make sure you have the following prerequisites inst
 To use this project, run the following command:
 
 ```bash
-python3 app.py path/to/your_img
+python3 app.py BASE_64_IMAGE_BLOB_STRING
 ```

app.py

@@ -7,14 +7,17 @@ import onnxruntime as rt
 from torchvision import transforms as T
 from tokenizer_base import Tokenizer
 from PIL import Image
+from huggingface_hub import hf_hub_download, try_to_load_from_cache
 
 
 class DocumentParserModel:
-    def __init__(self, model_path, img_size, charset):
-        self.charset = charset
-        self.tokenizer_base = Tokenizer(self.charset)
+    def __init__(self):
+        charset = r"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
+        img_size = (32, 128)
+
+        self.tokenizer_base = Tokenizer(charset)
         self.transform = self.create_transform_pipeline(img_size)
-        self.ort_session = self.initialize_onnx_model(str(model_path))
+        self.ort_session = self.initialize_onnx_model()
 
     def create_transform_pipeline(self, img_size):
         transforms = [
@@ -24,46 +27,50 @@ class DocumentParserModel:
         ]
         return T.Compose(transforms)
 
-    def initialize_onnx_model(self, model_path):
-        onnx_model = onnx.load(model_path)
+    def initialize_onnx_model(self):
+        repo_id = "stevenchang/captcha"
+        filename = "captcha.onnx"
+
+        filepath = try_to_load_from_cache(repo_id, filename)
+
+        if isinstance(filepath, str):
+            model_file = filepath
+        else:
+            model_file = result = hf_hub_download(repo_id, filename)
+
+        onnx_model = onnx.load(model_file)
         onnx.checker.check_model(onnx_model)
-        return rt.InferenceSession(model_path)
+        return rt.InferenceSession(model_file)
 
     def load_image_from_base64(self, base64_string):
         img_data = base64.b64decode(base64_string)
         image_buffer = io.BytesIO(img_data)
-        image = Image.open(image_buffer)
-        return image
 
-    def predict_text(self, image_path):
         try:
-            with self.load_image_from_base64(image_path) as img_org:
-                x = self.transform(img_org.convert("RGB")).unsqueeze(0)
-                ort_inputs = {self.ort_session.get_inputs()[0].name: x.cpu().numpy()}
-                logits = self.ort_session.run(None, ort_inputs)[0]
-                probs = torch.tensor(logits).softmax(-1)
-                preds, _ = self.tokenizer_base.decode(probs)
-                return preds[0]
+            image = Image.open(image_buffer)
+            return image
         except IOError:
-            print(f"Error: Cannot open image {image_path}")
+            print(f"Error: Cannot open image {image_blob}")
             return None
 
+    def predict_text(self, image_blob):
+        with self.load_image_from_base64(image_blob) as img_org:
+            x = self.transform(img_org.convert("RGB")).unsqueeze(0)
+            ort_inputs = {self.ort_session.get_inputs()[0].name: x.cpu().numpy()}
+            logits = self.ort_session.run(None, ort_inputs)[0]
+            probs = torch.tensor(logits).softmax(-1)
+            preds, _ = self.tokenizer_base.decode(probs)
+            return preds[0]
+
 
 if __name__ == "__main__":
     import sys
 
-    model_path = "captcha.onnx"
-    img_size = (32, 128)
-    charset = r"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
+    doc_parser = DocumentParserModel()
 
-    doc_parser = DocumentParserModel(
-        model_path=model_path,
-        img_size=img_size,
-        charset=charset,
-    )
     if len(sys.argv) > 1:
-        image_path = sys.argv[1]
-        result = doc_parser.predict_text(image_path)
+        image_blob = sys.argv[1]
+        result = doc_parser.predict_text(image_blob)
         print(result)
     else:
-        print("Please provide an image path.")
+        print("Please provide an image blob.")

captcha.onnx

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:2a672587ee82eb010dbef54dd0a38e99625293608ed4068c4bd20ebe467fede4
-size 95401304

handler.py

@@ -6,14 +6,7 @@ from app import DocumentParserModel
 LOGGER = logging.getLogger()
 LOGGER.setLevel(logging.INFO)
 
-model_path = "captcha.onnx"
-img_size = (32, 128)
-charset = r"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
-model = DocumentParserModel(
-    model_path=model_path,
-    img_size=img_size,
-    charset=charset,
-)
+model = DocumentParserModel()
 
 
 def lambda_handle(event, context):
@@ -22,4 +15,4 @@ def lambda_handle(event, context):
         LOGGER.info("No ML work to do. Just staying warm...")
         return "Keeping Lambda warm"
 
-    return {"statusCode": 200, "vc": model.predict_text(image_path=event["image_path"])}
+    return {"statusCode": 200, "result": model.predict_text(image_blob=event["image_blob"])}

requirements.txt

@@ -3,3 +3,4 @@ torchvision==0.12.0
 onnx==1.16.0
 onnxruntime==1.16.*
 Pillow==10.0.0
+huggingface_hub==0.21.4

serverless.yml

@@ -29,7 +29,7 @@ functions:
     environment:
       # On Lambda, the default location is not writable. Only the "/tmp" folder is writable. Therefore, we need to set the cache location inside "/tmp".
       TORCH_HOME: /tmp/.ml_cache
-      # TRANSFORMERS_CACHE: /tmp/.ml_cache/huggingface
+      HF_HUB_CACHE: /tmp/.ml_cache/huggingface
 custom:
   warmup:
     MLModelWarmer: