ColabでDiffusersモデルをGoogle Driveから読み出す

Colabでいちいち毎回StableDiffusionPipelineのためのモデルをHugging Faceからダウンロードしてくるのは馬鹿みたいだ。これはGoogle DriveにCloneしておいてそこから読み出すべきものだと思う。絶対に。
ちなみに最近気に入ってるモデルはNeverEnding Dream (Lykon/NeverEnding-Dream) です。

実際にやってみる

Colabを立ち上げてGoogle Driveをマウントしたら、まずは準備としてgit lfsを有効化する。これがないと大容量ファイルのgit cloneができない。

!curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh | sudo bash
!sudo apt-get install git-lfs
!git lfs install

できたら、次にGoogle Drive内にモデルの置き場所を作る。Drive側で作ってもいいけど面倒なのでColab上でやってしまう。今回はGoogle Drive上のStableDiffusion/models（Colab側から見たフルパスは/content/drive/MyDrive/StableDiffusion/models/） に置くことにする。
clone前にcdでディレクトリ移動するので、事前に必要なディレクトリを作るなり下のコードのパスを書き換えるなりする。ちなみにcdは%使わないとちゃんと移動してくれない。

%cd /content/drive/MyDrive/StableDiffusion/models/
!pwd
!git clone https://huggingface.co/stabilityai/stable-diffusion-2-1
!git clone https://huggingface.co/Lykon/NeverEnding-Dream

今回はsd2.1とNEDを落とすようにしてあるけど、単純にhuggingface上のリポジトリのURLを指定しているだけであることをわかってほしい。

モデルをローカルに落とせたら（もちろん我々にとって厳密にはローカルではないのだけれどもColabからみたらローカルみたいなもんだよね）、モデルを実際に読み込めるか試していく。

とりあえず必要なもの（diffusersとか）をざっと入れる。ぼくは毎回text-to-imageとimage-to-imageの両方ともをまとめてセットアップするので、どちらもimportする。スケジューラーは別のライブラリの都合で別で呼び出してるけど今回は正直どっちでもいい。

!pip install -q diffusers==0.11.1
!pip install -q transformers scipy ftfy accelerate
import torch
from diffusers import StableDiffusionPipeline
from diffusers import StableDiffusionImg2ImgPipeline
from diffusers import DPMSolverMultistepScheduler

diffusersとかまでcloneしといてもいいんだけどね。どうせバージョンとか全部指定して使ってるわけだし。ほんといつの間にか動かなくなるからね。ただColab側のアップデートで動かなくなることもあるので本当はColabの環境ごとコンテナ化できるようになってほしい。

できたらPipelineの呼び出し。だいたいの人はfrom_pretrainedで呼ぶと思うし、ここで一般的にはmodel_idにHuggingface上の識別子を指定するわけだけど、これをローカルのディレクトリに設定しておくことでそこをリポジトリとしてモデルを呼び出してくれる。
たとえば、

t2i = StableDiffusionPipeline.from_pretrained('/content/drive/MyDrive/StableDiffusion/models/NeverEnding-Dream')

こんな感じで、t2iとi2i両方立ち上げる。

model_id = '/content/drive/MyDrive/StableDiffusion/models/NeverEnding-Dream'
t2i = StableDiffusionPipeline.from_pretrained(
    model_id,
    feature_extractor=None,
    safety_checker=None,
    torch_dtype=torch.float16
    ).to('cuda')
t2i.scheduler=DPMSolverMultistepScheduler.from_config(t2i.scheduler.config)

i2i = StableDiffusionImg2ImgPipeline.from_pretrained(
    model_id,
    feature_extractor=None,
    safety_checker=None,
    torch_dtype=torch.float16
    ).to('cuda')
i2i.scheduler=DPMSolverMultistepScheduler.from_config(i2i.scheduler.config)

みてわかると思いますが、私は（判定と責任は人間の仕事だという立場なので）基本的にsafety_checkerは指定せずに使います。

あとはgenerate_imageっていう関数を作って、そこの引数の指定の仕方でtext-to-imageなのかimage-to-imageなのか判定させる。基本的にはgenerate_imageの引数にurlの指定があったらそのURLに画像取りにいってi2i走らせる感じ。あとはリネームとか移動とか。

import os
import datetime as dt
import numpy.random as rnd
from random import sample
import requests
from PIL import Image
from io import BytesIO

def get_now():
    now = dt.datetime.now(dt.timezone(dt.timedelta(hours=9)))
    now = str(now.year)+str(now.month).zfill(2)+str(now.day).zfill(2)+str(now.hour).zfill(2)+str(now.minute).zfill(2)+str(now.second).zfill(2)+'_'+str(rnd.randint(100))
    return now

def get_image(url):
    response = requests.get(url)
    init_image = Image.open(BytesIO(response.content)).convert("RGB")
    init_image = init_image.resize((768, 512))
    return init_image

def generate_image(prompt, url=None, categ='graffiti', seed=None, ninf=50, nbatch=1):
    dir = '/content/drive/MyDrive/StableDiffusion/images/'
    ngp = 'Bad quality'

    gen = seed if seed else rnd.randint(10000000)
    gen = torch.Generator("cuda").manual_seed(gen)

    for i in range(nbatch):
        if url:
            init = get_image(url=url)
            img = i2i(prompt, negative_prompt=ngp, generator=gen, num_inference_steps=ninf, image=init, strength=0.75, guidance_scale=7.5).images[0]
        else:
            img = t2i(prompt, negative_prompt=ngp, generator=gen, num_inference_steps=ninf).images[0]
        now  = get_now()
        newdir = dir + categ
        if not os.path.isdir(newdir):
            os.mkdir(newdir)
        newf = newdir + '/' + now +'.png'
        img.save(newf)
        print('Saved as', newf)

あと、実際に使ってるコード上では、指定したモデルがローカルにあれば読み込むし、なければHuggingface上に取りにいってsave_pretrainedでローカルに保存するっていうロジックを組んであるので、一旦使ったモデルは全部Drive上に保管される。今回は省略。

最後にGoogle Driveに移動させるようにしているのは意図的なものです。他の自作関数で何百枚とか生成させた後にまとめて移動させてる名残りです。

↓ちなみにNEDで生成するとこんな感じ↓

この辺の生成の挙動は基本的に確率的なものだし、そういう意味では一期一会(笑)なのでseedは記録していませんが、必要ならしたければファイル名を指定するnewfにでもstr(seed)みたいな感じで付与してあげればいいんじゃないかと思います。Stable-Diffusion-Videosとか使う時にはseedの指定によっては切れ目ない連続した静止画が作れるからまたちょっと話変わってくるんだけど。