Colabでいちいち毎回StableDiffusionPipelineのためのモデルをHugging Faceからダウンロードしてくるのは馬鹿みたいだ。これはGoogle DriveにCloneしておいてそこから読み出すべきものだ。絶対に。
モデルとして最近気に入ってるのはNeverEnding Dream (Lykon/NeverEnding-Dream) です。
Colabを立ち上げてGoogle Driveをマウントしたら、大容量ファイルをgit cloneするためにgit lfsを有効化する。
!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の、今回はMyDrive/StableDiffusion/models/内に落としていく。clone前に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を落とすようにしてあるけどその辺はHugging Face見てくれればいい。認証が必要なものは使わないので知らない。
モデルをローカルに落とせたら（厳密には我々にとってローカルではないが、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しといていいんだけどね。どうせバージョンとか全部指定して使ってるわけだし。
できたらPipelineの呼び出し。だいたいの人はfrom_pretrainedで呼ぶと思うし、ここで一般的にはmodel_idにHugging Face上の識別子を指定するわけだけど、ここをローカルのディレクトリに設定しておくことでそこをリポジトリとしてモデルを呼び出してくれる。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)
この辺の生成の挙動は基本的に確率的なものだし、そういう意味では一期一会(笑)なのでseedは記録しません。したければファイル名を指定するnewfにでもつけてあげてください。
最後にGoogle Driveに移動させるようにしているのは意図的なものです。他の自作関数で何百枚とか生成させた後にまとめて移動させてる名残りです。
↓ちなみにNEDで生成するとこんな感じ↓
