Update Dockerfile

Dockerfile

@@ -22,6 +22,7 @@ RUN pip install --no-cache-dir --upgrade pip && \
     ftfy \
     jinja2 \
     torchvision \
+    bitsandbytes \
     torch \
     --extra-index-url https://download.pytorch.org/whl/cu121
 
@@ -39,20 +40,4 @@ RUN mkdir /app/${OUTPUT_DIR} && chmod 777 /app/${OUTPUT_DIR}
 
 RUN accelerate config default
 
-CMD bash -c "accelerate launch train.py \
-  --mixed_precision=bf16 \
-  --gradient_accumulation_steps=1 \
-  --learning_rate=3e-5 \
-  --gradient_checkpointing \
-  --max_train_steps=2000 \
-  --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"
+CMD bash -c "accelerate launch --multi_gpu train.py"