CPU/Cuda versions

Author: OskarLiew

.github/workflows/cd.yaml

@@ -9,7 +9,7 @@ env:
   IMAGE_ID: ghcr.io/${{ github.repository_owner }}/docling-inference
 
 jobs:
-  push:
+  push-cuda:
     runs-on: ubuntu-latest
 
     steps:
@@ -23,6 +23,7 @@ jobs:
     - name: Set up Docker Buildx
       uses: docker/setup-buildx-action@v3
 
+    # Build and push cuda version
     - name: Build and push Docker image
       uses: docker/build-push-action@v6
       with:
@@ -31,3 +32,28 @@ jobs:
         tags: ${{ env.IMAGE_ID }}:latest,${{ env.IMAGE_ID }}:${{ github.sha }}
         cache-from: type=registry,ref=${{ env.IMAGE_ID }}:buildcache
         cache-to: type=registry,ref=${{ env.IMAGE_ID }}:buildcache,mode=max
+
+  push-cpu:
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v3
+
+    - name: Docker login
+      run: echo ${{ secrets.GITHUB_TOKEN }} | docker login ghcr.io -u ${{ github.actor }} --password-stdin
+
+    # Use buildx which can cache docker layers between runs
+    - name: Set up Docker Buildx
+      uses: docker/setup-buildx-action@v3
+
+    # Build and push cpu version
+    - name: Build and push Docker image
+      uses: docker/build-push-action@v6
+      with:
+        context: .
+        push: true
+        file: Dockerfile.cpu
+        tags: ${{ env.IMAGE_ID }}:cpu-latest,${{ env.IMAGE_ID }}:cpu-${{ github.sha }}
+        cache-from: type=registry,ref=${{ env.IMAGE_ID }}:cpu-buildcache
+        cache-to: type=registry,ref=${{ env.IMAGE_ID }}:cpu-buildcache,mode=max

Dockerfile

@@ -6,7 +6,7 @@ RUN pip install --upgrade pip && \
     poetry config virtualenvs.in-project true
 
 COPY pyproject.toml poetry.lock /app/
-RUN poetry install --no-ansi
+RUN poetry install --no-ansi --with=cuda
 
 FROM python:3.11-slim-bookworm
 WORKDIR /app

Dockerfile.cpu

@@ -0,0 +1,17 @@
+FROM python:3.11-slim-bookworm AS builder
+WORKDIR /app
+
+RUN pip install --upgrade pip && \
+    pip install poetry && \
+    poetry config virtualenvs.in-project true
+
+COPY pyproject.toml poetry.lock /app/
+RUN poetry install --no-ansi --with=cpu
+
+FROM python:3.11-slim-bookworm
+WORKDIR /app
+
+COPY --from=builder /app /app
+COPY src /app/src
+
+CMD [ "/app/.venv/bin/python", "-m", "src.main" ]

README.md

@@ -20,15 +20,15 @@ docker run -d \
   -e NUM_WORKERS=8 \
   -v hf_cache:/root/.cache/huggingface \
   -v ocr_cache:/root/.EasyOCR \
-  docling-inference
+  ghcr.io/aidotse/docling-inference:latest
 ```
 
 ### Docker compose
 
 ```yaml
 services:
   docling-inference:
-    image: docling-inference
+    image: ghcr.io/aidotse/docling-inference:latest
     ports:
       - 8080:8080
     environment:

docker-compose.yaml

@@ -1,6 +1,6 @@
 services:
   docling-inference:
-    image: docling-inference:dev
+    image: ghcr.io/aidotse/docling-inference:dev
     ports:
       - 8080:8080
     environment:

justfile

@@ -1,14 +1,16 @@
 compose := "docker compose"
-image := "docling-inference"
+image := "ghcr.io/aidotse/docling-inference"
 rev := shell("git rev-parse --short HEAD")
 
-
 default:
 	just --list
 
 build:
 	docker build -t {{image}}:dev -t {{image}}:{{rev}} .
 
+build-cpu:
+	docker build -f Dockerfile.cpu -t {{image}}:cpu-dev -t {{image}}:cpu-{{rev}} .
+
 up:
 	{{compose}} up --remove-orphans
 

poetry.lock

@@ -10,14 +10,28 @@ packages = [{include = "src"}]
 
 [tool.poetry.dependencies]
 python = "^3.11"
-torch = "^2.5.1"
 docling = "^2.10.0"
 fastapi = "^0.115.6"
 uvicorn = "^0.32.1"
 python-multipart = "^0.0.19"
 
 
 
+[[tool.poetry.source]]
+name = "pytorch_cpu"
+url = "https://download.pytorch.org/whl/cpu"
+priority = "explicit"
+
+
+[tool.poetry.group.cpu.dependencies]
+torch = {version = "^2.5.1+cpu", source = "pytorch_cpu"}
+torchvision = {version = "^0.20.1+cpu", source = "pytorch_cpu"}
+
+
+[tool.poetry.group.cuda.dependencies]
+torch = "^2.5.1"
+torchvision = "^0.20.1"
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"