tirex-mirror 2025.11.14__py3-none-any.whl → 2025.11.15__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
tirex/base.py CHANGED
@@ -3,6 +3,7 @@
3
3
 
4
4
  import logging
5
5
  import os
6
+ import warnings
6
7
  from abc import ABC, abstractmethod
7
8
  from typing import Literal, TypeVar
8
9
 
@@ -60,7 +61,7 @@ class PretrainedModel(ABC):
60
61
  if ckp_kwargs is None:
61
62
  ckp_kwargs = {}
62
63
  if device is None:
63
- device = "cuda:0" if backend == "cuda" else "cpu"
64
+ device = "cuda:0" if torch.cuda.is_available() else "cpu"
64
65
  if os.path.exists(path):
65
66
  print("Loading weights from local directory")
66
67
  checkpoint_path = path
@@ -93,7 +94,7 @@ class PretrainedModel(ABC):
93
94
  def load_model(
94
95
  path: str,
95
96
  device: str | None = None,
96
- backend: Literal["torch", "cuda"] | None = None,
97
+ backend: Literal["torch", "cuda"] = "torch",
97
98
  compile: bool = False,
98
99
  hf_kwargs=None,
99
100
  ckp_kwargs=None,
@@ -115,8 +116,22 @@ def load_model(
115
116
  model: ForecastModel = load_model("NX-AI/TiRex")
116
117
  """
117
118
 
118
- if backend is None:
119
- backend = "torch" if skip_cuda() or not xlstm_available() else "cuda"
119
+ if backend == "cuda" and not xlstm_available() and torch.cuda.is_available():
120
+ backend = "torch"
121
+ warnings.warn(
122
+ "Switching to 'torch' backend with device='gpu'. In oder to use a CUDA backend, please make sure that xlstm package is installed.",
123
+ UserWarning,
124
+ stacklevel=2,
125
+ )
126
+
127
+ if device is not None and device.startswith("cuda") and not torch.cuda.is_available():
128
+ raise ValueError(
129
+ "CUDA is not available! This could be because:\n"
130
+ " - No GPU is present on this machine\n"
131
+ " - GPU drivers are not installed or not functioning\n"
132
+ " - PyTorch is installed without CUDA support (CPU-only version)\n"
133
+ "To resolve: use device='cpu' for CPU inference, or install CUDA-enabled PyTorch if a GPU is available."
134
+ )
120
135
 
121
136
  try:
122
137
  _, model_string = parse_hf_repo_id(path).split("/")
@@ -70,8 +70,8 @@ class sLSTMCell(nn.Module):
70
70
  if input.device.type != "cuda":
71
71
  warnings.warn(
72
72
  f"You use TiRex with sLSTM CUDA kernels BUT DO NOT LOAD THE DEVICE ON A CUDA DEVICE (device type is {input.device.type})!"
73
- "This is not supported and calls to the model will likely lead to an error if you dont move your model to a CUDA device!"
74
- "If you want to run TiRex on CPU you need to disable sLSTM CUDA kernels but be aware of the downsides (see FAQ)"
73
+ "This is not supported and calls to the model will likely lead to an error if you don't move your model to a CUDA device!"
74
+ "If you want to run TiRex on CPU, please select backend='torch' and device='cpu'"
75
75
  )
76
76
 
77
77
  if not hasattr(self, "func"):
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: tirex-mirror
3
- Version: 2025.11.14
3
+ Version: 2025.11.15
4
4
  Summary: Unofficial mirror of NX-AI/tirex for packaging
5
5
  Author-email: Arpad Rozsas <rozsasarpi@gmail.com>
6
6
  License: NXAI COMMUNITY LICENSE AGREEMENT
@@ -1,5 +1,5 @@
1
1
  tirex/__init__.py,sha256=rfsOeCJ7eRqU3K3TOhfN5-4XUuZFqt11wBRxk5SoAWA,292
2
- tirex/base.py,sha256=cwmvqgvqHDFXkVZQ4qczPa6FPsp1c2SWvuByDR2Ecb8,4331
2
+ tirex/base.py,sha256=pGShyI3LwPAwkl0rIYqik5qb8uBAWqbQrVoVjq7UH_8,5096
3
3
  tirex/util.py,sha256=ggfETirJ589Dr5o3QThxTnjhgNnMCk10bNoJghnoeoA,31672
4
4
  tirex/api_adapter/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
5
5
  tirex/api_adapter/forecast.py,sha256=FnRgdI_4vJ7iFTqyWwxPE_C7MBB3nhLTnVMzTr5pihc,17581
@@ -10,12 +10,12 @@ tirex/models/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
10
10
  tirex/models/patcher.py,sha256=8T4c3PZnOAsEpahhrjtt7S7405WUjN6g3cV33E55PD4,1911
11
11
  tirex/models/tirex.py,sha256=URt-MClXu0zdUHACQ96Zu3Ytdb52vbeG_SXnj5C4tI8,7522
12
12
  tirex/models/slstm/block.py,sha256=V91Amgz8WAOOHo4fK1UZxd4Dgbx4-X6kUBS6X4m0tKQ,2006
13
- tirex/models/slstm/cell.py,sha256=Q0eRInybFx7Mj_pmlCl3CW9SShT2Qxegrr9KZzy7nrU,7320
13
+ tirex/models/slstm/cell.py,sha256=Otyil_AjpJbUckkINWGHxlqP14J5epm_J_zdWPzvD2g,7290
14
14
  tirex/models/slstm/layer.py,sha256=hrDydQJIAHf5W0A0Rt0hXG4yKXrOSY-HPL0UbigR6Q8,2867
15
- tirex_mirror-2025.11.14.dist-info/licenses/LICENSE,sha256=HlwHKnGTlE2oNm6734V-Vy62zlkWohnuZpYXSdkqDk4,7362
16
- tirex_mirror-2025.11.14.dist-info/licenses/LICENSE_MIRROR.txt,sha256=ulPZMcOZdN7JvISjiID3KUwovTjrPwiMv5ku9dM7nls,496
17
- tirex_mirror-2025.11.14.dist-info/licenses/NOTICE.txt,sha256=rcgDscFHb-uuZO3L0_vIxYhTYl-a2Rm0lBpp3_kKdFQ,147
18
- tirex_mirror-2025.11.14.dist-info/METADATA,sha256=aUeYkFk0Sg6aeXGAfodSaeMHfWuvTx5LuZqRmv_k5s8,11516
19
- tirex_mirror-2025.11.14.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
20
- tirex_mirror-2025.11.14.dist-info/top_level.txt,sha256=AOLDhfv0F_7nn3pFq0Kapg6Ky_28I_cGDXzQX3w9eO4,6
21
- tirex_mirror-2025.11.14.dist-info/RECORD,,
15
+ tirex_mirror-2025.11.15.dist-info/licenses/LICENSE,sha256=HlwHKnGTlE2oNm6734V-Vy62zlkWohnuZpYXSdkqDk4,7362
16
+ tirex_mirror-2025.11.15.dist-info/licenses/LICENSE_MIRROR.txt,sha256=ulPZMcOZdN7JvISjiID3KUwovTjrPwiMv5ku9dM7nls,496
17
+ tirex_mirror-2025.11.15.dist-info/licenses/NOTICE.txt,sha256=rcgDscFHb-uuZO3L0_vIxYhTYl-a2Rm0lBpp3_kKdFQ,147
18
+ tirex_mirror-2025.11.15.dist-info/METADATA,sha256=eI99-lkrGyhCPfQwFXik7QwZNgsX-qoRR0URpNCdmuM,11516
19
+ tirex_mirror-2025.11.15.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
20
+ tirex_mirror-2025.11.15.dist-info/top_level.txt,sha256=AOLDhfv0F_7nn3pFq0Kapg6Ky_28I_cGDXzQX3w9eO4,6
21
+ tirex_mirror-2025.11.15.dist-info/RECORD,,