Add cache_dir in from_pretrained.

app.py

@@ -19,10 +19,11 @@ import zipfile
 
 
 transformers.utils.move_cache()
-os.environ["HUGGINGFACE_HUB_CACHE"] = "/tmp/hf_cache"
-os.environ["HF_HOME"] = "/tmp/hf_cache"
-os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
-os.makedirs("/tmp/hf_cache", exist_ok=True)
+hf_cache_path = '/tmp/hf_cache'
+os.environ["HUGGINGFACE_HUB_CACHE"] = hf_cache_path
+os.environ["HF_HOME"] = hf_cache_path
+os.environ["TRANSFORMERS_CACHE"] = hf_cache_path
+os.makedirs(hf_cache_path, exist_ok=True)
 
 
 torch.set_float32_matmul_precision('high')
@@ -96,7 +97,7 @@ usage_to_weights_file = {
     'General-dynamic': 'BiRefNet_dynamic',
 }
 
-birefnet = transformers.AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', usage_to_weights_file['General'])), trust_remote_code=True)
+birefnet = transformers.AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', usage_to_weights_file['General'])), trust_remote_code=True, cache_dir=hf_cache_path)
 birefnet.to(device)
 birefnet.eval(); birefnet.half()
 
@@ -109,7 +110,7 @@ def predict(images, resolution, weights_file):
     # Load BiRefNet with chosen weights
     _weights_file = '/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else usage_to_weights_file['General']))
     print('Using weights: {}.'.format(_weights_file))
-    birefnet = transformers.AutoModelForImageSegmentation.from_pretrained(_weights_file, trust_remote_code=True)
+    birefnet = transformers.AutoModelForImageSegmentation.from_pretrained(_weights_file, trust_remote_code=True, cache_dir=hf_cache_path)
     birefnet.to(device)
     birefnet.eval(); birefnet.half()