CUDA / NVIDIA Setup

Getting Docker to talk to your NVIDIA GPU on Linux.

A note before we start: CUDA and GPU driver setup is a bit outside the scope of picoDiffusion itself. Every system is slightly different and things can vary between kernel versions, driver versions, and distributions. This guide covers the general steps for recent versions of Ubuntu, Linux Mint, and Debian. If you run into trouble, your distribution's documentation and the NVIDIA forums are the best places to look.

What about Mac or CPU?

picoDiffusion was built to run Stable Diffusion on Linux without the usual pain. That was the goal from the start. But we want to be upfront about what happens on other setups.

Apple Silicon (M1, M2, M3, M4)

The good news: PyTorch works really well on Apple Silicon. It supports MPS (Metal Performance Shaders), which is Apple's GPU acceleration, and the picoDiffusion code already detects and uses it.

The bad news: Docker on Mac runs a Linux virtual machine under the hood, and that VM cannot pass through the Metal GPU the way Docker on Linux passes through CUDA. So if you run picoDiffusion in Docker on a Mac, it falls back to CPU mode — which is very slow.

Running picoDiffusion natively on Mac (without Docker, just Python and pip) would give you MPS acceleration and reasonable performance. There is an at-your-own-risk native install guide, but we have not tested it on a Mac. We do not have one. We do not know anyone who will let us borrow one. Really, do not do any of this unless you enjoy suffering.

CPU only (any platform)

It works, but a single 512x512 image can take 5 to 15 minutes depending on your processor. During that time, the inference will use all available CPU cores at 100%. Your system may feel like it is locking up — fans will spin, other applications will be sluggish, and it will generally not be a fun time. We do not recommend it for regular use.

Our recommendation

A used NVIDIA GPU with 4GB or more VRAM can be found for very little money and will transform the experience. A Quadro P2000, GTX 1050 Ti, or anything in that range will generate 512x512 images in under a minute on Linux. It is worth it.

Step 1: Check if you already have NVIDIA drivers

Open a terminal and run:

nvidia-smi

If you see a table showing your GPU name, driver version, and CUDA version, your drivers are already installed. Skip to Step 3.

If you get "command not found" or an error, you need to install the drivers.

Step 2: Install NVIDIA drivers

Select your distribution:

You need the proprietary NVIDIA drivers. The open-source nouveau drivers will not work with CUDA.

Ubuntu and Mint include NVIDIA drivers in their package repositories. At the time of writing, the OS recommended nvidia-driver-535 by default. That one did not work for us. We installed nvidia-driver-580 instead and it worked. We did not test every version in between. We just picked the newest available and it was fine. We have not tested drivers beyond 580 yet.

sudo apt update
sudo apt install nvidia-driver-580

If 580 is not available in your repositories, see what is available and pick the newest one:

apt-cache search nvidia-driver-

Pick the plain one like nvidia-driver-580. Not the ones that say -open, -server, or -server-open. Just the bare package. Note: at the time of writing, nvidia-driver-590 was available but did not support our old P2000 card. Newer is not always better with NVIDIA drivers.

Avoid sudo ubuntu-drivers install as it will likely pick 535 again.

After installation, reboot:

sudo reboot

After rebooting, run nvidia-smi to confirm the driver is working.

Linux Mint note: Mint also has a Driver Manager in the system settings. You can use that to select driver version 580 instead of the command line if you prefer.

You need the proprietary NVIDIA drivers. The open-source nouveau drivers will not work with CUDA.

Debian requires enabling the non-free repository first:

sudo apt install software-properties-common
sudo add-apt-repository non-free
sudo apt update

Then install a driver. Same advice as Ubuntu/Mint: the default nvidia-driver package may pull an older version. Try to get something newer than 535. Check what is available:

apt-cache search nvidia-driver-

Pick the plain package, not the -open or -server variants. If a specific version like 580 is available:

sudo apt install nvidia-driver-580

If not, the generic package will get you the newest Debian has:

sudo apt install nvidia-driver

After installation, reboot:

sudo reboot

After rebooting, run nvidia-smi to confirm the driver is working.

Step 3: Install Docker

If you do not have Docker installed yet, the official install script is the quickest way:

curl -fsSL https://get.docker.com | sh
sudo usermod -aG docker $USER

Log out and back in (or reboot) for the group change to take effect. Then verify Docker works:

docker run hello-world

Step 4: Install the NVIDIA Container Toolkit

This is the piece that lets Docker containers access your GPU. Without it, Docker cannot see the GPU at all. These commands are the same for Ubuntu, Linux Mint, and Debian.

# Add the NVIDIA container toolkit repository
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | \
  sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg

curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
  sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
  sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

# Install it
sudo apt update
sudo apt install -y nvidia-container-toolkit

# Configure Docker to use it
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

Step 5: Verify everything works

Run this command to check that Docker can see your GPU:

docker run --rm --gpus all nvidia/cuda:12.6.0-base-ubuntu24.04 nvidia-smi

You should see the same nvidia-smi output as before, but this time it is running inside a container. If this works, you, Docker and your fancy GPU are on speaking terms. You are ready to run picoDiffusion.

If it does not work, the most common issues are:

The NVIDIA Container Toolkit is not installed or Docker was not restarted after installing it
Your user is not in the docker group (run groups to check)
The NVIDIA driver version is too old for the CUDA version — update your drivers

Next steps

Once docker run --gpus all nvidia/cuda:... shows your GPU, head back to the Getting Started guide and pick up from Step 1.

Still stuck? GPU driver issues can be frustrating and very system-specific. The NVIDIA Developer Forums and your distribution's community forums are good places to ask for help. Include the output of lspci | grep -i nvidia and uname -r when asking — it helps people diagnose the problem.