🎃

生成AIをGoogle Colaboratoryで簡単に 【Part9 画像生成AI SD3+ControlNet編】

2024/07/29に公開

はじめに

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

Part9の今回は、Part8の続きとなっております。
特にPart8を読んでいただく必要はないですが、中身を理解したい方は、そちらも合わせて読んでいただけますと幸いです。
なお、今回Part9ですが、Part1~Part7の内容を読まなくても、わかるように記載しています。

画像生成AIということで、直接的には音声対話システムとは無関係な技術のように見えますが、AIの発話に合わせて適切な画像を表示するみたいなことに使えるかなと思い、勉強しています。
(過去に私が実装した音声対話システムは下記で紹介しています。興味があれば一読いただけますと嬉しいです)
https://zenn.dev/asap/articles/5b1b7553fcaa76

初心者向けに無料版のGoogle Colaboratoryで動作可能な範囲内で実施しておりますので、気軽に試してみてください。
(SD3とControlnetを動かす際に、無料枠に収まるように無理やり処理しているところもあるので、大きな計算資源がある方や、課金ができる方は、自身の計算資源に合わせて、実装を効率化した方が良いと思います。あくまで、無料枠のGPUメモリの範囲内で動かせることを意識して実装しています)

本記事では、T5 Text Encoderをはずした状態で、ControlNetを動作させることと、T5 Text Encoderを量子化して導入した状態で、ControlNetを動作させることに取り組みました。
(T5 Text Encoderを量子化した場合、通常の方法でControlNetを動作させることは難しいため、本記事の内容が役に立つと幸いです)

今回は下記の記事を参考に記載しています。
こちらは公式でSD3のControlNetの使い方が記載されている記事になります。
https://huggingface.co/docs/diffusers/api/pipelines/controlnet_sd3

こちらはSD3用ではありませんが、SDXLなどでControlNetを利用している記事になります。非常にわかりやすくまとまっている記事なのでぜひ一読ください。
https://note.com/npaka/n/n06b9ca7994a4#246b7e90-f372-453e-a13d-1c33d8def8e6

本記事の内容を理解することで、
ポーズを指定して下記のような画像を生成できるようになります。

ぜひお試しください。

Stable Diffusion 3(SD3)とは

Part8の記事の解説に譲ります。

ControlNetとは

「Adding Conditional Control to Text-to-Image Diffusion Models」
という論文で導入された技術で、Stable Diffusionで生成される画像をプロンプト以外の方法で制御することができます。

記事で紹介されている論文のAbstractを引用(Google翻訳)します。

大規模な事前トレーニング済みのテキストから画像への拡散モデルに空間調整コントロールを追加するニューラル ネットワーク アーキテクチャである ControlNet を紹介します。ControlNet は、生産準備が整った大規模な拡散モデルをロックし、数十億枚の画像で事前トレーニングされたディープで堅牢なエンコーディング レイヤーを強力なバックボーンとして再利用して、さまざまな条件付きコントロールのセットを学習します。ニューラル アーキテクチャは、パラメーターをゼロから徐々に増やし、有害なノイズが微調整に影響しないようにする「ゼロ畳み込み」(ゼロ初期化畳み込みレイヤー) に接続されています。エッジ、深度、セグメンテーション、人間のポーズなど、さまざまな調整コントロールを、プロンプトありまたはなしで単一または複数の条件を使用して、Stable Diffusion でテストします。ControlNet のトレーニングは、小規模 (<50k) および大規模 (>1m) のデータセットで堅牢であることを示します。広範な結果から、ControlNet は画像拡散モデルを制御するためのより幅広いアプリケーションに役立つ可能性があることが示されています。

論文で説明されている通り、以下のようなposeを参照画像として入力することで、同じようなポーズの画像を生成できたり、線画を入力することで、その線画に合わせたような画像を出力できるようになります。

OpenPose

線画

SD3では、InstantXさんがControlNetを実装して提供してくださっています。
今回はそちらの実装を利用させていただきます。
https://huggingface.co/InstantX

2024年7月17日時点(執筆日)時点では下記のモデルが公開されています。

  • InstantX/SD3-Controlnet-Pose
  • InstantX/SD3-Controlnet-Canny
  • InstantX/SD3-Controlnet-Tile

SD3-Controlnet-Poseはposeを参照画像として入力することで、そのposeに合わせた姿勢の画像を出力してくれます。
SD3-Controlnet-Cannyは線画を参照画像として入力することで、その線画に合わせて画像を生成してくれます。
SD3-Controlnet-Tileは通常の画像を参照画像として入力することで、その画像の特徴を大まか残しながら、より質の高い画像を生成したり、画像の一部部分を変更したような画像を生成してくれます。

その他、Stable Diffusion XLなどでは、深度マップを参照画像にするようなControlNetやセグメンテーション画像を参照画像として、構図を固定するようなControlNetもありますが、まだSD3(Diffuser)では実装されていないように思います。

成果物

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

今回の実験

今回は、私自身がSD3とControlNetについて初めて触るため、好奇心にしたがって様々な実験をしています。

下記に実施した実験の内容を記載します。実験結果については最後にご紹介しています。

  • 実験1-3
    • T5のText Encoderをドロップした状態で、ControlNetを試す
      • SD3-Controlnet-Poseで実験
      • SD3-Controlnet-Cannyで実験
      • SD3-Controlnet-Tileで実験
  • 実験4-6
    • T5のText Encoderを量子化して導入した状態で、ControlNetを試す
      • SD3-Controlnet-Poseで実験
        • 生成途中の表示もあり。
      • SD3-Controlnet-Cannyで実験
      • SD3-Controlnet-Tileで実験
  • 実験7
    • 少し設定を変えて実験

事前準備

基本的にはPart8での事前準備と同様ですが、再掲します。

Hugging Faceのlogin tokenの取得と登録

SD3のモデルをローカルで利用可能にするために、Huggingfaceからログイン用のtokenを取得する必要があります。
ログインtokenの取得方法は下記の記事を参考にしてください。
https://zenn.dev/protoout/articles/73-hugging-face-setup

また、取得したログインtokenをGoogle Colabに登録する必要があります。
下記の記事を参考に登録してください。
https://note.com/npaka/n/n79bb63e17685

HF_LOGINという名前で登録してください。

SD3重みへのアクセス準備

続いて、HuggingfaceでStable Diffusionの重みにアクセスできるようにします。

https://huggingface.co/stabilityai/stable-diffusion-3-medium
上記のURLにアクセスします。

初めてアクセスした場合、上記のような画面になると思うので、作成したアカウントでログインしてください。

その後、入力フォームが表示されるかと思います。
そのフォームをすべて埋めて、提出することで、モデル重みにアクセスできるようになります。

解説

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

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

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

ディレクトリ構造

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

MyDrive/
    └ colab_AI_sample/
          └ colab_SD3ControlNet_sample/
                  ├ configs/
                  |    └ config.ini
                  ├ inputs/
                  |    └ refer.webp
                  ├ outputs/
                  ├ module/
                  |    └ module_sd3c.py
                  └ SD3ControlNet_sample.ipynb.ipynb

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

使い方解説

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

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

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

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

コード解説

主に、重要なSD3ControlNet_sample.ipynbmodule/module_sd3c.pyについて解説します。

SD3ControlNet_sample.ipynb

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

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

1セル目

./colab_AI_sample/colab_SD3ControlNet_sample/SD3ControlNet_sample.ipynb
#SD3 で必要なモジュールのインストール
!pip install -Uqq diffusers transformers ftfy accelerate bitsandbytes controlnet_aux omegaconf GPUtil

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

2セル目

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


from huggingface_hub import login
from google.colab import userdata
HF_LOGIN = userdata.get('HF_LOGIN')
login(HF_LOGIN)

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

%cd $pwd
!pwd

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

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

また、Google Colabに登録されたHF_LOGINtokenを読み込み、HuggingFaceにログインします。

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

3セル目

./colab_AI_sample/colab_SD3ControlNet_sample/SD3ControlNet_sample.ipynb
#モジュールをimportする
from module.module_sd3c import SD3C
import time

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

また、実行時間の計測のためtimeモジュールも読み込んでいます。

4セル目

./colab_AI_sample/colab_SD3ControlNet_sample/SD3ControlNet_sample.ipynb
#モデルの設定を行う。


config_text = """
[SD3C]
device = auto
n_steps=28
seed=42
shift = 3.0

model_path = stabilityai/stable-diffusion-3-medium-diffusers
controlnet_path = InstantX/SD3-Controlnet-Pose
;controlnet_path = InstantX/SD3-Controlnet-Canny
;controlnet_path = InstantX/SD3-Controlnet-Tile

guided_scale = 4.0
width = 1024
height = 1024

use_cpu_offload = False
use_text_encoder_3 = True
use_t5_quantization = True
save_latent = True

"""

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

基本的にはpart8と同様なので、差分について解説します。

  • controlnet_path = InstantX/SD3-Controlnet-Pose
    • ここで利用したいControlNetのモデルを指定します。

5セル目

./colab_AI_sample/colab_SD3ControlNet_sample/SD3ControlNet_sample.ipynb

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


main_prompt = """
Anime style, An anime digital illustration of young woman with striking red eyes, standing outside surrounded by falling snow and cherry blossoms. Her long, flowing silver hair is intricately braided, adorned an accessory made of lace, white tripetal flowers, and small, white pearls, cascading around her determined face. She wears a high-collared, form-fitting white dress with intricate lace details and cut-out sections on the sleeves, adding to the air of elegance and strength. The background features out-of-focus branches dotted with vibrant pink blossoms, set against a clear blue sky speckled with snow and bokeh, creating a dreamlike, serene atmosphere. The lighting casts a soft, almost divine glow on her, emphasizing her resolute expression. The artwork captures her from the waist up. Image is an anime style digital artwork. trending on artstation, pixiv, anime-style, anime screenshot
"""

negative_prompt=""

input_refer_image_path = "./inputs/refer.webp"
output_refer_image_path = "./inputs/refer.png"

ここでは生成されるpromptを指定しています。

下記のポストにて、公開してくださっているプロンプトを使わせていただいております。ありがとうございます。
https://x.com/Lykon4072/status/1801036331973275843

SD3では、ネガティブプロンプトを指定する必要はないため、空文字に設定しています。

また、main_promptは非常に長いプロンプトを指定しています。
Text_Encoderの最初の2つで使われているCLIPというモデルは、77tokenしか処理できないため、プロンプトの後半を切り捨てます。
一方で、3つ目のText_EncoderであるT5は、非常に長いプロンプトを処理できるため、T5を利用する場合はプロンプトのすべてが問題なく処理できます。

加えて参照画像のpathをinput_refer_image_pathで指定しています。
controlNetを利用する際に、OpenPoseの骨格や、線画を用意する必要がありますが、普通のユーザはそんなの持っていないと思いますので、ここでは普通の画像を指定すると、output_refer_image_pathに指定したpathにOpenPoseの骨格や、線画など、今回の実行で利用する参照画像を、普通の画像から生成して保存するようにしています。

すなわち、以下のような画像を入力することで、自動的にOpenPoseの骨格や、線画に変換してControlNetにて利用できます。(画像はSD3 Mediumで作成しました)

参照画像

OpenPose

線画

6セル目

./colab_AI_sample/colab_SD3ControlNet_sample/SD3ControlNet_sample.ipynb

sd = SD3C()
sd.prepare_referimage(input_refer_image_path = input_refer_image_path, output_refer_image_path = output_refer_image_path, low_threshold = 100, high_threshold = 200)

ここで、SD3Cクラスのインスタンスを作成します。

その上で、上記5セル目の説明でもしたように、通常の参照画像から、OpenPoseの骨格や、線画の画像に変換するprepare_referimageメソッドを実行しています。
こちらを実行するとoutput_refer_image_path = "./inputs/refer.png"のpathにControlNetに入力される画像が保存されます。

low_threshold = 100, high_threshold = 200の値は、線画用に画像からcanny edgeを取得する際の閾値になります。特に変更しなくても多くの画像でそこそこ良い結果になります。

7セル目

./colab_AI_sample/colab_SD3ControlNet_sample/SD3ControlNet_sample.ipynb

for i in range(3):
      start = time.time()
      image = sd.generate_image(main_prompt, neg_prompt = negative_prompt,image_path = output_refer_image_path, controlnet_conditioning_scale = 1.0)
      print("generate image time: ", time.time()-start)
      image.save("./outputs/SD3C_result_{}.png".format(i))

ここではSD3+ControlNetで画像を生成します。
3回for文を回すことで、3枚の画像を生成し、outputsフォルダに保存します。
また、1枚生成するのにかかる時間を計測して表示しています。

実行結果(生成される画像)などは、のちの実験結果の章で提示します。

画像はoutputsフォルダに保存されますが、実行のたびに同じ名前で保存されるので、過去に保存した画像は上書きされるようになりますので注意してください。
(seed値という乱数表の値を変更しない限り、実行のたびに必ず同じ画像が生成されます。seedは設定ファイル(4セル目)で変更できます。)

また、引数のcontrolnet_conditioning_scale = 1.0の値を変更することで、どの程度、参照画像が生成される画像に影響を与えるかを指定できます。
1.0はMAXの値で、0が最小の値です。0を指定すると全く参照画像の影響がなくなり、プロンプトのみによって生成されます。1.0の場合は、参照画像の影響が強くなります。
いい感じにプロンプトも反映させたい場合は、0.5-0.7付近で設定することをおすすめします。

module/module_sd3c.py

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

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

コード全文
./colab_AI_sample/colab_SD3ControlNet_sample/module/module_sd3c.py
import torch
from diffusers import StableDiffusion3ControlNetPipeline, AutoencoderTiny , FlowMatchEulerDiscreteScheduler
from diffusers.models import SD3ControlNetModel, SD3MultiControlNetModel
from diffusers.utils import load_image
from diffusers.pipelines.stable_diffusion_3.pipeline_output import StableDiffusion3PipelineOutput
from transformers import T5EncoderModel, BitsAndBytesConfig
from PIL import Image

import os
import configparser
# ファイルの存在チェック用モジュール
import errno
import time
import numpy as np
import GPUtil
import gc

class SD3Cconfig:
    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')
        SD3C_items = self.config_ini.items('SD3C')
        self.SD3C_config_dict = dict(SD3C_items)

class SD3C:
    def __init__(self,device = None, config_ini_path = './configs/config.ini'):

        SD3C_config = SD3Cconfig(config_ini_path = config_ini_path)
        config_dict = SD3C_config.SD3C_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.prompt_embeds = None
        self.negative_prompt_embeds = None
        self.pooled_prompt_embeds = None
        self.negative_pooled_prompt_embeds = None
        self.last_prompt = None
        self.last_prompt_2 = None
        self.last_prompt_3 = None
        self.last_neg_prompt = None
        self.last_neg_prompt_2 = None
        self.last_neg_prompt_3 = None

        self.n_steps = int(config_dict["n_steps"])
        self.seed = int(config_dict["seed"])
        self.generator = torch.Generator(device=self.device).manual_seed(self.seed)
        self.width = int(config_dict["width"])
        self.height = int(config_dict["height"])
        self.guided_scale = float(config_dict["guided_scale"])
        self.shift = float(config_dict["shift"])

        self.use_cpu_offload = config_dict["use_cpu_offload"]
        if self.use_cpu_offload == "True":
            self.use_cpu_offload = True
        else:
            self.use_cpu_offload = False

        self.use_text_encoder_3 = config_dict["use_text_encoder_3"]
        if self.use_text_encoder_3 == "True":
            self.use_text_encoder_3 = True
        else:
            self.use_text_encoder_3 = False

        self.use_T5_quantization = config_dict["use_t5_quantization"]
        if self.use_T5_quantization == "True":
            self.use_T5_quantization = True
        else:
            self.use_T5_quantization = False

        self.save_latent = config_dict["save_latent"]
        if self.save_latent == "True":
            self.save_latent = True
        else:
            self.save_latent = False

        self.model_path = config_dict["model_path"]
        self.controlnet_path = config_dict["controlnet_path"]


        self.pipe  = self.preprepare_model()


    def preprepare_model(self):

        pipe = None

        controlnet = SD3ControlNetModel.from_pretrained(self.controlnet_path, torch_dtype=torch.float16)
        sampler = FlowMatchEulerDiscreteScheduler(
                    shift = self.shift
                    )

        if self.use_text_encoder_3:
            if self.use_T5_quantization:
                quantization_config = BitsAndBytesConfig(load_in_8bit=True)

                self.text_encoder = T5EncoderModel.from_pretrained(
                    self.model_path,
                    subfolder="text_encoder_3",
                    quantization_config=quantization_config,
                )
                self.prepipe = StableDiffusion3ControlNetPipeline.from_pretrained(
                    self.model_path,
                    controlnet=None,
                    scheduler = None,
                    text_encoder_3=self.text_encoder,
                    transformer=None,
                    vae=None,
                    device_map="balanced"
                )

                pipe = StableDiffusion3ControlNetPipeline.from_pretrained(
                  self.model_path,
                  controlnet=controlnet,
                  scheduler = sampler,
                  text_encoder=None,
                  text_encoder_2=None,
                  text_encoder_3=None,
                  tokenizer=None,
                  tokenizer_2=None,
                  tokenizer_3=None,
                  torch_dtype=torch.float16
                )


            else:
                pipe = StableDiffusion3ControlNetPipeline.from_pretrained(
                    self.model_path,
                    controlnet=controlnet,
                    scheduler = sampler,
                    torch_dtype=torch.float16)
        else:
            pipe = StableDiffusion3ControlNetPipeline.from_pretrained(
                        self.model_path,
                        controlnet=controlnet,
                        scheduler = sampler,
                        text_encoder_3=None,
                        tokenizer_3=None,
                        torch_dtype=torch.float16)


        if self.use_T5_quantization:
            pipe = pipe.to("cpu")
        elif self.use_cpu_offload:
            pipe.enable_model_cpu_offload()
        else:
            pipe = pipe.to("cuda")

        print(pipe.scheduler.config)

        return pipe


    def prepare_referimage(self,input_refer_image_path,output_refer_image_path, low_threshold = 100, high_threshold = 200):
        def extract_last_segment(text):
            segments = text.split('-')
            return segments[-1]

        mode = extract_last_segment(self.controlnet_path)

        def prepare_openpose(input_refer_image_path,output_refer_image_path):
            from diffusers.utils import load_image
            from controlnet_aux import OpenposeDetector

            # 初期画像の準備
            init_image = load_image(input_refer_image_path)
            init_image = init_image.resize((self.width, self.height))

            # コントロール画像の準備
            openpose_detector = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
            openpose_image = openpose_detector(init_image)

            openpose_image.save(output_refer_image_path)

        def prepare_canny(input_refer_image_path,output_refer_image_path, low_threshold = 100, high_threshold = 200):
            init_image = load_image(input_refer_image_path)
            init_image = init_image.resize((self.width, self.height))
            import cv2


            # コントロールイメージを作成するメソッド
            def make_canny_condition(image, low_threshold = 100, high_threshold = 200):
                image = np.array(image)
                image = cv2.Canny(image, low_threshold, high_threshold)
                image = image[:, :, None]
                image = np.concatenate([image, image, image], axis=2)
                return Image.fromarray(image)

            control_image = make_canny_condition(init_image, low_threshold, high_threshold)
            control_image.save(output_refer_image_path)

        if mode == "Pose":
            prepare_openpose(input_refer_image_path,output_refer_image_path)
        elif mode == "Canny":
            prepare_canny(input_refer_image_path,output_refer_image_path, low_threshold = low_threshold, high_threshold = high_threshold)
        else:
            init_image = load_image(input_refer_image_path)
            init_image.save(output_refer_image_path)
            
    def check_prepipe(self):
        if hasattr(self, 'prepipe'):
            print("prepipe exists.")
        else:
            print("prepipe does not exist.")
            self.pipe = self.pipe.to("cpu")
            quantization_config = BitsAndBytesConfig(load_in_8bit=True)

            self.text_encoder = T5EncoderModel.from_pretrained(
                self.model_path,
                subfolder="text_encoder_3",
                quantization_config=quantization_config,
            )
            self.prepipe = StableDiffusion3ControlNetPipeline.from_pretrained(
                self.model_path,
                controlnet=None,
                scheduler = None,
                text_encoder_3=self.text_encoder,
                transformer=None,
                vae=None,
                device_map="balanced"
            )
            


    def generate_image(self, prompt, prompt_2 = None, prompt_3 = None, neg_prompt = "", neg_prompt_2 = None, neg_prompt_3 = None, image_path = None,  seed = None, controlnet_conditioning_scale = 1.0):

        def decode_tensors(pipe, step, timestep, callback_kwargs):
            latents = callback_kwargs["latents"]

            image = latents_to_rgb(latents,pipe)
            gettime = time.time()
            formatted_time_human_readable = time.strftime("%Y%m%d_%H%M%S", time.localtime(gettime))
            image.save(f"./outputs/latent_{formatted_time_human_readable}_{step}_{timestep}.png")


            return callback_kwargs

        def latents_to_rgb(latents,pipe):

            latents = (latents / pipe.vae.config.scaling_factor) + pipe.vae.config.shift_factor

            img = pipe.vae.decode(latents, return_dict=False)[0]
            img = pipe.image_processor.postprocess(img, output_type="pil")

            return StableDiffusion3PipelineOutput(images=img).images[0]

        if seed is not None:
            self.generator = torch.Generator(device=self.device).manual_seed(seed)

        def check_prompts(prompt, prompt_2, prompt_3, neg_prompt, neg_prompt_2, neg_prompt_3):
            # タプルに変数を格納
            last_prompts = (self.last_prompt, self.last_prompt_2, self.last_prompt_3,
                            self.last_neg_prompt, self.last_neg_prompt_2, self.last_neg_prompt_3)

            current_prompts = (prompt, prompt_2, prompt_3,
                            neg_prompt, neg_prompt_2, neg_prompt_3)

            # 全ての変数が一致しているか判定
            return last_prompts == current_prompts

        if image_path is None:
            raise ValueError("ControlNetを利用する場合は、画像のパスをimage_path引数に提示してください。")
        control_image = load_image(image_path)

        if prompt_2 is None:
            prompt_2 = prompt
        if prompt_3 is None:
            prompt_3 = prompt
        if neg_prompt_2 is None:
            neg_prompt_2 = neg_prompt
        if neg_prompt_3 is None:
            neg_prompt_3 = neg_prompt

        image = None
        #T5を利用して、かつ、量子化する場合
        if (self.use_T5_quantization) and (self.use_text_encoder_3):
            #GPUメモリを解放する都合上、繰り返しメソッドを実行するために、T5の量子化を行う場合はプロンプトの埋め込みを保存しておく必要がある。
            #前回の実行とプロンプトが一つでも異なる場合は、そのプロンプトをインスタンス変数に一旦保存した上で、埋め込みを計算し、インスタンス変数に保存する。
            if not check_prompts(prompt, prompt_2, prompt_3, neg_prompt, neg_prompt_2, neg_prompt_3):
                self.check_prepipe()
                
                self.last_prompt = prompt
                self.last_prompt_2 = prompt_2
                self.last_prompt_3 = prompt_3
                self.last_neg_prompt = neg_prompt
                self.last_neg_prompt_2 = neg_prompt_2
                self.last_neg_prompt_3 = neg_prompt_3

                with torch.no_grad():
                    (
                    self.prompt_embeds,
                    self.negative_prompt_embeds,
                    self.pooled_prompt_embeds,
                    self.negative_pooled_prompt_embeds,
                    ) = self.prepipe.encode_prompt(
                        prompt=prompt,
                        prompt_2=prompt_2,
                        prompt_3=prompt_3,
                        negative_prompt=neg_prompt,
                        negative_prompt_2 = neg_prompt_2,
                        negative_prompt_3 = neg_prompt_3,
                        )

                #無理やりGPUメモリを解放する。GoogleColabだと何故かtorch.cuda.empty_cache()を2回実行しないと解放されない。
                GPUtil.showUtilization()
                del self.prepipe
                del self.text_encoder
                torch.cuda.empty_cache()
                gc.collect()
                torch.cuda.empty_cache()
                GPUtil.showUtilization()
                ### ここまで

                #メモリが解放されたので、次のモデルをGPUでロード
                self.pipe = self.pipe.to("cuda")

            #プロンプトが一致している場合は、埋め込みを再利用する。
            else:
                pass

            if self.save_latent:
                image = self.pipe(
                    prompt_embeds=self.prompt_embeds.half(),
                    negative_prompt_embeds=self.negative_prompt_embeds.half(),
                    pooled_prompt_embeds=self.pooled_prompt_embeds.half(),
                    negative_pooled_prompt_embeds=self.negative_pooled_prompt_embeds.half(),
                    control_image = control_image,
                    controlnet_conditioning_scale=controlnet_conditioning_scale,
                    height = self.height,
                    width = self.width,
                    num_inference_steps=self.n_steps,
                    guidance_scale=self.guided_scale,
                    generator=self.generator,
                    callback_on_step_end=decode_tensors,
                    callback_on_step_end_tensor_inputs=["latents"],
                    ).images[0]

            else:
                image = self.pipe(
                        prompt_embeds=self.prompt_embeds.half(),
                        negative_prompt_embeds=self.negative_prompt_embeds.half(),
                        pooled_prompt_embeds=self.pooled_prompt_embeds.half(),
                        negative_pooled_prompt_embeds=self.negative_pooled_prompt_embeds.half(),
                        control_image = control_image,
                        controlnet_conditioning_scale=controlnet_conditioning_scale,
                        height = self.height,
                        width = self.width,
                        num_inference_steps=self.n_steps,
                        guidance_scale=self.guided_scale,
                        generator=self.generator
                        ).images[0]

        #T5を利用しない、もしくは、量子化しない場合
        else:
            if self.save_latent:
                image = self.pipe(
                    prompt=prompt,
                    prompt_2=prompt_2,
                    prompt_3=prompt_3,
                    negative_prompt=neg_prompt,
                    negative_prompt_2 = neg_prompt_2,
                    negative_prompt_3 = neg_prompt_3,
                    control_image = control_image,
                    controlnet_conditioning_scale=controlnet_conditioning_scale,
                    height = self.height,
                    width = self.width,
                    num_inference_steps=self.n_steps,
                    guidance_scale=self.guided_scale,
                    generator=self.generator,
                    callback_on_step_end=decode_tensors,
                    callback_on_step_end_tensor_inputs=["latents"],
                    ).images[0]
            else:
                image = self.pipe(
                    prompt=prompt,
                    prompt_2=prompt_2,
                    prompt_3=prompt_3,
                    negative_prompt=neg_prompt,
                    negative_prompt_2 = neg_prompt_2,
                    negative_prompt_3 = neg_prompt_3,
                    control_image = control_image,
                    controlnet_conditioning_scale=controlnet_conditioning_scale,
                    height = self.height,
                    width = self.width,
                    num_inference_steps=self.n_steps,
                    guidance_scale=self.guided_scale,
                    generator=self.generator
                    ).images[0]


        return image




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

SD3Cconfigクラス

./colab_AI_sample/colab_SD3ControlNet_sample/module/module_sd3c.py

class SD3Cconfig:
    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')
        SD3C_items = self.config_ini.items('SD3C')
        self.SD3C_config_dict = dict(SD3C_items)

基本的にPart8の内容と同じです。

SD3Cクラスのinitメソッド

./colab_AI_sample/colab_SD3ControlNet_sample/module/module_sd3c.py
class SD3C:
    def __init__(self,device = None, config_ini_path = './configs/config.ini'):

        SD3C_config = SD3Cconfig(config_ini_path = config_ini_path)
        config_dict = SD3C_config.SD3C_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.prompt_embeds = None
        self.negative_prompt_embeds = None
        self.pooled_prompt_embeds = None
        self.negative_pooled_prompt_embeds = None
        self.last_prompt = None
        self.last_prompt_2 = None
        self.last_prompt_3 = None
        self.last_neg_prompt = None
        self.last_neg_prompt_2 = None
        self.last_neg_prompt_3 = None

        self.n_steps = int(config_dict["n_steps"])
        self.seed = int(config_dict["seed"])
        self.generator = torch.Generator(device=self.device).manual_seed(self.seed)
        self.width = int(config_dict["width"])
        self.height = int(config_dict["height"])
        self.guided_scale = float(config_dict["guided_scale"])
        self.shift = float(config_dict["shift"])

        self.use_cpu_offload = config_dict["use_cpu_offload"]
        if self.use_cpu_offload == "True":
            self.use_cpu_offload = True
        else:
            self.use_cpu_offload = False

        self.use_text_encoder_3 = config_dict["use_text_encoder_3"]
        if self.use_text_encoder_3 == "True":
            self.use_text_encoder_3 = True
        else:
            self.use_text_encoder_3 = False

        self.use_T5_quantization = config_dict["use_t5_quantization"]
        if self.use_T5_quantization == "True":
            self.use_T5_quantization = True
        else:
            self.use_T5_quantization = False

        self.save_latent = config_dict["save_latent"]
        if self.save_latent == "True":
            self.save_latent = True
        else:
            self.save_latent = False

        self.model_path = config_dict["model_path"]
        self.controlnet_path = config_dict["controlnet_path"]


        self.pipe  = self.preprepare_model()

基本的にpart8の内容と同じですが、私の怠慢で、モデルのcompileを消去したのと、後述するプロンプトの埋め込みなどの初期化を行なっています。

モデルのcompileに関しては、Part8と同様な形で実装すれば問題なく実装できるかと思います。

SD3Cクラスのpreprepare_modelメソッド

./colab_AI_sample/colab_SD3ControlNet_sample/module/module_sd3c.py
class SD3:
   ・・・
    def preprepare_model(self):

        pipe = None
        controlnet = SD3ControlNetModel.from_pretrained(self.controlnet_path, torch_dtype=torch.float16)
        sampler = FlowMatchEulerDiscreteScheduler(
                    shift = self.shift
                    )

        if self.use_text_encoder_3:
            if self.use_T5_quantization:
                quantization_config = BitsAndBytesConfig(load_in_8bit=True)

                self.text_encoder = T5EncoderModel.from_pretrained(
                    self.model_path,
                    subfolder="text_encoder_3",
                    quantization_config=quantization_config,
                )
                self.prepipe = StableDiffusion3ControlNetPipeline.from_pretrained(
                    self.model_path,
                    controlnet=None,
                    scheduler = None,
                    text_encoder_3=self.text_encoder,
                    transformer=None,
                    vae=None,
                    device_map="balanced"
                )

                pipe = StableDiffusion3ControlNetPipeline.from_pretrained(
                  self.model_path,
                  controlnet=controlnet,
                  scheduler = sampler,
                  text_encoder=None,
                  text_encoder_2=None,
                  text_encoder_3=None,
                  tokenizer=None,
                  tokenizer_2=None,
                  tokenizer_3=None,
                  torch_dtype=torch.float16
                )


            else:
                pipe = StableDiffusion3ControlNetPipeline.from_pretrained(
                    self.model_path,
                    controlnet=controlnet,
                    scheduler = sampler,
                    torch_dtype=torch.float16)
        else:
            pipe = StableDiffusion3ControlNetPipeline.from_pretrained(
                        self.model_path,
                        controlnet=controlnet,
                        scheduler = sampler,
                        text_encoder_3=None,
                        tokenizer_3=None,
                        torch_dtype=torch.float16)


        if self.use_T5_quantization:
            pipe = pipe.to("cpu")
        elif self.use_cpu_offload:
            pipe.enable_model_cpu_offload()
        else:
            pipe = pipe.to("cuda")

        print(pipe.scheduler.config)

        return pipe

基本的にPart8の記事の内容と同じように、設定ファイルの内容に合わせて、必要なモデルを読み込んでいます。
前回記事と変わっている部分としては、まず下記の通りControlNetのモデルを読み込んでいます。

controlnet = SD3ControlNetModel.from_pretrained(self.controlnet_path, torch_dtype=torch.float16)

さらに使うPipelineもStableDiffusion3ControlNetPipelineに変更になっております。こちらのPipelineを利用することで、ControlNetをDiffuserで利用することができます。

そして、特に前回記事と大きく変わっている部分は、T5 Text Encoderを量子化した場合です。
少し込み入った話になるため、興味のある方は下記を展開してご覧ください。

T5 Text Encoderを量子化した場合の実装

GPUのモデルが潤沢にある場合は、量子化などもせずシンプルに下記のような形で利用できます。

pipe = StableDiffusion3ControlNetPipeline.from_pretrained(
                    self.model_path,
                    controlnet=controlnet,
                    scheduler = sampler,
                    torch_dtype=torch.float16)

実際に、公式の記事でも、このような使い方を紹介しています。
しかしながら、無料版のGoogle Colabでは16GBメモリのGPUしか利用できないため、T5を量子化する必要があります。

量子化した場合、シンプルな実装は下記のような形です。

quantization_config = BitsAndBytesConfig(load_in_8bit=True)

self.text_encoder = T5EncoderModel.from_pretrained(
    self.model_path,
    subfolder="text_encoder_3",
    quantization_config=quantization_config,
)

pipe = StableDiffusion3ControlNetPipeline.from_pretrained(
  self.model_path,
  controlnet=controlnet,
  scheduler = sampler,
  text_encoder_3=self.text_encoder,
  torch_dtype=torch.float16
)

しかしながら、上記のような形で実装をすると下記のようなエラーが表示されます。

RuntimeError: Input type (torch.cuda.HalfTensor) and weight type (torch.HalfTensor) should be the same

これは、入力テンソルと重みテンソルの型が一致していないことを示しています。
具体的には、入力テンソルがCUDAのハーフ精度テンソル(torch.cuda.HalfTensor)である一方で、重みテンソルがCPUのハーフ精度テンソル(torch.HalfTensor)であるために、計算が実行できない状況ということを示しています。

つまり、現時点ではStableDiffusion3ControlNetPipelineにおいて、T5の量子化がサポートされていない状況です。
したがって下記のようのモデルを定義しています。

quantization_config = BitsAndBytesConfig(load_in_8bit=True)

self.text_encoder = T5EncoderModel.from_pretrained(
    self.model_path,
    subfolder="text_encoder_3",
    quantization_config=quantization_config,
)
self.prepipe = StableDiffusion3ControlNetPipeline.from_pretrained(
    self.model_path,
    controlnet=None,
    scheduler = None,
    text_encoder_3=self.text_encoder,
    transformer=None,
    vae=None,
    device_map="balanced"
)

pipe = StableDiffusion3ControlNetPipeline.from_pretrained(
  self.model_path,
  controlnet=controlnet,
  scheduler = sampler,
  text_encoder=None,
  text_encoder_2=None,
  text_encoder_3=None,
  tokenizer=None,
  tokenizer_2=None,
  tokenizer_3=None,
  torch_dtype=torch.float16
)

上記では、
self.text_encoderBitsAndBytesにより量子化されます。
それをself.prepipeにて、Text Encoderだけを読み込み、transfomerやVAEは無視して、StableDiffusion3ControlNetPipelineを定義しています。

最後に、pipeにて、Text Encoderだけを無視してStableDiffusion3ControlNetPipelineを定義しています。

何をしているかというと、後述のコードで触れますが、Text EncoderとTransfomerやVAEを分けることで、テキストの埋め込みの生成と、画像の生成を分けて処理を行おうとしています。

量子化しているので、あくまでText Encoderの部分であり、ControlNetを利用しているのは、その後の部分(Transfomerの部分)になるので、分離して定義して処理を行うことで、Text Encoderの量子化の影響をControlNetが受けないように実装しています。

また、3つのモデルをすべてGPUに置くことが、16GBのVRAMではできないので、pipeに関しては、ここではcpuメモリにモデルをおいています。

SD3Cクラスのprepare_referimageメソッド

./colab_AI_sample/colab_SD3ControlNet_sample/module/module_sd3c.py
class SD3:
   ・・・
    def prepare_referimage(self,input_refer_image_path,output_refer_image_path, low_threshold = 100, high_threshold = 200):
        def extract_last_segment(text):
            segments = text.split('-')
            return segments[-1]

        mode = extract_last_segment(self.controlnet_path)

        def prepare_openpose(input_refer_image_path,output_refer_image_path):
            from diffusers.utils import load_image
            from controlnet_aux import OpenposeDetector

            # 初期画像の準備
            init_image = load_image(input_refer_image_path)
            init_image = init_image.resize((self.width, self.height))

            # コントロール画像の準備
            openpose_detector = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
            openpose_image = openpose_detector(init_image)

            openpose_image.save(output_refer_image_path)

        def prepare_canny(input_refer_image_path,output_refer_image_path, low_threshold = 100, high_threshold = 200):
            init_image = load_image(input_refer_image_path)
            init_image = init_image.resize((self.width, self.height))
            import cv2


            # コントロールイメージを作成するメソッド
            def make_canny_condition(image, low_threshold = 100, high_threshold = 200):
                image = np.array(image)
                image = cv2.Canny(image, low_threshold, high_threshold)
                image = image[:, :, None]
                image = np.concatenate([image, image, image], axis=2)
                return Image.fromarray(image)

            control_image = make_canny_condition(init_image, low_threshold, high_threshold)
            control_image.save(output_refer_image_path)

        if mode == "Pose":
            prepare_openpose(input_refer_image_path,output_refer_image_path)
        elif mode == "Canny":
            prepare_canny(input_refer_image_path,output_refer_image_path, low_threshold = low_threshold, high_threshold = high_threshold)
        else:
            init_image = load_image(input_refer_image_path)
            init_image.save(output_refer_image_path)

これは、今回新しく追加されたメソッドです。
下記のような普通の画像を、OpenPoseの骨格や線画に変換するメソッドです。

参照画像

OpenPose

線画

ここではまず下記で、参照画像をどのように変換するべきかを設定します。

mode = extract_last_segment(self.controlnet_path)

ControlNetのモデルの名前(path)は下記のようになっています。

controlnet_path = InstantX/SD3-Controlnet-Pose
controlnet_path = InstantX/SD3-Controlnet-Canny
controlnet_path = InstantX/SD3-Controlnet-Tile

したがって、文末を取得することで、どのようなControlNetを利用して、どのような参照画像が必要かがわかります。
そしてmodeの値によって下記のように、参照画像を生成しています。

if mode == "Pose":
    prepare_openpose(input_refer_image_path,output_refer_image_path)
elif mode == "Canny":
    prepare_canny(input_refer_image_path,output_refer_image_path, low_threshold = low_threshold, high_threshold = high_threshold)
else:
    init_image = load_image(input_refer_image_path)
    init_image.save(output_refer_image_path)

Poseの場合はOpenPoseの骨格画像をoutput_refer_image_pathに保存します
Cannyの場合はCanny Edgeを取得して同様に保存します。
それ以外の場合(おそらくTileの場合は、同じ画像を利用するため、そのまま保存します。)

SD3Cクラスのgenerate_imageメソッド

./colab_AI_sample/colab_SD3ControlNet_sample/module/module_sd3c.py
class SD3:
   ・・・
    def generate_image(self, prompt, prompt_2 = None, prompt_3 = None, neg_prompt = "", neg_prompt_2 = None, neg_prompt_3 = None, image_path = None,  seed = None, controlnet_conditioning_scale = 1.0):

        def decode_tensors(pipe, step, timestep, callback_kwargs):
            latents = callback_kwargs["latents"]

            image = latents_to_rgb(latents,pipe)
            gettime = time.time()
            formatted_time_human_readable = time.strftime("%Y%m%d_%H%M%S", time.localtime(gettime))
            image.save(f"./outputs/latent_{formatted_time_human_readable}_{step}_{timestep}.png")


            return callback_kwargs

        def latents_to_rgb(latents,pipe):

            latents = (latents / pipe.vae.config.scaling_factor) + pipe.vae.config.shift_factor

            img = pipe.vae.decode(latents, return_dict=False)[0]
            img = pipe.image_processor.postprocess(img, output_type="pil")

            return StableDiffusion3PipelineOutput(images=img).images[0]

        if seed is not None:
            self.generator = torch.Generator(device=self.device).manual_seed(seed)

        def check_prompts(prompt, prompt_2, prompt_3, neg_prompt, neg_prompt_2, neg_prompt_3):
            # タプルに変数を格納
            last_prompts = (self.last_prompt, self.last_prompt_2, self.last_prompt_3,
                            self.last_neg_prompt, self.last_neg_prompt_2, self.last_neg_prompt_3)

            current_prompts = (prompt, prompt_2, prompt_3,
                            neg_prompt, neg_prompt_2, neg_prompt_3)

            # 全ての変数が一致しているか判定
            return last_prompts == current_prompts

        if image_path is None:
            raise ValueError("ControlNetを利用する場合は、画像のパスをimage_path引数に提示してください。")
        control_image = load_image(image_path)

        if prompt_2 is None:
            prompt_2 = prompt
        if prompt_3 is None:
            prompt_3 = prompt
        if neg_prompt_2 is None:
            neg_prompt_2 = neg_prompt
        if neg_prompt_3 is None:
            neg_prompt_3 = neg_prompt

        image = None
        #T5を利用して、かつ、量子化する場合
        if (self.use_T5_quantization) and (self.use_text_encoder_3):
            #GPUメモリを解放する都合上、繰り返しメソッドを実行するために、T5の量子化を行う場合はプロンプトの埋め込みを保存しておく必要がある。
            #前回の実行とプロンプトが一つでも異なる場合は、そのプロンプトをインスタンス変数に一旦保存した上で、埋め込みを計算し、インスタンス変数に保存する。
            if not check_prompts(prompt, prompt_2, prompt_3, neg_prompt, neg_prompt_2, neg_prompt_3):
                self.check_prepipe()

                self.last_prompt = prompt
                self.last_prompt_2 = prompt_2
                self.last_prompt_3 = prompt_3
                self.last_neg_prompt = neg_prompt
                self.last_neg_prompt_2 = neg_prompt_2
                self.last_neg_prompt_3 = neg_prompt_3

                with torch.no_grad():
                    (
                    self.prompt_embeds,
                    self.negative_prompt_embeds,
                    self.pooled_prompt_embeds,
                    self.negative_pooled_prompt_embeds,
                    ) = self.prepipe.encode_prompt(
                        prompt=prompt,
                        prompt_2=prompt_2,
                        prompt_3=prompt_3,
                        negative_prompt=neg_prompt,
                        negative_prompt_2 = neg_prompt_2,
                        negative_prompt_3 = neg_prompt_3,
                        )

                #無理やりGPUメモリを解放する。GoogleColabだと何故かtorch.cuda.empty_cache()を2回実行しないと解放されない。
                GPUtil.showUtilization()
                del self.prepipe
                del self.text_encoder
                torch.cuda.empty_cache()
                gc.collect()
                torch.cuda.empty_cache()
                GPUtil.showUtilization()
                ### ここまで

                #メモリが解放されたので、次のモデルをGPUでロード
                self.pipe = self.pipe.to("cuda")

            #プロンプトが一致している場合は、埋め込みを再利用する。
            else:
                pass

            if self.save_latent:
                image = self.pipe(
                    prompt_embeds=self.prompt_embeds.half(),
                    negative_prompt_embeds=self.negative_prompt_embeds.half(),
                    pooled_prompt_embeds=self.pooled_prompt_embeds.half(),
                    negative_pooled_prompt_embeds=self.negative_pooled_prompt_embeds.half(),
                    control_image = control_image,
                    controlnet_conditioning_scale=controlnet_conditioning_scale,
                    height = self.height,
                    width = self.width,
                    num_inference_steps=self.n_steps,
                    guidance_scale=self.guided_scale,
                    generator=self.generator,
                    callback_on_step_end=decode_tensors,
                    callback_on_step_end_tensor_inputs=["latents"],
                    ).images[0]

            else:
                image = self.pipe(
                        prompt_embeds=self.prompt_embeds.half(),
                        negative_prompt_embeds=self.negative_prompt_embeds.half(),
                        pooled_prompt_embeds=self.pooled_prompt_embeds.half(),
                        negative_pooled_prompt_embeds=self.negative_pooled_prompt_embeds.half(),
                        control_image = control_image,
                        controlnet_conditioning_scale=controlnet_conditioning_scale,
                        height = self.height,
                        width = self.width,
                        num_inference_steps=self.n_steps,
                        guidance_scale=self.guided_scale,
                        generator=self.generator
                        ).images[0]

        #T5を利用しない、もしくは、量子化しない場合
        else:
            if self.save_latent:
                image = self.pipe(
                    prompt=prompt,
                    prompt_2=prompt_2,
                    prompt_3=prompt_3,
                    negative_prompt=neg_prompt,
                    negative_prompt_2 = neg_prompt_2,
                    negative_prompt_3 = neg_prompt_3,
                    control_image = control_image,
                    controlnet_conditioning_scale=controlnet_conditioning_scale,
                    height = self.height,
                    width = self.width,
                    num_inference_steps=self.n_steps,
                    guidance_scale=self.guided_scale,
                    generator=self.generator,
                    callback_on_step_end=decode_tensors,
                    callback_on_step_end_tensor_inputs=["latents"],
                    ).images[0]
            else:
                image = self.pipe(
                    prompt=prompt,
                    prompt_2=prompt_2,
                    prompt_3=prompt_3,
                    negative_prompt=neg_prompt,
                    negative_prompt_2 = neg_prompt_2,
                    negative_prompt_3 = neg_prompt_3,
                    control_image = control_image,
                    controlnet_conditioning_scale=controlnet_conditioning_scale,
                    height = self.height,
                    width = self.width,
                    num_inference_steps=self.n_steps,
                    guidance_scale=self.guided_scale,
                    generator=self.generator
                    ).images[0]


        return image

このメソッドは、ここまでで読み込んだモデルと設定を利用して、実際に画像を生成するメソッドです。
基本的には前回のPart8の記事で実装した内容と同じです。
self.pipe(StableDiffusion3ControlNetPipeline)のcallの使い方は下記の公式ドキュメントを参考にしています。(control_imageなど)
https://huggingface.co/docs/diffusers/api/pipelines/controlnet_sd3

今回の実装で大きく前回から変わっている部分は下記の部分です。
preprepare_modelメソッドの説明の中でも、トグルの中で記載していますが、T5を量子化した場合は、通常の方法で処理するとエラーが発生するため、プロンプトの埋め込み取得と画像生成の部分を分けて処理する必要があります。

加えて、GoogleColabの無料版のGPUで動作させるために、不要になったメモリを無理やり解放するなどの苦肉の実装をしています。内容の詳細に興味のある方だけ下記をご覧ください。

T5を量子化して利用する場合

該当部分のコードを改めて下記に示します。

#T5を利用して、かつ、量子化する場合
        if (self.use_T5_quantization) and (self.use_text_encoder_3):
            #GPUメモリを解放する都合上、繰り返しメソッドを実行するために、T5の量子化を行う場合はプロンプトの埋め込みを保存しておく必要がある。
            #前回の実行とプロンプトが一つでも異なる場合は、そのプロンプトをインスタンス変数に一旦保存した上で、埋め込みを計算し、インスタンス変数に保存する。
            if not check_prompts(prompt, prompt_2, prompt_3, neg_prompt, neg_prompt_2, neg_prompt_3):
                self.check_prepipe()

                self.last_prompt = prompt
                self.last_prompt_2 = prompt_2
                self.last_prompt_3 = prompt_3
                self.last_neg_prompt = neg_prompt
                self.last_neg_prompt_2 = neg_prompt_2
                self.last_neg_prompt_3 = neg_prompt_3

                with torch.no_grad():
                    (
                    self.prompt_embeds,
                    self.negative_prompt_embeds,
                    self.pooled_prompt_embeds,
                    self.negative_pooled_prompt_embeds,
                    ) = self.prepipe.encode_prompt(
                        prompt=prompt,
                        prompt_2=prompt_2,
                        prompt_3=prompt_3,
                        negative_prompt=neg_prompt,
                        negative_prompt_2 = neg_prompt_2,
                        negative_prompt_3 = neg_prompt_3,
                        )

                #無理やりGPUメモリを解放する。GoogleColabだと何故かtorch.cuda.empty_cache()を2回実行しないと解放されない。
                GPUtil.showUtilization()
                del self.prepipe
                del self.text_encoder
                torch.cuda.empty_cache()
                gc.collect()
                torch.cuda.empty_cache()
                GPUtil.showUtilization()
                ### ここまで

                #メモリが解放されたので、次のモデルをGPUでロード
                self.pipe = self.pipe.to("cuda")

            #プロンプトが一致している場合は、埋め込みを再利用する。
            else:
                pass

            if self.save_latent:
                image = self.pipe(
                    prompt_embeds=self.prompt_embeds.half(),
                    negative_prompt_embeds=self.negative_prompt_embeds.half(),
                    pooled_prompt_embeds=self.pooled_prompt_embeds.half(),
                    negative_pooled_prompt_embeds=self.negative_pooled_prompt_embeds.half(),
                    control_image = control_image,
                    controlnet_conditioning_scale=controlnet_conditioning_scale,
                    height = self.height,
                    width = self.width,
                    num_inference_steps=self.n_steps,
                    guidance_scale=self.guided_scale,
                    generator=self.generator,
                    callback_on_step_end=decode_tensors,
                    callback_on_step_end_tensor_inputs=["latents"],
                    ).images[0]

            else:
                image = self.pipe(
                        prompt_embeds=self.prompt_embeds.half(),
                        negative_prompt_embeds=self.negative_prompt_embeds.half(),
                        pooled_prompt_embeds=self.pooled_prompt_embeds.half(),
                        negative_pooled_prompt_embeds=self.negative_pooled_prompt_embeds.half(),
                        control_image = control_image,
                        controlnet_conditioning_scale=controlnet_conditioning_scale,
                        height = self.height,
                        width = self.width,
                        num_inference_steps=self.n_steps,
                        guidance_scale=self.guided_scale,
                        generator=self.generator
                        ).images[0]

preprepare_modelメソッドで2つに分けてモデルを定義しています。
一つが、Text Encoderのみを持つself.prepipe
もう一つがTransfomerとVAEを持つself.pipeです。

まずは、今回使うモデルの確認と、プロンプトを保存します。

if not check_prompts(prompt, prompt_2, prompt_3, neg_prompt, neg_prompt_2, neg_prompt_3):
    self.check_prepipe()

    self.last_prompt = prompt
    self.last_prompt_2 = prompt_2
    self.last_prompt_3 = prompt_3
    self.last_neg_prompt = neg_prompt
    self.last_neg_prompt_2 = neg_prompt_2
    self.last_neg_prompt_3 = neg_prompt_3

上記では、クラスインスタンスのプロンプトと、引数として利用したプロンプトが一致しない場合は、今回のプロンプトを保存して、下記の埋め込み処理を実施するif文になります。
(つまり、2回目以降同じプロンプトで実行した場合は、TextEncoderをスキップする)

またself.check_prepipe()メソッドは下記のように実装しています。

    def check_prepipe(self):
        if hasattr(self, 'prepipe'):
            print("prepipe exists.")
        else:
            print("prepipe does not exist.")
            self.pipe = self.pipe.to("cpu")
            quantization_config = BitsAndBytesConfig(load_in_8bit=True)

            self.text_encoder = T5EncoderModel.from_pretrained(
                self.model_path,
                subfolder="text_encoder_3",
                quantization_config=quantization_config,
            )
            self.prepipe = StableDiffusion3ControlNetPipeline.from_pretrained(
                self.model_path,
                controlnet=None,
                scheduler = None,
                text_encoder_3=self.text_encoder,
                transformer=None,
                vae=None,
                device_map="balanced"
            )

これは、self.prepipeが存在するかどうかを確認します。通常はinitメソッドで定義しているため存在しますが、後述する流れ中でモデルをメモリ容量を削減するために、削除することになります。その後に、再びプロンプトを変えて本メソッドを繰り返し実行した場合、改めてモデルを定義するメソッドです。

続いて、上記のifがTrueの場合、プロンプトから埋め込みを取得します。

with torch.no_grad():
    (
    self.prompt_embeds,
    self.negative_prompt_embeds,
    self.pooled_prompt_embeds,
    self.negative_pooled_prompt_embeds,
    ) = self.prepipe.encode_prompt(
        prompt=prompt,
        prompt_2=prompt_2,
        prompt_3=prompt_3,
        negative_prompt=neg_prompt,
        negative_prompt_2 = neg_prompt_2,
        negative_prompt_3 = neg_prompt_3,
        )

その後、不要になったText Encoderなどのメモリを破棄して、VRAMに余裕を作ります。

#無理やりGPUメモリを解放する。GoogleColabだと何故かtorch.cuda.empty_cache()を2回実行しないと解放されない。
GPUtil.showUtilization()
del self.prepipe
del self.text_encoder
torch.cuda.empty_cache()
gc.collect()
torch.cuda.empty_cache()
GPUtil.showUtilization()
### ここまで

#メモリが解放されたので、次のモデルをGPUでロード
self.pipe = self.pipe.to("cuda")

本来であればtorch.cuda.empty_cache()は一回呼べば十分なのですが、何回実験しても、2回呼ばないとメモリが空かなかったので、2回読んでいます。原因わかる人いたら教えてください。
さらに、その後の画像生成部をGPUメモリにおいて、処理に備えます。

その後は下記のような形で、プロンプトの埋め込みを入力して画像を生成します。

image = self.pipe(
    prompt_embeds=self.prompt_embeds.half(),
    negative_prompt_embeds=self.negative_prompt_embeds.half(),
    pooled_prompt_embeds=self.pooled_prompt_embeds.half(),
    negative_pooled_prompt_embeds=self.negative_pooled_prompt_embeds.half(),
    control_image = control_image,
    controlnet_conditioning_scale=controlnet_conditioning_scale,
    height = self.height,
    width = self.width,
    num_inference_steps=self.n_steps,
    guidance_scale=self.guided_scale,
    generator=self.generator
    ).images[0]

実験結果

ここからは、上記のコードによってGoogle Colabでパラメータを変更して、様々な実験を実施したため、その詳細を記載します。

実験の前提

今後、明示的に言及しない場合は、下記の前提に則って実験を行なっております。

PoseとCannyに関して

プロンプトはPart8と同様に下記を利用します。

Anime style, An anime digital illustration of young woman with striking red eyes, standing outside surrounded by falling snow and cherry blossoms. Her long, flowing silver hair is intricately braided, adorned an accessory made of lace, white tripetal flowers, and small, white pearls, cascading around her determined face. She wears a high-collared, form-fitting white dress with intricate lace details and cut-out sections on the sleeves, adding to the air of elegance and strength. The background features out-of-focus branches dotted with vibrant pink blossoms, set against a clear blue sky speckled with snow and bokeh, creating a dreamlike, serene atmosphere. The lighting casts a soft, almost divine glow on her, emphasizing her resolute expression. The artwork captures her from the waist up. Image is an anime style digital artwork. trending on artstation, pixiv, anime-style, anime screenshot

また、どの程度ControlNetによる条件付けを反映させるかのパラメータであるcontrolnet_conditioning_scaleを最大値の1.0にして実験します。

Tileに関して

Tileは参照画像をもとに、高解像にしたり、一部を変換させたりする用途で利用するため、今回は下記のプロンプトを利用します。

main_prompt = """
Anime style, pixiv
"""

上記により、元々実写感の強い画像を、アニメ画像風に変換できればと思います。

また、controlnet_conditioning_scaleは0.7に設定します。これは1.0だとあまりにも生成される画像が微妙だったので、恣意的に変更しました。

実験1

まずは、T5 Text EncoderをDropした状態で、ControlNetのPoseを利用してみます。

参照画像

普通の画像をOpenPoseの骨格画像に変換して、利用します。

参照画像の変換前

参照画像の変換後

設定

[SD3C]
device = auto
n_steps=28
seed=42
shift = 3.0

model_path = stabilityai/stable-diffusion-3-medium-diffusers
controlnet_path = InstantX/SD3-Controlnet-Pose

guided_scale = 4.0
width = 1024
height = 1024

use_cpu_offload = False
use_text_encoder_3 = False
use_t5_quantization = False
save_latent = False
controlnet_conditioning_scale = 1.0

結果

実行時間

generate image time:  39.25575375556946
generate image time:  38.35134029388428
generate image time:  39.60095191001892

生成された画像は下記に示します。


おおよそ、ControlNetを利用しない方が、破綻が少ない画像が生成されているように見えます。
(ただ、これは実験のためにcontrolnet_conditioning_scale = 1.0の最大値にしていることも大きいです。後述する実験の通り、この数値を少し下げたほうが破綻が減る気がします)

一方で、入力した骨格画像にはかなり忠実に画像を生成できているように見えます。

実験2

まずは、T5 Text EncoderをDropした状態で、ControlNetのCannyを利用してみます。

参照画像

普通の画像を線画画像に変換して、利用します。

参照画像の変換前

参照画像の変換後

設定

[SD3C]
device = auto
n_steps=28
seed=42
shift = 3.0

model_path = stabilityai/stable-diffusion-3-medium-diffusers
controlnet_path = InstantX/SD3-Controlnet-Canny

guided_scale = 4.0
width = 1024
height = 1024

use_cpu_offload = False
use_text_encoder_3 = False
use_t5_quantization = False
save_latent = False
controlnet_conditioning_scale = 1.0

結果

実行時間

generate image time:  43.07749342918396
generate image time:  40.56032943725586
generate image time:  40.25832676887512

実行時間はPoseと比較して若干かかりましたが、大きくは変わりません。

生成された画像は下記に示します。


こちらもcontrolnet_conditioning_scale = 1.0のせいなのか、画質には違和感を感じますが、元々の線画に忠実に絵を生成しているように見えます。

実験3

まずは、T5 Text EncoderをDropした状態で、ControlNetのTileを利用してみます。

参照画像

参照画像の変換は行いません。普通の画像をそのまま入力します。

参照画像の変換前

参照画像の変換後

設定

[SD3C]
device = auto
n_steps=28
seed=42
shift = 3.0

model_path = stabilityai/stable-diffusion-3-medium-diffusers
controlnet_path = InstantX/SD3-Controlnet-Tile

guided_scale = 4.0
width = 1024
height = 1024

use_cpu_offload = False
use_text_encoder_3 = False
use_t5_quantization = False
save_latent = False
controlnet_conditioning_scale = 0.7

結果

実行時間

generate image time: 41.53155183792114
generate image time: 41.0320258140564
generate image time: 40.09854006767273

実行時間はPoseと比較して若干かかりましたが、大きくは変わりません。

生成された画像は下記に示します。


一枚目など、結構綺麗に絵風に実写画像を変換できたかなと思います。

実験4

まずは、T5 Text Encoderを量子化した状態で導入し、ControlNetのPoseを利用してみます。

参照画像

普通の画像をOpenPoseの骨格画像に変換して、利用します。

参照画像の変換前

参照画像の変換後

設定

[SD3C]
device = auto
n_steps=28
seed=42
shift = 3.0

model_path = stabilityai/stable-diffusion-3-medium-diffusers
controlnet_path = InstantX/SD3-Controlnet-Pose

guided_scale = 4.0
width = 1024
height = 1024

use_cpu_offload = False
use_text_encoder_3 = True
use_t5_quantization = True
save_latent = False
controlnet_conditioning_scale = 1.0

結果

実行時間

generate image time:  70.73976230621338
generate image time:  42.649738073349
generate image time:  42.33554172515869

1回目はプロンプトの埋め込み処理や、モデルの処理などが入るため、若干時間がかかり、2回目以降はプロンプトの埋め込み処理をスキップできるので、その分処理が早くなっている

生成された画像は下記に示します。


前回のPart8の記事でもそうでしたが、T5のText Encoderを導入することで、生成される画像の質が上がるように見えます。
また、T5を導入して、処理が分割されましたが、問題なく参照画像によって、生成される画像がコントロールされていることがわかります。

実験5

まずは、T5 Text Encoderを量子化した状態で導入して、ControlNetのCannyを利用してみます。

参照画像

普通の画像を線画画像に変換して、利用します。

参照画像の変換前

参照画像の変換後

設定

[SD3C]
device = auto
n_steps=28
seed=42
shift = 3.0

model_path = stabilityai/stable-diffusion-3-medium-diffusers
controlnet_path = InstantX/SD3-Controlnet-Canny

guided_scale = 4.0
width = 1024
height = 1024

use_cpu_offload = False
use_text_encoder_3 = True
use_t5_quantization = True
save_latent = False
controlnet_conditioning_scale = 1.0

結果

実行時間

generate image time:  70.17689919471741
generate image time:  42.68410515785217
generate image time:  42.38533401489258

実行時間はT5導入後のPoseとほぼ変わりません。

生成された画像は下記に示します。


こちらもcontrolnet_conditioning_scale = 1.0のせいなのか、画質には違和感を感じますが、元々の線画に忠実に絵を生成しているように見えます。

実験6

まずは、T5 Text Encoderを量子化した状態で導入して、ControlNetのTileを利用してみます。

参照画像

参照画像の変換は行いません。普通の画像をそのまま入力します。

参照画像の変換前

参照画像の変換後

設定

[SD3C]
device = auto
n_steps=28
seed=42
shift = 3.0

model_path = stabilityai/stable-diffusion-3-medium-diffusers
controlnet_path = InstantX/SD3-Controlnet-Tile

guided_scale = 4.0
width = 1024
height = 1024

use_cpu_offload = False
use_text_encoder_3 = True
use_t5_quantization = True
save_latent = False
controlnet_conditioning_scale = 0.7

結果

実行時間

generate image time: 71.40609192848206
generate image time: 42.3540825843811
generate image time: 42.60647749900818

実行時間はT5導入後のPoseとほぼ変わりません。

生成された画像は下記に示します。


あまり変わっていないようにも見えますが、ちゃんと実写画像を絵風に変換できています。

実験7

まずは、T5 Text Encoderを量子化した状態で導入し、ControlNetのPoseを利用してみます。
その上で、controlnet_conditioning_scaleを0.5に下げて実験してみます

この値を下げることで、参照画像に忠実な画像が生成される確率は下がりますが、その代わり画像の破綻を減らせるのではないかと期待しています。

参照画像

普通の画像をOpenPoseの骨格画像に変換して、利用します。

参照画像の変換前

参照画像の変換後

設定

[SD3C]
device = auto
n_steps=28
seed=42
shift = 3.0

model_path = stabilityai/stable-diffusion-3-medium-diffusers
controlnet_path = InstantX/SD3-Controlnet-Pose

guided_scale = 4.0
width = 1024
height = 1024

use_cpu_offload = False
use_text_encoder_3 = True
use_t5_quantization = True
save_latent = False
controlnet_conditioning_scale = 0.5

結果

生成された画像は下記に示します。


おお!1.0だった時と比較して、かなり画像の質が上がりつつ、ちゃんと参照画像のポーズに近いポーズで生成できているっぽいです!
controlnet_conditioning_scaleの値は色々試すと良さそうですね!

まとめ

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

Part8の今回は、生成AIの一つ、画像生成AIのStable Diffusion 3 MediumモデルにControlNetを導入して使えるようにしました。

今回は、初めて使うモデルだったこともあり、色々実験してみました。

色々実験してみてわかったことは

  • controlnet_conditioning_scaleは色々試してみると良いと思いました
  • T5量子化をしても、ControlNetを利用できます
  • Cannyに関しては、かなり忠実に線画を利用するので、プロンプトには色合いや雰囲気を指定した方がよかったかもしれないです

記述内容などで間違っている部分があればご教授ください。よろしくお願いします!

Discussion