[DALL·E 3]で使われているConsistency Decoderの速度を測ってみた

2023年11月9日 03:48

OpenAIのDevDayリリースを読んでいたら、しれっと
[DALL·E 3]で[Consistency Models]が使われていることを発見しました。

Whisper v3 と一貫性デコーダー
オープンソース自動音声認識モデル（ASR）の次のバージョンである「Whisper large-v3」をリリースします。また、近い将来、APIでもWhisper v3をサポートする予定です。
また、Stable Diffusion VAEデコーダーのドロップインリプレースメントである「一貫性デコーダー (Consistency Decoder)」もオープンソース化しています。このデコーダーは、Stable Diffusion 1.0+ VAEと互換性のあるすべての画像を改善し、テキスト、顔、直線を大幅に改善します。

ChatGPTのOpenAI DevDayアナウンス (AICU社による翻訳)

Consistency Modelとは

めちゃめちゃ高速です

100枚が…7分…!? 1枚4秒だよ！！？

Real-Time Latent Consistency Model (LCM)から掘り下げる最近の画像生成AIの高速化研究

｜しらいはかせ(Hacker作家) @o_ob https://t.co/GWmoO4IjyN #note pic.twitter.com/1iwt6tBUF2
— Dr.(Shirai)Hakase しらいはかせCEO 技術書典15 (@o_ob) November 1, 2023

100枚が…7分…!? 1枚4秒だよ！！？ Real-Time Latent Consistency Model (LCM)から掘り下げる最近の画像生成AIの高速化研究

DALL-E3が公開されているわけではないのですが、論文がDALL-E3の論文（長い！）なので、これは使われているという理解で良さそう。
[Improving Image Generation with Better Captions]
Improved decoding for stable diffusion vaes.

論文より：良いキャプションで画像生成を改善する

要旨

我々は、高度に説明的な生成画像キャプションを保持することによって、テキスト-画像モードのプロンプトフォローが大幅に改善されることを示す。既存の外部画像モデルは、詳細な画像説明に従うことに苦労しており、しばしば単語を無視したり、ロムの意味を混同したりする。我々は、この[問題]は、訓練データセットに含まれるノイズの多い不正確な画像キャプションに起因すると仮定する。
我々はこの問題に対処するため、特注の画像キャプション作成ツールを開発し、それを訓練データセットの受信に利用する、
次に、いくつかの外部画像モデルを訓練し、これらの合成画像キャプションを使用することで、プロンプトフォローが確実に改善されることを示す。最後に、これらの結果をDALL-E 3：新しいテキスト画像生成システムに適用し、プロンプトフォローイングを混乱させるように設計された回避に対する性能のベンチマークを行う。
その結果、DALL-E 3は、競争相手と比較しても遜色ないことがわかりました：我々は、これらの評価のサンプルとコードを公開する。

公開されているコードを読み解く

VAEだけでも公開されているので遊んでみます。

Installation

$ pip install git+https://github.com/openai/consistencydecoder.git

Usage

import torch
from diffusers import StableDiffusionPipeline
from consistencydecoder import ConsistencyDecoder, save_image, load_image

# encode with stable diffusion vae
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, device="cuda:0"
)
pipe.vae.cuda()
decoder_consistency = ConsistencyDecoder(device="cuda:0") # Model size: 2.49 GB

image = load_image("assets/gt1.png", size=(256, 256), center_crop=True)
latent = pipe.vae.encode(image.half().cuda()).latent_dist.mean

# decode with gan
sample_gan = pipe.vae.decode(latent).sample.detach()
save_image(sample_gan, "gan.png")

# decode with vae
sample_consistency = decoder_consistency(latent)
save_image(sample_consistency, "con.png")

試してみたコード

import torch
from diffusers import StableDiffusionPipeline
from consistencydecoder import ConsistencyDecoder, save_image, load_image

# encode with stable diffusion vae
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, device="cuda:0"
)
pipe.vae.cuda()
decoder_consistency = ConsistencyDecoder(device="cuda:0") # Model size: 2.49 GB

image = load_image("assets/shirai-256.png", size=(256, 256), center_crop=True)
latent = pipe.vae.encode(image.half().cuda()).latent_dist.mean

# decode with gan
sample_gan = pipe.vae.decode(latent).sample.detach()
save_image(sample_gan, "gan-shirai.png")

# decode with vae
sample_consistency = decoder_consistency(latent)
save_image(sample_consistency, "con-shirai.png")

画像を突っ込んでみます。
メモリも6GBだとちょっと足りないかも。
256x256なら動作しますので画像を小さくします。

わくわく…

ぎゃー！
実行時間0.3032283000065945秒

実行時間24.49278059997596秒
・・・ん・・・？遅すぎませんか？？

順番逆にしても同じでした。
[CON] 実行時間26.30387265299987秒
[GAN] 実行時間0.14000741000018024秒
Google Colabにしてみてもあまり代わりませんでした

GoogleColab のT4 GPUでセルを分割したらかなり改善されました
[CON]実行時間1.1482076719998986秒
[GAN]実行時間0.08754536800006463秒

でもVAEだけだとそこまで速くはない印象ですね、中間のステップがぜんぜん違うしLatentを直接持ってきているので比較できないんだろうけど…。

計測用コード

# https://zenn.dev/timoneko/articles/16f9ee7113f3cd
import time
import torch
from diffusers import StableDiffusionPipeline
from consistencydecoder import ConsistencyDecoder, save_image, load_image

def execution_speed(func):
    """
    実行速度計測用のデコレータ
    """
    def wrapper(*args, **kwargs):
        start_time = time.perf_counter()
        func(*args, **kwargs)
        end_time = time.perf_counter()
        run_time = end_time - start_time
        print("実行時間" + str(run_time) + "秒")
    return wrapper

# encode with stable diffusion vae
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, device="cuda:0"
)
pipe.vae.cuda()
decoder_consistency = ConsistencyDecoder(device="cuda:0") # Model size: 2.49 GB

image = load_image("assets/shirai-256.png", size=(256, 256), center_crop=True)
latent = pipe.vae.encode(image.half().cuda()).latent_dist.mean

# decode with vae
@execution_speed
def CON():
    sample_consistency = decoder_consistency(latent)
    save_image(sample_consistency, "con-shirai.png")
CON()

#decode with gan
@execution_speed
def GAN():
    sample_gan = pipe.vae.decode(latent).sample.detach()
    save_image(sample_gan, "gan-shirai.png")
GAN()

画像を生成してみたい

そもそもVAEだけなので無意味かもですが、最後にこのコードを足してみる

# text to image
# pipe.enable_attention_slicing()
prompt = "1girl"
image = pipe(
    prompt, 
    width=512, 
    height=512, 
    num_inference_steps=10
).images[0]
image.save("1girl.png")

LayerNormKernelImplにHalfがないってことでコケます。

/usr/local/lib/python3.10/dist-packages/torch/nn/functional.py in layer_norm(input, normalized_shape, weight, bias, eps) 2541 layer_norm, (input, weight, bias), input, normalized_shape, weight=weight, bias=bias, eps=eps 2542 ) -> 2543 return torch.layer_norm(input, normalized_shape, weight, bias, eps, torch.backends.cudnn.enabled) 2544 2545

RuntimeError: "LayerNormKernelImpl" not implemented for 'Half'

Transformersのバージョンとかかなあ…

どなたか動作実績あれば教えてください。

まあ冒頭に紹介したツールを使えばいいんですけどモデルが限定されているので。