Wuerstchen

🚀 Würstchen - 概要

Würstchenは拡散モデルで、そのテキスト条件付きモデルは画像の高度に圧縮された潜在空間で動作します。なぜこれが重要なのでしょうか？データを圧縮することで、トレーニングと推論の両方の計算コストを大幅に削減することができます。1024x1024の画像でトレーニングするのは、32x32でトレーニングするよりもはるかにコストがかかります。通常、他の研究では4x - 8xの空間圧縮といった比較的小さな圧縮を利用しています。Würstchenはこれを極限まで追求し、独自の設計により42xの空間圧縮を達成しました。これは、一般的な方法では16xの空間圧縮を超えると詳細な画像を忠実に再構築できないため、これまでに例がないものです。Würstchenは2段階の圧縮、つまりステージAとステージBを採用しています。ステージAはVQGANで、ステージBは拡散オートエンコーダーです（詳細は論文を参照）。第3のモデルであるステージCは、この高度に圧縮された潜在空間で学習されます。このトレーニングに必要な計算量は、現在の最高性能のモデルに比べてごくわずかであり、推論も安価かつ高速に行えます。

✨ 主な機能

Würstchen - デコーダー

デコーダーは「ステージA」と「ステージB」のことを指します。デコーダーは、Prior（ステージC）によって生成された、または実際の画像から抽出された画像埋め込みを受け取り、それらの潜在表現をピクセル空間に戻します。具体的には、ステージBがまず画像埋め込みをVQGAN空間にデコードし、ステージA（VQGAN）が潜在表現をピクセル空間にデコードします。これにより、42倍の空間圧縮が達成されます。

注意: 再構築は損失を伴い、画像の情報が失われます。現在のステージBでは、再構築において細部が欠けることが多く、特に顔や手などを見るときに人間には顕著に感じられます。将来的にはこれらの再構築をさらに改善するために取り組んでいます！

画像サイズ

Würstchenは1024x1024から1536x1536の画像解像度でトレーニングされました。1024x2048のような解像度でも良好な出力が得られることがあります。ぜひ試してみてください。また、Prior（ステージC）は新しい解像度に非常に速く適応することがわかっています。したがって、2048x2048での微調整は計算コストが低くて済みます。

📦 インストール

このパイプラインは、事前学習モデル warp-ai/wuerstchen-prior と一緒に実行する必要があります。

💻 使用例

基本的な使用法

import torch
from diffusers import AutoPipelineForText2Image

device = "cuda"
dtype = torch.float16

pipeline =  AutoPipelineForText2Image.from_pretrained(
    "warp-diffusion/wuerstchen", torch_dtype=dtype
).to(device)

caption = "Anthropomorphic cat dressed as a fire fighter"

output = pipeline(
    prompt=caption,
    height=1024,
    width=1024,
    prior_guidance_scale=4.0,
    decoder_guidance_scale=0.0,
).images

画像サンプリング時間

下の図は、異なるバッチサイズ (num_images_per_prompt) におけるWürstchenの推論時間（A100上）を、Stable Diffusion XL（リファイナーなし）と比較したものです。左の図は推論時間（torch > 2.0を使用）を示しており、右の図は両方のパイプラインに事前に torch.compile を適用した場合を示しています。

📚 詳細ドキュメント

モデルの詳細

@inproceedings{
      pernias2024wrstchen,
            title={W"urstchen: An Efficient Architecture for Large-Scale Text-to-Image Diffusion Models},
            author={Pablo Pernias and Dominic Rampas and Mats Leon Richter and Christopher Pal and Marc Aubreville},
            booktitle={The Twelfth International Conference on Learning Representations},
            year={2024},
            url={https://openreview.net/forum?id=gU58d5QeGv}
      }

環境への影響

Würstchen v2 推定排出量 この情報に基づいて、Lacoste et al. (2019) で提示された Machine Learning Impact calculator を使用して、以下のCO2排出量を推定します。ハードウェア、実行時間、クラウドプロバイダー、およびコンピュートリージョンを利用して、炭素排出量を推定しました。

• ハードウェアタイプ: A100 PCIe 40GB
• 使用時間: 24602時間
• クラウドプロバイダー: AWS
• コンピュートリージョン: US-east
• 排出された炭素量 (電力消費量 x 時間 x 電力網の所在地に基づく炭素排出量): 2275.68 kg CO2 eq.

📄 ライセンス

このモデルはMITライセンスの下で提供されています。