本記事では、clip2latentと呼ばれる機械学習手法を用いて、テキストから顔画像や風景画像を生成する方法をご紹介します。

出典: justinpinkney/clip2latent

clip2latent

概要

clip2latentは、事前トレーニング済みCLIPとStyleGANを用いてテキストから画像を生成する際に、拡散モデルによるサンプリングを適用した画像生成手法です。

CLIP embeddingsを入力とする拡散モデルをトレーニングし、事前トレーニング済みのStyleGANの潜在ベクトルをサンプリングすることによりテキストに応じた画像を生成します。

この構成により、条件付き拡散モデルを画像とキャプションのペアデータセットでトレーニングする必要性をなくしトレーニングコストを低減しています。
論文中では、clip2latentはテキストプロンプトに基づいて1024x1024ピクセルの高解像画像を高速にサンプリングでき、かつ、低コストでトレーニングができることが示されています。

出典: clip2latent: Text driven sampling of a pre-trained StyleGAN using denoising diffusion and CLIP

詳細はこちらの論文をご参照ください。

本記事では上記手法を用いて、任意のテキストから顔画像や風景画像を生成してみます。

デモ(Colaboratory)

それでは、実際に動かしながらテキストから画像生成を行います。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

また、下記から直接Google Colaboratoryで開くこともできます。

なお、このデモはPythonで実装しています。
Pythonの実装に不安がある方、Pythonを使った機械学習について詳しく勉強したい方は、以下の書籍やオンライン講座などがおすすめです。

環境セットアップ

それではセットアップしていきます。 Colaboratoryを開いたら下記を設定しGPUを使用するようにしてください。

初めにGithubからソースコードを取得します。

%cd /content

!git clone https://github.com/justinpinkney/clip2latent.git

# Commits on Oct 7, 2022
%cd /content/clip2latent
!git checkout fc18ab19ec03f8a4fd4d6ac0e04f2b34db696946

次にライブラリをインストールします。

%cd /content/clip2latent

# install requirements
!pip install -r requirements-colab.txt

最後にライブラリをインポートします。

import torch
from clip2latent import models
from PIL import Image

device = 'cuda' if torch.cuda.is_available() else "cpu"
print("using device is", device)

以上で環境セットアップは完了です。

学習済みモデルのセットアップ

続いて、GANモデルをセットアップしていきます。
学習済みモデルは2種類用意されており、顔画像を生成するfacesと風景画像を生成するlandscapeです。

model_choices = {
    "faces": {
        "checkpoint": "https://huggingface.co/lambdalabs/clip2latent/resolve/main/ffhq-sg2-510.ckpt",
        "config": "https://huggingface.co/lambdalabs/clip2latent/resolve/main/ffhq-sg2-510.yaml",
        },
    "landscape": {
        "checkpoint": "https://huggingface.co/lambdalabs/clip2latent/resolve/main/lhq-sg3-410.ckpt",
        "config": "https://huggingface.co/lambdalabs/clip2latent/resolve/main/lhq-sg3-410.yaml",
    }
}

chosen_model = "faces" #@param ['faces', 'landscape']

checkpoint = model_choices[chosen_model]["checkpoint"]
cfg_file = model_choices[chosen_model]["config"]
model = models.Clip2StyleGAN(cfg_file, device, checkpoint)

Text to Image

それでは、テキストプロンプトを入力し画像を生成してみます。

prompt = "a photo of happy pierced blonde woman at mountain" #@param {type:"string"}

@torch.no_grad()
def infer(model, prompt, n_samples, scale, skips=250):
    images, clip_score = model(prompt, n_samples_per_txt=n_samples, cond_scale=scale, skips=skips, clip_sort=True)
    images = images.cpu()
    make_im = lambda x: (255*x.clamp(-1, 1)/2 + 127.5).to(torch.uint8).permute(1,2,0).numpy()
    images = [Image.fromarray(make_im(x)) for x in images]
    return images, clip_score
    
outputs, clip_score = infer(model, prompt, 5, 5)

for o, s in zip(outputs, clip_score):
  display(o.resize((256, 256)))
  print(f"CLIP similarity score = {s.item()}")

出力結果は以下の通りです。
テキストプロンプトに応じた画像が生成されています。注目すべきはモデルのロードさえ完了してしまえば画像の生成が非常に速い点です。
Tesla T4のGoogle Colaboratoryでさえ、5枚の画像生成が600ミリ秒前後で完了します。

参考までにlandscapeモデルにテキストプロンプトthe sun setting over the seaを入力した結果は以下の通りです。

まとめ

本記事では、clip2latentと呼ばれる機械学習手法を用いて、テキストから顔画像や風景画像を生成する方法をご紹介しました。
Stable Diffusionなど拡散モデルで画像を生成する手法が注目を集めますが、本手法では、拡散モデルを潜在空間の制御に用いており目の付け所が面白いです。

また本記事では、機械学習を動かすことにフォーカスしてご紹介しました。
もう少し学術的に体系立てて学びたいという方には以下の書籍などがお勧めです。ぜひご一読下さい。

また動かせるだけから理解して応用できるエンジニアの足掛かりに下記のUdemyなどもお勧めです。

参考文献

1.  論文 - clip2latent: Text driven sampling of a pre-trained StyleGAN using denoising diffusion and CLIP

2. GitHub - justinpinkney/clip2latent