holm/easydiffusion
1
0
Fork 0
forked from extern/easydiffusion
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Include onnx and polygraphy for TRT, and allow skipping the wheels for TRT

Browse Source
This commit is contained in:
cmdr2 2023-08-04 19:46:37 +05:30
parent b6e512e65f
commit 20b06db359
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
ui/easydiffusion/package_manager.py
View File

@ -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)