Update Dockerfile

Dockerfile

@@ -8,7 +8,8 @@ RUN apt-get update && apt-get install -y git && rm -rf /var/lib/apt/lists/*
 
 RUN pip install --no-cache-dir --upgrade pip && \
   pip install --no-cache-dir \
-    "git+https://github.com/huggingface/diffusers.git@main#egg=diffusers" \
+    "git+https://github.com/huggingface/diffusers.git@main#egg=diffusers" && \
+  pip install --no-cache-dir --prefer-binary \
     transformers \
     accelerate \
     safetensors \
@@ -24,18 +25,12 @@ RUN pip install --no-cache-dir --upgrade pip && \
     torchvision \
     bitsandbytes \
     torch \
-    --extra-index-url https://download.pytorch.org/whl/cu121
+    --extra-index-url https://download.pytorch.org/whl/cu121 && \
+  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
 
-ARG OUTPUT_DIR=fry-diffusion
-ARG INPUT_DS=nroggendorff/fries
-ARG BASE=glides/illustriousxl
-
-ENV OUTPUT_DIR=${OUTPUT_DIR} \
-  INPUT_DS=${INPUT_DS} \
-  BASE=${BASE}
-
 RUN mkdir /app/${OUTPUT_DIR} && chmod 777 /app/${OUTPUT_DIR}
 
 RUN accelerate config default