change docker os

Dockerfile

@@ -1,15 +1,50 @@
-# Use the official Python 3.9 image
-FROM python:3.9
+FROM nvidia/cuda:11.7.1-cudnn8-runtime-ubuntu20.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt update && \
+    apt install -y bash \
+                   build-essential \
+                   git \
+                   git-lfs \
+                   curl \
+                   ca-certificates \
+                   libsndfile1-dev \
+                   python3.8 \
+                   python3-pip \
+                   python3.8-venv && \
+    rm -rf /var/lib/apt/lists
+
+# make sure to use venv
+RUN python3 -m venv /opt/venv
+ENV PATH="/opt/venv/bin:$PATH"
+
+# pre-install the heavy dependencies (these can later be overridden by the deps from setup.py)
+RUN python3 -m pip install --no-cache-dir --upgrade pip && \
+    python3 -m pip install --no-cache-dir \
+        torch \
+        torchvision \
+        torchaudio \
+        --extra-index-url https://download.pytorch.org/whl/cu117 && \
+    python3 -m pip install --no-cache-dir \
+        accelerate \
+        datasets \
+        hf-doc-builder \
+        huggingface-hub \
+        Jinja2 \
+        librosa \
+        numpy \
+        scipy \
+        tensorboard \
+        transformers
 
 # Set the working directory to /code
 WORKDIR /code
 
 # Copy the current directory contents into the container at /code
-COPY ./requirements_cuda.txt /code/requirements_cuda.txt
 COPY ./requirements.txt /code/requirements.txt
 
 # Install requirements.txt 
-RUN pip install --no-cache-dir --upgrade -r /code/requirements_cuda.txt
 RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
 
 # Set up a new user named "user" with user ID 1000

requirements.txt

@@ -1,6 +1,5 @@
 fastapi==0.74.*
 requests==2.27.*
 sentencepiece==0.1.*
-transformers==4.*
 uvicorn[standard]==0.17.*
 diffusers[torch]==0.13.*

requirements_cuda.txt

@@ -1,4 +0,0 @@
---extra-index-url https://download.pytorch.org/whl/cu117
-torch
-torchvision
-torchaudio