From c465b61fb5af0e7c1d710fd37a3946f0abdb9b52 Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Sat, 18 Mar 2023 13:21:17 -0500
Subject: [PATCH] fix(scripts): check for ORT modules

---
 api/scripts/check-env.py | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/api/scripts/check-env.py b/api/scripts/check-env.py
index 774ffd6d7..cf7b8b164 100644
--- a/api/scripts/check-env.py
+++ b/api/scripts/check-env.py
@@ -3,7 +3,7 @@
 
 REQUIRED_MODULES = ["onnx", "diffusers", "safetensors", "torch"]
 
-REQUIRED_RUNTIME = [
+RUNTIME_MODULES = [
     "onnxruntime",
     "onnxruntime_gpu",
     "onnxruntime_rocm",
@@ -26,6 +26,21 @@ def check_modules() -> List[str]:
     return results
 
 
+def check_runtimes() -> List[str]:
+    results = []
+    for name in RUNTIME_MODULES:
+        try:
+            __import__(name)
+            module_version = version(name)
+            results.append(
+                f"runtime module {name} is present at version {module_version}"
+            )
+        except ImportError as e:
+            results.append(f"unable to import runtime module {name}: {e}")
+
+    return results
+
+
 def check_providers() -> List[str]:
     results = []
     try:
@@ -46,6 +61,7 @@ def check_providers() -> List[str]:
 
 ALL_CHECKS = [
     check_modules,
+    check_runtimes,
     check_providers,
 ]