Update src/streamlit_app.py

src/streamlit_app.py

@@ -5,8 +5,13 @@ from PIL import Image
 import numpy as np
 import io
 import requests
+import os
 from typing import List, Tuple
 
+# Set cache directories to writable locations
+os.environ['TORCH_HOME'] = '/tmp/torch_cache'
+os.environ['HF_HOME'] = '/tmp/hf_cache'
+
 # Configure page
 st.set_page_config(
     page_title="CLIP Classifier",
@@ -18,8 +23,12 @@ st.set_page_config(
 def load_clip_model():
     """Load CLIP model and preprocessing function"""
     try:
+        # Ensure cache directories exist
+        os.makedirs('/tmp/torch_cache', exist_ok=True)
+        os.makedirs('/tmp/clip_models', exist_ok=True)
+        
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        model, preprocess = clip.load("ViT-B/32", device=device)
+        model, preprocess = clip.load("ViT-B/32", device=device, download_root="/tmp/clip_models")
         return model, preprocess, device
     except Exception as e:
         st.error(f"Error loading CLIP model: {e}")