cINAWGD commited on
Commit
e8a949b
·
verified ·
1 Parent(s): a67c7e5

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +14 -0
main.py CHANGED
@@ -8,6 +8,20 @@ from fastapi import FastAPI, Query, HTTPException
8
  from transformers import AutoModelForImageClassification, ViTImageProcessor
9
  from typing import Optional
10
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
11
  # Initialize the model and processor globally to avoid reloading for each request
12
  model = AutoModelForImageClassification.from_pretrained("Falconsai/nsfw_image_detection")
13
  processor = ViTImageProcessor.from_pretrained('Falconsai/nsfw_image_detection')
 
8
  from transformers import AutoModelForImageClassification, ViTImageProcessor
9
  from typing import Optional
10
 
11
+ # Determine a writable cache directory
12
+ default_cache_dir = os.path.join(os.path.expanduser('~'), '.cache', 'huggingface', 'hub')
13
+
14
+ # Ensure the directory exists
15
+ try:
16
+ os.makedirs(default_cache_dir, exist_ok=True)
17
+ except PermissionError:
18
+ # Fallback to a temporary directory if user's home directory is not writable
19
+ default_cache_dir = os.path.join('/tmp', 'huggingface_cache')
20
+ os.makedirs(default_cache_dir, exist_ok=True)
21
+
22
+ # Set the environment variable to the created directory
23
+ os.environ['TRANSFORMERS_CACHE'] = default_cache_dir
24
+ os.environ['HF_HOME'] = default_cache_dir
25
  # Initialize the model and processor globally to avoid reloading for each request
26
  model = AutoModelForImageClassification.from_pretrained("Falconsai/nsfw_image_detection")
27
  processor = ViTImageProcessor.from_pretrained('Falconsai/nsfw_image_detection')