Update Dockerfile

Dockerfile

@@ -1,41 +1,24 @@
-FROM pytorch/pytorch:2.4.0-cuda12.1-cudnn9-runtime
+FROM nvidia/cuda:12.3.2-cudnn9-devel-ubuntu22.04
 
-WORKDIR /app
+RUN apt-get update && \
+    apt-get install -y python3.10 python3.10-dev python3.10-distutils curl git && \
+    update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 && \
+    curl -sS https://bootstrap.pypa.io/get-pip.py | python3.10
 
+RUN python3 -m pip install --upgrade pip
+
+WORKDIR /app
 COPY . .
 
-RUN apt-get update && apt-get install -y git && rm -rf /var/lib/apt/lists/*
-RUN apt-get install -y python3.10 python3.10-dev python3.10-distutils && \
-    ln -sf /usr/bin/python3.10 /usr/bin/python3 && \
-    curl -sS https://bootstrap.pypa.io/get-pip.py | python3.10
+RUN pip install --no-cache-dir torch==2.4.0+cu123 torchvision==0.19.0 --extra-index-url https://download.pytorch.org/whl/cu121
 
-RUN pip install --no-cache-dir --upgrade pip && \
-  pip install --no-cache-dir \
-    "git+https://github.com/huggingface/diffusers.git@main#egg=diffusers" && \
-  pip install --no-cache-dir --prefer-binary \
-    transformers \
-    accelerate \
-    safetensors \
-    opencv-python \
-    datasets \
-    dataclasses \
-    scikit-learn \
-    huggingface-hub \
-    tensorboard \
-    peft \
-    ftfy \
-    jinja2 \
-    torchvision \
-    bitsandbytes \
-    torch \
-    --extra-index-url https://download.pytorch.org/whl/cu121 && \
-  pip install --no-cache-dir \
+RUN pip install --no-cache-dir \
     "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.3.11/flash_attn-2.8.0+cu121torch2.4-cp310-cp310-linux_x86_64.whl"
 
-RUN mkdir /.cache && chmod 777 /.cache && chmod 777 /app
-
-RUN mkdir /app/${OUTPUT_DIR} && chmod 777 /app/${OUTPUT_DIR}
+RUN pip install --no-cache-dir \
+    diffusers transformers accelerate safetensors opencv-python datasets dataclasses \
+    scikit-learn huggingface-hub tensorboard peft ftfy jinja2 bitsandbytes
 
-RUN accelerate config default
+RUN mkdir /.cache && chmod 777 /.cache && chmod 777 /app
 
-CMD bash -c "python train.py"
+CMD ["python3", "train.py"]