Spaces:
Running
on
Zero
Compile CUDA kernels (SAM2) on HuggingFace Spaces (ZeroGPU)
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!