From 20f558584ac2f3f3ff8ceef5c0fb75b673f9598b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Felix=20Sp=C3=B6ttel?=
 <1682504+fspoettel@users.noreply.github.com>
Date: Tue, 28 Feb 2023 22:05:09 +0100
Subject: [PATCH] feat: allow multiple models to be installed in docker
 container

---
 docker/dev/docker-compose.yml | 4 ++--
 docker/dev/worker.Dockerfile  | 8 ++++----
 docker/prod/.env.example      | 2 +-
 docker/prod/worker.Dockerfile | 6 +++---
 scripts/download_model.py     | 8 --------
 scripts/download_models.py    | 8 ++++++++
 6 files changed, 18 insertions(+), 18 deletions(-)
 delete mode 100644 scripts/download_model.py
 create mode 100644 scripts/download_models.py

diff --git a/docker/dev/docker-compose.yml b/docker/dev/docker-compose.yml
index ab92d92..1f7efdd 100644
--- a/docker/dev/docker-compose.yml
+++ b/docker/dev/docker-compose.yml
@@ -5,7 +5,7 @@ x-app-variables: &app-variables
   DATABASE_URI: sqlite:///./whisperbox.sqlite
   ENVIRONMENT: development
   BROKER_URL: redis://redis:6379/0
-  WHISPER_MODEL: small
+  WHISPER_MODELS: small
 
 services:
   redis:
@@ -41,7 +41,7 @@ services:
       context: ../../
       dockerfile: docker/dev/worker.Dockerfile
       args:
-        WHISPER_MODEL: small
+        WHISPER_MODELS: small
     container_name: whisperbox_worker_dev
     volumes:
       - ../../:/code
diff --git a/docker/dev/worker.Dockerfile b/docker/dev/worker.Dockerfile
index 8ad134f..7549c8f 100644
--- a/docker/dev/worker.Dockerfile
+++ b/docker/dev/worker.Dockerfile
@@ -1,6 +1,6 @@
 FROM python:3.10
 
-ARG WHISPER_MODEL
+ARG WHISPER_MODELS
 
 WORKDIR /code
 
@@ -13,7 +13,7 @@ COPY --from=mwader/static-ffmpeg:5.1.2 /ffprobe /usr/local/bin/
 COPY pyproject.toml .
 RUN pip install --no-cache-dir --user .[worker,worker_dev]
 
-COPY scripts/download_model.py .
-RUN chmod +x download_model.py && python download_model.py ${WHISPER_MODEL}
+COPY scripts/download_models.py .
+RUN python download_models.py ${WHISPER_MODELS}
 
-ENTRYPOINT ["watchmedo", "auto-restart", "-d" , "app/worker", "-p", "*.py", "--recursive", "celery", "--", "--app=app.worker.main.celery", "worker", "--loglevel=info", "--concurrency=1"]
+CMD watchmedo auto-restart -d app/worker -p *.py --recursive celery -- --app=app.worker.main.celery worker --loglevel=info --concurrency=1
diff --git a/docker/prod/.env.example b/docker/prod/.env.example
index 801b483..654c500 100644
--- a/docker/prod/.env.example
+++ b/docker/prod/.env.example
@@ -1,2 +1,2 @@
 API_SECRET="change_me"
-WHISPER_MODEL="small"
+WHISPER_MODELS="small"
diff --git a/docker/prod/worker.Dockerfile b/docker/prod/worker.Dockerfile
index b0e61e6..d76f524 100644
--- a/docker/prod/worker.Dockerfile
+++ b/docker/prod/worker.Dockerfile
@@ -10,7 +10,7 @@ RUN python -m venv /opt/venv && \
 
 FROM python:3.10 as python-deploy
 
-ARG WHISPER_MODEL
+ARG WHISPER_MODELS
 
 WORKDIR /etc/whisperbox
 
@@ -24,7 +24,7 @@ COPY app ./app
 ENV VIRTUAL_ENV /opt/venv
 ENV PATH /opt/venv/bin:$PATH
 
-COPY scripts/download_model.py .
-RUN chmod +x download_model.py && python download_model.py ${WHISPER_MODEL:-small}
+COPY scripts/download_models.py .
+RUN python download_models.py ${WHISPER_MODELS}
 
 CMD celery --app=app.worker.main.celery worker --loglevel=info --concurrency=1
diff --git a/scripts/download_model.py b/scripts/download_model.py
deleted file mode 100644
index d4d6f79..0000000
--- a/scripts/download_model.py
+++ /dev/null
@@ -1,8 +0,0 @@
-import sys
-from whisper import _download, _MODELS # type: ignore
-
-if __name__ == "__main__":
-    model_name = sys.argv[1]
-    _download(_MODELS[model_name], "/models/", False)
-    if model_name != "large":
-        _download(_MODELS[f"{model_name}.en"], "/models/", False)
diff --git a/scripts/download_models.py b/scripts/download_models.py
new file mode 100644
index 0000000..ee92627
--- /dev/null
+++ b/scripts/download_models.py
@@ -0,0 +1,8 @@
+import os
+from whisper import _download, _MODELS  # type: ignore
+
+if __name__ == "__main__":
+    for name in os.environ["WHISPER_MODELS"].split(","):
+        _download(_MODELS[name], "/models/", False)
+        if name != "large":
+            _download(_MODELS[f"{name}.en"], "/models/", False)