fix: optimize docker build and reduce image size

Dockerfile

@@ -6,28 +6,24 @@ WORKDIR /
 
 # Update and upgrade the system packages (Worker Template)
 RUN apt-get update && \
-    apt-get upgrade -y
-
-# Install System Packages
-RUN apt-get install ffmpeg -y
-
-# Download Models
-COPY builder/download_models.sh /download_models.sh
-RUN chmod +x /download_models.sh && \
-    /download_models.sh
-RUN rm /download_models.sh
+    apt-get install -y --no-install-recommends ffmpeg && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
 
 # Install Python dependencies (Worker Template)
 COPY builder/requirements.txt /requirements.txt
-RUN pip install --upgrade pip && \
-    pip install -r /requirements.txt && \
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir -r /requirements.txt && \
     rm /requirements.txt
 
-ADD src .
+# Copy source code
+COPY src .
 
-# Cleanup section (Worker Template)
-RUN apt-get autoremove -y && \
-    apt-get clean -y && \
-    rm -rf /var/lib/apt/lists/*
+# Add test_input.json
+COPY test_input.json /
+
+# Download only the tiny model to save space (moved after copying code to leverage layer caching)
+RUN mkdir -p ./weights && \
+    wget -q https://openaipublic.azureedge.net/main/whisper/models/65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/tiny.pt -P ./weights
 
-CMD [ "python", "-u", "/rp_handler.py" ]
+CMD [ "python", "-u", "/rp_handler.py" ]

README.md

@@ -41,3 +41,9 @@ The following inputs can be used for testing the model:
     }
 }
 ```
+
+## Space Optimization
+
+To reduce the Docker image size and prevent "no space left on device" errors during builds, this worker only includes the "tiny" Whisper model by default. If you need additional models (base, small, medium, large), you can uncomment the relevant sections in the `builder/download_models.sh` script or modify the Dockerfile to download them at build time.
+
+The test_input.json file is configured to use the "tiny" model by default.

builder/download_models.sh

@@ -1,8 +1,24 @@
-mkdir ./weights
+#!/bin/bash
+
+# Create weights directory
+mkdir -p ./weights
+
+# Download only tiny model by default to save space
+echo "Downloading tiny model..."
 wget https://openaipublic.azureedge.net/main/whisper/models/65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/tiny.pt -P ./weights
-wget https://openaipublic.azureedge.net/main/whisper/models/ed3a0b6b1c0edf879ad9b11b1af5a0e6ab5db9205f891f668f8b0e6c6326e34e/base.pt -P ./weights
-wget https://openaipublic.azureedge.net/main/whisper/models/9ecf779972d90ba49c06d968637d720dd632c55bbf19d441fb42bf17a411e794/small.pt -P ./weights
-wget https://openaipublic.azureedge.net/main/whisper/models/345ae4da62f9b3d59415adc60127b97c714f32e89e936602e85993674d08dcb1/medium.pt -P ./weights
-wget https://openaipublic.azureedge.net/main/whisper/models/e4b87e7e0bf463eb8e6956e646f1e277e901512310def2c24bf0e11bd3c28e9a/large-v1.pt  -P ./weights
-wget https://openaipublic.azureedge.net/main/whisper/models/81f7c96c852ee8fc832187b0132e569d6c3065a3252ed18e56effd0b6a73e524/large-v2.pt  -P ./weights
-wget https://openaipublic.azureedge.net/main/whisper/models/e5b1a55b89c1367dacf97e3e19bfd829a01529dbfdeefa8caeb59b3f1b81dadb/large-v3.pt -P ./weights
+
+# Other models are commented out to save space
+# Uncomment if you need these models and have enough disk space
+# echo "Downloading base model..."
+# wget https://openaipublic.azureedge.net/main/whisper/models/ed3a0b6b1c0edf879ad9b11b1af5a0e6ab5db9205f891f668f8b0e6c6326e34e/base.pt -P ./weights
+# 
+# echo "Downloading small model..."
+# wget https://openaipublic.azureedge.net/main/whisper/models/9ecf779972d90ba49c06d968637d720dd632c55bbf19d441fb42bf17a411e794/small.pt -P ./weights
+# 
+# echo "Downloading medium model..."
+# wget https://openaipublic.azureedge.net/main/whisper/models/345ae4da62f9b3d59415adc60127b97c714f32e89e936602e85993674d08dcb1/medium.pt -P ./weights
+# 
+# echo "Downloading large models..."
+# wget https://openaipublic.azureedge.net/main/whisper/models/e4b87e7e0bf463eb8e6956e646f1e277e901512310def2c24bf0e11bd3c28e9a/large-v1.pt -P ./weights
+# wget https://openaipublic.azureedge.net/main/whisper/models/81f7c96c852ee8fc832187b0132e569d6c3065a3252ed18e56effd0b6a73e524/large-v2.pt -P ./weights
+# wget https://openaipublic.azureedge.net/main/whisper/models/e5b1a55b89c1367dacf97e3e19bfd829a01529dbfdeefa8caeb59b3f1b81dadb/large-v3.pt -P ./weights

builder/requirements.txt

@@ -1,12 +1,10 @@
 # Required Python packages get listed here, one per line.
 # Reccomended to lock the version number to avoid unexpected changes.
 
-runpod~=1.7.0
+runpod==1.7.13
 tqdm==4.64.1
-torch==1.13.0
-torchaudio==0.13.0
 numpy==1.23.5
 transformers==4.25.1
 ffmpeg-python==0.2.0
 more_itertools==9.0.0
-git+https://github.com/openai/whisper.git
+git+https://github.com/openai/whisper.git

test_input.json

@@ -0,0 +1,6 @@
+{
+  "input": {
+    "audio_url": "https://github.com/openai/whisper/raw/main/tests/jfk.flac",
+    "model": "tiny"
+  }
+}