text2img の簡単なウェブアプリを作った

TL;DR

最近 text2img の発展が著しいので、自分でもモデルを触って遊んでみようとかやっていた。
自分で遊ぶ用に簡単な web アプリを作ったりした https://github.com/yoheikikuta/text2img_webapp ので、そのあたりの話を書いておくエントリ。

text2img とは

入力した文章に基づく画像を生成する過程を指す言葉であり、最近 diffusion model を使いつつ巨大なモデルを構築することによって飛躍的にその性能が向上している。

この手の話は昔からやられているが、なかなか難易度が高いので人間が描いたものと遜色がないというレベルに結構難しいなと思っていた。 それが 2021 年 初頭に DALL·E https://arxiv.org/abs/2102.12092 によって驚くような性能を発揮したことで、ワクワクして論文を読んだり podcast で取り上げたりした。 podcast は以下の3回のシリーズものとして出している。

この後に画像生成で diffusion model が流行りを見せ、それを text2img に取り入れて更に性能を引き上げた unCLIP (DALL·E2) https://arxiv.org/abs/2204.06125 や Imagen https://arxiv.org/abs/2205.11487 などが出てきて、これはもう人間が描いたものと遜色がないというレベルまで来ている。

自分で遊べる pretrained model が欲しいなと思っていたら、GLIDE https://arxiv.org/abs/2112.10741 という手法の GitHub repository https://github.com/openai/glide-text2im で簡易版ではあるが pretrained model が MIT ライセンスで公開されていたので、それを使って遊んでみようとなった。

web アプリを作る

アプリとしてはまったく簡素なもので、以下のように 日本語文章 → (Cloud Translation で英訳) → 対応する英語文章を作成 → (text2img モデルで画像生成) → 生成画像を表示 するだけのものであり、https://github.com/yoheikikuta/text2img_webapp にコードを残していてこれがほぼ全ての情報である。

構成は frontend を streamlit で実装して Cloud Run で動かし、backend を FastAPI で実装して GPU enabled な GCE で動かしている、という特に何も面白いこともないものである。 結果的には特に面白いことはないが、ここに至るまでいくつか試したものがあったりするので備忘録がてらそれをメモしておく。

GPU を使う

GPU を使わないと（diffusion step の数をかなり減らしても）画像生成が数分掛かるので GPU を使いたいが、GPU はそこそこお金が掛かるので使う時だけ動かすようにしたい。 Cloud Run が GPU をサポートしてくれるようになると嬉しいけど、今のところはサポートされてない（将来的に使えるようになるかもという話は聞いたり聞かなかったりする）。

Cloud Run for Anthos では GPU を使える https://cloud.google.com/anthos/run/docs/configuring/compute-power-gpu という話を聞いてちょっと試してみたりした。 これは GPU enabled な node pool を有する (Cloud Run for Anthos を有効化した) GKE cluster 準備して、その node に NVIDIA driver 準備して、とかでとりあえずまともに起動するだけでもちょっと手間取った。 そして node 数が最低 3 個じゃないとそもそも cluster 作れなかったりして、あれこれ試してるうちにちょっとした金額になってしまったという… 個人で GPU を使ったちょっとしたサービスを作りたい、という場合には殆ど選択肢に入らなそう。

ということで大人しく Computer Engine を使うことにしたのだが、Terraform 使って resource 作成時に docker container が起動していてあとは text2img の endpoint を叩くだけにしたい、というところでもちょっと大変だった。

Compute Engine では Container-Optimized OS https://cloud.google.com/container-optimized-os/docs という Docker container に最適化された VM image が提供されているが、これで GPU 使おうとすると cloud-init 使って NVIDIA driver と CUDA インストールをしないといけない。 いまだにそんなことしないといけないのは面倒だなと思って、deeplearning-platform-release プロジェクトの PyTorch が使える image を以下から探して使うことにした。

gcloud compute images list --project deeplearning-platform-release --no-standard-images

そしたら今度は起動時に docker container を走らせておくのがどうすべきかというところで困って、軽く調べてもよさそうな解決策が見つからなかったので、結局 systemd の service として docker container を起動するものを登録して、それを起動時に metadata_startup_script で実行するという形になった。 これなら COS image で cloud-init 書いてるのと一緒やん…ということになってしまった。

これくらいの用途ならごちゃごちゃやらずにスッと使えるものがあって欲しい気がするが、自分には見つけられなかったので知ってる人がいたら教えてください。 ちなみに現状では T4 使うようにしてるけど、A100 使うように変えると動かなかったりして、まだまだイマイチである。

適当なところ

とりあえず自分が遊ぶように作ったので、適当なところが多い。 ライブラリのバージョンを fix してないとか、backend の API が公開されてるとか、local の google cloud の credential で resource 作ってて適切な権限持たせた service account 作ったりしてないとか、Container Registry 使ってて Artifact Registry 使ってないのと雑に latest タグ使ってるとか、その他もろもろ。

もう少し真面目に作っていくか〜という気持ちになりたかったが、後述のように pretrained model が非力で十分なやる気が喚起されなかった（という言い訳）。 今後もっと良いモデルが使えるようになれば改めて頑張るかもしれない。

pretarined model の非力さ

今回遊んでみて少しは楽しかったのだが、pretrained model を使った生成画像が論文で報告されている画像と比べてイマイチだったというのが残念なところであった。

公開されている pretrained model は小さいモデルになっており、学習では（悪用を防ぐためと思われるが）人物のデータとか暴力的な object とかが除かれたデータが使われているとか、upscale のモデルは二段階ではなくて一段階目のみ提供されている。

最近の text2img ではインターネット上になさそうな画像（例えば テディベアがバタフライで泳いでる とか）を生成できるところに凄さがあるのだが、それが殆ど再現できてないと感じられた。 これだと最近の本質的な発展が取り込まれていないので、遊びがいがないな〜という結果になってしまった。 dalle-mini https://huggingface.co/spaces/dalle-mini/dalle-mini の結果の方がかなりよさそうに感じるので、お手軽に遊ぶならこちらの方がいいかなと思う。

もっとよいモデルが使えるようになったらまた改めて着手してみるとするかな、ということで今回は終了。

まとめ

text2img の web アプリを作ったというエントリ。
text2img の発展がすごいのでフルのモデルを自分で動かしてみたいなぁ。