diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..c60f026 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,4 @@ +./venv +./danbooru-aesthetic +./logs +*.ckpt diff --git a/.gitignore b/.gitignore index 89d4032..680b96c 100644 --- a/.gitignore +++ b/.gitignore @@ -40,6 +40,7 @@ lib64/ parts/ sdist/ var/ +venv/ wheels/ share/python-wheels/ *.egg-info/ @@ -54,4 +55,4 @@ MANIFEST /src/ #Obsidian -.obsidian/ \ No newline at end of file +.obsidian/ diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..8ad4af7 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,10 @@ +FROM pytorch/pytorch:latest + +RUN apt update && \ + apt install -y git curl unzip vim && \ + pip install git+https://github.com/derfred/lightning.git@waifu-1.6.0#egg=pytorch-lightning +RUN mkdir /waifu +COPY . /waifu/ +WORKDIR /waifu +RUN grep -v pytorch-lightning requirements.txt > requirements-waifu.txt && \ + pip install -r requirements-waifu.txt diff --git a/docs/en/training/executing.md b/docs/en/training/executing.md index 7ccb57f..4bfafac 100644 --- a/docs/en/training/executing.md +++ b/docs/en/training/executing.md @@ -1,13 +1,51 @@ -# 1. Executing +# 3. Executing -## Installation +There are two modes of executing the training: +1. Using docker image. This is the fastest way to get started. +2. Using system python install. Allows more customization. + +Note: You will need to provide the initial checkpoint for resuming the training. This must be a version with the full EMA. Otherwise you will get this error: +``` +RuntimeError: Error(s) in loading state_dict for LatentDiffusion: + Missing key(s) in state_dict: "model_ema.diffusion_modeltime_embed0weight", "model_ema.diffusion_modeltime_embed0bias".... (Many lines of similar outputs) +``` + +## 1. Using docker image + +An image is provided at `ghcr.io/derfred/waifu-diffusion`. Execute it using by adjusting the NUM_GPU variable: +``` +docker run -it -e NUM_GPU=x ghcr.io/derfred/waifu-diffusion +``` + +Next you will want to download the starting checkpoint into the file `model.ckpt` and copy the training data in the directory `/waifu/danbooru-aesthetic`. + +Finally execute the training using: +``` +sh train.sh -t -n "aesthetic" --resume_from_checkpoint model.ckpt --base ./configs/stable-diffusion/v1-finetune-4gpu.yaml --no-test --seed 25 --scale_lr False --data_root "./danbooru-aesthetic" +``` + +## 2. system python install First install the dependencies: ```bash pip install -r requirements.txt ``` -## Executing -```bash -sh train.sh +Next you will want to download the starting checkpoint into the file `model.ckpt` and copy the training data in the directory `/waifu/danbooru-aesthetic`. + +Also you will need to edit the configuration in `./configs/stable-diffusion/v1-finetune-4gpu.yaml`. In the `data` section (around line 70) change the `batch_size` and `num_workers` to the number of GPUs you are using: ``` +data: + target: main.DataModuleFromConfig + params: + batch_size: 4 + num_workers: 4 + wrap: false +``` + +Finally execute the training using the following command. You need to adjust the `--gpu` parameter according to your GPU settings. +```bash +sh train.sh -t -n "aesthetic" --resume_from_checkpoint model.ckpt --base ./configs/stable-diffusion/v1-finetune-4gpu.yaml --no-test --seed 25 --scale_lr False --data_root "./danbooru-aesthetic" --gpu=0,1,2,3, +``` + +In case you get an error stating `KeyError: 'Trying to restore optimizer state but checkpoint contains only the model. This is probably due to ModelCheckpoint.save_weights_only being set to True.'` follow these instructions: https://discord.com/channels/930499730843250783/953132470528798811/1018668937052962908 diff --git a/requirements.txt b/requirements.txt index 090b52d..5e1c8ab 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,10 +1,10 @@ -numpy==1.19.2 +numpy==1.21.6 albumentations==0.4.3 -opencv-python==4.1.2.30 +opencv-python pudb==2019.2 imageio==2.9.0 imageio-ffmpeg==0.4.2 -pytorch-lightning==1.6.5 +pytorch-lightning==1.6.0 omegaconf==2.1.1 test-tube>=0.7.5 streamlit>=0.73.1 diff --git a/train.sh b/train.sh index 13144e9..32314a4 100644 --- a/train.sh +++ b/train.sh @@ -1 +1,15 @@ -python3 main.py --train --resume model.ckpt --base ./configs/stable-diffusion/v1-finetune-4gpu.yaml --no-test --seed 25 --scale_lr False --gpus 0,1,2,3 +#!/bin/bash + +ARGS="" +if [ ! -z "$NUM_GPU" ]; then + ARGS="--gpu=" + for i in $(seq 0 $((NUM_GPU-1))) + do + ARGS="$ARGS$i," + done + + sed -i "s/batch_size: 4/batch_size: $NUM_GPU/g" ./configs/stable-diffusion/v1-finetune-4gpu.yaml + sed -i "s/num_workers: 4/num_workers: $NUM_GPU/g" ./configs/stable-diffusion/v1-finetune-4gpu.yaml +fi + +python3 main.py $ARGS "$@"