use lru cache for simplicity

app.py

@@ -7,6 +7,7 @@ Any new model should implement the following functions:
 
 import os
 import glob
+import functools
 import spaces
 import gradio as gr
 from huggingface_hub import get_token
@@ -17,7 +18,6 @@ from utils import (
     FlaggedCounter,
 )
 from flagging import HuggingFaceDatasetSaver
-import numpy as np
 
 import install_private_repos
 from seavision import load_model, AHOY
@@ -48,29 +48,16 @@ h1 {
 }
 """
 
-rng = np.random.default_rng(0xDEADBEEF)
-
-class SeaVisionModel:
-    def __init__(self, model_name="ahoy-RGB-b2"):
-        self.model_name = model_name
-        self.model = None
-
-    def get_model(self):
-        if self.model is None:
-            self.model = load_model(self.model_name)
-        return self.model
-
-    def run_inference(self, image):
-        model = self.get_model()
-        results = model(image)
-        return results.draw(image, diameter=4)
-
-vision_model = SeaVisionModel()
+@functools.lru_cache(maxsize=1)
+def get_model():
+    return load_model("ahoy-RGB-b2")
 
 @spaces.GPU
 def inference(image):
     """Run inference on image and return annotated image."""
-    return vision_model.run_inference(image)
+    model = get_model()
+    results = model(image)
+    return results.draw(image, diameter=4)
 
 # Flagging
 dataset_name = "SEA-AI/crowdsourced-sea-images"
@@ -106,7 +93,7 @@ with gr.Blocks(theme=theme, css=css, title="SEA.AI Vision Demo") as demo:
         examples=glob.glob("examples/*.jpg"),
         inputs=img_input,
         outputs=img_output,
-        fn=lambda image: inference(model, image),
+        fn=inference,
         cache_examples=True,
     )