Google Colab ではじめる Trinart Stable Diffusion

2022年9月11日 09:32

「Trinart Stable Diffusion」が二次元キャラの画像生成が得意らしいということで試してみました。

・Trinart Stable Diffusion v2
・Stable Diffusion v1.4

1. Trinart Stable Diffusion

「Trinart Stable Diffusion」は、「Stable Diffusion」を二次元イラストのデータセットでファインチューニングして作られた画像生成AIです。Twitterボット「とりんさまAI」@trinsama ( https://twitter.com/trinsama ) で学習しています。

2. Colabでの実行

Colabでの実行手順は、次のとおりです。

(1) 新規のColabのノートブックを開き、メニュー「編集 → ノートブックの設定」で「GPU」を選択。

# 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   35C    P0    27W / 250W |      0MiB / 16280MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

(3) Googleドライブのマウント。
あらかじめGoogleドライブのマイドライブ直下に「work」フォルダを作成しておきます。

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

(4) パッケージのインストール

# パッケージのインストール
!pip install -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
!pip install pytorch_lightning tensorboard==2.8 omegaconf einops taming-transformers==0.0.1 clip transformers kornia test-tube
!pip install diffusers invisible-watermark

(5) StableDiffusionのインストール。

# StableDiffusionのインストール
!git clone https://github.com/CompVis/stable-diffusion.git
%cd stable-diffusion
!pip install -e .
%mkdir outputs

(6) Trinartのチェックポイントのダウンロード。
ファイルが大きいので、「git lfs install」も必要でした。

# モデルのダウンロード
!git lfs install
!git clone https://huggingface.co/naclbit/trinart_stable_diffusion_v2

・trinart_stable_diffusion_v2
　・trinart2_step60000.ckpt
　・trinart2_step95000.ckpt
　・trinart2_step115000.ckpt

(7) テキストから画像の生成。
今回は、猫耳メイド(cute cat ear maid)で画像生成してみました。

# テキストからの画像生成
!python scripts/txt2img.py \
    --plms \
    --ckpt ./trinart_stable_diffusion_v2/trinart2_step115000.ckpt \
    --skip_grid \
    --n_samples 1  \
    --n_iter 1 \
    --outdir outputs \
    --ddim_steps 100 \
    --prompt "cute cat ear maid"

(8) 生成した画像の確認。
左端のフォルダアイコンでファイル一覧を表示し、「work/stable-diffusion/outputs」フォルダ内の画像ファイルをダブルクリックします。

3. 関連

【おまけ】 APIリファレンス

txt2img.pyのパラメータは、次のとおりです。

・-h : ヘルプ
・--prompt [PROMPT] : プロンプト
・--outdir [OUTDIR] : 出力フォルダ
・--skip_grid : グリッドを保存せず、個々の画像のみを保存
・--skip_save : 個々の画像を保存せず。速度測定で利用
・--plms : palmsサンプリングの利用
・--ddim_eta [ETA]: ddim eta (eta=0.0 は決定論的サンプリング)
・--ddim_steps [STEPS] : ddim サンプリングステップ数 (50〜100程度)
・--H [H] : 画像の高さ
・--W [W] : 画像の幅
・--C [C] : 潜在チャネル
・--f [F] : ダウンサンプリング係数
・--n_samples [N_SAMPLES] : バッチサイズ
・--n_iter [N_ITER]: 出力画像数
・--n_rows [ROWS]: 行数
・--scale [SCALE] : プロンプトに従う度合い (7〜11程度)
・--from-file [FROM_FILE] : プロンプトファイル
・--config [CONFIG] : コンフィグファイルのパス (v1-inference.yaml)
・--ckpt [CKPT] : 事前学習モデルのパス (sd-v1-4-full-ema.ckpt)
・--seed [SEED] : 乱数シード
・--precision {full,autocast} : 精度評価