Fix the deployment space problem

api/app.py

@@ -1,29 +1,45 @@
 # api/app.py
+import io
+import json
+import traceback
+from pathlib import Path
+
 from fastapi import FastAPI, File, UploadFile
-from fastapi.responses import JSONResponse, FileResponse, Response
 from fastapi.middleware.cors import CORSMiddleware
-from pathlib import Path
+from fastapi.responses import FileResponse, JSONResponse, Response
 from PIL import Image
-import io, traceback, json
 
-app = FastAPI(title="VizRef API")
+import numpy as np
+import sys
+
+# make sure project root is importable (scripts/inference.py lives at repo root)
+ROOT = Path(__file__).resolve().parent.parent
+sys.path.append(str(ROOT))
 
+# ---- YOUR ORIGINAL MODEL LOADER (unchanged in spirit) ----
+from scripts.inference import ModelInference
+
+# create model ONCE (like your original code)
+# if your ModelInference supports map_location/device, set to CPU for Spaces
+model = ModelInference(
+    checkpoint_path=str(ROOT / "checkpoints" / "best_checkpoint.pth"),
+    multi_task=True
+)
+
+# -------- FastAPI app & CORS ----------
+app = FastAPI(title="VizRef API")
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
+    allow_origins=["*"], allow_credentials=True,
+    allow_methods=["*"], allow_headers=["*"],
 )
 
-# --- serve the frontend ---
+# -------- Serve your UI ----------
 @app.get("/")
 def root():
-    # robust absolute path to ui/index.html
-    ui_path = Path(__file__).resolve().parent.parent / "ui" / "index.html"
-    return FileResponse(str(ui_path))
+    return FileResponse(str(ROOT / "ui" / "index.html"))
 
-# --- API: health (frontend expects /api/health) ---
+# -------- API: health ----------
 @app.get("/api/health")
 def health_check():
     return {
@@ -32,14 +48,14 @@ def health_check():
         "model_info": {"model_name": "resnet50", "multi_task": True},
     }
 
-# --- API: training data (frontend expects /api/training_data) ---
+# -------- API: training data ----------
 @app.get("/api/training_data")
 def get_training_data():
-    with open("data/splits/train.json", "r") as f:
+    with open(ROOT / "data" / "splits" / "train.json", "r") as f:
         data = json.load(f)
     return {"status": "success", "data": data}
 
-# --- API: proxy image (frontend calls /api/proxy_image) ---
+# -------- API: proxy image ----------
 @app.get("/api/proxy_image")
 def proxy_image(url: str):
     import requests
@@ -57,9 +73,8 @@ def proxy_image(url: str):
     except:
         return Response(status_code=404)
 
-# --- API: predict (frontend posts to /api/predict) ---
+# --- helper to make numpy types JSON-safe ---
 def _pythonify(obj):
-    import numpy as np
     if isinstance(obj, dict):
         return {k: _pythonify(v) for k, v in obj.items()}
     if isinstance(obj, list):
@@ -72,23 +87,20 @@ def _pythonify(obj):
         return bool(obj)
     return obj
 
+# -------- API: predict (your original logic, adapted to JSON) ----------
 @app.post("/api/predict")
 async def predict(file: UploadFile = File(...)):
     try:
         raw = await file.read()
-        # decode image
         try:
             img = Image.open(io.BytesIO(raw)).convert("RGB")
         except Exception as e:
-            return JSONResponse(
-                {"status": "error", "message": f"Cannot read image: {e}"},
-                status_code=400,
-            )
+            return JSONResponse({"status": "error", "message": f"Cannot read image: {e}"}, status_code=400)
 
         tmp = Path("/tmp/vizref_upload.jpg")
         img.save(tmp)
 
-        # IMPORTANT: if your model loader expects CPU, ensure map_location="cpu" in your ModelInference
+        # your original call
         result = model.predict_image(str(tmp))
 
         return JSONResponse({