From 1219b408b5778eb2095b80b6590631a706b3edcc Mon Sep 17 00:00:00 2001
From: JB <jb-apps@outlook.com>
Date: Mon, 22 Dec 2025 08:54:03 -0800
Subject: [PATCH] feat: Update Dockerfile and entrypoint script for CUDA
 diagnostics and web server startup; adjust PyTorch version and URL handling

---
 abogen/webui/Dockerfile           |  9 ++++++--
 abogen/webui/conversion_runner.py | 10 ++++++++-
 abogen/webui/entrypoint.sh        | 36 +++++++++++++++++++++++++++++++
 docker-compose.yaml               |  2 +-
 4 files changed, 53 insertions(+), 4 deletions(-)
 create mode 100755 abogen/webui/entrypoint.sh

diff --git a/abogen/webui/Dockerfile b/abogen/webui/Dockerfile
index cd4070a..4c65733 100644
--- a/abogen/webui/Dockerfile
+++ b/abogen/webui/Dockerfile
@@ -1,4 +1,4 @@
-FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
+FROM nvidia/cuda:12.6.3-cudnn-runtime-ubuntu22.04
 
 ENV PYTHONDONTWRITEBYTECODE=1 \
   PYTHONUNBUFFERED=1 \
@@ -6,7 +6,7 @@ ENV PYTHONDONTWRITEBYTECODE=1 \
   VIRTUAL_ENV=/opt/venv \
   PATH=/opt/venv/bin:$PATH
 
-ARG TORCH_INDEX_URL=https://download.pytorch.org/whl/cu124
+ARG TORCH_INDEX_URL=https://download.pytorch.org/whl/cu126
 ARG TORCH_VERSION=
 ARG USE_GPU=true
 
@@ -59,6 +59,10 @@ ENV ABOGEN_UPLOAD_ROOT=/data/uploads \
   HF_HOME=/data/huggingface \
   HUGGINGFACE_HUB_CACHE=/data/huggingface/hub
 
+# Copy and setup entrypoint script
+COPY abogen/webui/entrypoint.sh /entrypoint.sh
+RUN chmod +x /entrypoint.sh
+
 # Create non-root user and setup permissions
 RUN useradd -m -u 1000 abogen \
     && mkdir -p /data/uploads /data/outputs /data/cache /data/voice-cache /data/huggingface \
@@ -66,4 +70,5 @@ RUN useradd -m -u 1000 abogen \
 
 USER abogen
 
+ENTRYPOINT ["/entrypoint.sh"]
 CMD ["abogen-web"]
diff --git a/abogen/webui/conversion_runner.py b/abogen/webui/conversion_runner.py
index 0b996ea..f41e624 100644
--- a/abogen/webui/conversion_runner.py
+++ b/abogen/webui/conversion_runner.py
@@ -1595,7 +1595,15 @@ def run_conversion_job(job: Job) -> None:
             if not disable_gpu:
                 device = _select_device()
             _np, KPipeline = load_numpy_kpipeline()
-            pipelines[provider_norm] = KPipeline(lang_code=job.language, repo_id="hexgrad/Kokoro-82M", device=device)
+            # Try to initialize with the selected device; fall back to CPU if CUDA fails
+            try:
+                pipelines[provider_norm] = KPipeline(lang_code=job.language, repo_id="hexgrad/Kokoro-82M", device=device)
+            except RuntimeError as e:
+                if "CUDA" in str(e) and device != "cpu":
+                    job.add_log(f"CUDA initialization failed, falling back to CPU: {e}", level="warning")
+                    pipelines[provider_norm] = KPipeline(lang_code=job.language, repo_id="hexgrad/Kokoro-82M", device="cpu")
+                else:
+                    raise
             if not kokoro_cache_ready:
                 _initialize_voice_cache(job)
                 kokoro_cache_ready = True
diff --git a/abogen/webui/entrypoint.sh b/abogen/webui/entrypoint.sh
new file mode 100755
index 0000000..b11305f
--- /dev/null
+++ b/abogen/webui/entrypoint.sh
@@ -0,0 +1,36 @@
+#!/bin/bash
+# Entrypoint script for abogen container
+# Performs CUDA diagnostics and starts the web server
+
+set -e
+
+echo "=== Abogen Container Starting ==="
+
+# Check CUDA availability
+if command -v nvidia-smi &> /dev/null; then
+    echo "NVIDIA Driver detected:"
+    nvidia-smi --query-gpu=name,driver_version,memory.total,memory.free --format=csv,noheader 2>/dev/null || echo "  (nvidia-smi query failed)"
+    
+    # Check PyTorch CUDA support
+    python3 -c "
+import torch
+print(f'PyTorch version: {torch.__version__}')
+print(f'CUDA available: {torch.cuda.is_available()}')
+if torch.cuda.is_available():
+    print(f'CUDA version (PyTorch): {torch.version.cuda}')
+    print(f'GPU count: {torch.cuda.device_count()}')
+    for i in range(torch.cuda.device_count()):
+        props = torch.cuda.get_device_properties(i)
+        print(f'  GPU {i}: {props.name} ({props.total_memory // 1024**2} MB)')
+else:
+    print('WARNING: PyTorch cannot access CUDA. Running on CPU.')
+" 2>&1 || echo "PyTorch CUDA check failed"
+else
+    echo "No NVIDIA driver detected. Running on CPU."
+fi
+
+echo "================================="
+echo ""
+
+# Start the application
+exec "$@"
diff --git a/docker-compose.yaml b/docker-compose.yaml
index 256b11d..1c92069 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -19,7 +19,7 @@ services:
       context: .
       dockerfile: abogen/webui/Dockerfile
       args:
-        TORCH_INDEX_URL: ${TORCH_INDEX_URL:-https://download.pytorch.org/whl/cu124}
+        TORCH_INDEX_URL: ${TORCH_INDEX_URL:-https://download.pytorch.org/whl/cu126}
         TORCH_VERSION: ${TORCH_VERSION:-}
         USE_GPU: ${USE_GPU:-true}
     image: abogen:latest