From 20b06db35985a77107a02ee76d30c5c4db226c9e Mon Sep 17 00:00:00 2001 From: cmdr2 Date: Fri, 4 Aug 2023 19:46:37 +0530 Subject: [PATCH] Include onnx and polygraphy for TRT, and allow skipping the wheels for TRT --- ui/easydiffusion/package_manager.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/ui/easydiffusion/package_manager.py b/ui/easydiffusion/package_manager.py index c246c54d..6a0521b2 100644 --- a/ui/easydiffusion/package_manager.py +++ b/ui/easydiffusion/package_manager.py @@ -25,7 +25,7 @@ installing = [] # remove this once TRT releases on pypi if platform.system() == "Windows": trt_dir = os.path.join(app.ROOT_DIR, "tensorrt") - if os.path.exists(trt_dir): + if os.path.exists(trt_dir) and os.path.isdir(trt_dir) and len(os.listdir(trt_dir)) > 0: files = os.listdir(trt_dir) packages = manifest["tensorrt"]["install"] @@ -61,6 +61,10 @@ def install(module_name): raise RuntimeError(f"Can't install unknown package: {module_name}!") commands = manifest[module_name]["install"] + if module_name == "tensorrt": + commands += [ + "protobuf==3.20.3 polygraphy==0.47.1 onnx==1.14.0 --extra-index-url=https://pypi.ngc.nvidia.com --trusted-host pypi.ngc.nvidia.com" + ] commands = [f"python -m pip install --upgrade {cmd}" for cmd in commands] installing.append(module_name)