在免费版 Google Colab 上使用 Diffusers 运行 DeepFloyd IF

TL;DR

DeepFloyd 的 IF 模型可以在免费版 Google Colab 上运行，只需将 T5‑XXL 文本编码器量化为 8 位，使用 🤗 Diffusers 模块化加载模型组件，并将权重卸载到 CPU 或磁盘，即使笔记本的 13 GB RAM 和 15 GB GPU 限制，也能实现完整的文本到图像、图像变体和修复（inpainting）管道。

介绍 – 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 RAM 都无法满足。

为内存受限硬件进行优化

关键工具

  • 🤗 Accelerate – 用于大模型设备放置的工具。
  • bitsandbytes – 用于 PyTorch 模型的 8 位量化。
  • 🤗 safetensors – 快速、安全的检查点加载。
  • 🤗 Diffusers – 将上述功能集成的高级扩散管道。

节省内存的策略

  1. 将 T5‑XXL 量化为 8 位 – 将编码器检查点从约 20 GB（fp32）降低到约 8 GB（8 位 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. 为变体管道加载 Stage 1 UNet（unet=None），传入原始图像并使用 strength 控制添加的噪声量。
  3. 使用 IFImg2ImgSuperResolutionPipeline（或 Stable Diffusion 放大器）进行放大。
  4. 对于修复，还需提供二值掩码图像；同样的模块化加载和内存释放模式适用。

这三条管道共享相同的内存优化技巧：8 位文本编码器、fp16 UNet、自动设备映射以及显式垃圾回收。

实际考虑因素与性能权衡

  • 速度 vs. 内存 – 8 位量化以及反复加载/卸载会增加推理延迟。生产环境下，建议使用 ≥40 GB VRAM（如 A100）的 GPU，将所有组件保持在设备上以获得最大吞吐量。
  • 质量 – 官方 IF 演示（托管于 Hugging Face Spaces）使用全精度模型，能够提供略高的保真度，尤其在大图像时更为明显。
  • 许可证 – 用户必须在模型卡页面接受 DeepFloyd IF 许可证后才能加载任何检查点。

结论 – 让大型扩散模型大众化

通过将开源检查点（DeepFloyd IF、StabilityAI 放大器）与社区驱动的库（Diffusers、Transformers、Accelerate、bitsandbytes）相结合，Hugging Face 证明了 >10 B 参数的扩散模型可以在免费 Google Colab 实例上运行。这展示了模块化管道和量化技术的力量，能够让最前沿的生成式 AI 更加普及。

Sources