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

ARG OUTPUT_DIR=fries-style
ARG INPUT_DS=nroggendorff/fries
ARG BASE=glides/illustriousxl

ENV OUTPUT_DIR=${OUTPUT_DIR} \
  INPUT_DS=${INPUT_DS} \
  BASE=${BASE}

RUN mkdir -p /.cache /app/${OUTPUT_DIR} && chmod -R 777 /.cache /app/${OUTPUT_DIR}

RUN accelerate config default

CMD bash -c "accelerate launch \
  --multi_gpu \
  --mixed_precision=bf16 \
  train.py \
    --gradient_accumulation_steps=1 \
    --learning_rate=3e-5 \
    --gradient_checkpointing \
    --max_train_steps=7000 \
    --lr_scheduler=cosine \
    --lr_warmup_steps=200 \
    --center_crop \
    --random_flip \
    --pretrained_model_name_or_path=${BASE} \
    --dataset_name=${INPUT_DS} \
    --resolution=1024 \
    --train_batch_size=4 \
    --output_dir=${OUTPUT_DIR} \
    --checkpointing_steps=100000 \
    --validation_prompt='art' \
    --num_validation_images=4 \
    --validation_epochs=100 \
    --push_to_hub"