Allow selection of the onnx runtime

README.md

@@ -4,14 +4,27 @@
 
 -
 
-
 Pangu-Weather: A 3D High-Resolution Model for Fast and Accurate Global Weather Forecast, arXiv preprint: 2211.02556, 2022.
-https://arxiv.org/abs/2211.02556
+<https://arxiv.org/abs/2211.02556>
 
 Pangu-Weather was created by Kaifeng Bi, Lingxi Xie, Hengheng Zhang, Xin Chen, Xiaotao Gu and Qi Tian. It is released by Huawei Cloud.
 
 The trained parameters of Pangu-Weather are made available under the terms of the BY-NC-SA 4.0 license.
 
 The commercial use of these models is forbidden.
 
-See https://github.com/198808xc/Pangu-Weather for further details.
+See <https://github.com/198808xc/Pangu-Weather> for further details.
+
+### Installation
+
+To install the package, run:
+
+```bash
+pip install ai-models-panguweather
+```
+
+This will install the package and its dependencies, in particular the ONNX runtime. The installation script will attempt to guess which runtime to install. You can force a given runtime by specifying the the `ONNXRUNTIME` variable, e.g.:
+
+```bash
+ONNXRUNTIME=onnxruntime-cpu pip install ai-models-panguweather
+```

ai_models_panguweather/__init__.py

@@ -5,4 +5,4 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-__version__ = "0.0.2"
+__version__ = "0.0.3"

setup.py

@@ -55,14 +55,16 @@ def has_gpu():
         return check_gpus() > 0
 
 
-onnxruntime = "onnxruntime"
-if sys.platform == "darwin":
-    if platform.machine() == "arm64":
-        onnxruntime = "onnxruntime-silicon"
-
-
-if has_gpu():
-    onnxruntime = "onnxruntime-gpu"
+if "ONNXRUNTIME" in os.environ:
+    onnxruntime = os.environ["ONNXRUNTIME"]
+else:
+    onnxruntime = "onnxruntime"
+    if sys.platform == "darwin":
+        if platform.machine() == "arm64":
+            onnxruntime = "onnxruntime-silicon"
+
+    if has_gpu():
+        onnxruntime = "onnxruntime-gpu"
 
 
 setuptools.setup(