Stanislav commited on
Commit
3d9f0a5
·
1 Parent(s): 3164b96

feat: HF cache in dino class

Browse files
Files changed (2) hide show
  1. models/dino.py +7 -0
  2. run_fastapi.py +0 -4
models/dino.py CHANGED
@@ -3,6 +3,13 @@ from transformers import GroundingDinoProcessor, GroundingDinoForObjectDetection
3
 
4
  import cv2
5
 
 
 
 
 
 
 
 
6
 
7
  class DinoWrapper:
8
  """
 
3
 
4
  import cv2
5
 
6
+ import os
7
+
8
+ HF_CACHE = "/tmp/hf_cache"
9
+ os.makedirs(HF_CACHE, exist_ok=True)
10
+ os.environ["TRANSFORMERS_CACHE"] = HF_CACHE
11
+
12
+
13
 
14
  class DinoWrapper:
15
  """
run_fastapi.py CHANGED
@@ -15,10 +15,6 @@ from models.dino import DinoWrapper
15
 
16
  from huggingface_hub import hf_hub_download
17
 
18
- # HF Cache
19
- HF_CACHE_DIR = os.path.join("/tmp", "hf_cache")
20
- os.environ["TRANSFORMERS_CACHE"] = HF_CACHE_DIR
21
- os.makedirs(HF_CACHE_DIR, exist_ok=True)
22
 
23
  # Check write permissions
24
  print("WRITE to ./weights:", os.access("weights", os.W_OK))
 
15
 
16
  from huggingface_hub import hf_hub_download
17
 
 
 
 
 
18
 
19
  # Check write permissions
20
  print("WRITE to ./weights:", os.access("weights", os.W_OK))