File size: 1,391 Bytes
96e5830
e5432d0
2f99d3e
 
 
 
41f9ba5
 
e5432d0
2addcb5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2f99d3e
 
 
61dc47b
 
2addcb5
 
 
 
 
 
65b4e82
2f99d3e
 
 
e5432d0
 
 
 
 
61dc47b
e5432d0
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
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"