本記事では、RQ-VAE Transformerと呼ばれる機械学習手法を用いて、テキストから画像を生成(Image to Text)する方法をご紹介します。

RQ-VAE Transformer

概要

RQ-VAE Transformerは高解像度画像生成技術です。

RQ-VAE Transformerは、高解像度画像を効果的に生成するため、残差量子化VAE(Residual-Quantized VAE)とRQ-Transformerで構成される2段階のフレームワークです。
この構成は、計算コストを効率的に削減しつつ、画像生成タスクにおいて様々なベンチマークで既存の自己回帰モデルより優れていると示されています。

出典: Autoregressive Image Generation using Residual Quantization

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

本記事では、上記手法を用いて、任意のテキストを入力して画像を生成する方法をご紹介します。

デモ(Colaboratory)

それでは、実際に動かしながらText to Imageを行っていきます。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

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

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

環境セットアップ

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

はじめに、GitHubからソースコードを取得します。

%cd /content

!git clone https://github.com/kakaobrain/rq-vae-transformer.git

続いて、ライブラリのインストールします。
注意点として、「RESTART RUNTIME」が表示された場合「ランタイム」→「ランタイムを再起動」を実行してください。

%cd /content/rq-vae-transformer

# もし RESTART RUNTIMEが表示されたら「ランタイム」→「ランタイムを再起動」
!pip install -r requirements.txt

続いて、ライブラリをインポートします。

%cd /content/rq-vae-transformer

import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

from PIL import Image
import yaml
import torch
import torchvision
import clip
import torch.nn.functional as F

from notebooks.notebook_utils import TextEncoder, load_model, get_generated_images_by_texts

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

環境セットアップが完了したので、続いて学習済みモデルを論文発表元から取得します。

%cd /content/rq-vae-transformer
!mkdir pretrained
%cd pretrained

!wget https://arena.kakaocdn.net/brainrepo/models/RQVAE/dcd95e8f08408e113aab6451fae895f5/cc3m.tar.gz
!tar -xvf cc3m.tar.gz

/content/rq-vae-transformer/pretrainedにモデルがダウンロードされます。

先ほどダウンロードしたモデルをロードしておきます。
RQ-VAEとRQ-Transferの2ステージ構成のため、2つのモデルをダウンロードしています。

vqvae_path = '/content/rq-vae-transformer/pretrained/cc3m/stage1/model.pt'
model_path = '/content/rq-vae-transformer/pretrained/cc3m/stage2/model.pt'

# load stage 1 model: RQ-VAE
model_vqvae, _ = load_model(vqvae_path)

# load stage 2 model: RQ-Transformer
model_ar, config = load_model(model_path, ema=False)

# GPUにセット
model_ar = model_ar.cuda().eval()
model_vqvae = model_vqvae.cuda().eval()

# CLIPモデルのダウンロード
model_clip, preprocess_clip = clip.load("ViT-B/32", device='cpu')
model_clip = model_clip.cuda().eval()

# prepare text encoder to tokenize natual languages
text_encoder = TextEncoder(tokenizer_name=config.dataset.txt_tok_name, 
                           context_length=config.dataset.context_length)

Text to Image

それでは、Text to Imageを実行します。
初めに任意のテキストを英語で設定します。

#@markdown 入力テキストを設定してください。
text_prompts = "'a photo of \"Cherry blossoms in the snow\"'" #@param {type:"string"}

num_samples = 16
temperature= 1.0
top_k=1024
top_p=0.95

今回は雪の中の桜の写真を生成してみます。

pixels = get_generated_images_by_texts(model_ar,
                                       model_vqvae,
                                       text_encoder,
                                       model_clip,
                                       preprocess_clip,
                                       text_prompts,
                                       num_samples,
                                       temperature,
                                       top_k,
                                       top_p,
                                      )

生成結果を出力します。

num_visualize_samples = 16
images = [pixel.cpu().numpy() * 0.5 + 0.5 for pixel in pixels]
images = torch.from_numpy(np.array(images[:num_visualize_samples]))
images = torch.clamp(images, 0, 1)
grid = torchvision.utils.make_grid(images, nrow=4)

img = Image.fromarray(np.uint8(grid.numpy().transpose([1,2,0])*255))
display(img)

出力結果は以下の通りです。

高解像度画像が4x4の16枚表示されます。
精細な画像が生成されていますね。

まとめ

本記事では、RQ VAE Transformerを用いたText to Imageを行いました。
テキスト次第で様々な画像が生成されるため是非いろいろ試してみてください。

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

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

参考文献

1.  論文 - Autoregressive Image Generation using Residual Quantization

2. GitHub - kakaobrain/rq-vae-transformer