prepare transcription API for GPU support (#34)

transcription-api/Dockerfile.gpu

@@ -0,0 +1,32 @@
+FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04
+
+ENV PYTHON_VERSION=3.11
+
+RUN export DEBIAN_FRONTEND=noninteractive \
+    && apt-get -qq update \
+    && apt-get -qq install --no-install-recommends \
+    python${PYTHON_VERSION} \
+    python${PYTHON_VERSION}-venv \
+    python3-pip \
+    ffmpeg \
+    curl \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN ln -s -f /usr/bin/python${PYTHON_VERSION} /usr/bin/python3 && \
+    ln -s -f /usr/bin/python${PYTHON_VERSION} /usr/bin/python && \
+    ln -s -f /usr/bin/pip3 /usr/bin/pip
+
+COPY requirements.txt requirements.txt
+RUN pip3 install -r requirements.txt \
+    && pip3 install python-multipart \
+    && pip3 install torch==1.13.0+cu117 -f https://download.pytorch.org/whl/torch
+
+WORKDIR /app
+
+COPY . .
+
+HEALTHCHECK --interval=30s --timeout=10s --start-period=15s \
+CMD ["python3", "healthcheck.py"]
+
+EXPOSE 8000
+CMD ["python3", "main.py"]

transcription-api/main.py

@@ -15,15 +15,18 @@ async def transcribe_endpoint(file: UploadFile = File(None),
                               filename: str = None,
                               model_size: ModelSize = ModelSize.small, 
                               language: Languages = Languages.auto,
-                              device: DeviceType = DeviceType.cpu):
+                              device: str = "cpu"):
     
-    print(f"Transcribing with model {model_size.value} on device {device.value}...")
+    if device != "cpu" and device != "cuda":
+        return {"detail": "Device must be either cpu or cuda"}
+    
+    print(f"Transcribing with model {model_size.value} on device {device}...")
     if file is not None:
         # if a file is uploaded, use it
-        return await transcribe_file(file, model_size.value, language.value, device.value)
+        return await transcribe_file(file, model_size.value, language.value, device)
     elif filename is not None:
         # if a filename is provided, use it
-        return await transcribe_from_filename(filename, model_size.value, language.value, device.value)
+        return await transcribe_from_filename(filename, model_size.value, language.value, device)
     else:
         return {"detail": "No file uploaded and no filename provided"}
 

transcription-api/requirements.txt

@@ -3,4 +3,6 @@ faster_whisper==0.9.0
 pydantic==2.1.1
 python-dotenv==1.0.0
 uvicorn==0.23.2
-python-multipart
+python-multipart
+ffmpeg-python
+soundfile