Implement model caching in app.py and initialize default model at startup. Update test.py to use variable for app URL.

app.py

@@ -9,8 +9,14 @@ from ultralytics import YOLO
 import requests
 import supervision as sv
 
+# Global variable to store model instances
+MODEL_CACHE = {}
 
 def init_model(model_id: str):
+    # Return cached model if it exists
+    if model_id in MODEL_CACHE:
+        return MODEL_CACHE[model_id]
+    
     # Define models
     MODEL_OPTIONS = {
         "YOLOv11-Nano": "medieval-yolov11n.pt",
@@ -19,19 +25,26 @@ def init_model(model_id: str):
         "YOLOv11-Large": "medieval-yolov11l.pt",
         "YOLOv11-XLarge": "medieval-yolov11x.pt"
     }
+    
     if model_id in MODEL_OPTIONS:
         os.makedirs("models", exist_ok=True)
         model_path = hf_hub_download(
             repo_id="biglam/medieval-manuscript-yolov11",
-            filename=MODEL_OPTIONS[model_id]
+            filename=MODEL_OPTIONS[model_id],
+            cache_dir="models"  # Specify cache directory
         )
-        local_path = os.path.join("models", model_path)
-        return YOLO(local_path)
+        model = YOLO(model_path)
+        MODEL_CACHE[model_id] = model
+        return model
     else:
         raise ValueError(f"Model {model_id} not found")
 
 app = FastAPI()
 
+# Initialize default model at startup
+@app.on_event("startup")
+async def startup_event():
+    init_model("YOLOv11-XLarge")  # Initialize default model
 
 @app.post("/predict")
 async def predict(image: UploadFile,
@@ -39,7 +52,7 @@ async def predict(image: UploadFile,
                   conf: float = 0.25,
                   iou: float = 0.7
                   ):
-    # Initialize model at startup
+    # Get model from cache or initialize it
     model = init_model(model_id)
     
     # Download and open image from URL

test.py

@@ -18,7 +18,7 @@ with open(image_path, 'rb') as f:
     }
 
     # Send POST request to the endpoint
-    response = requests.post(local_app + ':7860/predict', 
+    response = requests.post(hf_app + ':7860/predict', 
                            files=files,
                            params=params)