Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

simply change __init__.py, add onnxruntime with dml support. #3045

Open
wants to merge 7 commits into
base: main
Choose a base branch
from
Open
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