FROM pytorch/pytorch:2.4.0-cuda12.1-cudnn9-runtime

WORKDIR /app

COPY . .

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" \
    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

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

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"