Update app.py

app.py

@@ -17,6 +17,17 @@ logger = logging.getLogger(__name__)
 os.environ['MPLCONFIGDIR'] = '/tmp/matplotlib'
 os.environ['HOME'] = '/tmp'
 
+# Set cache directories explicitly
+cache_dir = os.environ.get('POINT_E_CACHE_DIR', '/tmp/point_e_model_cache')
+clip_cache_dir = os.environ.get('CLIP_MODEL_DIR', '/tmp/clip_models')
+
+# Ensure cache directories exist and are writable
+for directory in [cache_dir, clip_cache_dir, '/tmp/point_e_models']:
+    if not os.path.exists(directory):
+        os.makedirs(directory, exist_ok=True)
+    # Make sure permissions are correct
+    os.chmod(directory, 0o777)
+
 app = Flask(__name__)
 CORS(app)
 
@@ -25,9 +36,6 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 logger.info(f"Using device: {device}")
 
 # Set up model directory for downloads
-os.makedirs("/tmp/point_e_models", exist_ok=True)
-
-# Download model weights or use cached version
 model_path = "/tmp/point_e_models/base40M-textvec.pt"
 if not os.path.exists(model_path):
     logger.info("Model weights not found. Downloading...")
@@ -45,7 +53,11 @@ if not os.path.exists(model_path):
 # Load the model
 logger.info("Loading base model...")
 try:
-    base_model = model_from_config(MODEL_CONFIGS["base40M-textvec"], device=device)
+    base_model = model_from_config(
+        MODEL_CONFIGS["base40M-textvec"], 
+        device=device,
+        cache_dir=cache_dir  # Pass cache directory explicitly
+    )
     base_model.load_state_dict(torch.load(model_path, map_location=device))
     base_model.eval()
     logger.info("Base model loaded successfully")