add absolute path to cpp files in setup

bcmi · Jan 29, 2024 · 016776e · 016776e
1 parent b6bd5c7
commit 016776e
Showing 1 changed file with 8 additions and 7 deletions.
diff --git a/setup.py b/setup.py
@@ -5,6 +5,7 @@
 import torch
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension
 import os
+from os.path import join
 
 
 def readme():
@@ -91,17 +92,17 @@ def gen_packages_items():
     packages = list(gen_packages_items())
     return packages
 
-def get_ext_modules():
+def get_ext_modules(cur_dir):
     # if encounter compilation issues, please refer to https://github.com/HuiZeng/Image-Adaptive-3DLUT?tab=readme-ov-file#build.
     if torch.cuda.is_available():
         return CUDAExtension('trilinear', 
-                             ['libcom/image_harmonization/source/trilinear_cpp/src/trilinear_cuda.cpp', 
-                              'libcom/image_harmonization/source/trilinear_cpp/src/trilinear_kernel.cu'],
+                             [join(cur_dir, 'libcom/image_harmonization/source/trilinear_cpp/src/trilinear_cuda.cpp'), 
+                              join(cur_dir, 'libcom/image_harmonization/source/trilinear_cpp/src/trilinear_kernel.cu')],
                              )
     else:
         return CppExtension('trilinear',  
-                            ['libcom/image_harmonization/source/trilinear_cpp/src/trilinear.cpp'], 
-                            include_dirs=['libcom/image_harmonization/source/trilinear_cpp/src']
+                            [join(cur_dir, 'libcom/image_harmonization/source/trilinear_cpp/src/trilinear.cpp')], 
+                            include_dirs=[join(cur_dir, 'libcom/image_harmonization/source/trilinear_cpp/src')]
                             )
 
 if __name__ == '__main__':
@@ -128,11 +129,11 @@ def get_ext_modules():
         setup_requires=parse_requirements('requirements/build.txt'),
         tests_require=parse_requirements('requirements/tests.txt'),
         install_requires=parse_requirements('requirements/requirements.txt'),
-        ext_modules=[get_ext_modules()],
+        ext_modules=[get_ext_modules(cur_dir)],
         cmdclass={'build_ext': BuildExtension},
         extras_require={
             'all': parse_requirements('requirements.txt'),
             'build': parse_requirements('requirements/build.txt'),
             'test': parse_requirements('requirements.txt'),
         },
-    )
+    )