Browse Source

infer download_root from XDG_CACHE_HOME if avail (#257)

David Marx 2 years ago
parent
commit
82725cea9c
1 changed files with 4 additions and 1 deletions
  1. 4 1
      whisper/__init__.py

+ 4 - 1
whisper/__init__.py

@@ -90,7 +90,10 @@ def load_model(name: str, device: Optional[Union[str, torch.device]] = None, dow
     if device is None:
     if device is None:
         device = "cuda" if torch.cuda.is_available() else "cpu"
         device = "cuda" if torch.cuda.is_available() else "cpu"
     if download_root is None:
     if download_root is None:
-        download_root = os.path.join(os.path.expanduser("~"), ".cache", "whisper")
+        download_root = os.getenv(
+            "XDG_CACHE_HOME", 
+            os.path.join(os.path.expanduser("~"), ".cache", "whisper")
+        )
 
 
     if name in _MODELS:
     if name in _MODELS:
         checkpoint_file = _download(_MODELS[name], download_root, in_memory)
         checkpoint_file = _download(_MODELS[name], download_root, in_memory)