Getting Started with Textual Inversion in Google Colab

I have summarized my experience trying out "Textual Inversion" using "textual_inversion.py" from "Diffusers".

1. Textual Inversion

"Textual Inversion" is a method for fine-tuning using 3 to 5 images. It allows you to teach the "Stable Diffusion" model unique objects or art styles.

In this guide, we will use "textual_inversion.py" from "Diffusers". It is quite nice because it requires little code and consumes minimal GPU memory.

2. License Verification

Access the model card below, verify the license, click "Access Repository," log in to "Hugging Face" (create an account if you don't have one), and agree to the terms.

3. Preparing Training Images

Fine-tuning requires 3 to 5 training images.

4. Executing Fine-tuning

The steps to perform fine-tuning in Colab are as follows:

(1) Open a new Colab notebook and select "GPU" (High-RAM is not required).
Confirm that the GPU memory is 16GB or more (the following shows 16280MiB).

# GPUの確認 
!nvidia-smi

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.32.03    Driver Version: 460.32.03    CUDA Version: 11.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  Tesla P100-PCIE...  Off  | 00000000:00:04.0 Off |                    0 |
| N/A   33C    P0    27W / 250W |      0MiB / 16280MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

(2) Mount Google Drive and move to the working folder.
Create a "work" folder directly under your Google Drive root in advance.

# Googleドライブのマウントと作業フォルダへの移動
from google.colab import drive
drive.mount('/content/drive')
%cd '/content/drive/My Drive/work'

(3) Install packages.

# パッケージのインストール 
!pip install diffusers[training] accelerate transformers

(4) Log in to HuggingFace.
Copy the token from the link, paste it into the text field, and click the Login button.

# HuggingFaceにログイン
from huggingface_hub import notebook_login
notebook_login()

(5) Set up "HuggingFace Accelerate".
Configure for single GPU (0, 0, NO, NO, NO).

# HuggingFace Accelerateの初期化
!accelerate config

In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 0
Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU [4] MPS): 0
Do you want to run your training on CPU only (even if a GPU is available)? [yes/NO]:NO
Do you want to use DeepSpeed? [yes/NO]: NO
Do you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16]: NO

I have briefly explained "HuggingFace Accelerate" below.

(6) Place training images in the "inputs" folder directly under the "work" folder.

(7) Clone the diffusers repository.
We will use the examples within the diffusers repository.

# diffusersリポジトリのクローン
!git clone https://github.com/huggingface/diffusers.git

(8) Executing fine-tuning.
In this instance, we will train an "art style." The placeholder token is <sorami>, and the initializer token is set to comic. It took about 2 hours on a P100.

# 学習の実行
!accelerate launch ./diffusers/examples/textual_inversion/textual_inversion.py \
  --pretrained_model_name_or_path="CompVis/stable-diffusion-v1-4" \
  --use_auth_token \
  --train_data_dir="inputs" \
  --learnable_property="style" \
  --placeholder_token="<sorami>" \
  --initializer_token="comic" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --max_train_steps=3000 \
  --learning_rate=5.0e-04 \
  --scale_lr \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --output_dir="textual_inversion_sorami"

"--learnable_property" is used to specify whether to train an art style (style) or an object (object).
"--placeholder_token" is used to specify the placeholder token, which is the word representing the new concept learned. It is recommended to add <> to avoid conflicts with other words.
"--initializer_token" is used to specify the initializer token, which serves as the initial value for the concept being learned at the start of training.

5. Executing inference

The steps to perform inference in Colab are as follows.

(1) Preparing the StableDiffusion pipeline.

import torch
from diffusers import StableDiffusionPipeline

# StableDiffusionパイプラインの準備
pipe = StableDiffusionPipeline.from_pretrained(
    "./textual_inversion_sorami",
    torch_dtype=torch.float16
    ).to("cuda")

(2) Executing inference.
This time, I generated an image of a cute cat ear maid of <sorami> style.

from torch import autocast

# プロンプト
prompt = "cute cat ear maid of <sorami> style"

# 推論の実行
with autocast("cuda"):
    image = pipe(
        prompt, 
        num_inference_steps=50, 
        guidance_scale=7.5
        ).images[0]
image.save("output.png")

(3) Checking the generated image.
Display the file list using the folder icon on the far left and double-click output.png.

6. API Reference

The parameters for "textual_inversion.py" are as follows.

7. References

・Resurrecting Ugly Sonic with Textual Inversion

・Comparison between DreamBooth and Textual Inversion

・09/18/22 Survey on additional training for Stable Diffusion

・09/26/22 Record of additional learning for Stable Diffusion

8. Related

The training method in the "Textual Inversion" repository is introduced below.