本記事では、Lambda Diffusersを用いてテキストからポケモン風の画像を生成する方法をご紹介します。

出典: LambdaLabsML/lambda-diffusers

Stable Diffusion fine tuned on Pokémon

概要

近年良くも悪くも注目を集めるStable Diffusionですが、Stable Diffusion fine tuned on PokémonはLambdaLabsがポケモンの画像とキャプションからなるデータセットを用いてStable DiffusionをFine Tuningしたモデルです。

このモデルはHuggingFaceから取得することができ、任意のテキストを入力すると上記の画像のようにポケモン風画像が生成されます。

なおStable Diffusionの技術解説は以下の記事をご参照ください。

本記事では上記手法を用いて、任意のテキストからポケモン風画像を生成するとともに、LambdaLabsが公開するその他のモデルを動かしていきます。

デモ(Colaboratory)

それでは、実際に動かしながらLambdaLabsが公開するStable Diffusionモデルを試していきます。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

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

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

環境セットアップ

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

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

%cd /content

!git clone https://github.com/LambdaLabsML/lambda-diffusers.git

# Commits on Dec 8, 2022
%cd /content/lambda-diffusers
!git checkout ef843a02cf0bc4e96c44a95fb3f65fa1939125b6

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

%cd /content/lambda-diffusers

# for lambdalabs/sd-pokemon-diffusers
!pip install diffusers==0.3.0
!pip install transformers==4.22.1 scipy==1.7.3 ftfy==6.1.1

# for lambda-diffusers
!pip install -r requirements.txt

インストール完了後、一度ランタイムを再起動します。

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

%cd /content/lambda-diffusers

from pathlib import Path
from lambda_diffusers import StableDiffusionImageEmbedPipeline
from PIL import Image
import torch
from torch import autocast

from diffusers import StableDiffusionPipeline


device = "cuda" if torch.cuda.is_available() else "cpu"

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

Image to Image

まずテキストからポケモン風画像を生成する前に、LambdaLabsが公開する別のモデルを用いて、Image to Imageを動かしていきます。
このモデルは入力画像を変化させいくつかのバリエーション画像を出力します。

まずは学習済みモデルをロードします。

pipe = StableDiffusionImageEmbedPipeline.from_pretrained(
    "lambdalabs/sd-image-variations-diffusers",
    revision="273115e88df42350019ef4d628265b8c29ef4af5",)
pipe = pipe.to(device)

次に、入力画像をセットアップし、Image to Imageを実行します。

!rm -rf outputs
!rm -rf inputs
!mkdir outputs inputs

!wget -c https://casie.jp/media/wp-content/uploads/2020/09/Van_Gogh_Self-Portrait_with_Straw_Hat_1887-Detroit.jpg \
      -O inputs/test_01.jpg

input_img_path = "inputs/test_01.jpg"

im = Image.open(input_img_path)
num_samples = 2
image = pipe(num_samples*[im], guidance_scale=3.0)
image = image["sample"]

base_path = Path("outputs")
base_path.mkdir(exist_ok=True, parents=True)
for idx, im in enumerate(image):
  im.save(base_path/f"{idx:06}.jpg")

出力結果は以下の通りです。
左の画像が出力結果で、中央と右側がモデルによって変化させた画像です。

Pokemon Text to Image

続いて本題のPokemon Text to Imageを動かしていきます。
先ほどと同様にHuggingFaceから学習済みモデルをセットアップします。

pipe = StableDiffusionPipeline.from_pretrained("lambdalabs/sd-pokemon-diffusers", torch_dtype=torch.float16)  
pipe = pipe.to("cuda")

続いて、テキスト(Prompt)を設定します。

prompt = "Oda Nobunaga" #@param {type:"string"}

# ガイダンススケール
scale = 7.5
# 生成枚数
n_samples = 3
disable_safety = True

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

!rm -rf outputs_poke
!mkdir outputs_poke

if disable_safety:
  def null_safety(images, **kwargs):
      return images, False
  pipe.safety_checker = null_safety

with autocast("cuda"):
  images = pipe(n_samples*[prompt], guidance_scale=scale).images

for idx, im in enumerate(images):
  im.save(f"outputs_poke/{idx:06}.png")

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

Oda Nobunaga

その他の結果も見てみます。

Birth of Venus

Abraham Lincoln

まとめ

本記事では、FineTuneされたStable Diffusionでテキストからポケモン風画像を生成する方法をご紹介しました。
Stable Diffusion関連技術の発達は目覚ましく日々新しいモデルやサービスが生まれています。
楽しみながら追いかけていきたいです。

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

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

参考文献

2. GitHub - LambdaLabsML/lambda-diffusers