Fix model loading: use /data for HF cache directory instead of /app/data

Dockerfile

@@ -31,11 +31,11 @@ RUN mkdir -p /data/uploads /data/results && chmod -R 777 /data
 # This allows the credentials to be passed as a secret and written to file at runtime
 RUN echo '#!/bin/sh' > /entrypoint.sh && \
     echo 'set -e' >> /entrypoint.sh && \
+    echo 'mkdir -p /data/uploads /data/results' >> /entrypoint.sh && \
+    echo 'chmod -R 777 /data' >> /entrypoint.sh && \
     echo 'if [ -n "$FIREBASE_CREDENTIALS" ]; then' >> /entrypoint.sh && \
-    echo '  mkdir -p /data' >> /entrypoint.sh && \
-    echo '  touch /data/firebase-adminsdk.json' >> /entrypoint.sh && \
-    echo '  chmod 600 /data/firebase-adminsdk.json' >> /entrypoint.sh && \
     echo '  printf "%s" "$FIREBASE_CREDENTIALS" > /data/firebase-adminsdk.json' >> /entrypoint.sh && \
+    echo '  chmod 600 /data/firebase-adminsdk.json' >> /entrypoint.sh && \
     echo 'fi' >> /entrypoint.sh && \
     echo 'exec "$@"' >> /entrypoint.sh && \
     chmod +x /entrypoint.sh

app/colorize_model.py

@@ -30,25 +30,36 @@ class ColorizeModel:
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         logger.info("Using device: %s", self.device)
         self.dtype = torch.float16 if self.device == "cuda" else torch.float32
-        self.hf_token = os.getenv("HF_TOKEN") or None
+        # Check for Hugging Face token (try both environment variable names)
+        self.hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN") or None
 
         # Configure writable cache to avoid permission issues on Spaces
-        # Prefer user home cache: ~/.cache/huggingface
-        default_home_cache = os.path.join(os.path.expanduser("~"), ".cache", "huggingface")
-        hf_cache_dir = os.getenv("HF_HOME", default_home_cache)
-        os.environ.setdefault("HF_HOME", hf_cache_dir)
-        os.environ.setdefault("HUGGINGFACE_HUB_CACHE", hf_cache_dir)
-        os.environ.setdefault("TRANSFORMERS_CACHE", hf_cache_dir)
+        # Use /data directory which is writable in Hugging Face Spaces
+        data_dir = os.getenv("DATA_DIR", "/data")
+        hf_cache_dir = os.path.join(data_dir, "hf_cache")
+        
+        # Set cache environment variables
+        os.environ["HF_HOME"] = hf_cache_dir
+        os.environ["HUGGINGFACE_HUB_CACHE"] = hf_cache_dir
+        os.environ["TRANSFORMERS_CACHE"] = hf_cache_dir
+        
         try:
             os.makedirs(hf_cache_dir, exist_ok=True)
-        except Exception:
-            # Fallback to a local data dir if home is not writable
-            hf_cache_dir = os.path.abspath(os.path.join(".", "data", "hf_cache"))
+            logger.info("HF cache directory: %s", hf_cache_dir)
+        except Exception as e:
+            # Fallback to user home if /data is not available (local dev)
+            logger.warning("Failed to create cache in /data: %s, trying home directory", str(e))
+            default_home_cache = os.path.join(os.path.expanduser("~"), ".cache", "huggingface")
+            hf_cache_dir = os.getenv("HF_HOME", default_home_cache)
             os.environ["HF_HOME"] = hf_cache_dir
             os.environ["HUGGINGFACE_HUB_CACHE"] = hf_cache_dir
             os.environ["TRANSFORMERS_CACHE"] = hf_cache_dir
-            os.makedirs(hf_cache_dir, exist_ok=True)
-        logger.info("HF cache directory: %s", hf_cache_dir)
+            try:
+                os.makedirs(hf_cache_dir, exist_ok=True)
+                logger.info("HF cache directory (fallback): %s", hf_cache_dir)
+            except Exception as e2:
+                logger.error("Failed to create cache directory: %s", str(e2))
+                raise RuntimeError(f"Cannot create Hugging Face cache directory: {str(e2)}")
 
         # Avoid libgomp warning by setting a valid integer
         os.environ.setdefault("OMP_NUM_THREADS", "1")

app/main.py

@@ -75,6 +75,7 @@ app.mount("/uploads", StaticFiles(directory=str(UPLOAD_DIR)), name="uploads")
 
 # Initialize ColorizeNet model
 colorize_model = None
+model_load_error: Optional[str] = None
 
 @app.get("/")
 async def root():
@@ -89,13 +90,17 @@ async def root():
 @app.on_event("startup")
 async def startup_event():
     """Initialize the colorization model on startup"""
-    global colorize_model
+    global colorize_model, model_load_error
     try:
-        logger.info("Loading ColorizeNet model...")
+        logger.info("Loading ColorizeNet model with MODEL_ID: %s", settings.MODEL_ID)
+        logger.info("HF_TOKEN present: %s", "Yes" if os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN") else "No")
         colorize_model = ColorizeModel(settings.MODEL_ID)
         logger.info("ColorizeNet model loaded successfully")
+        model_load_error = None
     except Exception as e:
-        logger.error("Failed to load ColorizeNet model: %s", str(e))
+        error_msg = str(e)
+        logger.error("Failed to load ColorizeNet model: %s", error_msg)
+        model_load_error = error_msg
         # Don't raise - allow health check to work even if model fails
 
 @app.on_event("shutdown")
@@ -156,10 +161,14 @@ async def verify_request(request: Request):
 @app.get("/health")
 async def health_check():
     """Health check endpoint"""
-    return {
+    response = {
         "status": "healthy",
-        "model_loaded": colorize_model is not None
+        "model_loaded": colorize_model is not None,
+        "model_id": settings.MODEL_ID
     }
+    if model_load_error:
+        response["model_error"] = model_load_error
+    return response
 
 @app.post("/upload")
 async def upload_image(