Reduce container size (m-bain#1)

* Smaller dockerfile

* Small dockerfile no model

* README no need for -- with new docker

* Revert "README no need for -- with new docker"

This reverts commit e9816fb2df019148e148e656381b44f1f044446d.

---------

Co-authored-by: 陳鈞 <jim60105@gmail.com>

Dockerfile

@@ -1,21 +1,34 @@
 ARG WHISPER_MODEL=base
 ARG LANG=en
+ARG TORCH_HOME=/cache/torch
+ARG HF_HOME=/cache/huggingface
 
-FROM nvcr.io/nvidia/pytorch:23.05-py3 as base
+FROM python:3.10-slim as dependencies
 
-WORKDIR /app
-ENV TORCH_HOME=/cache/torch
-ENV HF_HOME=/cache/huggingface
+# Setup venv
+RUN python3 -m venv /venv
+ARG PATH="/venv/bin:$PATH"
+RUN python3 -m pip install --upgrade pip setuptools
+# Install requirements
+RUN python3 -m pip install torch torchaudio --extra-index-url https://download.pytorch.org/whl/cu118
 
-# ffmpeg
-COPY --link --from=mwader/static-ffmpeg:6.0 /ffmpeg /usr/local/bin/
-COPY --link --from=mwader/static-ffmpeg:6.0 /ffprobe /usr/local/bin/
+# Add git
+ARG DEBIAN_FRONTEND=noninteractive
+RUN apt-get update && apt-get install -y --no-install-recommends git
 
-# Install requirements
-COPY ./whisperX/requirements.txt .
-RUN python3 -m pip install --no-cache-dir -r ./requirements.txt ujson torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2
+# Install whisperX
+COPY ./whisperX /code
+RUN python3 -m pip install /code
+
+
+FROM dependencies as load_model
 
-FROM base AS load_model
+ARG TORCH_HOME
+ARG HF_HOME
+ARG PATH="/venv/bin:$PATH"
+
+# Preload vad model
+RUN python3 -c 'from whisperx.vad import load_vad_model; load_vad_model("cpu");'
 
 # Preload fast-whisper
 ARG WHISPER_MODEL
@@ -24,24 +37,36 @@ RUN python3 -c 'import faster_whisper; model = faster_whisper.WhisperModel("'${W
 # Preload align model
 ARG LANG
 COPY load_align_model.py .
-RUN python load_align_model.py ${LANG}
+RUN python3 load_align_model.py ${LANG}
 
-FROM base AS final
 
-# Install whisperX
-COPY ./whisperX/ .
-RUN python3 -m pip install --no-cache-dir .
+FROM python:3.10-slim
+
+# Copy and use venv
+COPY --from=dependencies /venv /venv
+ARG PATH="/venv/bin:$PATH"
+ENV PATH=${PATH}
 
 # Non-root user
 RUN useradd -m -s /bin/bash appuser
 USER appuser
 
 COPY --chown=appuser --from=load_model /cache /cache
 
+WORKDIR /app
+ARG TORCH_HOME
+ARG HF_HOME
+ENV TORCH_HOME=${TORCH_HOME}
+ENV HF_HOME=${HF_HOME}
+
+# ffmpeg
+COPY --link --from=mwader/static-ffmpeg:6.0 /ffmpeg /usr/local/bin/
+COPY --link --from=mwader/static-ffmpeg:6.0 /ffprobe /usr/local/bin/
+
 ARG WHISPER_MODEL
 ENV WHISPER_MODEL=${WHISPER_MODEL}
 ARG LANG
 ENV LANG=${LANG}
 
 STOPSIGNAL SIGINT
-ENTRYPOINT whisperx --model ${WHISPER_MODEL} --language ${LANG} $*
+ENTRYPOINT whisperx --model ${WHISPER_MODEL} --language ${LANG} $@

Dockerfile.no_model

@@ -1,31 +1,51 @@
-FROM nvcr.io/nvidia/pytorch:23.05-py3 as base
+ARG TORCH_HOME=/cache/torch
+ARG HF_HOME=/cache/huggingface
 
-WORKDIR /app
-ENV TORCH_HOME=/cache/torch
-ENV HF_HOME=/cache/huggingface
-
-# ffmpeg
-COPY --link --from=mwader/static-ffmpeg:6.0 /ffmpeg /usr/local/bin/
-COPY --link --from=mwader/static-ffmpeg:6.0 /ffprobe /usr/local/bin/
+FROM python:3.10-slim as dependencies
 
+# Setup venv
+RUN python3 -m venv /venv
+ARG PATH="/venv/bin:$PATH"
+RUN python3 -m pip install --upgrade pip setuptools
 # Install requirements
-COPY ./whisperX/requirements.txt .
-RUN python3 -m pip install --no-cache-dir -r ./requirements.txt ujson torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2
+RUN python3 -m pip install torch torchaudio --extra-index-url https://download.pytorch.org/whl/cu118
 
-FROM base AS final
+# Add git
+ARG DEBIAN_FRONTEND=noninteractive
+RUN apt-get update && apt-get install -y --no-install-recommends git
 
 # Install whisperX
-COPY ./whisperX/ .
-RUN python3 -m pip install --no-cache-dir .
+COPY ./whisperX /code
+RUN python3 -m pip install /code
+
+
+FROM python:3.10-slim
+
+# Copy and use venv
+COPY --from=dependencies /venv /venv
+ARG PATH="/venv/bin:$PATH"
+ENV PATH=${PATH}
 
 # Non-root user
 RUN useradd -m -s /bin/bash appuser \
     && mkdir -p /cache \
     && chown -R appuser:appuser /cache
 USER appuser
 
-ENV WHISPER_MODEL=
-ENV LANG=
+WORKDIR /app
+ARG TORCH_HOME
+ARG HF_HOME
+ENV TORCH_HOME=${TORCH_HOME}
+ENV HF_HOME=${HF_HOME}
+
+# ffmpeg
+COPY --link --from=mwader/static-ffmpeg:6.0 /ffmpeg /usr/local/bin/
+COPY --link --from=mwader/static-ffmpeg:6.0 /ffprobe /usr/local/bin/
+
+ARG WHISPER_MODEL
+ENV WHISPER_MODEL=${WHISPER_MODEL}
+ARG LANG
+ENV LANG=${LANG}
 
 STOPSIGNAL SIGINT
-ENTRYPOINT whisperx $*
+ENTRYPOINT whisperx $@

README.md

@@ -44,7 +44,7 @@ docker build --build-arg LANG=ja --build-arg WHISPER_MODEL=large-v2 -t whisperx:
 Mount the current directory as `/app` and run WhisperX with additional input arguments:
 
 ```bash
-docker run --gpus all -it -v "$(PWD):/app" whisperx:large-v2-ja -- --output_format srt audio.mp3
+docker run --gpus all -it -v ".:/app" whisperx:large-v2-ja -- --output_format srt audio.mp3
 ```
 
 > [!NOTE]