Enable model snapshot downloads to be saved in a specified local dire…

…ctory.
google-research · Oct 30, 2024 · c1e2094 · c1e2094
1 parent e3d45d6
commit c1e2094
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/src/timesfm/timesfm_base.py b/src/timesfm/timesfm_base.py
@@ -130,6 +130,7 @@ class TimesFmCheckpoint:
   huggingface_repo_id: str | None = None
   type: Any = None
   step: int | None = None
+  local_dir: str | None = None
 
 
 class TimesFmBase:

diff --git a/src/timesfm/timesfm_torch.py b/src/timesfm/timesfm_torch.py
@@ -55,8 +55,9 @@ def load_from_checkpoint(
     checkpoint_path = checkpoint.path
     repo_id = checkpoint.huggingface_repo_id
     if checkpoint_path is None:
-      checkpoint_path = path.join(snapshot_download(repo_id),
-                                  "torch_model.ckpt")
+      checkpoint_path = path.join(
+                snapshot_download(repo_id, local_dir=checkpoint.local_dir),
+                "torch_model.ckpt")
     self._model = ppd.PatchedTimeSeriesDecoder(self._model_config)
     loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)
     logging.info("Loading checkpoint from %s", checkpoint_path)