Skip to content

Commit

Permalink
fix #5857
Browse files Browse the repository at this point in the history
  • Loading branch information
Baiyuetribe committed Dec 31, 2024
1 parent c8630f4 commit 2b83fdc
Showing 1 changed file with 11 additions and 3 deletions.
14 changes: 11 additions & 3 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ def find_version():
return version_major[0] + "." + version_minor[0] + "." + ncnn_version
raise RuntimeError("Unable to find version string.")


# Parse environment variables
Vulkan_LIBRARY = os.environ.get("Vulkan_LIBRARY", "")
CMAKE_TOOLCHAIN_FILE = os.environ.get("CMAKE_TOOLCHAIN_FILE", "")
Expand All @@ -38,11 +39,13 @@ def find_version():
ENABLE_ARC = os.environ.get("ENABLE_ARC", "")
ENABLE_VISIBILITY = os.environ.get("ENABLE_VISIBILITY", "")


# Parse variables from command line with setup.py install
class InstallCommand(install):
user_options = install.user_options + [
('vulkan=', None, 'Enable the usage of Vulkan.'),
("vulkan=", None, "Enable the usage of Vulkan."),
]

def initialize_options(self):
install.initialize_options(self)
self.vulkan = None
Expand All @@ -53,6 +56,7 @@ def finalize_options(self):
def run(self):
install.run(self)


# Convert distutils Windows platform specifiers to CMake -A arguments
PLAT_TO_CMAKE = {
"win32": "Win32",
Expand All @@ -61,6 +65,7 @@ def run(self):
"win-arm64": "ARM64",
}


# A CMakeExtension needs a sourcedir instead of a file list.
# The name must be the _single_ output extension from the CMake build.
# If you need multiple extensions, see scikit-build.
Expand Down Expand Up @@ -88,13 +93,16 @@ def build_extension(self, ext):
# Set Python_EXECUTABLE instead if you use PYBIND11_FINDPYTHON
# EXAMPLE_VERSION_INFO shows you how to pass a value into the C++ code
# from Python.
enable_vulkan = "ON"
if Vulkan_LIBRARY == "":
enable_vulkan = "OFF"
cmake_args = [
"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={}".format(extdir),
"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE={}".format(extdir),
"-DPYTHON_EXECUTABLE={}".format(sys.executable),
"-DCMAKE_BUILD_TYPE={}".format(cfg), # not used on MSVC, but no harm
"-DNCNN_PYTHON=ON",
"-DNCNN_VULKAN=ON",
"-DNCNN_VULKAN={}".format(enable_vulkan),
"-DNCNN_DISABLE_RTTI=OFF",
"-DNCNN_DISABLE_EXCEPTION=OFF",
"-DNCNN_BUILD_BENCHMARK=OFF",
Expand Down Expand Up @@ -208,5 +216,5 @@ def build_extension(self, ext):
package_dir={"": "python"},
install_requires=requirements,
ext_modules=[CMakeExtension("ncnn")],
cmdclass={'install': InstallCommand, "build_ext": CMakeBuild},
cmdclass={"install": InstallCommand, "build_ext": CMakeBuild},
)

0 comments on commit 2b83fdc

Please sign in to comment.