Fix CUDA/ROCm compatibility: default to CPU for AMD GPUs

- Change default device to CPU in docker-compose and main.py
- Set compute_type=int8 for CPU inference
- Auto-detect device if env not set

Dockerfile

@@ -1,5 +1,5 @@
-# Use ROCm PyTorch base image
+# Use ROCm PyTorch base image with compatible PyTorch
-FROM rocm/pytorch:latest
+FROM rocm/pytorch:rocm7.2.3_ubuntu22.04_py3.10_pytorch_release_2.10.0
 
 # Set environment variables for ROCm and Python
 ENV PYTHONDONTWRITEBYTECODE=1

docker-compose.yml

@@ -8,9 +8,10 @@ services:
     ports:
       - "8000:8000"
     environment:
-      - WHISPERX_MODEL=turbo
+      - WHISPERX_MODEL=large-v2
-      - WHISPERX_DEVICE=cuda
+      - WHISPERX_DEVICE=cpu
-      - WHISPERX_COMPUTE_TYPE=float16
+      - WHISPERX_COMPUTE_TYPE=int8
+      - HF_HUB_DISABLE_TELEMETRY=1
     volumes:
       # Mount Hugging Face cache if needed
       - hf_cache:/app/.cache/huggingface

pyproject.toml

@@ -19,10 +19,10 @@ dependencies = [
     "av<16.0.0",
     "numpy>=2.1.0,<2.3.0; python_version >='3.13'",
     "pyannote-audio>=3.3.2,<4.0.0",
-    "torch~=2.8.0",
+    "torch~=2.4.0",
-    "torchaudio~=2.8.0",
+    "torchaudio~=2.4.0",
-    "transformers>=4.48.0",
+    "torchvision~=0.19.0",
-    "triton>=3.3.0; sys_platform == 'linux' and platform_machine == 'x86_64'", # only install triton on x86_64 Linux
+    "transformers==4.42.0",
     "fastapi>=0.104.0",
     "uvicorn[standard]>=0.24.0",
     "python-multipart>=0.0.6",

whisperx/api/main.py

@@ -6,34 +6,29 @@ from fastapi import FastAPI, UploadFile, File, Form, HTTPException
 from fastapi.responses import JSONResponse
 import torch
 import whisperx
-from whisperx.schema import TranscriptionResult
 
 
 model = None
-align_model_metadata = None
 
 
-def load_transcription_model(model_name: str = "turbo", device: str = None, compute_type: str = "float16"):
+def load_transcription_model(model_name: str = "large-v2", device: str = None, compute_type: str = None):
-    global model, align_model_metadata
+    global model
     if device is None:
-        device = "cuda" if torch.cuda.is_available() else "cpu"
+        device = os.getenv("WHISPERX_DEVICE", "cuda" if torch.cuda.is_available() else "cpu")
+    if compute_type is None:
+        compute_type = "int8" if device == "cpu" else "float16"
     print(f"Loading WhisperX model: {model_name} on {device} with {compute_type}")
     model = whisperx.load_model(model_name, device, compute_type=compute_type)
-    # For alignment, load the metadata
-    align_model_metadata = whisperx.alignment.DEFAULT_ALIGN_MODELS_HF
     print("Model loaded and ready.")
 
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    # Load the model at startup
+    model_name = os.getenv("WHISPERX_MODEL", "large-v2")
-    model_name = os.getenv("WHISPERX_MODEL", "turbo")
+    device = os.getenv("WHISPERX_DEVICE", "cuda" if torch.cuda.is_available() else "cpu")
-    device = os.getenv("WHISPERX_DEVICE", "cuda")
+    compute_type = os.getenv("WHISPERX_COMPUTE_TYPE", "int8" if device == "cpu" else "float16")
-    compute_type = os.getenv("WHISPERX_COMPUTE_TYPE", "float16")
     load_transcription_model(model_name, device, compute_type)
     yield
-    # Cleanup if needed
-    print("Shutting down API")
 
 
 app = FastAPI(