💫
高品質動画生成AI【Pyramid Flow】をローカルで実行する

2024/10/22に公開
 Pyramid Flowとはhttps://github.com/jy0205/Pyramid-Flow

2024年10月10日に公開された、ローカルで動作させることができる非常に高品質な最新の動画生成AIです。
非常に最新の動画生成AIのため、まだベータ版のような安定性ですが、ポテンシャルは非常に感じました。
性能としては、最大768pの解像度で24FPSで最大10秒の動画を生成することができ、さらに、テキストから動画への変換（t2v）、画像から動画への変換（i2v）が利用できます。
生成されるデモ動画は下記をご覧ください。

https://pyramid-flow.github.io/
これらが、「既存のオープンソースデータセットのみ」で学習されたモデルにより生成されているようです。

（ここが一番驚きました）
動画生成や画像生成のモデル作成には大量のデータセットが必要と思っていたため、まさかオープンソースのデータセットだけ（とは言っても多いですが）でこの品質に到達できるとは・・・
今回は、RTX3060の安価なGPUで動作確認をしたため、防備録として記事を残します。
ちなみにデモは下記から実行できます。

https://huggingface.co/spaces/Pyramid-Flow/pyramid-flow
私は、WebUIを宗教上の理由で利用できないため、実際にPyhonコードで実行できるようにしていきたいともいます。

 注意記事執筆時点（2024年10月20日）で、今回のリポジトリは完全体ではありません。
たとえば、リポジトリ上ではmodel.enable_sequential_cpu_offload()を利用することで8GB未満のGPUメモリで推論が可能になると記載されていますが、768pの解像度で12GBのGPUでは、最後まで推論を完了することができません。
どうやら、内部は自己回帰モデルとなっているようで、生成されたフレーム数が増えるほど利用するGPUメモリが増えていくため、最初のstepは推論できますが、ある一定のstepを超えた段階でOOMが発生します。
また、現時点では初期モデル重みとしてSD3(Stable Diffusion 3)による初期化を歳用しているようですが、その場合、生成された人間の構造に問題が発生するようです。したがって著者は新規でPyramid Flow をゼロからトレーニングしており、近日中にその重みを公開すると宣言しています。

（おそらくその重みのほうが性能が高い）
したがって、今回の記事での動画や実行コードなどはあくまで現時点での実行コードであり、推論結果であることを留意してください。
また、特定のPythonバージョンでしか動かないなど、まだまだ使いやすいリポジトリではありません。

今後、Diffusersライブラリに統合されることが期待されますが、その場合とはまた構築する環境や実行コードなどが変わると思われます。

 環境構築あくまで2024年10月20日現在の環境構築手法です。

 実行環境OS:Ubuntu 20.04

GPU:RTX3060 12GB

CUDA:12.2

RAM:64GB (おそらく25GB以上は必要そうです)
Pythonはanacondaによる仮想環境を利用します。

（pyenv+venvでも問題ないと思います）

 リポジトリをクローンするgit clone https://github.com/jy0205/Pyramid-Flow
リポジトリをクローンして、カレントディレクトリを移動します。
cd Pyramid-Flow

 Anacondaによる仮想環境を構築するconda create -n pyramid python==3.8.10
リポジトリで指定されているPythonのバージョンで仮想環境を構築します

（Pythonのバージョンを揃えないと実行できません。なんとかしてほしい）
下記コマンドで仮想環境に入ります。
conda activate pyramid
下記コマンドで必要なモジュールをインストールします。
pip install -r requirements.txt

 モデルのダウンロードモデル自体は下記に存在します。

https://huggingface.co/rain1011/pyramid-flow-sd3

こちらをすべて自力でダウンロードしてきても良いですが、公式から便利なコードが提供されているため、今回はそちらを利用します。
Pyramid-Flowフォルダ内に、下記のmodel_download.pyを作成してください。
model_download.py
from huggingface_hub import snapshot_download

model_path = 'model_assets'   # The local directory to save downloaded checkpoint
snapshot_download("rain1011/pyramid-flow-sd3", local_dir=model_path, local_dir_use_symlinks=False, repo_type='model')
その後、下記のコマンドで実行して、Pyramid Flowのモデルをダウンロードしてください。
python model_download.py

 実行コードの作成Pyramid-Flowフォルダ内に、下記のmain.pyを作成してください。
main.py
import torch
from PIL import Image
from pyramid_dit import PyramidDiTForVideoGeneration
from diffusers.utils import load_image, export_to_video

torch.cuda.set_device(0)
model_dtype, torch_dtype = 'bf16', torch.bfloat16   
model_path = 'model_assets'   # The local directory to save downloaded checkpoint

#use_resolutions = "384p"
use_resolutions = "768p"
#use_mode = "t2v"
use_mode = "i2v"
image_path = "videoframe_2013.png"


# used for 384p model variant
if use_resolutions == "384p":
    width = 640
    height = 384
    guidance_scale=7.0 
    sequential_cpu_offload = False
    cpu_offloading = True
    model_variant='diffusion_transformer_384p'
    temp=31
# used for 768p model variant
else:
    width = 1280
    height = 768
    guidance_scale=9.0
    sequential_cpu_offload = True
    cpu_offloading = False
    model_variant='diffusion_transformer_768p'
    temp=8

image = None
if use_mode == "t2v":
    print("Using text to video mode")
else:
    image = load_image(image=image_path)
    print("Using image to video mode")



model = PyramidDiTForVideoGeneration(
    model_path,                                         
    model_dtype,
    model_variant=model_variant,   
)


model.vae.enable_tiling()
if sequential_cpu_offload:
    model.enable_sequential_cpu_offload()

prompt = "A side profile shot of a woman with fireworks exploding in the distance beyond her"

if use_mode == "t2v":
    with torch.no_grad(), torch.cuda.amp.autocast(enabled=True, dtype=torch_dtype):
        frames = model.generate(
            prompt=prompt,
            num_inference_steps=[20, 20, 20],
            video_num_inference_steps=[10, 10, 10],
            height=height,     
            width=width,
            temp=temp,                    # temp=16: 5s, temp=31: 10s
            guidance_scale=guidance_scale,         # The guidance for the first frame, set it to 7 for 384p variant
            video_guidance_scale=5.0,   # The guidance for the other video latent
            output_type="pil",
            cpu_offloading=cpu_offloading,
            save_memory=True,           # If you have enough GPU memory, set it to `False` to improve vae decoding speed
        )
else:
    with torch.no_grad(), torch.cuda.amp.autocast(enabled=True, dtype=torch_dtype):
        frames = model.generate_i2v(
            prompt=prompt,
            input_image=image,
            num_inference_steps=[10, 10, 10],
            temp=temp,
            video_guidance_scale=4.0,
            output_type="pil",
            save_memory=True,           # If you have enough GPU memory, set it to `False` to improve vae decoding speed
        )


export_to_video(frames, "./sample.mp4", fps=24)

 実行下記コマンドを実行することで、動画が生成できます。
python main.py

 生成された動画
 Text to Videoをためすまずは、t2vを試します。

プロンプトはデモ動画と同じプロンプトを利用します。
A side profile shot of a woman with fireworks exploding in the distance beyond her
では、まずは384pの解像度のモデルで24FPSで10秒の動画を再生してみます。
https://youtu.be/-4fIBK5JjpA
では、続いては768pの解像度のモデルで動画を生成しようと思います。

しかし、残念ながらRTX3060では、768pの解像度のモデルではOOMになってしまいます。

これは内部的に、自己回帰モデルが利用されているため、生成した動画のframe数が増えるほど、GPUメモリを逼迫します。
ただ、10秒の動画を作成するのは難しそうでしたが、2秒程度の動画であればギリギリ切り抜けることができそうだったので、768pの解像度で動画を作成する場合は2秒程度の動画が生成される設定になっています。

（より大きなGPUメモリが利用できる方はtempの値を増やして調整してください）
では、同じプロンプトで作成した動画は下記です。
https://youtu.be/qbUEdWWPpp8
やはり、まだ安定性は低く感じてしまいますが、それでもオープンソースのデータセットだけでこのレベルの性能の動画生成AIを作ることができるというのは、今後の進化が非常に期待できる結果だと思います。

 Image to Videoをためす続いて、i2vを試します。
これも、デモのものを流用します。

プロンプトは下記です。
FPV flying over the Great Wall
画像は下記のものを利用します。

では、まずは384pの解像度のモデルで24FPSで10秒の動画を再生してみます。
https://youtu.be/7tyflA4GkkI
では、続いては768pの解像度のモデルで動画を生成しようと思います。

（こちらも生成される動画時間は制限しています）
https://youtu.be/cV-EbkJepYs
やはり、384pの解像度においては、後半すこし安定性にかける部分はありますが、それでもかなりポテンシャルを感じます。

 コードの簡単な解説main.py
model_dtype, torch_dtype = 'bf16', torch.bfloat16   # Use bf16 (not support fp16 yet)
model_path = 'model_assets'   # The local directory to save downloaded checkpoint
基本的にはPyramid Flowはbf16の利用を想定されているようですので、fp16は利用できないそうなので、利用するGPUの世代によっては注意してください。

RTX2000番台の場合は実行が難しい可能性があります。
また、モデルの重みはmodel_assetsフォルダに格納されていることとしています。
main.py
use_resolutions = "768p" #"384p"
use_mode = "t2v" #"i2v"
image_path = "videoframe_2013.png"

# used for 384p model variant
if use_resolutions == "384p":
    print("Using 384p model variant")
    width = 640
    height = 384
    guidance_scale=7.0 
    sequential_cpu_offload = False
    cpu_offloading = True
    model_variant='diffusion_transformer_384p'
    temp=31
# used for 768p model variant
else:
    print("Using 768p model variant")
    width = 1280
    height = 768
    guidance_scale=9.0
    sequential_cpu_offload = True
    cpu_offloading = False
    model_variant='diffusion_transformer_768p'
    temp=8

image = None
if use_mode == "t2v":
    print("Using text to video mode")
else:
    image = load_image(image=image_path)
    print("Using image to video mode")
    if use_resolutions == "768p":
        image = image.resize((width, height), Image.LANCZOS)
        temp=7
    else:
        image = image.resize((width, height), Image.LANCZOS)
こちらでは、上から「生成する動画の解像度の指定」「動画を生成するモードの指定」「画像から動画を生成する場合の画像が保存されているpath」を指定しています。
また、生成する動画の解像度によって、（当然ですが）動画の縦横指定や、guidance_scale、temp（動画の長さを指定するパラメータ、3ごとに約1秒）などが変わるため、それを指定しています。
また、生成モードとしてi2vが指定された場合は、画像をパスから読み取ったあとに、リサイズをしています。

またi2vのほうが、GPUメモリが厳しいため、ひとつだけtempを小さくしています。
main.py
model = PyramidDiTForVideoGeneration(
    model_path,
    model_dtype,
    model_variant=model_variant,   
)

model.vae.enable_tiling()
if sequential_cpu_offload:
    model.enable_sequential_cpu_offload()

prompt = "A side profile shot of a woman with fireworks exploding in the distance beyond her"

つづいて。モデルの読み込みとインスタンスの作成、そして、CPUオフロードの設定と、プロンプトの指定を行っています。

model.enable_sequential_cpu_offload()を利用することで、生成速度が大幅に遅くなりますが、代わりに小さなVRAMでもモデルを動かすことができるようになる起死回生の技術です。
モデルの層ごとに、VRAMとRAMを行き来させ、必要な層だけGPUのメモリに置くようにしています。この移動に時間がかかるため生成速度は遅くなりますが、必要な層だけをGPUメモリに置くことができるため、必要なVRAMが小さくできます。
main.py
if use_mode == "t2v":
    with torch.no_grad(), torch.cuda.amp.autocast(enabled=True, dtype=torch_dtype):
        frames = model.generate(
            prompt=prompt,
            num_inference_steps=[20, 20, 20],
            video_num_inference_steps=[10, 10, 10],
            height=height,     
            width=width,
            temp=temp,                    # temp=16: 5s, temp=31: 10s
            guidance_scale=guidance_scale,         # The guidance for the first frame, set it to 7 for 384p variant
            video_guidance_scale=5.0,   # The guidance for the other video latent
            output_type="pil",
            cpu_offloading=cpu_offloading,
            save_memory=True,           # If you have enough GPU memory, set it to `False` to improve vae decoding speed
        )
else:
    with torch.no_grad(), torch.cuda.amp.autocast(enabled=True, dtype=torch_dtype):
        frames = model.generate_i2v(
            prompt=prompt,
            input_image=image,
            num_inference_steps=[10, 10, 10],
            temp=temp,
            video_guidance_scale=4.0,
            output_type="pil",
            save_memory=True,           # If you have enough GPU memory, set it to `False` to improve vae decoding speed
        )


export_to_video(frames, "./sample.mp4", fps=24)
最後に動画生成を実行した上で、動画を保存しています。

これまでの設定に応じて、どれが実行されるかが変わります。
cpu_offloadingは384pの解像度の動画を生成するときにのみ利用しています。

これは上述したmodel.enable_sequential_cpu_offload()と同様にVRAMを節約する技術ではありますが、model.enable_sequential_cpu_offload()よりは、節約量がちいさい設定になります。その代わり、生成速度は早くなります。
自身のVRAMの大きさに合わせて、使い分けることをおすすめします。

 まとめかなりポテンシャルの高い動画生成モデルがローカルで利用できるようになっていました！

近日中に新しい重みが公開されるようですので、とてもたのしみですね！
また、Diffusersライブラリにも、CogVideoXという動画生成AIが統合されています。こちらのモデルもいろいろ試してみて、結構気に入っているので、こちらのモデルについての記事も今後書きたいと思います。
ここまで読んでくださって、ありがとうございました。