見出し画像

Google Colab で Tune-A-Video による動画生成を試す

npaka

「Google Colab」で「Tune-A-Video」による動画生成を試したのでまとめました。

【注意】「Tune-A-Video」を動作させるには、「Google Colab Pro/Pro+」のプレミアム (A100 40GB) が必要です。

1. Tune-A-Video

「Tune-A-Video」は、テキストから動画を生成するための、Image DiffusionモデルのOne-Shotチューニングの手法になります。

2. Colabでの実行

Google Colabでの実行手順は、次のとおりです。

(1) 新規のColabのノートブックを開き、メニュー「編集 → ノートブックの設定」で「GPU」の「プレミアム」を選択

(2) xformersのインストール。

# xformersのインストール (A100用)
!pip install https://github.com/brian6091/xformers-wheels/releases/download/0.0.15.dev0%2B4c06c79/xformers-0.0.15.dev0+4c06c79.d20221205-cp38-cp38-linux_x86_64.whl

ビルドに時間がかかるため、Google ColabのA100とT4に対応している以下のWheelを使わせてもらいました。

(3) Tune-A-Videoのインストール。
PyTorchのバージョン調整も行っています。

# Tune-A-Videoのインストール
!git clone https://github.com/showlab/Tune-A-Video.git
%cd Tune-A-Video
!pip install -r requirements.txt
!pip install torch==1.13.1 torchtext==0.14.1 torchvision==0.14.1 torchaudio==0.13.1

(4) checkpointsフォルダを作成してStable Diffusionモデルをダウンロード。
今回は、「StableDiffusion v1.4」を使っています。DreamBoothしたモデルでもOKとのことです。

# checkpointsフォルダの作成と移動
import os
os.makedirs("./checkpoints", exist_ok=True)
%cd ./checkpoints

# Stable Diffusionモデルのダウンロード
!git lfs install
!git clone https://huggingface.co/CompVis/stable-diffusion-v1-4

# 作業フォルダに戻る
%cd ..

(5) 学習の実行。
学習に7分ほどかかりました。outputsに学習したモデルが出力されます。

%%time

# 学習
!accelerate launch train_tuneavideo.py --config="configs/man-surfing.yaml"

今回はサンプルの設定ファイルをそのまま使います。

・man-surfing.yaml

pretrained_model_path: "./checkpoints/stable-diffusion-v1-4"
output_dir: "./outputs/man-surfing_lr3e-5_seed33"

train_data:
  video_path: "data/man-surfing.mp4"
  prompt: "a man is surfing"
  n_sample_frames: 8
  width: 512
  height: 512
  sample_start_idx: 0
  sample_frame_rate: 1

validation_data:
  prompts:
    - "a panda is surfing"
    - "a boy, wearing a birthday hat, is surfing"
    - "a raccoon is surfing, cartoon style"
    - "Iron Man is surfing in the desert"
  video_length: 8
  width: 512
  height: 512
  num_inference_steps: 50
  guidance_scale: 7.5

learning_rate: 3e-5
train_batch_size: 1
max_train_steps: 300
checkpointing_steps: 1000
validation_steps: 100
trainable_modules:
  - "attn1.to_q"
  - "attn2.to_q"
  - "attn_temp"

seed: 33
mixed_precision: fp16
use_8bit_adam: False
gradient_checkpointing: True
enable_xformers_memory_efficient_attention: True

・man-surfing.mp4 (動画)と"a man is surfing" (男はサーフィンしている) という学習データから、以下の動画を生成できるモデルを学習しています。

・"a panda is surfing" (パンダがサーフィンをしている)
・"a boy, wearing a birthday hat, is surfing" (誕生日の帽子をかぶった少年がサーフィンをしている)
・"a raccoon is surfing, cartoon style" (アライグマがサーフィンをしている、漫画のスタイル)
・"Iron Man is surfing in the desert" (砂漠でサーフィンするアイアンマン)

(6) 動画の生成。
学習したモデルのパスは、左端のフォルダアイコンで確認できます。
(以下の画像では、"./outputs/man-surfing_lr3e-5_seed33/2023-01-30T14-02-30")

from tuneavideo.pipelines.pipeline_tuneavideo import TuneAVideoPipeline
from tuneavideo.models.unet import UNet3DConditionModel
from tuneavideo.util import save_videos_grid
import torch

# パイプラインの準備
model_id = "./outputs/man-surfing_lr3e-5_seed33/2023-01-30T14-02-30"
unet = UNet3DConditionModel.from_pretrained(
    model_id, 
    subfolder='unet', 
    torch_dtype=torch.float16
).to('cuda')
pipe = TuneAVideoPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4", 
    unet=unet, 
    torch_dtype=torch.float16
).to("cuda")

# 推論の実行
prompt = "a panda is surfing"
video = pipe(
    prompt, 
    video_length=8, 
    height=512, 
    width=512, 
    num_inference_steps=50, 
    guidance_scale=7.5
).videos

# 動画の保存
save_videos_grid(video, "outputs/video.gif")

(7) 生成された動画 (outputs/video.gif) を確認。



この記事が気に入ったらサポートをしてみませんか？