【Stable Diffusion】モデル読み込みの高速化

【Stable Diffusion】モデル読み込みの高速化 機械学習

「Stable Diffusionの処理を高速化したい」
「ダウンロード済みのモデルが再度ダウンロードされる事象が発生している・・・」

このような場合には、この記事の内容が参考になります。
この記事では、Stable Diffusionのモデル読み込みを高速にする方法を解説しています。

本記事の内容

  • Stable Diffusionにおけるモデル読み込みの高速化
  • モデル変換(Pickle形式→Safetensors形式)のためのシステム要件
  • モデル変換処理
  • モデル読み込みの検証

それでは、上記に沿って解説していきます。

Stable Diffusionにおけるモデル読み込みの高速化

Stable Diffusionでは、画像生成の作成時間についてよく言及されています。
「1枚あたり〇秒で生成可能」という記述を見たことがありませんか?

実際、画像生成時間の短縮化は重要なポイントです。
当ブログでも、次の記事で画像生成の高速化について解説しています。

もちろん、サンプラーの改善・改良は今後も重要なことに変わりありません。
それ以外でも、処理の高速化が可能だということを知っていますか?

そのうちの一つに、モデル読み込みの高速化があります。
私の環境では、次のようにモデル読み込みの時間を高速化できました。

モデル名処理前(Pickle形式)処理後(Safetensors形式)
Stable Diffusion v1-58.03秒4.47秒
Stable Diffusion v2-17.51秒4.17秒

処理前のPickle形式は、デフォルトのモデル読み込みとなります。
Pickle形式からSafetensors形式への変更により、モデル読み込み時間を短縮できています。

もちろん、この数字は各自の環境により異なります。
ちなみに、利用している環境は以下となります。

> python -m xformers.info
WARNING:root:A matching Triton is not available, some optimizations will not be enabled.
Error caught was: No module named 'triton'
xFormers 0.0.14.dev
memory_efficient_attention.flshatt:      available - requires GPU with compute capability 7.5+
memory_efficient_attention.cutlass:      available
memory_efficient_attention.small_k:      available
is_triton_available:                     False
is_functorch_available:                  False
pytorch.version:                         1.13.0+cu116
pytorch.cuda:                            available
gpu.compute_capability:                  8.6
gpu.name:                                NVIDIA GeForce RTX 3090

NVIDIA GeForce RTX 3090搭載のWindowsマシンを利用しています。
そして、xFormersによりPyTorchの高速化を実施済みです。

したがって、処理前でもそこそこ速いと言えます。
それをさらに速くすることが可能になったということです。

以上、Stable Diffusionにおけるモデル読み込みの高速化を説明しました。
次は、モデル変換(Pickle形式→Safetensors形式)のためのシステム要件を説明します。

モデル変換(Pickle形式→Safetensors形式)のためのシステム要件

モデル変換には、大きく次の二つが必要となります。

  • safetensors
  • Diffusers 0.10以降

それぞれを以下で説明します。

safetensors

safetensorsにより、Pickle形式からSafetensors形式のファイル変換が可能になります。
safetensorsについては、次の記事で説明しています。

safetensorsのインストールは、次のコマンドで対応可能です。

pip install safetensors

Diffusers 0.10以降

Diffusersについては、次の記事で解説しています。

現時点では、0.10.2が最新バージョンとなっています。
おそらく、すぐにバージョンはアップするでしょう。

とりあえず、モデルの変換に必要なのは0.10.0以降です。
インストール済みなら、次のコマンドで最新バージョンに更新しておきましょう。

pip install diffusers --upgrade

以上、モデル変換(Pickle形式→Safetensors形式)のためのシステム要件を説明しました。
次は、モデル変換処理を説明します。

モデル変換処理

モデルの変換処理自体は、簡単と言えます。
次のコードを実行するだけです。

まずは、Stable Diffusion v1-5をSafetensors形式のモデルに変換します。
あらかじめ「safe_model」ディレクトリは作成しておきましょう。

from diffusers import StableDiffusionPipeline

MODEL = "runwayml/stable-diffusion-v1-5"
SAFE_MODEL = "./safe_model/stable-diffusion-v1-5"

pipe = StableDiffusionPipeline.from_pretrained(MODEL)
pipe.save_pretrained(SAFE_MODEL, safe_serialization=True)

コードを実行すると、少しの時間待たされます。
処理が完了すると、「SAFE_MODEL」を確認します。

safetensorsを拡張子に持つファイルを確認できます。
同じように、Stable Diffusion v2-1も変換できます。

from diffusers import StableDiffusionPipeline

MODEL = "stabilityai/stable-diffusion-2-1"
SAFE_MODEL = "./safe_model/stable-diffusion-2-1"

pipe = StableDiffusionPipeline.from_pretrained(MODEL)
pipe.save_pretrained(SAFE_MODEL, safe_serialization=True)

なお、safetensorsでは安全なモデル読み込みが可能と説明されています。
もしかしたら、Pickle形式では安全ではないのかもしれません。

安全ではないから、Diffusersで同じモデルを何度もダウンロードすることがあるのかもしれません。
キャッシュに同じモデルが存在しているはずなのに、再度ダウンロードする事象のことです。

しばしば、この事象にあって無駄に5GBのダウンロード処理が行われていました。
今後は、このSafetensors形式をモデルに指定します。

そうなれば、無駄にダウンロードが始まる怪現象はなくなるはずです。

以上、モデル変換処理を説明しました。
次は、モデル読み込みの検証を説明します。

モデル読み込みの検証

モデルの読み込みの方法は、パスを指定するだけです。
次のように、パス指定に切り替えます。

MODEL = "stabilityai/stable-diffusion-2-1"
↓
MODEL = "./safe_model/stable-diffusion-2-1"

もちろん、フルパス指定でもOK。
ここでは、高速化できた結果として挙げた数値の裏付けを説明しておきます。

モデル名処理前(Pickle形式)処理後(Safetensors形式)
Stable Diffusion v1-58.03秒4.47秒
Stable Diffusion v2-17.51秒4.17秒

上記の数値は、次のコードを利用して算出しています。

from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
import time

# MODEL = "runwayml/stable-diffusion-v1-5"
# MODEL = "./safe_model/stable-diffusion-v1-5"
# MODEL = "stabilityai/stable-diffusion-2-1"
MODEL = "./safe_model/stable-diffusion-2-1"

def get_elapsed_time(model_id):
    # 時間計測開始
    time_start = time.perf_counter()

    pipe = StableDiffusionPipeline.from_pretrained(model_id)
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)

    pipe = pipe.to("cuda")

    # 時間計測終了
    time_end = time.perf_counter()

    # 経過時間(秒)
    elapsed_time = time_end - time_start

    return elapsed_time


all_time = 0
trial_count = 10

for i in range(trial_count):
    tmp_time = get_elapsed_time(MODEL)
    print("経過時間:{0}".format(tmp_time) + "[秒]")
    all_time = all_time + tmp_time

print("-" * 50)
print("合計時間:{0}".format(all_time) + "[秒]")
print("平均時間:{0}".format(all_time / trial_count) + "[秒]")

プログラムの内容は、モデル読み込みだけを10回試行した場合の読み込みにかかる時間を算出することです。
最後に、合計時間と平均時間を表示します。

各「MODEL」毎に実行した結果は、以下。

Stable Diffusion v1-5 処理前(Pickle形式)

「runwayml/stable-diffusion-v1-5」をモデルに指定した場合の結果です。

Fetching 15 files: 100%|██████████| 15/15 [00:00<00:00, 7522.07it/s]
経過時間:8.189880299993092[秒]
Fetching 15 files: 100%|██████████| 15/15 [00:00<00:00, 7498.76it/s]
経過時間:8.001528199994937[秒]
Fetching 15 files: 100%|██████████| 15/15 [00:00<00:00, 7583.72it/s]
経過時間:8.167271500002244[秒]
Fetching 15 files: 100%|██████████| 15/15 [00:00<00:00, 5003.94it/s]
経過時間:8.003240299993195[秒]
Fetching 15 files: 100%|██████████| 15/15 [00:00<00:00, 4909.45it/s]
経過時間:8.01132749998942[秒]
Fetching 15 files: 100%|██████████| 15/15 [00:00<00:00, 4836.61it/s]
経過時間:7.989443400001619[秒]
Fetching 15 files: 100%|██████████| 15/15 [00:00<00:00, 5936.46it/s]
経過時間:8.026545300002908[秒]
Fetching 15 files: 100%|██████████| 15/15 [00:00<00:00, 5005.53it/s]
経過時間:7.9812774999882095[秒]
Fetching 15 files: 100%|██████████| 15/15 [00:00<00:00, 14456.47it/s]
経過時間:7.976256900001317[秒]
Fetching 15 files: 100%|██████████| 15/15 [00:00<00:00, 7462.29it/s]
経過時間:8.000048600006266[秒]
--------------------------------------------------
合計時間:80.34681949997321[秒]
平均時間:8.03468194999732[秒]

Stable Diffusion v1-5 処理後(Safetensors形式)

「./safe_model/stable-diffusion-v1-5」をモデルに指定した場合の結果です。

経過時間:4.542423599996255[秒]
経過時間:4.34834700000647[秒]
経過時間:4.472572600003332[秒]
経過時間:4.4918803999898955[秒]
経過時間:4.558145800008788[秒]
経過時間:4.405960200005211[秒]
経過時間:4.476813500004937[秒]
経過時間:4.442639200002304[秒]
経過時間:4.45199689999572[秒]
経過時間:4.556190700008301[秒]
--------------------------------------------------
合計時間:44.74696990002121[秒]
平均時間:4.474696990002121[秒]

Stable Diffusion v2-1 処理前(Pickle形式)

「stabilityai/stable-diffusion-2-1」をモデルに指定した場合の結果です。

Fetching 12 files: 100%|██████████| 12/12 [00:00<00:00, 6008.31it/s]
経過時間:7.6619338000018615[秒]
Fetching 12 files: 100%|██████████| 12/12 [00:00<00:00, 2958.08it/s]
経過時間:7.569101599990972[秒]
Fetching 12 files: 100%|██████████| 12/12 [00:00<00:00, 3997.11it/s]
経過時間:7.555761399999028[秒]
Fetching 12 files: 100%|██████████| 12/12 [00:00<00:00, 12234.24it/s]
経過時間:7.531382200002554[秒]
Fetching 12 files: 100%|██████████| 12/12 [00:00<00:00, 7656.17it/s]
経過時間:7.493757200005348[秒]
Fetching 12 files: 100%|██████████| 12/12 [00:00<00:00, 6053.84it/s]
経過時間:7.594000000011874[秒]
Fetching 12 files: 100%|██████████| 12/12 [00:00<00:00, 12009.46it/s]
経過時間:7.563966399990022[秒]
Fetching 12 files: 100%|██████████| 12/12 [00:00<00:00, 4004.43it/s]
経過時間:7.460174399995594[秒]
Fetching 12 files: 100%|██████████| 12/12 [00:00<00:00, 2927.28it/s]
経過時間:7.387935299993842[秒]
Fetching 12 files: 100%|██████████| 12/12 [00:00<00:00, 11992.29it/s]
経過時間:7.34581060000346[秒]
--------------------------------------------------
合計時間:75.16382289999456[秒]
平均時間:7.516382289999456[秒]

Stable Diffusion v2-1 処理後(Safetensors形式)

「./safe_model/stable-diffusion-2-1」をモデルに指定した場合の結果です。

経過時間:4.36155279999366[秒]
経過時間:4.057511000006343[秒]
経過時間:4.19884950001142[秒]
経過時間:4.159023400003207[秒]
経過時間:4.211333900006139[秒]
経過時間:4.089154300003429[秒]
経過時間:4.1470158000011[秒]
経過時間:4.145440799999051[秒]
経過時間:4.230742499988992[秒]
経過時間:4.166927400001441[秒]
--------------------------------------------------
合計時間:41.76755140001478[秒]
平均時間:4.176755140001478[秒]

以上、モデル読み込みの検証を説明しました。

タイトルとURLをコピーしました