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)},
)