chore: better structure for cuda/cpu

README.md

@@ -64,7 +64,7 @@ run it
 
 ## Docker Setup
 
-Note: currently we only have Docker support for CUDA.
+Note: currently we only have Docker support for CUDA and CPU.
 
 To build and run the project using Docker, follow these steps:
 
@@ -75,14 +75,15 @@ To build and run the project using Docker, follow these steps:
 
 2. Build the Docker image using the following command:
    ```bash
-   DOCKER_BUILDKIT=1 docker build \
-     --secret id=ssh_key,src=$HOME/.ssh/your_ssh_key \
-     --secret id=known_hosts,src=$HOME/.ssh/known_hosts \
-     --secret id=hf_token,env=HF_TOKEN \
-     -t entropix \
-     -f docker/Dockerfile.entropix_cuda .
-   ```
-   Make sure to replace `your_ssh_key` with the path to your actual SSH key.
+  DOCKER_BUILDKIT=1 docker build \
+    --secret id=ssh_key,src=$HOME/.ssh/your_ssh_key \
+    --secret id=known_hosts,src=$HOME/.ssh/known_hosts \
+    --secret id=hf_token,env=HF_TOKEN \
+    --build-arg BACKEND=cuda \
+    -t entropix \
+    -f docker/Dockerfile.entropix .   
+    ```
+   Make sure to replace `your_ssh_key` with the path to your actual SSH key, and to set the `BACKEND` argument to `cpu` if you want to build the CPU image.
 
 3. Once the image is built, you can run it using:
    ```bash

docker/Dockerfile.entropix_cuda → docker/Dockerfile.entropix

@@ -1,4 +1,13 @@
-FROM nvidia/cuda:12.3.1-base-ubuntu22.04
+ARG BACKEND=cuda
+# CUDA base image
+FROM nvidia/cuda:12.3.1-base-ubuntu22.04 AS cuda_base
+# CPU base image
+FROM ubuntu:22.04 AS cpu_base
+
+FROM ${BACKEND}_base AS final
+ARG DEBIAN_FRONTEND=noninteractive
+ARG BACKEND
+
 # Required to disable interactive prompts when installing openssh.
 ARG DEBIAN_FRONTEND=noninteractive
 RUN apt-get update && apt-get install -y \