HuggingFace ブログ : SDXL のための単純な最適化の探求 (翻訳/解説)

翻訳 : クラスキャット セールスインフォメーション
作成日時 : 11/02/2023

* 本ページは、HuggingFace Blog の以下のドキュメントを翻訳した上で適宜、補足説明したものです：

• Exploring simple optimizations for SDXL (Authors : Sayak Paul, Steven Liu ; Published : 10/24/2023)

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

 

 

 

HuggingFace ブログ : SDXL のための単純な最適化の探求

Stable Diffusion XL (SDXL) は高品質な超リアルな画像を生成するための Stability AI による最新の潜在拡散モデルです。それは手やテキストそして空間的に正しい構成を適切に取得するような、以前の Stable Diffusion モデルの課題を乗り越えています。更に、SDXL はまたよりコンテキスト aware で、より良く見える画像を生成するためにプロンプトにおいて少ない単語を必要とするだけです。

けれども、これらの改良のすべては大幅に大規模なモデルという代償を払っています。どのくらい大きくなるのでしょう？ベース SDXL モデルは 3.5B (= 35 億) パラメータ (特に UNet) を持ちます、これは以前の Stable Diffusion モデルよりもおよそ 3 倍は大きいです。

SDXL を推論スピードやメモリ使用量のためにどのように最適化できるかを探求するため、私たちは幾つかのテストを A100 GPU (40 GB) で実行しました。各推論実行について、4 画像を生成してそれを 3 回繰り返します。推論遅延を計算する際、3 反復から最後の反復だけを考えます。

その結果、SDXL をそのまま完全精度で実行してデフォルトのアテンション機構を使用する場合、それは 28 GB メモリを消費して 72.2 秒かかります！

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0").to("cuda")
pipe.unet.set_default_attn_processor()

これはあまり現実的ではなく、4 画像以上の画像を生成する場合が多いので、更に遅くなる可能性があります。そしてより強力な GPU を持たない場合、メモリ不足エラーのメッセージに悩まされることになります。それでは推論スピードを増してメモリ使用量を減らすにはどのように SDXL を最適化できるでしょう？

🤗 Diffusers には、多くの最適化トリックやテクニックがあり SDXL のようなメモリ集約的なモデルを実行するのに役立ちます、そして私たちはその方法を示します！フォーカスする 2 つのものは推論スピードとメモリです。

🧠 この投稿で説明されるテクニックはすべての パイプライン に適用可能です。

 

推論スピード

拡散はランダムな過程ですので、好きな画像を得られる保証はありません。多くの場合、推論を複数回実行して繰り返す必要があり、そしてそれがスピードの最適化が重要である理由です。このセクションは低精度な重みを使用し、メモリ効率の良いアテンションと PyTorch 2.0 の torch.compile を組み込み、スピードをブーストして推論時間を削減することにフォーカスします。

 

低精度 (Lower precision)

モデル重みは浮動小数点データ型として表現される特定の精度で保存されます。標準の浮動小数点データ型は float32 (fp32) で、これは広範囲な浮動小数点を正確に表すことができます。推論については、必ずしも正確である必要はない場合が多いので float16 (fp16) を使用するべきです、これはより狭い範囲の浮動小数点を捉えます。つまり、fp16 は fp32 に比べてストアするメモリの総量が半分しかかからず、計算することがより簡単なので 2 倍高速です。更に、現代的な GPU カードは fp16 計算を実行するためにハードウェアを最適化していて、更に高速にします。

🤗 Diffusers では、モデルがロードされるとき重みを変換するために torch.dtype パラメータを指定することで推論のために fp16 を使用することができます :

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
).to("cuda")
pipe.unet.set_default_attn_processor()

まったく最適化されていない SDXL パイプラインと比較して、fp16 の使用はメモリ 21.7 GB を使用して 14.8 秒だけかかります。推論を殆ど丸々 1 分スピードアップしています！

 

メモリ効率の良いアテンション

transformers モジュールで使用されるアテンションブロックは大きなボトルネックになる可能性があります、メモリは入力シークエンスが長くなるにつれて二次関数的に増えるからです。これはすぐに多くのメモリを使用して、メモリ不足エラーメッセージを残す可能性があります 😬

メモリ効率の良いアテンション・アルゴリズムは、それが疎性、タイリングいずれによるものであれ、アテンションを計算するメモリ負荷を軽減することを求めます。これらの最適化されたアルゴリズムは個別にインストールが必要なサードパーティのライブラリとして多くの場合利用可能でした。しかし PyTorch 2.0 からは、これはもはや当てはまりません。PyTorch 2 は スケールされたドット積アテンション (SDPA) を導入しました、これは Flash アテンション、メモリ効率的なアテンション (xFormers)、そして C++ の PyTorch 実装の融合された実装を提供します。SDPA は多分推論を高速化する最も簡単な方法です : 貴方が 🤗 Diffusers で PyTorch ≥ 2.0 を使用していれば、それはデフォルトで自動的に有効にされます！

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
).to("cuda")

まったく最適化されていない SDXL パイプラインと比較して、fp16 と SDPA の使用は同じメモリ総量を取り、推論時間は 11.4 秒に向上します。これを他の最適化を比較する新しいベースラインとして使用しましょう。

 

torch.compile

PyTorch 2.0 はまた PyTorch コードを推論用により最適化されたカーネルに just-in-time (JIT) コンパイルするための torch.compile API も導入しました。他のコンパイラ・ソリューションとは異なり、torch.compile は既存のコードに最小限の変更を必要とするだけで、それはモデルを関数でラップするほどに簡単です。

mode パラメータを使用して、コンパイルの際にメモリオーバーヘッドや推論スピードを最適化することができます、これはより柔軟な方法を与えます。

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
).to("cuda")
pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)

前のベースライン (fp16 + SDPA) と比較して、UNet を torch.compile でラップすると推論時間が 10.2 秒に改善されます。

⚠️ The first time you compile a model is slower, but once the model is compiled, all subsequent calls to it are much faster!

 

モデルのメモリ使用量

今日のモデルはますます大規模化し、それらをメモリに収めることを困難にしています。このセクションはこれらの巨大なモデルのメモリ使用量を削減できる方法に焦点を当ててそれらをコンシューマー GPU で実行できるようにします。これらのテクニックは CPU オフロード、潜在データを一度にではなく幾つかのステップで画像にデコードすること、そしてオートエンコーダの蒸留バージョンの使用を含みます。

 

モデルの CPU オフロード

モデルオフロードは、拡散モデルの他のコンポーネント (テキストエンコーダ、VAE) を CPU にロードする一方で、UNet を GPU メモリにロードすることでメモリを節約します。このようにして、UNet はそれが必要でなくなるまで GPU 上で複数の反復を実行できます。

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
).to("cuda")
pipe.enable_model_cpu_offload()

ベースラインと比較して、それは 20.2 GB メモリを消費するようになり、1.5 GB のメモリを節約します。

 

シーケンシャル CPU オフロード

推論が遅くなる代わりに更にメモリを節約できる別のタイプのオフロードはシーケンシャル CPU オフロードです。UNet のようにモデル全体をオフロードするのではなく、様々な UNet サブモジュールにストアされているモデル重みが CPU にオフロードされて forward パスの直前にだけ GPU 上にロードされます。基本的には、毎回モデルの一部だけをロードしますのでより多くのメモリを節約することさえ可能にします。唯一の欠点は、サブモジュールを何回もロードしてオフロードするので、非常に遅くなることです。

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
).to("cuda")
pipe.enable_sequential_cpu_offload()

ベースラインと比較して、これは 19.9GB のメモリを必要としますが、推論時間は 67 秒に増加します。

 

スライシング

SDXL では、変分エンコーダ (VAE) は (UNet により予測される) refined (精製された・洗練された) 潜在データをリアルな画像にデコードします。このステップのメモリ要件は予測される画像の数 (バッチサイズ) によりスケールします。画像解像度と利用可能な GPU VRAM に依存して、それは非常にメモリ集約的である可能性があります。

ここが「スライシング」が役立つ場所です。デコードされる入力テンソルはスライスに分割され、それをデコードする計算は幾つかのステップに渡り完結します。これはメモリを節約し、より大きいバッチサイズを可能にします。

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
).to("cuda")
pipe.enable_vae_slicing()

スライスされた計算で、メモリは 15.4 GB に削減されます。シーケンシャル CPU オフロードを追加する場合、更に 11.45 GB に削減され、これはプロンプト毎に 4 画像 (1024×1024) を生成させます。ただし、シーケンシャル・オフロードで、推論遅延もまた増大します。

 

計算のキャッシュ

テキスト条件付き画像生成モデルは通常は入力プロンプトから埋め込みを計算するためにテキストエンコーダを使用します。SDXL は 2 つのテキストエンコーダを使用します！これは推論遅延にかなり影響します。けれども、これらの埋め込みは逆拡散過程を通して変更されませんので、それらを事前計算して再利用することができます。このようにして、テキスト埋め込みの計算後、メモリからテキストエンコーダを削除できます。

まず、テキストエンコーダと対応するトークナイザーをロードして入力プロンプトから埋め込みを計算します :

tokenizers = [tokenizer, tokenizer_2]
text_encoders = [text_encoder, text_encoder_2]

(
    prompt_embeds,
    negative_prompt_embeds,
    pooled_prompt_embeds,
    negative_pooled_prompt_embeds
) = encode_prompt(tokenizers, text_encoders, prompt)

次に、GPU メモリをフラッシュしてテキストエンコーダを削除します :

del text_encoder, text_encoder_2, tokenizer, tokenizer_2
flush()

これで埋め込みを SDXL パイプラインに直接送り込めるようになりました :

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    text_encoder=None,
    text_encoder_2=None,
    tokenizer=None,
    tokenizer_2=None,
    torch_dtype=torch.float16,
).to("cuda")

call_args = dict(
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        pooled_prompt_embeds=pooled_prompt_embeds,
        negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
        num_images_per_prompt=num_images_per_prompt,
        num_inference_steps=num_inference_steps,
)
image = pipe(**call_args).images[0]

SDPA と fp16 の組み合わせで、メモリを 21.9 GB まで削減できました。メモリ最適化のための上述の他のテクニックはまたキャッシュされた計算とともに使用できます。

 

Tiny Autoencoder

前述のように、VAE は潜在データを画像にデコードします。自然に、このステップは VAE のサイズにより直接ボトルネックになります。そこで、より小さいオートエンコーダだけを使用しましょう！ハブ で利用可能な madebyollin による Tiny Autoencoder は 10 MB だけで、SDXL によるオリジナル VAE から蒸留されています。

from diffusers import AutoencoderTiny

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
)
pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", torch_dtype=torch.float16)
pipe.to("cuda")

このセットアップで、メモリ要件を 15.6 GB に削減しながら同時に推論遅延も削減します。

⚠️ The Tiny Autoencoder can omit some of the more fine-grained details from images, which is why the Tiny AutoEncoder is more appropriate for image previews.

 

Conclusion

To conclude and summarize the savings from our optimizations:

⚠️ While profiling GPUs to measure the trade-off between inference latency and memory requirements, it is important to be aware of the hardware being used. The above findings may not translate equally from hardware to hardware. For example, `torch.compile` only seems to benefit modern GPUs, at least for SDXL.

We hope these optimizations make it a breeze to run your favorite pipelines. Try these techniques out and share your images with us! 🤗

以上