add ARCH_LIST and install package for cuda extensions

This commit is contained in:
Peter Schroedl
2024-12-05 23:37:43 +01:00
parent fa428ca4d9
commit 38723ebb07
2 changed files with 6 additions and 2 deletions
+1 -1
View File
@@ -4,4 +4,4 @@ tqdm>=4.66.1
hydra-core>=1.3.2
iopath>=0.1.10
pillow>=9.4.0
-e .
git+https://github.com/pschroedl/ComfyUI-SAM2-Realtime.git@main#egg=sam2_realtime
+5 -1
View File
@@ -4,6 +4,10 @@
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
import os
# Set the CUDA architecture list
os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0 8.6+PTX 8.7 9.0 9.0a"
from setuptools import find_packages, setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
@@ -53,7 +57,7 @@ setup(
license=LICENSE,
packages=find_packages(),
install_requires=REQUIRED_PACKAGES,
python_requires=">=3.10",
python_requires=">=3.10.15",
ext_modules=get_extensions(),
cmdclass={"build_ext": BuildExtension.with_options(no_python_abi_suffix=True)},
)