生成AIをGoogle Colaboratoryで簡単に 【Part7 画像生成AI SDXL Turbo編】

はじめに

今回は、初心者向けにGoogle Colaboratoryで、簡単に生成AIを使えるようにする環境を作ります。
（音声対話システムを作成するための個人的な忘備録の意味合いが強いです）

Part7の今回は、画像生成AIを使えるようにします。
今回利用する音声認識AIはSDXL Turboになります。

なお、今回Part7ですが、Part1~Part6の内容を読まなくても、わかるように記載しています。

画像生成AIということで、直接的には音声対話システムとは無関係な技術のように見えますが、AIの発話に合わせて適切な画像を表示するみたいなことに使えるかなと思い、勉強しています。
（Part6でSDXLを紹介したのは、今回SDXL Turboを紹介したかったからです。）

SDXL Turboとは

SDXL Turboとは、Stability AI社が開発した画像生成AIの一つで、Part6で取り扱ったSDXLよりも高速に画像を生成できることが特徴になっています。

短時間で、画像生成が可能であるという特徴から、リアルタイム音声対話システムに導入して、対話の場面に合わせた画像の生成ができないかと思い、勉強しています。

SDXL turboが高速に画像生成ができる理由というのは、（SDXLの）蒸留モデルであるということ、また蒸留する際に敵対的拡散蒸留という手法を用いて、SDKLでは50 Stepほどかかっていた生成処理を、1 Stepで処理できるようになったため、超高速で画像を生成できます。

ただし、蒸留モデルであるからか、ネガティブプロンプトが利用できないことや、生成される画像の解像度や質がSDXLよりも低いことがデメリットに挙げられます。
しかしながら、処理速度を優先したい場面においては有力な選択肢になるかと思います。
（前回と同様、Diffuserというライブラリを利用します）

成果物

下記のリポジトリをご覧ください。
https://github.com/personabb/colab_AI_sample/tree/main/colab_SDXLturbo_sample

解説

下記の通り、解説を行います。
まずは上記のリポジトリをcloneしてください。

./

git clone https://github.com/personabb/colab_AI_sample.git

その後、cloneしたフォルダ「colab_AI_sample」をマイドライブの適当な場所においてください。

ディレクトリ構造

Google Driveのディレクトリ構造は下記を想定します。

MyDrive/
    └ colab_AI_sample/
          └ colab_SDXLturbo_sample/
                  ├ configs/
                  |    └ config.ini
                  ├ outputs/
                  ├ module/
                  |    └ module_sd_turbo.py
                  └ StableDiffusionXL-turbo_sample.ipynb

• colab_AI_sampleフォルダは適当です。なんでも良いです。1階層である必要はなく下記のように複数階層になっていても良いです。
  • MyDrive/hogehoge/spamspam/hogespam/colab_SDXLturbo_sample
• outputsフォルダには、生成後の画像が格納されます。最初は空です。
  • 連続して生成を行う場合、過去の生成内容を上書きするため、ダウンロードするか、名前を変えておくことをオススメします。

使い方解説

StableDiffusionXL-turbo_sample.ipynbをGoogle Colabratoryアプリで開いてください。
ファイルを右クリックすると「アプリで開く」という項目が表示されるため、そこからGoogle Colabratoryアプリを選択してください。

もし、ない場合は、「アプリを追加」からアプリストアに行き、「Google Colabratory」で検索してインストールをしてください。

Google Colabratoryアプリで開いたら、StableDiffusionXL-turbo_sample.ipynbのメモを参考にして、一番上のセルから順番に実行していけば、問題なく最後まで動作して、画像生成をすることができると思います。

また、最後まで実行後、パラメータを変更して再度実行する場合は、「ランタイム」→「セッションを再起動して全て実行する」をクリックしてください。

コード解説

主に、重要なStableDiffusionXL-turbo_sample.ipynbとmodule/module_sd_turbo.pyについて解説します。

StableDiffusionXL-turbo_sample.ipynb

該当のコードは下記になります。
https://github.com/personabb/colab_AI_sample/blob/main/colab_SDXLturbo_sample/StableDiffusionXL-turbo_sample.ipynb

下記に１セルずつ解説します。

１セル目

./colab_AI_sample/colab_SDXLturbo_sample/StableDiffusionXL-turbo_sample.ipynb

#SDXL Turbo で必要なモジュールのインストール
!pip install diffusers
!pip install transformers
!pip install scikit-learn
!pip install ftfy
!pip install accelerate
!pip install invisible_watermark
!pip install safetensors

ここでは、必要なモジュールをインストールしています。
Google colabではpytorchなどの基本的な深層学習パッケージなどは、すでにインストール済みなため上記だけインストールすれば問題ありません。

2セル目

./colab_AI_sample/colab_SDXLturbo_sample/StableDiffusionXL-turbo_sample.ipynb

#Google Driveのフォルダをマウント（認証入る）
from google.colab import drive
drive.mount('/content/drive')

# カレントディレクトリを本ファイルが存在するディレクトリに変更する。
import glob
import os
pwd = os.path.dirname(glob.glob('/content/drive/MyDrive/**/colab_SDXLturbo_sample/StableDiffusionXL-turbo_sample.ipynb', recursive=True)[0])
print(pwd)

%cd $pwd
!pwd

ここでは、Googleドライブの中身をマウントしています。
マウントすることで、Googleドライブの中に入っているファイルを読み込んだり、書き込んだりすることが可能になります。

マウントをする際は、Colabから、マウントの許可を行う必要があります。
ポップアップが表示されるため、指示に従い、マウントの許可を行なってください。

また、続けて、カレントディレクトリを/から/content/drive/MyDrive/**/colab_SDXLturbo_sampleに変更しています。
（**はワイルドカードです。任意のディレクトリ（複数）が入ります）
カレントディレクトリは必ずしも変更する必要はないですが、カレントディレクトリを変更することで、これ以降のフォルダ指定が楽になります

3セル目

./colab_AI_sample/colab_SDXLturbo_sample/StableDiffusionXL-turbo_sample.ipynb

#モジュールをimportする
from module.module_sd_turbo import SDXL as SDXL_turbo

module/module_sd_turbo.pyのSDXLクラスをSDXL_turboモジュールとしてインポートします。
この中身の詳細は後の章で解説します。

4セル目

./colab_AI_sample/colab_SDXLturbo_sample/StableDiffusionXL-turbo_sample.ipynb

#モデルの設定を行う。

config_text = """
[SDXL-turbo]
device = auto
n_steps=4
seed=42
guidance_scale = 0.0

vae_model_path = stabilityai/sdxl-vae
base_model_path = stabilityai/sdxl-turbo

cfg_scale = 1
width = 512
height = 512

"""

with open("configs/config.ini", "w", encoding="utf-8") as f:
  f.write(config_text)

このセルでは、設定ファイルconfigs/config.iniの中身をconfig_textの内容で上書きしています。
SDXL Turboは上記の設定に併せて動作をします。
基本的にはこの設定ファイルを変更する必要はありません。
WebUIで主流に使われている設定をこちらでも採用しております。

強いて言えば、下記の部分でモデルを指定していますが、そちらを変更することでモデルを変更できます。ただし、質を犠牲に生成速度を上げるというのは、あまり需要がないのか、SDXL TurboかつDiffuserに対応しているモデルはほとんどないので、モデルを変更するというのは期待し内容が良いと思います。（むしろ使えるモデルがあるなら、私が教えて欲しいです）

vae_model_path = stabilityai/sdxl-vae
base_model_path = stabilityai/sdxl-turbo

また、生成される画像の解像度は下記で指定しています。
変更することで、縦長の画像にしたり、横長の画像にしたりできますが、大きく変更すると画像が崩れる可能性があるため、注意してください。

width = 512 #横幅
height = 512 #縦幅

さらに、今回のモデルでは下記のようにstep数を4に設定しています。

n_steps=4

理論上1 Stepで画像生成できるというのがSDXL Turboの売りではございますが、試行錯誤の結果step数を4までなら増やしたほうが、生成される画像の質が高いことがわかり、4を採用しています。
また4以上にすると、逆に画像が崩れ始めるそうですので、上限4の範囲内で、速度と質を考慮しながら設定いただければと思います。
（4 Stepでも十分高速な画像生成が可能です）

最後に、本プログラムではseedを42に固定しているため、同じpromptを入れて、同じ設定の場合、実行のたびに毎回同じ画像が生成されます。
一方で、このseedの値をランダムに変更することで、生成される画像はランダムに変化します。

5セル目

./colab_AI_sample/colab_SDXLturbo_sample/StableDiffusionXL-turbo_sample.ipynb

#読み上げるプロンプトを設定する。

main_prompt = """
1 girl ,Yellowish-white hair ,short hair ,red small ribbon,red eyes,red hat ,school uniform ,solo ,smile ,upper body ,Anime ,Japanese,best quality,high quality,ultra highres,ultra quality
"""

negative_prompt = ""

ここで、生成する画像をプロンプトで指定します。
main_promptには、生成したい画像の特徴を入力してください。
例えば、人数、性別、髪の毛や服装の特徴など

negative_promptはSDXL Turboでは使用できないため、空文字を設定しています。

6セル目

./colab_AI_sample/colab_SDXLturbo_sample/StableDiffusionXL-turbo_sample.ipynb

sd = SDXL_turbo()

SDXL_turboクラスのインスタンスを生成しています。
画像生成単体の速度感を体感して欲しく、インスタンスの作成（モデルのダウンロードや読み込み）は別のセルとして実行するようにしました。

7セル目

./colab_AI_sample/colab_SDXLturbo_sample/StableDiffusionXL-turbo_sample.ipynb

for i in range(3):
      image = sd.generate_image(main_prompt, negative_prompt)
      image.save("./outputs/SDXL_result_{}.png".format(i))

ここで、画像を生成しています。
for文を使い、同じpromptで3枚の画像を生成しています。
画像はoutputsフォルダに保存されますが、実行のたびに同じ名前で保存されるので、過去に保存した画像は上書きされるようになりますので注意してください。

前回のPart6で扱ったSDXLと比較して、非常に高速に画像生成できているかと思います。

実行すると下記のような画像が生成されます。（seed値という乱数表の値を変更しない限り、実行のたびに必ず同じ画像が生成されます。seedは設定ファイルで変更できます。）

一方でPart6の再掲ですが、SDXLでは下記のような画像が生成されていました。
（プロンプトが違うので参考程度にご覧ください）

SDXLと比較すると、大分質が低いように見えますが、3枚目のように可愛らしい画像を表示させるポテンシャルはあるかと思います。

module/module_sd_turbo.py

続いて、StableDiffusionXL-turbo_sample.ipynbから読み込まれるモジュールの中身を説明します。

下記にコード全文を示します。

コード全文

./colab_AI_sample/colab_SDXLturbo_sample/module/module_sd_turbo.py

from diffusers import AutoPipelineForText2Image, AutoencoderKL, EulerAncestralDiscreteScheduler
import torch

import os
import configparser
# ファイルの存在チェック用モジュール
import errno

class SDXLconfig:
    def __init__(self, config_ini_path = './configs/config.ini'):
        # iniファイルの読み込み
        self.config_ini = configparser.ConfigParser()
        
        # 指定したiniファイルが存在しない場合、エラー発生
        if not os.path.exists(config_ini_path):
            raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), config_ini_path)
        
        self.config_ini.read(config_ini_path, encoding='utf-8')
        SDXL_items = self.config_ini.items('SDXL-turbo')
        self.SDXL_config_dict = dict(SDXL_items)

class SDXL:
    def __init__(self,device = None, config_ini_path = './configs/config.ini'):
        
        SDXL_config = SDXLconfig(config_ini_path = config_ini_path)
        config_dict = SDXL_config.SDXL_config_dict

        if device is not None:
            self.device = device
        else:
            device = config_dict["device"]

            self.device = "cuda" if torch.cuda.is_available() else "cpu"
            if device != "auto":
                self.device = device
        
        self.n_steps = int(config_dict["n_steps"])
        self.seed = int(config_dict["seed"])
        self.guidance_scale = float(config_dict["guidance_scale"])
        self.generator = torch.Generator(device=self.device).manual_seed(self.seed)

        self.vae_model_path = config_dict["vae_model_path"]
        self.VAE_FLAG = True
        if self.vae_model_path == "None":
            self.vae_model_path = None
            self.VAE_FLAG = False
            
        self.base_model_path = config_dict["base_model_path"]
        self.cfg_scale = float(config_dict["cfg_scale"])
        self.width = int(config_dict["width"])
        self.height = int(config_dict["height"])

        self.base  = self.preprepare_model()
        

    def preprepare_model(self):
        if self.VAE_FLAG:
            print("load vae")
            vae = AutoencoderKL.from_pretrained(
                self.vae_model_path,
                torch_dtype=torch.float16)

            base = AutoPipelineForText2Image.from_pretrained(
                self.base_model_path, 
                vae=vae,
                torch_dtype=torch.float16,
                variant="fp16",
                use_safetensors=True
            )
            base.to(self.device)
    
            
        else:
            print("non vae")
            base = AutoPipelineForText2Image.from_pretrained(
                self.base_model_path, 
                torch_dtype=torch.float16,
                variant="fp16",
                use_safetensors=True
            )
            base.to(self.device)


        base.scheduler = EulerAncestralDiscreteScheduler.from_config(
                    base.scheduler.config, 
                    timestep_spacing= "trailing"
                    )
                    
        return base

            
    
    def generate_image(self, prompt, neg_prompt,seed = None):
        
        if seed is not None:
            self.generator = torch.Generator(device=self.device).manual_seed(seed)
        
        image = self.base(
            prompt=prompt,
            negative_prompt = neg_prompt,
            cfg_scale=self.cfg_scale,
            guidance_scale = self.guidance_scale,
            num_inference_steps=self.n_steps,
            width = self.width, 
            height = self.height,
            generator=self.generator
            ).images[0]

        return image
    
    def generate_image_display(self, prompt, neg_prompt,cnt = 0,seed = None):
        
        if seed is not None:
            self.generator = torch.Generator(device=self.device).manual_seed(seed)
        
        image = self.base(
            prompt=prompt,
            negative_prompt = neg_prompt,
            cfg_scale=self.cfg_scale,
            guidance_scale = self.guidance_scale,
            num_inference_steps=self.n_steps,
            width = self.width, 
            height = self.height,
            generator=self.generator
            ).images[0]
        
        image.show()
        image.save("output_images/sample{}.png".format(cnt))

        return image
    

では一つ一つ解説していきます。

SDXLconfigクラス

./colab_AI_sample/colab_SDXLturbo_sample/module/module_sd_turbo.py

class SDXLconfig:
    def __init__(self, config_ini_path = './configs/config.ini'):
        # iniファイルの読み込み
        self.config_ini = configparser.ConfigParser()
        
        # 指定したiniファイルが存在しない場合、エラー発生
        if not os.path.exists(config_ini_path):
            raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), config_ini_path)
        
        self.config_ini.read(config_ini_path, encoding='utf-8')
        SDXL_items = self.config_ini.items('SDXL-turbo')
        self.SDXL_config_dict = dict(SDXL_items)

ここではconfig_ini_path = './configs/config.ini'で指定されている設定ファイルをSDXL_config_dictとして読み込んでいます。
辞書型で読み込んでいるため、設定ファイルの中身をpythonの辞書として読み込むことが可能になります。

SDXLクラスのinitメソッド

./colab_AI_sample/colab_SDXLturbo_sample/module/module_sd_turbo.py

class SDXL:
    def __init__(self,device = None, config_ini_path = './configs/config.ini'):
        
        SDXL_config = SDXLconfig(config_ini_path = config_ini_path)
        config_dict = SDXL_config.SDXL_config_dict

        if device is not None:
            self.device = device
        else:
            device = config_dict["device"]

            self.device = "cuda" if torch.cuda.is_available() else "cpu"
            if device != "auto":
                self.device = device
        
        self.n_steps = int(config_dict["n_steps"])
        self.seed = int(config_dict["seed"])
        self.guidance_scale = float(config_dict["guidance_scale"])
        self.generator = torch.Generator(device=self.device).manual_seed(self.seed)

        self.vae_model_path = config_dict["vae_model_path"]
        self.VAE_FLAG = True
        if self.vae_model_path == "None":
            self.vae_model_path = None
            self.VAE_FLAG = False
            
        self.base_model_path = config_dict["base_model_path"]
        self.cfg_scale = float(config_dict["cfg_scale"])
        self.width = int(config_dict["width"])
        self.height = int(config_dict["height"])

        self.base  = self.preprepare_model()

本クラスはStableDiffusionXL-turbo_sample.ipynbからSDXL_turboという名前で呼び出されます。

まず、設定ファイルの内容をconfig_dictに格納しています。これは辞書型のため、config_dict["device"]のような形で設定ファイルの内容を文字列として取得することができます。
あくまで、すべての文字を文字列として取得するため、int型やbool型にしたい場合は、適宜型変更をする必要があることに注意してください。

続いて下記の順番で処理を行います。

• モデルを動作させるdeviceを指定する
• 設定ファイルの各種設定を取得する
• モデルを定義する。
  • 設定ファイルに合わせて、適切なモデルを定義する
  • self.preprepare_model()メソッドで定義する

SDXLクラスのpreprepare_modelメソッド

./colab_AI_sample/colab_SDXLturbo_sample/module/module_sd_turbo.py

class SDXL:
　　　・・・
    def preprepare_model(self):
        if self.VAE_FLAG:
            print("load vae")
            vae = AutoencoderKL.from_pretrained(
                self.vae_model_path,
                torch_dtype=torch.float16)

            base = AutoPipelineForText2Image.from_pretrained(
                self.base_model_path, 
                vae=vae,
                torch_dtype=torch.float16,
                variant="fp16",
                use_safetensors=True
            )
            base.to(self.device)
    
            
        else:
            print("non vae")
            base = AutoPipelineForText2Image.from_pretrained(
                self.base_model_path, 
                torch_dtype=torch.float16,
                variant="fp16",
                use_safetensors=True
            )
            base.to(self.device)


        base.scheduler = EulerAncestralDiscreteScheduler.from_config(
                    base.scheduler.config, 
                    timestep_spacing= "trailing"
                    )
                    
        return base

このメソッドはinitメソッドで呼ばれて、モデルを定義して読み込むメソッドになります。

SDXL Turboでは、baseモデルとvaeモデルの3つのモデルがある。（SDXLであったrefinerモデルはないです。高速化のためだと思います）
基本的にはbaseモデルのみが画像を生成するため、解像度が512x512に制限されていると考えられます。
またvaeモデルは、適切なものを設定することで、より綺麗な画像を生成することができるようになります。

今回の設定ファイルでは、2つすべてのモデルを利用することになっていますが、設定ファイル次第では、VAEを使わない設定にすることも可能なように、このメソッドは設計しています。

SDXLクラスのgenerate_imageメソッド

./colab_AI_sample/colab_SDXLturbo_sample/module/module_sd_turbo.py

class SDXL:
　　　・・・
    def generate_image(self, prompt, neg_prompt,seed = None):
        
        if seed is not None:
            self.generator = torch.Generator(device=self.device).manual_seed(seed)
        
        image = self.base(
            prompt=prompt,
            negative_prompt = neg_prompt,
            cfg_scale=self.cfg_scale,
            guidance_scale = self.guidance_scale,
            num_inference_steps=self.n_steps,
            width = self.width, 
            height = self.height,
            generator=self.generator
            ).images[0]

        return image

このメソッドは、ここまでで読み込んだモデルと設定を利用して、実際に画像を生成するメソッドです。
本メソッドの引数にてseedを指定すると、設定ファイルを上書きして、seedを指定できます。
したがって、ここで乱数を指定することで、ランダムな画像を生成できるようになります。

また、引数でseedを指定しない場合は、

self.generator = torch.Generator(device=self.device).manual_seed(seed)

が呼ばれます。したがって、生成のたびに一つ大きなseedに変更されて適用される（42→43→44）ため、生成のたびに画像を変化させることができます。

まとめ

今回は、初心者向けにGoogle Colaboratoryで、簡単に生成AIを使えるようにする環境を作りました。

Part7の今回は、生成AIの一つ、画像生成AIのSDXL Turboモデルを使えるようにしました。

画像生成AIを利用することで、音声対話システムで音声対話を行いながら、その対話に合わせた画像を生成することができるため、よりエンターテイメント性を高めることができます。

実際に、音声対話システムに組み込んでみたので、そちらに関しては後日記事にする予定です

（下記の対話システムに組み込んでいます）
https://zenn.dev/asap/articles/5b1b7553fcaa76

今回取り扱ったモデルは、SDXLの蒸留モデルであるため、生成される画像の質は正直SDXLより大分劣っていると思います。
一方で、SDXLでは十数秒かかっていた画像生成プロセスを1秒程度で処理できるリアルタイム性は唯一無二の価値かと思います。（特にリアルタイム音声対話システムに組み込む際は）
用途に応じて、使いこなしていこうと思いました。
（WebUIだけでなく、もっとDiffuserでも動かせるモデルが増えるといいなと思いました）

次回は、満を持して、2024年7月時点で最新の画像生成AiであるStable Diffusion 3モデルを試してみたいと思います。