在免費等級的 Google Colab 上使用 Diffusers 執行 DeepFloyd IF

TL;DR

DeepFloyd 的 IF 模型可以在免費等級的 Google Colab 上執行，方法是將 T5‑XXL 文字編碼器量化為 8 位元，使用 🤗 Diffusers 模組化載入模型組件，並將權重卸載至 CPU 或磁碟，從而即使在筆記本僅有 13 GB RAM 與 15 GB GPU 限制的情況下，也能支援完整的文字轉圖像、圖像變換與修補管線。

介紹 – IF 的功能與限制

DeepFloyd 在 2023 年 4 月底發布了 IF，這是一個受 Google Imagen 啟發的像素空間文字到圖像擴散模型。與 Stable Diffusion 相比，IF：

  • 直接在未壓縮的圖像上運作，保留臉部與手部等高頻細節。
  • 使用強大的 T5‑XXL 編碼器取代 CLIP，提升文字忠實度，成為首個能可靠產生可讀文字的開源模型。

代價在於模型尺寸：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 記憶體中皆無法滿足。

為記憶體受限硬體進行最佳化

主要工具

  • 🤗 Accelerate – 用於大型模型設備配置的工具。
  • bitsandbytes – 用於 PyTorch 模型的 8 位元量化。
  • 🤗 safetensors – 快速且安全的檢查點載入。
  • 🤗 Diffusers – 整合上述功能的高階擴散管線。

記憶體節省策略

  1. 將 T5‑XXL 量化為 8 位元 – 將編碼器檢查點從約 20 GB（fp32）縮減至約 8 GB（8‑bit safetensors）。
  2. 延遲載入組件 – Diffusers 允許一次只載入文字編碼器或 UNet，避免同時的記憶體峰值。
  3. 對 UNet 權重使用 fp16 – 第 1 階段與第 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. 第 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. 第 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. 第 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 載入第 1 階段 UNet，傳入原始圖像與 strength 以控制噪聲添加量。
  3. 使用 IFImg2ImgSuperResolutionPipeline（或 Stable Diffusion 放大器）進行升級。
  4. 對於修補，額外提供二值遮罩圖像；同樣的模組化載入與記憶體釋放模式適用。

這三條管線皆共享相同的記憶體最佳化技巧：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 展示了即使在免費的 Google Colab 環境中，也能執行超過 100 億參數的擴散模型。此案例凸顯了模組化管線與量化技術的威力，讓最先進的生成式 AI 更加普及。

Sources