Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions backend/src/packages/chaiNNer_onnx/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

from api import KB, MB, Dependency, add_package
from gpu import nvidia
from system import is_arm_mac
from system import is_arm_mac, is_windows

general = "ONNX uses .onnx models to upscale images."
conversion = "It also helps to convert between PyTorch and NCNN."
Expand All @@ -15,7 +15,7 @@
f"{general} {conversion} It is fastest when CUDA is supported. If TensorRT is"
" installed on the system, it can also be configured to use that."
)
inst_hint = f"{general} It does not support AMD GPUs."
inst_hint = f"{general} It does not support AMD GPUs, in linux."


def get_onnx_runtime():
Expand All @@ -28,6 +28,13 @@ def get_onnx_runtime():
import_name="onnxruntime",
extra_index_url="https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/",
)
elif is_windows:
return Dependency(
display_name="ONNX Runtime (DirectMl)",
pypi_name="onnxruntime-directml",
version="1.17.1",
size_estimate=15 * MB,
)
else:
return Dependency(
display_name="ONNX Runtime",
Expand Down