Compile CUDA kernels (SAM2) on HuggingFace Spaces (ZeroGPU)

#1
by wondervictor - opened

Hi @hysts , I've changed the code of SAM2 and need to re-compile the CUDA kernels of SAM2. However, I can not find the right path for CUDA_HOME. I'm concerned whether it's feasible to compile CUDA kernels.

@wondervictor
Are you trying to build CUDA kernels on ZeroGPU? ZeroGPU Spaces don't have CUDA dev tools, like nvcc, so it's just not possible to build packages that require CUDA to build. Instead, you can build a wheel for the package in your local environment with CUDA and install it at startup using subprocess.run.

I haven't tested this myself for SAM2, so I'm not 100% sure if this works without modification, but I usually use Dockerfile like

FROM nvidia/cuda:12.3.2-cudnn9-devel-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update && \
    apt-get upgrade -y && \
    apt-get install -y --no-install-recommends \
    git \
    git-lfs \
    wget \
    curl \
    # python build dependencies \
    build-essential \
    libssl-dev \
    zlib1g-dev \
    libbz2-dev \
    libreadline-dev \
    libsqlite3-dev \
    libncursesw5-dev \
    xz-utils \
    tk-dev \
    libxml2-dev \
    libxmlsec1-dev \
    libffi-dev \
    liblzma-dev && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

RUN useradd -m -u 1000 user
USER user
ENV HOME=/home/user \
    PATH=/home/user/.local/bin:${PATH}

RUN curl https://pyenv.run | bash
ENV PATH=${HOME}/.pyenv/shims:${HOME}/.pyenv/bin:${PATH}
ARG PYTHON_VERSION=3.10.13
RUN pyenv install ${PYTHON_VERSION} && \
    pyenv global ${PYTHON_VERSION} && \
    pyenv rehash && \
    pip install --no-cache-dir -U pip setuptools wheel ninja

RUN pip install torch==2.2.0

ENV TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6"
ENV TCNN_CUDA_ARCHITECTURES=86;80;75;70;61;60
ENV FORCE_CUDA=1
ENV CUDA_HOME=/usr/local/cuda

WORKDIR /work

CMD ["python", "setup.py", "bdist_wheel"]

to create a docker image to build wheels.

docker build . -t wheel-builder

And you can build a wheel for a package like this:

cd /path/to/your/package
docker run --rm -v `pwd`:/work --gpus all wheel-builder 

Once you've built a wheel, you can add it to your Space repo and install it with a command like

subprocess.run(shlex.split("pip install your-pre-built-package.whl"))

You can check out this Space as an example. It contains a wheel here and installs it like this.

Thank you, I'll try it!

Sign up or log in to comment