From 899c9f5aa556ef1cc72f9b06bfff738a4fd69e87 Mon Sep 17 00:00:00 2001
From: JB <jb-apps@outlook.com>
Date: Sun, 21 Dec 2025 08:31:51 -0800
Subject: [PATCH] feat: Add USE_GPU argument to Dockerfile and docker-compose
 for GPU support

---
 abogen/webui/Dockerfile | 7 +++++++
 docker-compose.yaml     | 1 +
 2 files changed, 8 insertions(+)

diff --git a/abogen/webui/Dockerfile b/abogen/webui/Dockerfile
index 6c5731e..037dc53 100644
--- a/abogen/webui/Dockerfile
+++ b/abogen/webui/Dockerfile
@@ -8,6 +8,7 @@ ENV PYTHONDONTWRITEBYTECODE=1 \
 
 ARG TORCH_INDEX_URL=https://download.pytorch.org/whl/cu124
 ARG TORCH_VERSION=
+ARG USE_GPU=true
 
 RUN apt-get update \
   && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
@@ -38,6 +39,12 @@ RUN pip install --upgrade pip \
     https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl \
   && pip install --no-cache-dir "mutagen>=1.47.0"
 
+# Install onnxruntime-gpu for CUDA acceleration (supertonic uses ONNX Runtime)
+# Set USE_GPU=false to skip this for CPU-only deployments
+RUN if [ "$USE_GPU" = "true" ]; then \
+      pip install --no-cache-dir onnxruntime-gpu; \
+    fi
+
 ENV ABOGEN_HOST=0.0.0.0 \
   ABOGEN_PORT=8808
 
diff --git a/docker-compose.yaml b/docker-compose.yaml
index 924ff10..e47e6ec 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -16,6 +16,7 @@ services:
       args:
         TORCH_INDEX_URL: ${TORCH_INDEX_URL:-https://download.pytorch.org/whl/cu124}
         TORCH_VERSION: ${TORCH_VERSION:-}
+        USE_GPU: ${USE_GPU:-true}
     image: abogen:latest
     user: "${ABOGEN_UID:-1000}:${ABOGEN_GID:-1000}"
     ports: