実装例：DreamBoothによるstable diffusion再学習を速攻で試す流れ

# はじめに

DreaBoothはgoogleらがCVPR2023で発表したdiffusion model系の再学習のしくみ。

[0] N. Ruiz, et. al. "DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation" CVPR2023. 

再学習させるデータに過適合するのを防ぎ、また構文の意味を損なわないように配慮して再学習させる。

- dreamboothのロジックは以下にざっくりまとめまています
https://qiita.com/masataka46/items/134e08af7fcb4d7a726c

既に多くの解説が出回っているので、ここではハマりどころを重点的にまとめる。

# 1. 環境構築
以下の２つを試した。
- 自分のRTX3090パソコンにDockerで環境構築
- colabで環境構築

colab系は既に多くの資料があるので、ここでは前者を記述。

## 1.1 Ubuntuに基本的なものを入れる
- Nvidia driver
- GPU対応のDocker
- その他、pycharmやvs-code、iterm など
ここでは省略。

## 1.2 Dockerを立ち上げる
例えば以下のようなDockerfileを用いる。

```Dockerfile
FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04

ENV TZ=Asia/Tokyo
RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime
RUN echo $TZ > /etc/timezone

RUN apt-get update
RUN apt-get install -y apt-file
RUN apt-file update

RUN apt-get install -y git
RUN apt-get install -y wget
RUN apt-get install -y curl
RUN apt-get install -y vim

# python
RUN apt-get install -y python3-pip

# alias
RUN echo alias python="python3" >> /root/.bashrc
RUN echo alias pip="pip3" >> /root/.bashrc

# Install python packages
RUN pip install --upgrade pip
RUN pip install opencv-python
RUN pip install scipy
RUN pip install scikit-learn
RUN pip install poetry
# Install VIM
RUN apt-get update
RUN apt-get install -y apt-file
RUN apt-file update
RUN apt-get install -y vim

# insatll langchain, and else
RUN pip install langchain==0.1.14
RUN pip install faiss-gpu tiktoken sentence_transformers
RUN pip install wandb
RUN CMAKE_ARGS="-DLLAMA_CUBLAS=on" FORCE_CMAKE=1 pip install llama-cpp-python
RUN pip install flash_attn
RUN pip install openai==0.28.0
RUN pip install pandas==2.2.2
RUN pip install -U langchain-community
RUN pip install openpyxl==3.1.2
RUN pip install python-docx==1.1.2
RUN pip install pypdf==5.1.0
RUN pip install git+https://github.com/huggingface/diffusers
RUN pip install accelerate>=0.16.0
RUN pip install torchvision==0.20.1
RUN pip install ftfy
RUN pip install tensorboard
RUN pip install peft==0.7.0

WORKDIR /root/share
EXPOSE 8080/TCP
```
必要ないものも多分に含まれています。
diffusersもinstallしています。

```bash
# in the case to use cuda
IMAGE_REPOSITORY=dreambooth
IMAGE_TAG=latest
IMAGE_FULLNAME=${IMAGE_REPOSITORY}:${IMAGE_TAG}

# docker build if changed.
docker build -t ${IMAGE_FULLNAME} .

# allow display connection
xhost local:

# run container
docker run \
--interactive \
--tty \
--rm \
--mount=type=bind,src="$(pwd)",dst=/root/share \
--mount=type=bind,src=/etc/group,dst=/etc/group,readonly \
--mount=type=bind,src=/etc/passwd,dst=/etc/passwd,readonly \
--mount=type=bind,src=/tmp/.X11-unix,dst=/tmp/.X11-unix \
$( if [ -e $HOME/.Xauthority ]; then echo "--mount=type=bind,src=$HOME/.Xauthority,dst=/root/.Xauthority"; fi ) \
--gpus=all \
--env=QT_X11_NO_MITSHM=1 \
--env=DISPLAY=${DISPLAY} \
--env=NVIDIA_DRIVER_CAPABILITIES=all \
--net=host \
--publish=8080:8080 \
--publish=8080:8080/udp \
--name=${IMAGE_REPOSITORY}_${IMAGE_TAG}_$(date "+%Y_%m%d_%H%M%S") \
${IMAGE_FULLNAME}
```

このDockerRun.shとDockerfileを同じディレクトリ内に配置し、Dockerをbuild, runする。

```bash
# DockerRun.sh, Dockerfileがある階層にて
bash DockerRun.sh
```

ポイントは、GPUで演算させる場合、`torch_xla` をインストールしないこと。このパッケージはTPU等利用する場合に必要のようだが、GPUで演算させる場合はエラーとなる。

# 2. accelerateの設定をする
```bash
accelerate config
```
でGPU周りなどの設定をする。ここが一番のハマりどころ。いろいろ聞かれるが、設定を間違えると動作しない。

私の場合、
- 自分のlocalパソコン
- GPUは１個

なので、

```shell
In which compute environment are you running?
Please select a choice using the arrow or number keys, and selecting with enter
 ➔  This machine
    AWS (Amazon SageMaker)
```
This machineでリターン。

```shell
Which type of machine are you using?
Please select a choice using the arrow or number keys, and selecting with enter
 ➔  No distributed training
    multi-CPU
    multi-XPU
    multi-GPU
    multi-NPU
    multi-MLU
    multi-MUSA
    TPU
```
私の場合、GPU１個なので、いちばん上の「No distributed training」でリターン。

```shell
Do you want to run your training on CPU only (even if a GPU / Apple Silicon / Ascend NPU device is available)? [yes/NO]:
```
GPU使うので、NO。

```shell
Do you wish to optimize your script with torch dynamo?[yes/NO]:
```
よくわからないのでNO。

```shell
Do you want to use DeepSpeed? [yes/NO]:
```
よくわからないのでNO。

```shell
What GPU(s) (by id) should be used for training on this machine as a comma-seperated list? [all]:
```
ここではGPUのIDを入力する。私の場合、
`nvidia-smi` で

```shell
nvidia-smi
Wed Jan  1 09:20:16 2025
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.154.05             Driver Version: 535.154.05   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  NVIDIA GeForce RTX 3090        On  | 00000000:05:00.0  On |                  N/A |
|  0%   43C    P8              28W / 350W |    310MiB / 24576MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+

+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
+---------------------------------------------------------------------------------------+
```
と表示されたので、GPU欄に記載の０を入力。


```shell
Would you like to enable numa efficiency? (Currently only supported on NVIDIA hardware). [yes/NO]:
```
よくわからないのでNO。

```shell
Do you wish to use mixed precision?
Please select a choice using the arrow or number keys, and selecting with enter
 ➔  no
    fp16
    bf16
    fp8
```
学習時のbit数のことか？float32で学習させる気なので最上部でリターン。
うまく設定できれば

```shell
accelerate configuration saved at /root/.cache/huggingface/accelerate/default_config.yaml
```
と表示される。

# 3. 学習させる画像を用意する
適当な画像を適当なディレクトリに配置する。私は特定の種類の犬画像を５種類用意した。

# 4. train_dreambooth.pyを持ってくる
下記リンクの `diffusers/examples/dreambooth/train_dreambooth.py` をコピペして適当な位置に配置。

https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth.py

# 5. outputのディレクトリを作成
（自動で作成してくれるのかもしれないが）outputのディレクトリを作成する。
例えば
```bash
mkdir ../output
```

# 6. 環境変数の設定
今回、diffusionのv1-4を使う。
```
export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export INSTANCE_DIR="PATH/TO/IMAGE_DIR"
export OUTPUT_DIR="PATH/TO/OUTPUT_DIR"
```

# 7. 学習開始
以下のようなコマンドで学習を開始。
```shell
accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a photo of sks dog" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=400
```
instance_promptは適当なものに設定する。その他、目的に合わせて諸々ハイパラを変更する。

以下のようにVRAM 19GB近く使ってました。減らすには上記の量子化すればよい？
```shell
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.154.05             Driver Version: 535.154.05   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  NVIDIA GeForce RTX 3090        On  | 00000000:05:00.0  On |                  N/A |
| 43%   61C    P2             303W / 350W |  18953MiB / 24576MiB |     92%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+

+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    0   N/A  N/A      5384      G   /usr/lib/xorg/Xorg                          168MiB |
|    0   N/A  N/A      5517      G   /usr/bin/gnome-shell                         75MiB |
|    0   N/A  N/A      6653      G   ...2509800,12662543674627072428,262144       55MiB |
|    0   N/A  N/A    762978      C   /usr/bin/python3                          18638MiB |
+---------------------------------------------------------------------------------------+
```
10分程度で学習は終了する。

outputしたディレクトリに
```
feature_extractor  logs  model_index.json  safety_checker  scheduler  text_encoder  tokenizer  unet  vae
```
のように諸々保存される。

# 8. 学習したモデルを用いて推論
今回、以下のリンクにある `diffusers/src/diffusers/pipelines/`のREADME.mdに従って推論させる。

https://github.com/huggingface/diffusers/tree/main/src/diffusers/pipelines

以下の
```python
# make sure you're logged in with `huggingface-cli login`
from diffusers import StableDiffusionPipeline, LMSDiscreteScheduler

pipe = StableDiffusionPipeline.from_pretrained("../output/")
pipe = pipe.to("cuda")

prompt = "a photo of sks dog"
image = pipe(prompt).images[0]

image.save("tmp_dog.png")
```
を記述した `infer.py`を作成する。
- 冒頭pipeのpretrain読み込む部分は学習時に `OUTPUT_DIR` で指定したディレクトリを指定
- promptは学習時に指定したpromptを指定
- 保存する画像名はテキトーに指定


```shell
python infer.py
```
で推論させる。それっぽい画像が生成される。