[RQ VAE Transformer] AIでテキストから画像生成 [Python]

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

1.RQ-VAE Transformer
- 1.1.概要
2.デモ(Colaboratory)
3.まとめ
4.参考文献

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を使用するようにしてください。

「ランタイムのタイプを変更」→「ハードウェアアクセラレータ」を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)