Python 生成AI：画像から類似画像を生成するためのテクニック「Stable Diffusion」

　本記事では、Pythonを用いた画像生成AIの雛形コードを載せました。例えば、下図のような手書きの画像があります。

　　　　　　　　　　　　図1. 手書きによる画像

本プログラムを実行すると、上図1を元にした下図５つのような生成AIにより類似の画像を生成します。

　　　　　　　　　　　　図2. 生成AIによる画像1

　　　　　　　　　　　　図3. 生成AIによる画像2

　　　　　　　　　　　　図4. 生成AIによる画像3

　　　　　　　　　　　　図5. 生成AIによる画像4

　　　　　　　　　　　　図6. 生成AIによる画像5

　使用するのは生成AIのPythonライブラリ「Stable Diffusion」です。基本機能は、プロンプト（指示、命令文）と呼ぶ単語や短い文章を指定することで、それに従った画像を生成することです。
　通常は画像から画像を生成したい場合、ネット上で容易に見つかるコードでテストすると、お世辞にも元の画像と似通わない画像が生成されます。それはプロンプトに指定するキーワードが少ないことや、文章の表現があいまいだったりすることが原因です。
　そこで、元となる画像からあらかじめプロンプトを抽出して指定することで、その問題を回避します。ライブラリは「clip-interrogator」を使用します。

例えば、下記は図1から抽出したプロンプトです。

cartoon drawing of a blue water bottle holding a blue umbrella, epic fantasy battle, excrement, fantasy game spell icon, poop, computer drawing, burnt armor, sheild, inspired by RHADS, smelly, photo of war, shortsword, korean mmorpg, gelatinous, turqouise, fart

　このプロンプトと共に元の画像を指定することで、見当違いな画像は生成されずに、類似の画像を生成することができるようになります。

■インストールするライブラリ

pytorchのインストール

https://pytorch.org/ のリンク先に従います。各PC環境を指定してそのコマンドを実行してください。

例）conda install pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia

私の場合は、anaconda環境でGPUにnvidiaのGeforce3060を使用しているため上記コマンドでした。

diffusersのインストール

anaconda環境の場合

conda install -c conda-forge diffusers

anaconda環境でない場合

pip install diffusers["torch"]

その他の必要なライブラリ

自然言語を処理するライブラリ

pip install transformers

pytorchの補助ライブラリ

pip install accelerate

画像からそれを構成するプロンプトを抽出するライブラリ

pip install clip-interrogator

■本プログラム

　注意すべきは、元となる画像の品質が高いほどビデオメモリを必要とすることです。私のPC環境は、Geforce3060の12MBのグラフィックボードを使用しています。風景写真はビデオメモリ容量のエラーで落ちました。

import os
from PIL import Image
from clip_interrogator import Config, Interrogator
import torch
from diffusers import StableDiffusionImg2ImgPipeline

# 読み込む画像のパス
image_path = './poop/class4/120.png'

# ファイルパスからファイル名の取得
file_name = os.path.basename(image_path)

# 画像オブジェクトを読み込む
image = Image.open(image_path).convert('RGB')

# 画像からプロンプトを生成する
ci = Interrogator(Config(clip_model_name="ViT-L-14/openai"))

# プロンプト（呪文）の取得
prompt = ci.interrogate(image)
print(prompt)

# プロンプトをファイルへ出力する
with open(f'clip_{file_name[:-4]}.txt', 'w') as f:
    f.write(prompt)

# モデルの指定
#model_id = "CompVis/stable-diffusion-v1-4"
model_id = "runwayml/stable-diffusion-v1-5"
#model_id = "stabilityai/stable-diffusion-2-1-base"

# 画像から画像を生成する場合
pipeline = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.float16)

# GPUで演算する場合、CUDA（Compute Unified Device Architecture）を指定（ビデオカードがある場合：GPUで処理する設定）
pipeline.to("cuda")

# 否定するプロンプト
negative_prompt = 'people'

# 画像を生成する
num = 5
for i in range(1, num+1, 1):
    with torch.autocast("cuda"):
        #torch.manual_seed(i) # ランダムシード
        image = pipeline(prompt,
                         image = image,
                         guidance_scale = 7.5,
                         negative_prompt = negative_prompt,
                         strength = 0.75).images[0]
    image.save(f"{file_name[:-4]}_{i}.png")

以上

<広告>

リンク

リンク

リンク