WSL2でAnimateDiff-Lightningを試してみる

「超高速のテキストからビデオへの生成モデル。オリジナルの AnimateDiff よりも 10 倍以上速くビデオを生成できる」らしいAnimateDiff-Lightningを試してみます。

ByteDance/AnimateDiff-Lightning · Hugging Face

使用するPCはドスパラさんの「GALLERIA UL9C-R49」。スペックは
・CPU: Intel® Core™ i9-13900HX Processor
・Mem: 64 GB
・GPU: NVIDIA® GeForce RTX™ 4090 Laptop GPU(16GB)
・GPU: NVIDIA® GeForce RTX™ 4090 (24GB)
・OS: Ubuntu22.04 on WSL2（Windows 11）
です。

1. 準備

venv構築。

python3 -m venv animatediff-lightning
cd $_
source bin/activate

デモサイトにある requirements.txt を参考にパッケージのインストールします。

pip install torch transformers accelerate diffusers

2. 流し込むコード

README.mdに記載のdiffusersのコードを参考にして、ちょこっと修正したコードを流し込みます。

import torch
from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler
from diffusers.utils import export_to_gif
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

device = "cuda"
dtype = torch.float16

step = 4  # Options: [1,2,4,8]
repo = "ByteDance/AnimateDiff-Lightning"
ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
base = "emilianJR/epiCRealism"  # Choose to your favorite base model.

adapter = MotionAdapter().to(device, dtype)
adapter.load_state_dict(load_file(hf_hub_download(repo ,ckpt), device=device))
pipe = AnimateDiffPipeline.from_pretrained(base, motion_adapter=adapter, torch_dtype=dtype).to(device)
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", beta_schedule="linear")

# 以下、コメントアウトと追加

# output = pipe(prompt="A girl smiling", guidance_scale=1.0, num_inference_steps=step)
# export_to_gif(output.frames[0], "animation.gif")

import os
import time

def i(
    prompt: str=None,
):
    start = time.process_time()
    if not prompt:
        prompt = "A girl smiling"
    # generate image
    image = pipe(
        prompt,
        guidance_scale=1.0,
        num_inference_steps=step
    )
    # save a generated image to "./outputs/image_%Y%m%d_%H%M%S.gif"
    image_basename = "animation"
    outdir="outputs"
    os.makedirs(outdir, exist_ok=True) 
    export_to_gif(
        image.frames[0],
        outdir + os.sep + image_basename + "_" + time.strftime('%Y%m%d_%H%M%S') + ".gif"
    )
    end = time.process_time()
    ##
    total_time = end - start
    print(f"   total time = {total_time:f} [s]")

関数 i を定義し、引数にプロンプトを指定できるようにしたのと、出力先ディレクトリとファイル名が重複しないようにしているだけです。

3. 試してみる

24GB、16GB、それぞれのGPUカードで試します。

(1) RTX 4090（24GB）

上記コードを流し込みます。そして、聞きましょう。

>>> i("Doraemon smiling")
100%|████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:01<00:00,  3.48it/s]
   total time = 3.343191 [s]
>>>

3.3秒、はやい！

Doraemon smiling

>>> i("Nobi Nobita smiling")
100%|████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:01<00:00,  3.78it/s]
   total time = 2.822016 [s]
>>>

2回目以降は、だいたい2.8秒付近です。

Nobi Nobita smiling

どことなく、ドラえもんのカラーが入っているような入っていないような…。

(2) RTX 4090 Laptop GPU（16GB）

残念ながらVRAMが溢れました。

このため、生成時間も45秒とかかります。

>>> i("Doraemon smiling")
100%|████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:01<00:00,  2.06it/s]
   total time = 45.047692 [s]
>>>

これは、、、なんだ？

4. まとめ

GPUのVRAMは 24GB 必要です。
そして、確かにAnimateDiffより速い！

関連