Diffusers を使用した無料ティアの Google Colab で DeepFloyd IF を実行する

TL;DR

DeepFloyd の IF モデルは、T5‑XXL テキストエンコーダを 8 ビットに量子化し、🤗 Diffusers を使ってモデルコンポーネントをモジュラーにロードし、重みを CPU またはディスクにオフロードすることで、無料ティアの Google Colab でも実行できます。これにより、ノートブックの 13 GB RAM と 15 GB GPU の制限があるにもかかわらず、テキスト‑ツー‑イメージ、画像バリエーション、インペインティングの全パイプラインが利用可能になります。

はじめに – IF の機能と制約

DeepFloyd は 2023 年 4 月下旬に、Google の Imagen に触発されたピクセル空間テキスト‑ツー‑イメージ拡散モデルとして IF をリリースしました。Stable Diffusion と比較すると、IF は次の点が異なります：

  • ・圧縮されていない画像上で直接動作し、顔や手などの高周波ディテールを保持します。
  • ・CLIP の代わりに強力な T5‑XXL エンコーダを使用し、テキストの忠実度を向上させ、読みやすいテキストを確実に生成できる初のオープンソースモデルとなります。

トレードオフはサイズです。T5‑XXL（45 億パラメータ）、IF‑I UNet（43 億パラメータ）、IF‑II アップスケーラ UNet（12 億パラメータ）を合わせると 100 億パラメータを超え、Stable Diffusion 2.1 の約 13 億総数と比べてはるかに大きくなります。フルモデルを float32 で実行すると 40 GB 超の GPU メモリが必要ですが、無料の Colab T4（15 GB VRAM）や 13 GB の CPU RAM では利用できません。

メモリ制約ハードウェア向けの最適化

主なツール

  • 🤗 Accelerate – 大規模モデルのデバイス配置用ユーティリティ。
  • bitsandbytes – PyTorch モデルの 8 ビット量子化。
  • 🤗 safetensors – 高速で安全なチェックポイントロード。
  • 🤗 Diffusers – 上記を統合したハイレベルな拡散パイプライン。

メモリ節約戦略

  1. T5‑XXL を 8 ビットに量子化 – エンコーダのチェックポイントを約 20 GB（fp32）から約 8 GB（8‑bit safetensors）に削減します。
  2. コンポーネントを遅延ロード – Diffusers は、テキストエンコーダだけ、または UNet だけをその時点でロードでき、同時のメモリピークを防ぎます。
  3. UNet の重みを fp16 で使用 – Stage 1 と Stage 2 の UNet は半精度でロードすれば GPU メモリに収まります。
  4. PyTorch オブジェクトを明示的に解放 – モデルオブジェクトを del し、各ステージ後に gc.collect() と torch.cuda.empty_cache() を実行します。

ステップバイステップ：テキスト‑ツー‑イメージ生成

1. 最新の依存関係をインストール

pip install --upgrade \
  diffusers~=0.16 \
  transformers~=4.28 \
  safetensors~=0.3 \
  sentencepiece~=0.1 \
  accelerate~=0.18 \
  bitsandbytes~=0.38 \
  torch~=2.0 -q

2. 8 ビット T5 エンコーダをロード

from transformers import T5EncoderModel
text_encoder = T5EncoderModel.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0",
    subfolder="text_encoder",
    device_map="auto",
    load_in_8bit=True,
    variant="8bit",
)

3. UNet をロードせずにプロンプト埋め込みを作成

from diffusers import DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0",
    text_encoder=text_encoder,
    unet=None,
    device_map="auto",
)
prompt = "a photograph of an astronaut riding a horse holding a sign that says Pixel's in space"
prompt_embeds, negative_embeds = pipe.encode_prompt(prompt)

4. エンコーダを解放して UNet 用のメモリを確保

import gc, torch

def flush():
    gc.collect()
    torch.cuda.empty_cache()

del text_encoder, pipe
flush()

5. Stage 1 拡散（64×64） – UNet のみをロード

pipe = DiffusionPipeline.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0",
    text_encoder=None,
    variant="fp16",
    torch_dtype=torch.float16,
    device_map="auto",
)
generator = torch.Generator().manual_seed(1)
image = pipe(
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=negative_embeds,
    output_type="pt",
    generator=generator,
).images

結果は 64 × 64 のテンソルで、pt_to_pil で可視化できます。

6. Stage 2 超解像（64→256） – IF‑II パイプライン

pipe = DiffusionPipeline.from_pretrained(
    "DeepFloyd/IF-II-L-v1.0",
    text_encoder=None,
    variant="fp16",
    torch_dtype=torch.float16,
    device_map="auto",
)
image = pipe(
    image=image,
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=negative_embeds,
    output_type="pt",
    generator=generator,
).images

7. Stage 3 超解像（256→1024） – StabilityAI x4 アップスケーラ

pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-x4-upscaler",
    torch_dtype=torch.float16,
    device_map="auto",
)
final_image = pipe(prompt, image=image, generator=generator).images[0]

必要に応じて IF のウォーターマークを手動で適用します：

from diffusers.pipelines.deepfloyd_if import IFWatermarker
watermarker = IFWatermarker.from_pretrained("DeepFloyd/IF-I-XL-v1.0", subfolder="watermarker")
watermarker.apply_watermark(final_image, pipe.unet.config.sample_size)

このパイプラインは、無料の Colab セッション内だけで 1024 × 1024 の画像を生成します。

画像バリエーションとインペインティング – 同じチェックポイントを再利用

IF のチェックポイントは IFImg2ImgPipeline（バリエーション）と IFInpaintingPipeline もサポートしています。ワークフローはテキスト‑ツー‑イメージの手順と同様です：

  1. 8 ビット T5 エンコーダをロードし、バリエーション用プロンプトをエンコードした後、エンコーダを解放します。
  2. バリエーションパイプライン用に unet=None で Stage 1 UNet をロードし、元画像と strength を渡してノイズ追加量を制御します。
  3. IFImg2ImgSuperResolutionPipeline（または Stable Diffusion アップスケーラ）でアップスケールします。
  4. インペインティングの場合は、バイナリマスク画像を追加で提供します。同じモジュラーなロードとメモリ解放のパターンが適用されます。

3 つのパイプラインはすべて同じメモリ最適化テクニックを共有します：8 ビットテキストエンコーダ、fp16 UNet、device‑map の自動配置、明示的なガーベジコレクション。

実用上の考慮点とパフォーマンスのトレードオフ

  • 速度 vs. メモリ – 8 ビット量子化と繰り返しのロード/アンロードにより推論レイテンシが増加します。実運用では、40 GB 以上の VRAM（例：A100）を持つ GPU がすべてのコンポーネントをデバイス上に保持し、最大スループットを実現すべきです。
  • 品質 – 公式の IF デモ（Hugging Face Spaces にホスト）ではフル精度モデルを実行しており、特に大きな画像で若干高い忠実度を示します。
  • ライセンス – ユーザーはチェックポイントをロードする前に、モデルカード上の DeepFloyd IF ライセンスに同意する必要があります。

結論 – 大規模拡散モデルの民主化

オープンソースのチェックポイント（DeepFloyd IF、StabilityAI アップスケーラ）とコミュニティ主導のライブラリ（Diffusers、Transformers、Accelerate、bitsandbytes）を組み合わせることで、Hugging Face は 100 億パラメータを超える拡散モデルを無料の Google Colab インスタンスで実行できることを示しました。これは、モジュラーパイプラインと量子化の力が最先端生成 AI へのアクセスを広げることを実証しています。

Sources