Update timesfm_backend.py

timesfm_backend.py

@@ -1,6 +1,7 @@
 import time
 import json
 import logging
+import os
 from typing import Any, Dict, List, Optional, Tuple, Sequence
 
 import numpy as np
@@ -108,11 +109,26 @@ class TimesFMBackend(ChatBackend):
         if self._model is not None:
             return
         try:
+            import os
             import timesfm  # 2.5 API
-            # instantiate 2.5 torch model
-            model = timesfm.TimesFM_2p5_200M_torch()
-            model.load_checkpoint()  # uses built-in 2.5 checkpoint loader
-            # compile with sane defaults; caller can override via payload
+
+            hf_token = getattr(settings, "HF_TOKEN", None) or os.environ.get("HF_TOKEN")
+            cache_dir = getattr(settings, "TIMESFM_CACHE_DIR", None)
+
+            model = timesfm.TimesFM_2p5_200M_torch.from_pretrained(
+                self.model_id,
+                token=hf_token,
+                cache_dir=cache_dir,
+                local_files_only=False,
+            )
+
+            try:
+                # .model holds the underlying nn.Module; fall back to instance if absent.
+                target = getattr(model, "model", model)
+                target.to(self.device)  # type: ignore[arg-type]
+            except Exception:
+                pass
+
             cfg = timesfm.ForecastConfig(
                 max_context=1024,
                 max_horizon=256,
@@ -123,11 +139,6 @@ class TimesFMBackend(ChatBackend):
                 fix_quantile_crossing=True,
             )
             model.compile(cfg)
-            try:
-                # .to(device) may be no-op for this wrapper; safe to try
-                model.to(self.device)  # type: ignore[attr-defined]
-            except Exception:
-                pass
             self._model = model
             logger.info("TimesFM 2.5 model loaded on %s", self.device)
         except Exception as e: