Update Dockerfile

Dockerfile

@@ -1,10 +1,36 @@
 FROM python:latest
 
+
+ARG OUTPUT_DIR=zelda-style
+ARG INPUT_DS=nroggendorff/zelda-style
+ARG BASE=nroggendorff/zelda-style-diffusion
+
+ENV OUTPUT_DIR=${OUTPUT_DIR} \
+  INPUT_DS=${INPUT_DS} \
+  BASE=${BASE}
+
 WORKDIR /app
 
 COPY . .
 
-RUN pip install --no-cache-dir -r requirements.txt
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir \
+      "git+https://github.com/huggingface/diffusers.git@main#egg=diffusers" \
+      transformers \
+      accelerate \
+      safetensors \
+      opencv-python \
+      datasets \
+      dataclasses \
+      scikit-learn \
+      huggingface-hub \
+      tensorboard \
+      peft \
+      ftfy \
+      jinja2 \
+      torchvision \
+      torch \
+      --extra-index-url https://download.pytorch.org/whl/cu121
 
 RUN mkdir /.cache && chmod 777 /.cache
 
@@ -12,4 +38,20 @@ RUN mkdir /app/zelda-xl && chmod 777 /app/zelda-xl
 
 RUN accelerate config default
 
-CMD ["accelerate", "launch", "train.py", "--mixed_precision=fp16", "--learning_rate=1e-5", "--gradient_checkpointing", "--max_train_steps=360", "--lr_scheduler=constant", "--lr_warmup_steps=0", "--use_ema", "--center_crop", "--random_flip", "--pretrained_model_name_or_path=nroggendorff/animexl", "--dataset_name=nroggendorff/zelda", "--resolution=1024", "--train_batch_size=1", "--output_dir=zelda-xl", "--push_to_hub"]
+CMD bash -c "accelerate launch train.py \
+  --mixed_precision=bf16 \
+  --gradient_accumulation_steps=1 \
+  --learning_rate=3e-5 \
+  --gradient_checkpointing \
+  --max_train_steps=3000 \
+  --lr_scheduler=cosine_with_restarts \
+  --lr_warmup_steps=200 \
+  --center_crop \
+  --random_flip \
+  --pretrained_model_name_or_path=${BASE} \
+  --dataset_name=${INPUT_DS} \
+  --resolution=1024 \
+  --train_batch_size=20 \
+  --output_dir=${OUTPUT_DIR} \
+  --checkpointing_steps=100000 \
+  --push_to_hub"