UniDiffusersをgoogle colabで試してみた

UniDiffusersとは

UniDiffusersとは、1つのstable diffusionのモデルでText2Image, Image2Text, Image2Image, Text2Textをやってしまおうというもの。すごく便利、、
https://github.com/thu-ml/unidiffuser

diffusersにもv0.17.0から追加されています。

リンク

Colab
github

準備

Google Colabを開き、メニューから「ランタイム→ランタイムのタイプを変更」でランタイムを「GPU」に変更します。

環境構築

インストール手順です。

!pip install diffusers transformers accelerate

推論

4つのタスクをやってみようと思います。
(1)Text2Image

import torch
from diffusers import UniDiffuserPipeline

model_id_or_path = "thu-ml/unidiffuser-v1"
pipe = UniDiffuserPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)
pipe.to("cuda")

# This mode can be inferred from the input provided to the `pipe`. 
pipe.set_text_to_image_mode()

prompt = "an elephant under the sea"
sample = pipe(prompt=prompt, num_inference_steps=20, guidance_scale=8.0).images[0]
display(sample)

(2) Image2Text
先ほどの画像に対してやってみます。

import torch

from diffusers import UniDiffuserPipeline

device = "cuda"
model_id_or_path = "thu-ml/unidiffuser-v1"
pipe = UniDiffuserPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)
pipe.to(device)

from PIL import Image
# Image-to-text generation
init_image = Image.open("/content/an elephant under the sea.jpg").resize((512, 512))

sample = pipe(image=init_image, num_inference_steps=20, guidance_scale=8.0)
i2t_text = sample.text[0]
print(i2t_text)

output

An elephant swimming in the water of a aquarium

おー確かにOK

(3)Image2Image
(1)の画像に対して行います。

import torch

from diffusers import UniDiffuserPipeline
from PIL import Image

device = "cuda"
model_id_or_path = "thu-ml/unidiffuser-v1"
pipe = UniDiffuserPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)
pipe.to(device)

# Image variation can be performed with a image-to-text generation followed by a text-to-image generation:
init_image = Image.open("/content/an elephant under the sea.jpg").resize((512, 512))

sample = pipe(image=init_image, num_inference_steps=20, guidance_scale=8.0)
i2t_text = sample.text[0]
print(i2t_text)

# 2. Text-to-image generation
sample = pipe(prompt=i2t_text, num_inference_steps=20, guidance_scale=8.0)
final_image = sample.images[0]
display(final_image)

海の底から出ちゃったけど象と海はある

(4)Text2Text

import torch

from diffusers import UniDiffuserPipeline

device = "cuda"
model_id_or_path = "thu-ml/unidiffuser-v1"
pipe = UniDiffuserPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)
pipe.to(device)

# Text variation can be performed with a text-to-image generation followed by a image-to-text generation:
# 1. Text-to-image generation
prompt = "an elephant under the sea"

sample = pipe(prompt=prompt, num_inference_steps=20, guidance_scale=8.0)
t2i_image = sample.images[0]
t2i_image.save("unidiffuser_text2img_sample_image.png")

# 2. Image-to-text generation
sample = pipe(image=t2i_image, num_inference_steps=20, guidance_scale=8.0)
final_prompt = sample.text[0]
print(final_prompt)

output

An elephant is standing in a aquarium
# An elephant swims in the water and its head
# An elephantine water in a aquarium

まあ想像通り。そんなに多くのバリエーションは作れなそう。でも答えは揺らぐ

最後に

今回はさまざまなタスクを1つのdiffusion modelでやろうというmultimodalなdiffusion model,UniDiffuserについてgoogle colabで試してみました。似たようなやつにCoDiというのがありますが、あっちはgoogle colab freeでは動かんので手軽に試したい場合はこちらを利用ください。

今後ともLLM, Diffusion model, Image Analysis, 3Dに関連する試した記事を投稿していく予定なのでよろしくお願いします。