見出し画像

100fps超え画像生成StreamDiffusionのデモに飽きたら次に進もう。stremオブジェクトで動かすコツ公開

めぐチャンネル

StreamDiffusionはリアルタイム画像生成ができる最新の素晴らしい生成AIです。日本人の学生チームが開発しOSSとして公開してくれています。詳しくは以下のリポジトリや論文を参照してください。今回はラッパーを使わずにStreamオブジェクトをそのまま使用してパラメータを変更しながら様々な生成を試みたので記事に残します。
Githubリポジトリは常にアップデートされるので最新版を使うようにしましょう。

解説を論文執筆者でプロジェクトの代表でもある「あき先生」が解説してくださっているYuTubeもあります。論文の拾い読みもしていただいており、とても参考になりますので、ぜひご覧ください。（3時間と長いですが飛ばせるような、無駄な部分はありません）

デモやサンプルコードも豊富にあるので、機能や性能を知りたい場合はそちらを動かしてください。大変興味深いデモもあります。アプリに組み込む場合もラッパーが準備されているのでそちらを使えば厄介なpipeやstreamオブジェクトの操作は必要ありません。しかし、性能を引きだしなが柔軟にアプリで使用するためには直接pipeやstreamを操作する方がよさそうです。この記事では後者のケースでStreamDiffusionを使用する時のコツについてこれまでにテストした手法を解説します。

環境構築

Readmeに記述されているとおりに構築を進めます。デモやexampleを動かすのであればfor userで大丈夫ですが、pipeやstreamを扱う場合はFor Developer（開発者向け）で構築します。最難関はTensorRTだと思います。すぐに動く環境もあれば、どうしてもうまく動かない環境もあるようです。エラーメッセージを注視しながら修正をしましょう。筆者の環境でもエラーで動きませんでしたが、原因は容易に特定できてすぐに動かすことができました。
筆者のハードウエア環境
CPU：Corei5-13400 F
MEM：32G
GPU：GforceRTX4090
OS
Ubuntu 20.04jp

基本コード

リポジトリのMainのReadmeにt2iとi2iの基本コードが記述されています。環境が上手く構築できればコピペで動きます。exampleには多くの有用な例があるので参考にしましょう。StreamDiffusionの処理の流れは基本コードを見ると良く理解できます。できればDiffuserについて事前に慣れておけば入りやすいと思います。
なおリポジトリのsrc以下がStreamDiffujinの核になるコードです。pipeline.pyが中核です。
以下Image-to-Imageのコードで作業を進めます。

コードの解説

あき先生の3時間に及ぶ解説配信の中の、コードの説明部分の講義メモです。

モデルの読込み

冒頭にはモジュールのimport関係があります。pipeの初期化が続きます。ここはDiffuserでの記述です。モデルの読み込などを行います。以下は最も単純な書き方です。

pipe = StableDiffusionPipeline.from_pretrained("KBlueLeaf/kohaku-v2.1").to(
    device=torch.device("cuda"),
    dtype=torch.float16,
)

streamオブジェクトを定義

次にStreamDiffusionでpipeをラップしstreamオブジェクトを定義します。
t_index_listが重要なパラメータですので後の作業で変更などの解説します。

# Diffusers pipelineをStreamDiffusionにラップ
stream = StreamDiffusion(
    pipe,
    t_index_list=[32, 45],
    torch_dtype=torch.float16,
)

LCM-LoRAのマージ

高速化には不可欠です。ここはDiffuserでも同じ操作があるので書き方の違いだけです。Style-LoRAを追加したい場合はここで行うと纏まりがあっていいでしょう。

# 読み込んだモデルがLCMでなければマージする
stream.load_lcm_lora()
stream.fuse_lora()

TinyVAEで高速化

ここは残念ながら追及していません。公式に従います。

# Tiny VAEで高速化
stream.vae = AutoencoderTiny.from_pretrained("madebyollin/taesd").to(device=pipe.device, dtype=pipe.dtype)通常はxformaersです。TensorRTを有効にすると大幅に高速化できますが、利用するためには様々な考慮が必要なので、今回は使いません。Demoやサンプルで効果を確認してください。

xformersで高速化

ここではxformersで高速化を行っています。

# xformersで高速化# xformersで高速化
pipe.enable_xformers_memory_efficient_attention()

TensorRTを使って高速化

TensorRTによる高速化の有効性はサンプルを動かすとよくわかると思います。TensorRTを使用する場合は上記コードを下記コードで置きかえます。GPUのTensorコアの性能次第では２〜３倍高速化出来ます。以下のコードで、"engines"は実行時にコンパイルされて生成するTensorRT用のengineを収納するディレクトリ名（へのPATHを含む）です。image2image,やtext2imageなどの生成方法やモデル、LoRAのウエイト等を変更したら新たにengineを生成させる必要があります。　"engines"を異なるディレクトリ名に変更して再コンパイルさせるようにします。変更をしないと以前に生成したengineを使用して生成されるので正しい画像が得られません。

from streamdiffusion.acceleration.tensorrt import accelerate_with_tensorrt

stream = accelerate_with_tensorrt(
    stream, "engines", max_batch_size=2,
)

以下の制約はv0.1.1で修正されました。
2023年12月30日現在、以下の制約が有ることを確認しています。仕様なのか、筆者の環境の問題か、バグなのかは不明です。」
　t2i　t_index_listの要素数　2stepまで、
　i2i　t_index_listの要素数　3stepまで
　max_batch_sizeとの関連はない。
TensorRTについては文末のレファレンスに詳しい説明へのLinkを置きます。

プロンプトとレファレンス画像の読み込み

i2iなので画像も読み込みます。プロンプトのembeddedingなどの処理も同時に行っています。embeddedingについてはDiffuserのDOCSやexampleを参照しましょう。
load_imagはpillow形式です。

prompt = "1girl with dog hair, thick frame glasses"
# streamを準備する
stream.prepare(prompt)
# 画像を読み込む
init_image = load_image("assets/img2img_example.png").resize((512, 512))

パイプラインの初期化

StreamDiffusionは生成過程でパイプライン処理を取り入れて高速化をしています。初期状態ではパイプラインの各段が1（だったと思います）で初期化されているので、画像を満たすためにウォームアップを行います。t_index_list=[32, 45]としてしているので2本(多分ここは4本では？)のパイプを使います。なのでサンプルでは2回のウォームアップになっています。ウォームアップを行わなくても特段大きな影響があるわけではありませんが、パイプラインが画像で満たされるまで数枚茶色の画像が生成されます。その後は正しい画像になります。

# Warmup >= len(t_index_list) x frame_buffer_size
for _ in range(2):
    stream(init_image)

画像の生成

重要な部分は
x_output = stream(init_image)
になります。ここで生成を行います。x_outputはtensor形式の画像です。以下でPillow形式に変換して表示しています。
postprocess_image(x_output, output_type="pil")[0].show()

# 実行
while True:
    x_output = stream(init_image)
    postprocess_image(x_output, output_type="pil")[0].show()
    input_response = input("Press Enter to continue or type 'stop' to exit: ")
    if input_response == "stop":
        break

1枚だけ生成するのであれば上記の代わりに以下のコードでOKです。out_imag.show()の前後でファイルへの書き出しを行うことができます。

x_output = stream(init_image)
out_image = postprocess_image(x_output, output_type="pil")[0]
out_image.show()#PILで表示

コードのカスタマイズ

上記でも触れましたが、StreamDiffusinのコードを自在にカスタマイズするには事前にDiffuuserに慣れておくことが賢明だと思います。用語やコードにDiffuserが出てきます。SD-webUIの知識では厳しいかもしれません。

モデルをファイルからロードする

参考コードでは以下を使用しています。
pipe = StableDiffusionPipeline.from_pretrained("KBlueLeaf/kohaku-v2.1")
このコードはHugingFaceに公開されているコードをダウロードし、GPUにロードします。一方でCivitaiに公開されているコードはsafetensor形式のファイルのみです。通常はダウンロードをしてmodelsディレクトリなどに纏めていると思います。このモデルを使うためには異なるコードが必要です。
Diffuserではファイルからロードするクラスも準備されているので上記の代わりに利用します。サンプルコードはGPUへのロードを直後に行っていますが、Diffuserのサンプルのように変更した方がわかりやすいかもしれません。
ファイルからのロードはi2iやt2iなど個別に設けられています。
StableDiffusionPipelineの代わりに（注-->そのままでも出来るかも）
StableDiffusionImg2ImgPipelineを使います。そして.from_pretrainedの代わりに.from_single_fileを用いいます。

from diffusers import StableDiffusionImg2ImgPipeline


pipe = StableDiffusionImg2ImgPipeline.from_single_file(
    "/models/Counterfeit-V3.0/Counterfeit-V3.0_fix_fp16.safetensors").to(
    device=torch.device("cuda"),
    dtype=torch.float16,
)

Style-LoRA(=独自キャラなど）の適用

入力イメージに対して独自キャラ化をする場合にはどうしても必要な機能です。実は入力にも独自キャラを用いて画風を変える「やってみた」は多いのですが、キャラを変えてしまう例は少ないように思います。筆者の方法がベストではないかもしれませんが、独自LoRAとパラメータの適切な設定でi2iでもキャラ変更が可能です。変更のための準備を説明します。

LoRA適用の方法の実際

既に独自LoRAを持っていることが前提です。ない場合はKohya_ssなどで学習させてください。モデルのバージョンには注意が必要です。現状のStreamDiffusionはSD1.5しか対応していません。

stream.load_loa()を用いる場合

stream.load_lora("./models/LoRA/megu_sports_v02.safetensors")
stream.fuse_lora(lora_scale=0.8)

この方法もsrc/streamdiffusion/pipeline.pyの中でpipeにバインドされています。なので、以下に示したDiffuserを使う方法も可能だと考えられます。
lora_scale=0.8は適応するLoRAの強さを指定しています。

DiffuserのLoRA適用方法を使用する

Diffuserでは複数のモデルのロードとマージ手段が準備されています。

pipe.load_lora_weights("latent-consistency/lcm-lora-sdv1-5", adapter_name="lcm") #Stable  Diffusion 1.5 のLCM LoRA
pipe.load_lora_weights("./models/LoRA/megu_sports_v02.safetensors", adapter_name="papercut")
pipe.set_adapters(["lcm", "papercut"], adapter_weights=[1.0, 0.8])

このコードを
pipe = StableDiffusionImg2ImgPipeline.from_single_file(
・・
）
のすぐ下に追加します。
pipe.set_adapters(["lcm", "papercut"], adapter_weights=[1.0, 0.8])は
LoRAの効き具合を設定していますが、この機能を使うためには追加のモジュールが必要です。 peftというロードモジュールです。

pip install -U peft
でインストールしておきます。

t_index_listの変更

t_index_listでは生成時の実行STEPを定義します。リストの要素は１～４個で
多い方が処理が多く精度の高い画像が生成されます。i2iの場合、レファレンスの種類によってこのリストの書き方が大きく違います。自分のキャラ絵をレファレンスに使用して雰囲気をプロンプトで変えるのであれば、サンプル通り、[32, 45]で妥当だと思います。一方で変えたいキャラと異なる画像をレファレンスに使用する場合は全く異なる設定が必要です。以下はt_index_listの例とレファレンス画像と生成画像です。

t_index_listの要素の影響を見る

実際にパイプラインの数で出力される画像がどうなるのか以下のような方法で確認できます。生成画像を入力に戻して再生成するのですが、パイプラインがあるので、全てのパイプラインを通過した後にしか画像が変化しないことがよくわかります。画像の生成部を以下で置換します。

from diffusers.utils import make_image_grid

# Run the stream infinitely
image_list=[]
image_list.append(init_image )
for i in tqdm(range(13)):
    seed=i,
    x_output = stream(image)
    image=postprocess_image(x_output, output_type="pil")[0]
    image_list.append(image)
images=make_image_grid(image_list, rows=2, cols=7)
images.show()

以下の生成結果で赤い四角で囲んだ部分が変化のない画像です。STEPに合った画像数になっています。この試験では独自Style-LoRAを有効にしています。STEP完了後の次の画像に大きなLoRAの効果が出ていることが良くわかります。STEPが大きいほどLoRAの効果は高くなります。

RCFGの変更

画質の向上に役立つパラメータ設定です。詳しくは論文を参照してください。Readmeでも簡単に述べられています。RCFGの設定は3種類準備されています。
cfg_type = "none"　　　指定なし　
cfg_type = "full"　　　　通常のCFG
cfg_type = "self"　　　　Self-Negative
cfg_type = "initialize"　　Onetime-Negative
有効化するためにはstreamオブジェクト作成時にcfg_typeを設定します。以下のコードでcfg_type=cfg_typeの部分です。事前に上記でcfg_typeを設定しておきます。

stream = StreamDiffusion(
    pipe,
    t_index_list=[35,40 ,42,45],
    torch_dtype=torch.float16,
    cfg_type=cfg_type,
)

生成画像の比較です。cfg_typeにより生成品質に差がでることが解ります。

"none"　　　　　　"full"　　　　　　"self"　　　　　　"initialize"

カスタマイズの例

ぼやけた同じキャラを精緻化

設定はオリジナルのままの
t_index_list=[32, 45]
で十分です。2stepなので生成も早く
25fps程度です。生成部分を以下のコードに入れ替えます。
Style-LoRAの強度は0.5程度で充分です。t_index_listを微調整するとさらに精緻な画像が得られます。

# Run the stream infinitely
image_list=[]
image_list.append(init_image )
for i in tqdm(range(1)):
    seed=i,
    x_output = stream(image)
    image=postprocess_image(x_output, output_type="pil")[0]
    image_list.append(image)
images=make_image_grid(image_list, rows=1, cols=2)
images.show()

リアル画像のアニメ変更

t_index_list=[38,40 ,42,45]と深く設定しています。
4ステップなので、16fps程度です。
上記コードのままです。レファレンスイメージとt_index_listが異なります。

棒人間をキャラに変換

t_index_list=[15,20,30]は浅く設定。通常はContorolNETを使うと思います。
20fps程度で生成できます。棒人間→キャラ変更はハードルが高いようです。大きな変化のある棒人間は変換できませんでした。
こちらも、上記コードのままです。レファレンスイメージとt_index_listが異なります。

異なるアニメキャラを変換

画風が全く異なるキャラ絵を使って独自キャラを生成することもできますが、スタイルを重視するか、キャラ特性を重視するかでパラメータの設定が異なります。以下の例は破綻が少なくなるよう調整した例です。独自キャラがレファレンスキャラと同じポースを取っていますが、服装はレファレンスキャラに近く変換出来ていません。
t_index_list=[35,40 ,42,45]と深めで4stepです。

破綻はありますが手を挙げています

ぼやけたキャラ動画をリアルタイムで精緻化

筆者はキャラを動かす手法としてTalking-Head-Animeface-3を利用しています。Live2Dのようなモデルを必要とせず、キャラ絵1枚あれば限られた範囲でリアルタイムに自在に動かすことができます。一方で動きの大きな画像は周辺がぼやけたり歪んだりします。このひずみをStreameDiffusionで精緻化できます。同じキャラのスタイル変更は容易です。
生成部分を動画の下のコードで置換します。

stream_img_list=glob.glob("./img_streem/danse/"+".")
でimg_streemディレクトリにある連番のフレーム画像のファイル名を取得

sorted_files = sorted(stream_img_list, key=natural_keys)
file番号通りにソート

imgCV_RGB1= np.array(init_image, dtype=np.uint8)
imgCV_BGR1 = np.array(imgCV_RGB1)[:, :, ::-1]
cv2.imshow("org",imgCV_BGR1)
cv2.waitKey(1)
では、OpenCV形式に変換してタグが同じイメージウインドウへ上書きしています。同様に生成後の画像も表示してディスプレー上では動画に見えるようにしています。

動画変換用コード

#Run the stream infinitely
#fオリジナル動画フレーム画像ファイル名を取得
import glob
stream_img_list=glob.glob("./img_streem/danse/"+"*.*")
sorted_files = sorted(stream_img_list, key=natural_keys)#file番号通りにソート

image_list=[]
for n in  range(10):
     for i in tqdm(range(len(sorted_files ))):
        seed=i,
        init_image = load_image(sorted_files [i]).resize((512, 512))
        x_output = stream(init_image )
        image=postprocess_image(x_output, output_type="pil")[0]
         #オリジナル画像の表示         imgCV_RGB1= np.array(init_image, dtype=np.uint8)
        imgCV_BGR1 = np.array(imgCV_RGB1)[:, :, ::-1]
        cv2.imshow("org",imgCV_BGR1)
        cv2.waitKey(1)    
         #変換画像の表示         imgCV_RGB = np.array(image, dtype=np.uint8)
        imgCV_BGR = np.array(imgCV_RGB)[:, :, ::-1]
        cv2.imshow("i2i",imgCV_BGR)
        cv2.waitKey(10)#10mSのウエイト
    i=0

ダンス動画をリアルタイムでキャラ変換

Style-LoRAを効かせて、リアルタイム変換を行った例です。キャラが小さいので腕の交差などで破綻はでていますが、リアルタイム変換できることが今までにないことで、改めてStreamDiffusionの実力を感じます。
t_index_list=[35,45]
LoRA強度=0.7です。

まとめ

ついに待望のリアルタイム画像生成が開発されました。いずれモデル無しでキャラをリアルタイムで自在に動かせるようになると考えていましたが、思ったよりも早かったと思います。今回の記事では先週金曜日深夜の発表から取り急ぎStreamDiffusionの特性を理解することを目的に重点的に作業をした結果をまとめました。TensorRTがStreamで有効に出来ていない事以外、ほぼ網羅できたと思います。公式のサンプルには工夫を凝らした例もありますので参考になると思います。

TensorRT関係リンク