Skip to content

Commit eb9dfbe

Browse files
committed
feat(depth-estimation): enhance ONNX inference and CUDA requirements
- transform.py: improve ONNX runtime execution provider handling - requirements_cuda.txt: add onnxruntime-gpu dependency
1 parent 5e4e528 commit eb9dfbe

2 files changed

Lines changed: 27 additions & 3 deletions

File tree

skills/transformation/depth-estimation/requirements_cuda.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
# onnxruntime-gpu includes both CUDA and TensorRT execution providers.
55

66
onnxruntime-gpu>=1.17.0
7+
nvidia-cudnn-cu12>=9.0
78

89
# ── Common dependencies ─────────────────────────────────────────────
910
huggingface_hub>=0.20.0

skills/transformation/depth-estimation/scripts/transform.py

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -236,8 +236,33 @@ def _download_coreml_model(self, variant_id: str):
236236

237237
# ── ONNX Runtime backend (Windows/Linux — all GPUs) ────────────────
238238

239+
@staticmethod
240+
def _add_nvidia_dll_paths():
241+
"""Add pip-installed NVIDIA DLL directories to PATH so ORT finds cudnn, cublas, etc."""
242+
import site
243+
import glob
244+
245+
for sp in site.getsitepackages():
246+
nvidia_dir = os.path.join(sp, "nvidia")
247+
if not os.path.isdir(nvidia_dir):
248+
continue
249+
for bin_dir in glob.glob(os.path.join(nvidia_dir, "*", "bin")):
250+
if bin_dir not in os.environ.get("PATH", ""):
251+
os.environ["PATH"] = bin_dir + os.pathsep + os.environ.get("PATH", "")
252+
# Python 3.8+ on Windows: also register via os.add_dll_directory
253+
if hasattr(os, "add_dll_directory"):
254+
try:
255+
os.add_dll_directory(bin_dir)
256+
except OSError:
257+
pass
258+
_log(f"Added NVIDIA DLL path: {bin_dir}", "DepthEstimation")
259+
260+
239261
def _load_onnx(self, model_name: str, config: dict) -> dict:
240262
"""Load ONNX model with best available EP: CUDA → TRT → DirectML → CPU."""
263+
# Add pip-installed NVIDIA DLL dirs to PATH (cudnn, cublas, etc.)
264+
self._add_nvidia_dll_paths()
265+
241266
import onnxruntime as ort
242267
from huggingface_hub import hf_hub_download
243268

@@ -625,9 +650,7 @@ def on_config_update(self, config: dict):
625650
self.blend_mode = config["blend_mode"]
626651
_log(f"Blend mode updated: {self.blend_mode}", self._tag)
627652

628-
def get_output_mode(self) -> str:
629-
"""Use base64 for privacy transforms — avoids temp file cleanup issues."""
630-
return "base64"
653+
631654

632655

633656
if __name__ == "__main__":

0 commit comments

Comments
 (0)