Improve HF cache directory fallback to use /tmp when /data not writable

app/colorize_model.py

@@ -34,8 +34,10 @@ class ColorizeModel:
         self.hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN") or None
 
         # Configure writable cache to avoid permission issues on Spaces
-        # Use /data directory which is writable in Hugging Face Spaces
-        data_dir = os.getenv("DATA_DIR", "/data")
+        # Prefer DATA_DIR if available, otherwise fallback to /tmp
+        data_dir = os.getenv("DATA_DIR")
+        if not data_dir:
+            data_dir = "/tmp"
         hf_cache_dir = os.path.join(data_dir, "hf_cache")
         
         # Set cache environment variables
@@ -47,20 +49,36 @@ class ColorizeModel:
             os.makedirs(hf_cache_dir, exist_ok=True)
             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)
+            # Fallback to /tmp/hf_cache if DATA_DIR was set but not writable
+            tmp_cache_dir = os.path.join("/tmp", "hf_cache")
+            logger.warning("Failed to create cache in %s: %s, trying %s", data_dir, str(e), tmp_cache_dir)
+            hf_cache_dir = tmp_cache_dir
             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)
-                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)}")
+                logger.info("HF cache directory (tmp): %s", hf_cache_dir)
+            except Exception as e_tmp:
+                # Final fallback to user home (local dev)
+                logger.warning("Failed to create cache in /tmp: %s, trying user home", str(e_tmp))
+                default_home_cache = os.path.join(os.path.expanduser("~"), ".cache", "huggingface")
+                hf_cache_dir = 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
+                try:
+                    os.makedirs(hf_cache_dir, exist_ok=True)
+                    logger.info("HF cache directory (home): %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)}")
 
+        else:
+            # Ensure environment variables reflect the final cache dir
+            os.environ["HF_HOME"] = hf_cache_dir
+            os.environ["HUGGINGFACE_HUB_CACHE"] = hf_cache_dir
+            os.environ["TRANSFORMERS_CACHE"] = hf_cache_dir
         # Avoid libgomp warning by setting a valid integer
         os.environ.setdefault("OMP_NUM_THREADS", "1")