🧪 Google Colaboratory+Stable Diffusionで作成した画像をSlackに飛ばす

2022-10-19 /Development #Stable Diffusion

Stable Diffusion、手元のノートPCでは画像の生成にすごく時間がかかってしまうので、Google Colaboratoryを利用して遊んでいます。

Google Colaboratoryで画像を作ると、クラウド上のプロジェクトのディレクトリに画像が生成されるので、確認したりダウンロードするのに左メニューのファイルブラウザを開いてポチポチと押してまわる事になるのですが、面倒だし一覧性も良く無いので、Slackに飛ばすようにして使っています。普段Pythonはあんまり書かないし、後からまた使いたくなった時に忘れているかもしれないのでメモ。

準備

まず、Google ColaboratoryでStable Diffusion他、必要なライブラリをインストールします。

コピーしました
pip install diffusers==0.3.0 transformers scipy ftfy torch requests

Hugging Face で取得してきたトークンも変数として設定しておきます。

コピーしました
YOUR_TOKEN="<huggingfaceで取得したtoken>"

次にStable Diffusionパイプラインが使えるように準備します。

コピーしました
from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=YOUR_TOKEN).to("cuda")

あとSlackに投稿したいので、Slackのtokenも取得してきて変数を設定しておきます。

コピーしました
SLACK_TOKEN="<slackで取得したtoken>"

実装

あとは画像の生成〜保存〜Slackへの投稿〜投稿済みの画像を削除、という処理を記述します。

コピーしました
import os
import time
import requests
import torch
from datetime import datetime as dt

def file_name(num):
	tstr = dt.now().strftime('%Y-%m-%d-%H-%M-%S')
	return f'{tstr}_{num}.png'

def send_image(file_path):
	files = {'file': open(file_path, 'rb')}
	params = {"token": SLACK_TOKEN, 'channels': '#channel_name', 'initial_comment': 'from google colab'}
	r = requests.post('https://slack.com/api/files.upload', params=params, files=files)

for num in range(5):
  prompt = "photo of a beautiful cat, expulsion thunder from hands, colorful vibrant thunder sky, night, beautiful soft purple blue lighting, beautiful composition, dramatic lighting"
  size = 512
  image = pipe(prompt,
             height=size,
             width=size,
             guidance_scale=7.5,
             num_inference_steps=50,
             generator=torch.Generator("cuda").manual_seed(num),
             )["sample"][0]
  filepath = file_name(num)
  image.save(filepath)
  time.sleep(1)
  send_image(filepath)
  time.sleep(1)
  os.remove(filepath)
  time.sleep(1)

書いてある通りなのですが、send_image の部分でSlackに投稿しています。

コツといえるほどのものか微妙ですが、file_name でファイル名を指定する際に、日付とあと乱数のシードIDをファイル名にメモしておくと、後々良いのができた時に再現できるのでオススメです。

あとはこれらの関数を使いつつ画像を生成するのですが、1つのプロンプトからたくさん画像を生成してその中から選びたいので、 for 文で生成したい画像の数だけ繰り返して画像を生成するようにしています。一斉に動かすのも何か良く無い気がしたのでsleepを挟んでちょっとづつ実行するようにしています。

という処理を書いておいたので、一度実行すればあとはできたものから順にSlackに流れてくるのを他の事でもしながらのんびり待つ事ができるようになりました。

たまに思いついた文章で生成してみたりして楽しく遊んではいるものの、実用的な使用方法が全然思い当たらない…

参考サイト

Google Colab ではじめる Stable Diffusion v1.4

Comment