Update setup.py

setup.py

@@ -1,73 +1,30 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.
-
-from setuptools import find_packages, setup
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
-
-# Package metadata
-NAME = "SAM 2"
-VERSION = "1.0"
-DESCRIPTION = "SAM 2: Segment Anything in Images and Videos"
-URL = "https://github.com/facebookresearch/segment-anything-2"
-AUTHOR = "Meta AI"
-AUTHOR_EMAIL = "[email protected]"
-LICENSE = "Apache 2.0"
-
-# Read the contents of README file
-with open("README.md", "r") as f:
-    LONG_DESCRIPTION = f.read()
-
-# Required dependencies
-REQUIRED_PACKAGES = [
-    "torch>=2.3.1",
-    "torchvision>=0.18.1",
-    "numpy>=1.24.4",
-    "tqdm>=4.66.1",
-    "hydra-core>=1.3.2",
-    "iopath>=0.1.10",
-    "pillow>=9.4.0",
-]
-
-EXTRA_PACKAGES = {
-    "demo": ["matplotlib>=3.9.1", "jupyter>=1.0.0", "opencv-python>=4.7.0"],
-    "dev": ["black==24.2.0", "usort==1.0.2", "ufmt==2.0.0b2"],
-}
-
-
-def get_extensions():
-    srcs = ["sam2/csrc/connected_components.cu"]
-    compile_args = {
-        "cxx": [],
-        "nvcc": [
-            "-DCUDA_HAS_FP16=1",
-            "-D__CUDA_NO_HALF_OPERATORS__",
-            "-D__CUDA_NO_HALF_CONVERSIONS__",
-            "-D__CUDA_NO_HALF2_OPERATORS__",
-            "-allow-unsupported-compiler",
-        ],
-    }
-    ext_modules = [CUDAExtension("sam2._C", srcs, extra_compile_args=compile_args)]
-    return ext_modules
-
-
-# Setup configuration
-setup(
-    name=NAME,
-    version=VERSION,
-    description=DESCRIPTION,
-    long_description=LONG_DESCRIPTION,
-    long_description_content_type="text/markdown",
-    url=URL,
-    author=AUTHOR,
-    author_email=AUTHOR_EMAIL,
-    license=LICENSE,
-    packages=find_packages(exclude="notebooks"),
-    install_requires=REQUIRED_PACKAGES,
-    extras_require=EXTRA_PACKAGES,
-    python_requires=">=3.10.0",
-    ext_modules=get_extensions(),
-    cmdclass={"build_ext": BuildExtension.with_options(no_python_abi_suffix=True)},
-)
+from setuptools import find_packages, setup
+from torch.utils.cpp_extension import BuildExtension, CppExtension
+
+def get_extensions():
+    # 使用 CPU 的 C++ 扩展
+    srcs = ["sam2/csrc/connected_components.cpp"]  # 使用 `.cpp` 文件而不是 `.cu` 文件
+    compile_args = {
+        "cxx": [],
+        # nvcc 编译器标志可以删除或注释掉
+        # "nvcc": [
+        #     "-DCUDA_HAS_FP16=1",
+        #     "-D__CUDA_NO_HALF_OPERATORS__",
+        #     "-D__CUDA_NO_HALF_CONVERSIONS__",
+        #     "-D__CUDA_NO_HALF2_OPERATORS__",
+        #     "-allow-unsupported-compiler",
+        # ],
+    }
+    # 使用 CppExtension 而不是 CUDAExtension
+    ext_modules = [CppExtension("sam2._C", srcs, extra_compile_args=compile_args["cxx"])]
+    return ext_modules
+
+# Setup configuration
+setup(
+    name="SAM2",
+    version="1.0",
+    description="SAM 2: Segment Anything in Images and Videos",
+    packages=find_packages(),
+    ext_modules=get_extensions(),
+    cmdclass={"build_ext": BuildExtension.with_options(no_python_abi_suffix=True)},
+)