画像生成 AI 入門: Python による拡散モデルの理論と実践

画像生成 AI 入門: Python による拡散モデルの理論と実践#

Section 07. Play with Diffusion Model#

Stable Diffusion を中心とした拡散モデルを用いて、最先端の画像生成技術を実際に動かして実践していきます。

Lecture 26. unCLIP#

OpenAI の DALL-E 2 のベースとなっている unCLIP [Ramesh+ CoRR'22] を再現したもモデルや Stable Diffusion を fine-tuning した Stable unCLIP を用いて、画像生成を実現します。

セットアップ#

GPU が使用できるか確認#

本 Colab ノートブックを実行するために GPU ランタイムを使用していることを確認します。CPU ランタイムと比べて画像生成がより早くなります。以下の nvidia-smi コマンドが失敗する場合は再度講義資料の GPU 使用設定 のスライド説明や Google Colab の FAQ 等を参考にランタイムタイプが正しく変更されているか確認してください。

!nvidia-smi

Sat Jul 22 09:56:34 2023       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.105.17   Driver Version: 525.105.17   CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| 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 T4            Off  | 00000000:00:04.0 Off |                    0 |
| N/A   44C    P8    12W /  70W |      0MiB / 15360MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

利用する Python ライブラリをインストール#

diffusers ライブラリをインストールすることで拡散モデルを簡単に使用できるようにします。diffusers ライブラリを動かす上で必要となるライブラリも追加でインストールします:

transformers: 拡散モデルにおいて核となる Transformer モデルが定義されているライブラリ
accelerate: transformers と連携してより高速な画像生成をサポートするライブラリ

!pip install diffusers==0.16.1
!pip install transformers accelerate

Collecting diffusers==0.16.1
  Downloading diffusers-0.16.1-py3-none-any.whl (934 kB)
?25l     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 0.0/934.9 kB ? eta -:--:--
     ━━━━╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 102.4/934.9 kB 2.9 MB/s eta 0:00:01
     ━━━━━━━━━━━━━━━━━━━━━━━╺━━━━━━━━━━━━━━━ 553.0/934.9 kB 8.1 MB/s eta 0:00:01
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 934.9/934.9 kB 9.0 MB/s eta 0:00:00
?25hRequirement already satisfied: Pillow in /usr/local/lib/python3.10/dist-packages (from diffusers==0.16.1) (8.4.0)
Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from diffusers==0.16.1) (3.12.2)
Collecting huggingface-hub>=0.13.2 (from diffusers==0.16.1)
  Downloading huggingface_hub-0.16.4-py3-none-any.whl (268 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 268.8/268.8 kB 13.9 MB/s eta 0:00:00
?25hRequirement already satisfied: importlib-metadata in /usr/lib/python3/dist-packages (from diffusers==0.16.1) (4.6.4)
Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from diffusers==0.16.1) (1.22.4)
Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from diffusers==0.16.1) (2022.10.31)
Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from diffusers==0.16.1) (2.27.1)
Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.13.2->diffusers==0.16.1) (2023.6.0)
Requirement already satisfied: tqdm>=4.42.1 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.13.2->diffusers==0.16.1) (4.65.0)
Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.13.2->diffusers==0.16.1) (6.0.1)
Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.13.2->diffusers==0.16.1) (4.7.1)
Requirement already satisfied: packaging>=20.9 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.13.2->diffusers==0.16.1) (23.1)
Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->diffusers==0.16.1) (1.26.16)
Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->diffusers==0.16.1) (2023.5.7)
Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.10/dist-packages (from requests->diffusers==0.16.1) (2.0.12)
Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->diffusers==0.16.1) (3.4)
Installing collected packages: huggingface-hub, diffusers
Successfully installed diffusers-0.16.1 huggingface-hub-0.16.4
Collecting transformers
  Downloading transformers-4.31.0-py3-none-any.whl (7.4 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 7.4/7.4 MB 74.0 MB/s eta 0:00:00
?25hCollecting accelerate
  Downloading accelerate-0.21.0-py3-none-any.whl (244 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 244.2/244.2 kB 32.1 MB/s eta 0:00:00
?25hRequirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from transformers) (3.12.2)
Requirement already satisfied: huggingface-hub<1.0,>=0.14.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.16.4)
Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (1.22.4)
Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from transformers) (23.1)
Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (6.0.1)
Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (2022.10.31)
Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from transformers) (2.27.1)
Collecting tokenizers!=0.11.3,<0.14,>=0.11.1 (from transformers)
  Downloading tokenizers-0.13.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (7.8 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 7.8/7.8 MB 113.9 MB/s eta 0:00:00
?25hCollecting safetensors>=0.3.1 (from transformers)
  Downloading safetensors-0.3.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.3 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.3/1.3 MB 61.5 MB/s eta 0:00:00
?25hRequirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.10/dist-packages (from transformers) (4.65.0)
Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from accelerate) (5.9.5)
Requirement already satisfied: torch>=1.10.0 in /usr/local/lib/python3.10/dist-packages (from accelerate) (2.0.1+cu118)
Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.14.1->transformers) (2023.6.0)
Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.14.1->transformers) (4.7.1)
Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch>=1.10.0->accelerate) (1.11.1)
Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=1.10.0->accelerate) (3.1)
Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=1.10.0->accelerate) (3.1.2)
Requirement already satisfied: triton==2.0.0 in /usr/local/lib/python3.10/dist-packages (from torch>=1.10.0->accelerate) (2.0.0)
Requirement already satisfied: cmake in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch>=1.10.0->accelerate) (3.25.2)
Requirement already satisfied: lit in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch>=1.10.0->accelerate) (16.0.6)
Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (1.26.16)
Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2023.5.7)
Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2.0.12)
Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (3.4)
Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=1.10.0->accelerate) (2.1.3)
Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch>=1.10.0->accelerate) (1.3.0)
Installing collected packages: tokenizers, safetensors, transformers, accelerate
Successfully installed accelerate-0.21.0 safetensors-0.3.1 tokenizers-0.13.3 transformers-4.31.0

unCLIP を扱うパイプラインを構築#

DALL-E 2 を始め、unCLIP 自体の事前学習済みモデルは現状公開されておりません。

代わりに有志が unCLIP をベースに学習済みモデルを構築し、UnCLIPPipeline や StableUnCLIPPipeline で簡単に使用できるように整備しています。

まず準備として画像を複数生成した場合に結果を確認しやすいように、画像をグリッド上に表示する関数を以下のように定義します。この関数は 🤗 Hugging Face Stable Diffusion のブログ記事のものを利用しています。

from typing import List
from PIL import Image
from PIL.Image import Image as PilImage

def image_grid(imgs: List[PilImage], rows: int, cols: int) -> PilImage:
    assert len(imgs) == rows * cols

    w, h = imgs[0].size
    grid = Image.new('RGB', size=(cols * w, rows * h))
    grid_w, grid_h = grid.size

    for i, img in enumerate(imgs):
        grid.paste(img, box=(i%cols*w, i//cols*h))
    return grid

今回は unCLIP の論文で DALL-E 2 の生成画像の例としてあげられている以下の vibrant portrait painting of Salvador Dalí with a robotic half face というプロンプトから生成された画像と、unCLIP および Stable unCLIP による生成画像を比較していきます。

from diffusers.utils import load_image

prompt = "vibrant portrait painting of Salvador Dalí with a robotic half face"

image_dalle2 = load_image("https://ar5iv.labs.arxiv.org/html/2204.06125/assets/figures/banner_samples/Vibrant_portrait_painting_of_salvador_dali_with_a_robotic_half_face.jpg")

print(f"Prompt: {prompt}")
print(f"Original image size: {image_dalle2.size}")
image_dalle2 = image_dalle2.resize((512, 512))
image_dalle2

Prompt: vibrant portrait painting of Salvador Dalí with a robotic half face
Original image size: (1024, 1024)

../_images/2bf567d7c2981658f9cc1308549168d9610a73d0c46025903fa09215d94c2d77.png

UnCLIPPipeline による画像生成#

前述の通り、DALL-E 2 自体の学習済みの重みは公開されていません。

かわりに、unCLIP をベースとした Karlo というモデルが kakaobrain という企業から公開されています。本セクションでは kakaobrain が公開している kakaobrain/karlo-v1-alpha を使用して UnCLIPPipeline によるパイプラインの動作を確認します。

import torch
from diffusers import UnCLIPPipeline

model_id = "kakaobrain/karlo-v1-alpha"
data_type = torch.float16

pipe = UnCLIPPipeline.from_pretrained(
    model_id, torch_dtype=data_type,
)
pipe = pipe.to("cuda")

UnCLIPPiepline を用いて画像を生成します。unCLIP は prior と decoder の 2 つのモジュールの組み合わせからなるモデルでした。ここでは prior_num_inference_steps と decoder_num_inference_steps を指定して、それぞれのモジュールにおける diffusion step 数を指定可能です。更には super_res_num_inference_steps によって超解像用のモジュールにおける diffusion step も指定可能です。

generator = torch.Generator().manual_seed(19950815)

image_kalro = pipe(
    prompt,
    prior_num_inference_steps=25,
    decoder_num_inference_steps=25,
    super_res_num_inference_steps=7,
    generator=generator,
).images[0]

print(f"Prompt: {prompt}")
print(f"Original image size: {image_kalro.size}")
image_kalro = image_kalro.resize((512, 512))
image_kalro

Prompt: vibrant portrait painting of Salvador Dalí with a robotic half face
Original image size: (256, 256)

../_images/41e05e4c04f96655a9d412ffb63a2f443420e2e2d81be09f1d35a87e9e420435.png

StableUnCLIPPipeline による画像生成#

StableDiffusion の学習済みパラメータを unCLIP へ転用する試みがあります。

Stable unCLIP は Stable Diffusion v2.1 を fine-tuning したもので、CLIP による画像・テキスト両者で条件付けが可能になっています。本セクションでは stabilityai/stable-diffusion-2-1-unclip-small を使用して StableUnCLIPPipeline によるパイプラインの動作を確認します。

Stable unCLIP は推論時に画像埋め込みベクトルにどの程度ノイズを加えるかを決定する noise_level を設定可能です。この noise_level が高いほど、最終的な生成画像の多様性が高まります。デフォルトではノイズを追加しない noise_level = 0 が設定されています。

Stable unCLIP は上述の KakaoBrain の Karlo をベースに利用しています。以下のようにして prior となるモデル群を読み込んでいきます。

from diffusers import UnCLIPScheduler, DDPMScheduler, StableUnCLIPPipeline
from diffusers.models import PriorTransformer
from transformers import CLIPTokenizer, CLIPTextModelWithProjection

prior_model_id = "kakaobrain/karlo-v1-alpha"
prior = PriorTransformer.from_pretrained(prior_model_id, subfolder="prior", torch_dtype=data_type)

prior_text_model_id = "openai/clip-vit-large-patch14"
prior_tokenizer = CLIPTokenizer.from_pretrained(prior_text_model_id)
prior_text_model = CLIPTextModelWithProjection.from_pretrained(prior_text_model_id, torch_dtype=data_type)
prior_scheduler = UnCLIPScheduler.from_pretrained(prior_model_id, subfolder="prior_scheduler")
prior_scheduler = DDPMScheduler.from_config(prior_scheduler.config)

Karlo では prior として OpenAI の CLIP ViT-L/14 が使用されていますが、Stable unCLIP の学習済みモデル stabilityai/stable-diffusion-2-1-unclip-small でも同様に使用しています。似た学習済みモデルとして stabilityai/stable-diffusion-2-1-unclip がありますが、こちらは OpenCLIP ViT-H で学習されているため、今回 Karlo に合わせて前者の学習済みモデルを使用しています。

stable_unclip_model_id = "stabilityai/stable-diffusion-2-1-unclip-small"

pipe = StableUnCLIPPipeline.from_pretrained(
    stable_unclip_model_id,
    torch_dtype=data_type,
    variant="fp16",
    prior_tokenizer=prior_tokenizer,
    prior_text_encoder=prior_text_model,
    prior=prior,
    prior_scheduler=prior_scheduler,
)
pipe = pipe.to("cuda")

StableDiffusionUnCLIPPipeline を用いて画像を生成します。このパイプラインでは 2023/07/20 現在、prior に対する指定のみが存在しています（参考）。今回は以下のようにして prior_num_inference_steps を設定して画像を生成させてみます。

generator = torch.Generator().manual_seed(19950815)

image_sd_unclip = pipe(
    prompt,
    prior_num_inference_steps=25,
    generator=generator,
).images[0]

print(f"Prompt: {prompt}")
print(f"Original image size: {image_sd_unclip.size}")
image_sd_unclip = image_sd_unclip.resize((512, 512))
image_sd_unclip

Prompt: vibrant portrait painting of Salvador Dalí with a robotic half face
Original image size: (768, 768)

../_images/39b23ed8c78ed7a048adca433363a411866e2467df3e359cd89263db3cf408f5.png

では、DALL-E 2 による生成画像と、これまで unCLIP Karlo と Stable unCLIP で生成した画像を比較してみます。

image_grid([image_dalle2, image_kalro, image_sd_unclip], rows=1, cols=3)

../_images/ae571307f274015beeda1948a90ed422ca041ff7107f4de991589ab1a9d2aa0c.png

それぞれ違った良さがありそうです。解像度の観点だと、DALL-E 2 は 1024 x 1024、Karlo は 256 x 256、Stable unCLIP は 768 x 768 でした。DALL-E 2 以外はモデルの重みが公開されているので、研究的な意味でも扱いやすくなっています。