prepare transcription API for GPU support (#34)

.github/workflows/docker-image-gpu.yml

@@ -0,0 +1,38 @@
+name: Docker Transcription-API GPU
+
+on:
+  release:
+    types: [published]
+
+jobs:
+  build:
+    name: Build
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v2
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+
+      - name: Login to DockerHub
+        uses: docker/login-action@v2
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+
+      - name: Build and push Docker images
+        uses: docker/build-push-action@v4
+        with:
+          context: transcription-api
+          file: transcription-api/Dockerfile.gpu
+          platforms: linux/amd64,linux/arm64
+          push: true
+          tags: |
+            ${{ format('{0}/whishper-transcription-api:{1}-gpu', secrets.DOCKERHUB_USERNAME, github.event.release.tag_name) }}
+            ${{ format('{0}/whishper-transcription-api:latest-gpu', secrets.DOCKERHUB_USERNAME) }}
+

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

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