Update main.py

main.py

@@ -8,6 +8,20 @@ from fastapi import FastAPI, Query, HTTPException
 from transformers import AutoModelForImageClassification, ViTImageProcessor
 from typing import Optional
 
+# Determine a writable cache directory
+default_cache_dir = os.path.join(os.path.expanduser('~'), '.cache', 'huggingface', 'hub')
+
+# Ensure the directory exists
+try:
+    os.makedirs(default_cache_dir, exist_ok=True)
+except PermissionError:
+    # Fallback to a temporary directory if user's home directory is not writable
+    default_cache_dir = os.path.join('/tmp', 'huggingface_cache')
+    os.makedirs(default_cache_dir, exist_ok=True)
+
+# Set the environment variable to the created directory
+os.environ['TRANSFORMERS_CACHE'] = default_cache_dir
+os.environ['HF_HOME'] = default_cache_dir
 # Initialize the model and processor globally to avoid reloading for each request
 model = AutoModelForImageClassification.from_pretrained("Falconsai/nsfw_image_detection")
 processor = ViTImageProcessor.from_pretrained('Falconsai/nsfw_image_detection')