feat: Update Dockerfile and entrypoint script for CUDA diagnostics and web server startup; adjust PyTorch version and URL handling

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"]

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()
+            # 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

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 "$@"

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