diffusers0.25.0で実装された aMUSEd を試してみる

2023年12月28日 23:59

npakaさんのnoteで0.25.0のリリースを知る。

aMUSEd というアルゴリズムによる画像生成もあるようなので実験してみる。

Muse: Text-To-Image Generation via Masked Generative Transformers

ミューズ: マスクされた生成変換器によるテキストから画像への生成

Huiwen Chang、Han Zhang、Jarred Barber、AJ Maschinot、Jose Lezama、Lu Jiang、Ming-Hsuan Yang、Kevin Murphy、William T. Freeman、Michael Rubinstein、Yuanzhen Li、Dilip Krishnan
Muse は、拡散モデルや自己回帰モデルよりも大幅に効率的でありながら、最先端の画像生成パフォーマンスを実現するテキストから画像への Transformer モデルです。 Muse は、離散トークン空間のマスクされたモデリングタスクでトレーニングされます。事前トレーニングされた大規模言語モデル (LLM) から抽出されたテキスト埋め込みを考慮して、Muse はランダムにマスクされた画像トークンを予測するようにトレーニングされます。 Imagen や DALL-E 2 などのピクセル空間拡散モデルと比較して、Muse は離散トークンを使用し、必要なサンプリング反復が少ないため、大幅に効率的です。 Parti などの自己回帰モデルと比較して、Muse は並列デコードを使用するため効率的です。事前トレーニングされた LLM の使用により、きめ細かい言語理解が可能になり、高忠実度の画像生成に変換され、オブジェクト、その空間関係、ポーズ、カーディナリティなどの視覚概念を理解できるようになります。当社の 900M パラメータモデルは、新しい SOTA を実現します。 CC3M、FID スコアは 6.06。 Muse 3B パラメータモデルは、ゼロショット COCO 評価で 7.88 の FID と、0.32 の CLIP スコアを達成しました。また、Muse は、モデルの微調整や反転を必要とせずに、インペイント、アウトペイント、マスクフリー編集などの多くの画像編集アプリケーションを直接有効にします。詳細な結果は、この https URL でご覧いただけます

https://arxiv.org/abs/2301.00704

Hugging Faceのリリースより

aMUSEd は、MUSE アーキテクチャに基づいた軽量のテキストから画像へのモデルです。 aMUSEd は、一度に多くの画像を素早く生成するなど、軽量で高速なモデルを必要とするアプリケーションで特に役立ちます。 aMUSEd は現在研究リリースです。

aMUSEd は、多くの拡散モデルよりも少ない順方向パスでイメージを生成できる VQVAE トークンベースのトランスフォーマーです。 MUSE とは対照的に、T5-XXL の代わりに小型のテキストエンコーダー CLIP-L/14 を使用します。 ammused はパラメータ数が少なく、フォワードパス生成プロセスが少ないため、多くの画像を迅速に生成できます。この利点は、特に大きなバッチサイズで見られます。

Google Colab T4環境で実行してみる

!pip install diffusers

これで 0.25.0がインストールできる。パッケージの依存関係から他のパッケージのバージョンも見ておこう

Requirement already satisfied: diffusers in /usr/local/lib/python3.10/dist-packages (0.25.0)
Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.10/dist-packages (from diffusers) (7.0.0)
Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from diffusers) (3.13.1)
Requirement already satisfied: huggingface-hub>=0.19.4 in /usr/local/lib/python3.10/dist-packages (from diffusers) (0.19.4)
Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from diffusers) (1.23.5)
Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from diffusers) (2023.6.3)
Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from diffusers) (2.31.0)
Requirement already satisfied: safetensors>=0.3.1 in /usr/local/lib/python3.10/dist-packages (from diffusers) (0.4.1)
Requirement already satisfied: Pillow in /usr/local/lib/python3.10/dist-packages (from diffusers) (9.4.0)
Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.19.4->diffusers) (2023.6.0)
Requirement already satisfied: tqdm>=4.42.1 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.19.4->diffusers) (4.66.1)
Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.19.4->diffusers) (6.0.1)
Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.19.4->diffusers) (4.5.0)
Requirement already satisfied: packaging>=20.9 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.19.4->diffusers) (23.2)
Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.10/dist-packages (from importlib-metadata->diffusers) (3.17.0)
Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->diffusers) (3.3.2)
Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->diffusers) (3.6)
Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->diffusers) (2.0.7)
Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->diffusers) (2023.11.17)

!pip install diffusers
!pip install accelerate

一度に9枚生成してみた

import torch
from diffusers import AmusedPipeline

# Initialize the pipeline
pipe = AmusedPipeline.from_pretrained(
    "amused/amused-512", variant="fp16", torch_dtype=torch.float16
)
pipe.vqvae.to(torch.float32)  # Fix for vqvae producing NaNs in fp16
pipe = pipe.to("cuda")

# Define your prompt
num_images = 9
prompt = "miku in dancing"
for i in range(num_images):
    image = pipe(prompt, generator=torch.manual_seed(i)).images[0]
    image.save(f"miku_in_dancing{i}.png")

ちょっと結果は酷かったので公開するのはやめておきます

image2imageも試してみた。プロンプトは「apple watercolor」。

import torch
from diffusers import AmusedImg2ImgPipeline
from diffusers.utils import load_image

pipe = AmusedImg2ImgPipeline.from_pretrained(
    "amused/amused-512", variant="fp16", torch_dtype=torch.float16
)
pipe.vqvae.to(torch.float32)  # vqvae is producing nans in fp16
pipe = pipe.to("cuda")

prompt = "apple watercolor"
input_image = (
    load_image(
        "miku_39.png"
    )
    .resize((512, 512))
    .convert("RGB")
)

image = pipe(prompt, input_image, strength=0.7, generator=torch.manual_seed(39)).images[0]
image.save("image2image_512.png")