From 1d5c787480d5d4ba570d3f266c8553786ca2b268 Mon Sep 17 00:00:00 2001
From: pluja <github_pluja@r3d.red>
Date: Thu, 26 Oct 2023 22:49:37 +0200
Subject: [PATCH] prepare transcription API for GPU support (#34)

---
 .github/workflows/docker-image-gpu.yml | 38 ++++++++++++++++++++++++++
 transcription-api/Dockerfile.gpu       | 32 ++++++++++++++++++++++
 transcription-api/main.py              | 11 +++++---
 transcription-api/requirements.txt     |  4 ++-
 4 files changed, 80 insertions(+), 5 deletions(-)
 create mode 100644 .github/workflows/docker-image-gpu.yml
 create mode 100644 transcription-api/Dockerfile.gpu

diff --git a/.github/workflows/docker-image-gpu.yml b/.github/workflows/docker-image-gpu.yml
new file mode 100644
index 0000000..3c20323
--- /dev/null
+++ b/.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) }}
+
diff --git a/transcription-api/Dockerfile.gpu b/transcription-api/Dockerfile.gpu
new file mode 100644
index 0000000..0126ac8
--- /dev/null
+++ b/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"]
\ No newline at end of file
diff --git a/transcription-api/main.py b/transcription-api/main.py
index 4332fbb..cc0ca13 100644
--- a/transcription-api/main.py
+++ b/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"}
 
diff --git a/transcription-api/requirements.txt b/transcription-api/requirements.txt
index e30b3d6..551050f 100644
--- a/transcription-api/requirements.txt
+++ b/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
\ No newline at end of file
+python-multipart
+ffmpeg-python
+soundfile
\ No newline at end of file