在免费版 Google Colab 上用 🧨 diffusers 运行 IF
摘要：我们展示了如何使用 🧨 diffusers 在免费版的 Google Colab 上运行最强大的开源文本到图像模型之一 IF。
您也可以直接在 Hugging Face Space 中探索该模型的功能。
图片压缩自官方 IF GitHub 仓库。
简介
IF 是一种基于像素的文本到图像生成模型，由 DeepFloyd 于 2023 年 4 月下旬发布。该模型架构深受 Google 闭源模型 Imagen 的启发。
与现有的文本到图像模型（如 Stable Diffusion）相比，IF 有两个明显的优势
- 该模型直接在“像素空间”（即未压缩的图像）上操作，而不是像 Stable Diffusion 那样在潜空间中运行去噪过程。
- 该模型在 T5-XXL 的输出上进行训练，这是一个比 Stable Diffusion 用作文本编码器的 CLIP 更强大的文本编码器。
因此，IF 更擅长生成具有高频细节（例如人脸和手）的图像，并且是第一个能够可靠地生成带文本图像的开源图像生成模型。
在像素空间操作和使用更强大的文本编码器的缺点是 IF 的参数数量明显更多。T5、IF 的文本到图像 UNet 和 IF 的上采样器 UNet 分别具有 45 亿、43 亿和 12 亿个参数。相比之下，Stable Diffusion 2.1 的文本编码器和 UNet 分别只有 4 亿和 9 亿个参数。
尽管如此，如果针对低内存使用对模型进行优化，仍然可以在消费级硬件上运行 IF。在这篇博文中，我们将向您展示如何使用 🧨 diffusers 来实现这一点。
在第一部分中，我们解释了如何使用 IF 进行文本到图像的生成，在第二和第三部分中，我们介绍了 IF 的图像变体和图像修复功能。
💡 注意：我们在这里通过牺牲速度来换取内存的提升，以便能够在免费版的 Google Colab 中运行 IF。如果您有 A100 等高端 GPU，我们建议将所有模型组件都保留在 GPU 上以获得最大速度，就像官方 IF 演示中那样。
💡 注意：一些较大的图像已被压缩，以便在博客格式中加载得更快。使用官方模型时，它们的质量应该会更好！
让我们开始吧 🚀！
IF 的文本生成能力
目录
接受许可证
在使用 IF 之前，您需要接受其使用条件。为此：
- 确保拥有一个 Hugging Face 帐户并已登录
- 在 DeepFloyd/IF-I-XL-v1.0 的模型卡上接受许可证。在第一阶段模型卡上接受许可证将自动为其他 IF 模型接受。
- 确保在本地登录。安装
huggingface_hub
- 确保在本地登录。安装
pip install huggingface_hub --upgrade
在 Python shell 中运行登录函数
from huggingface_hub import login
login()
并输入您的 Hugging Face Hub 访问令牌。
优化 IF 以在内存受限的硬件上运行
最先进的机器学习不应该只掌握在少数精英手中。机器学习的民主化意味着让模型能够在不仅仅是最先进的硬件上运行。
深度学习社区已经创造了世界一流的工具，用于在消费级硬件上运行资源密集型模型
- 🤗 accelerate 提供了处理大型模型的实用程序。
- bitsandbytes 使所有 PyTorch 模型都能进行8位量化。
- 🤗 safetensors 不仅确保执行安全的代码，还显著加快了大型模型的加载时间。
Diffusers 无缝集成了上述库，以便在优化大型模型时提供简单的 API。
免费版的 Google Colab 在 CPU RAM（13 GB RAM）和 GPU VRAM（T4 为 15 GB RAM）方面都受到限制，这使得运行整个超过 10B 参数的 IF 模型具有挑战性！
让我们来计算一下 IF 模型组件在全 float32 精度下的大小
- T5-XXL 文本编码器：20GB
- 阶段 1 UNet：17.2 GB
- 阶段 2 超分辨率 UNet：2.5 GB
- 阶段 3 超分辨率模型：3.4 GB
我们无法在 float32 精度下运行该模型，因为 T5 和阶段 1 UNet 的权重都比可用的 CPU RAM 大。
在 float16 精度下，T5、阶段1 和阶段2 UNet 的组件大小分别为 11GB、8.6GB 和 1.25GB，这对于 GPU 来说是可行的，但我们在加载 T5 时仍然会遇到 CPU 内存溢出错误（部分 CPU 被其他进程占用）。
因此，我们通过使用 bitsandbytes 8位量化进一步降低 T5 的精度，这使得 T5 检查点可以仅用 8 GB 的空间来保存。
既然每个组件都能单独装入 CPU 和 GPU 内存，我们需要确保在需要时，组件能独占所有的 CPU 和 GPU 内存。
Diffusers 支持模块化加载单个组件，即我们可以加载文本编码器而不加载 UNet。这种模块化加载将确保我们只在管道的特定步骤中加载我们需要的组件，以避免耗尽可用的 CPU RAM 和 GPU VRAM。
让我们试一试 🚀
可用资源
免费版的 Google Colab 提供了大约 13 GB 的 CPU RAM
!grep MemTotal /proc/meminfo
MemTotal: 13297192 kB
以及一个带有 15 GB VRAM 的 NVIDIA T4
!nvidia-smi
Sun Apr 23 23:14:19 2023
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.85.12 Driver Version: 525.85.12 CUDA Version: 12.0 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|===============================+======================+======================|
| 0 Tesla T4 Off | 00000000:00:04.0 Off | 0 |
| N/A 72C P0 32W / 70W | 1335MiB / 15360MiB | 0% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=============================================================================|
+-----------------------------------------------------------------------------+
安装依赖项
一些优化可能需要最新版本的依赖项。如果您遇到问题，请仔细检查并升级版本。
! 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
1. 文本到图像生成
我们将逐步介绍如何使用 Diffusers 通过 IF 进行文本到图像的生成。我们将简要解释 API 和优化，但更深入的解释可以在 Diffusers、Transformers、Accelerate 和 bitsandbytes 的官方文档中找到。
1.1 加载文本编码器
我们将使用 8 位量化加载 T5。Transformers 直接通过 load_in_8bit 标志支持 bitsandbytes。
标志 variant="8bit" 将下载预量化的权重。
我们还使用 device_map 标志，以允许 transformers 将模型层卸载到 CPU 或磁盘。Transformers 的大型模型支持任意设备映射，可用于将模型参数直接加载到可用设备上。传递 "auto" 将自动创建设备映射。有关更多信息，请参见 transformers 文档。
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"
)
1.2 创建文本嵌入
用于访问扩散模型的 Diffusers API 是 DiffusionPipeline 类及其子类。DiffusionPipeline 的每个实例都是一个完全自包含的方法和模型集合，用于运行扩散网络。我们可以通过将替代实例作为关键字参数传递给 from_pretrained 来覆盖它使用的模型。
在这种情况下，我们为 unet 参数传递 None，因此不会加载 UNet。这使我们能够在不将 UNet 加载到内存中的情况下运行扩散过程的文本嵌入部分。
from diffusers import DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained(
"DeepFloyd/IF-I-XL-v1.0",
text_encoder=text_encoder, # pass the previously instantiated 8bit text encoder
unet=None,
device_map="auto"
)
IF 还带有一个超分辨率管道。我们将保存提示嵌入，以便稍后可以直接将它们传递给超分辨率管道。这将允许超分辨率管道在没有文本编码器的情况下加载。
不仅仅是宇航员骑马，我们还给他们一个牌子吧！
让我们定义一个合适的提示
prompt = "a photograph of an astronaut riding a horse holding a sign that says Pixel's in space"
并通过 8 位量化的 T5 模型运行它
prompt_embeds, negative_embeds = pipe.encode_prompt(prompt)
1.3 释放内存
一旦创建了提示嵌入，我们就不再需要文本编码器了。但是，它仍然在 GPU 的内存中。我们需要移除它，以便加载 UNet。
释放 PyTorch 内存并非易事。我们必须垃圾回收指向 GPU 上实际分配内存的 Python 对象。
首先，使用 Python 关键字 del 删除所有引用已分配 GPU 内存的 Python 对象
del text_encoder
del pipe
删除 python 对象不足以释放 GPU 内存。垃圾回收才是真正释放 GPU 内存的时候。
此外，我们将调用 torch.cuda.empty_cache()。这个方法并非严格必要，因为缓存的 cuda 内存将立即可用于进一步的分配。清空缓存可以让我们在 Colab UI 中验证内存是否可用。
我们将使用一个辅助函数 flush() 来刷新内存。
import gc
import torch
def flush():
gc.collect()
torch.cuda.empty_cache()
然后运行它
flush()
1.4 阶段 1：主扩散过程
利用我们现在可用的 GPU 内存，我们可以重新加载只包含 UNet 的 DiffusionPipeline 来运行主扩散过程。
variant 和 torch_dtype 标志由 Diffusers 用于以 16 位浮点格式下载和加载权重。
pipe = DiffusionPipeline.from_pretrained(
"DeepFloyd/IF-I-XL-v1.0",
text_encoder=None,
variant="fp16",
torch_dtype=torch.float16,
device_map="auto"
)
通常，我们直接将文本提示传递给 DiffusionPipeline.__call__。然而，我们之前计算了文本嵌入，我们可以将其传递进去。
IF 还带有一个超分辨率扩散过程。设置 output_type="pt" 将返回原始的 PyTorch 张量而不是 PIL 图像。这样，我们可以将 PyTorch 张量保留在 GPU 上，并直接将它们传递给阶段 2 的超分辨率管道。
让我们定义一个随机生成器并运行阶段 1 扩散过程。
generator = torch.Generator().manual_seed(1)
image = pipe(
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_embeds,
output_type="pt",
generator=generator,
).images
让我们手动将原始张量转换为 PIL 图像，并先睹为快最终结果。阶段 1 的输出是一张 64x64 的图像。
from diffusers.utils import pt_to_pil
pil_image = pt_to_pil(image)
pipe.watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)
pil_image[0]
再次，我们移除 Python 指针并释放 CPU 和 GPU 内存
del pipe
flush()
1.5 阶段 2：超分辨率 64x64 到 256x256
IF 带有一个用于上采样的独立扩散过程。
我们用一个独立的管道运行每个扩散过程。
超分辨率管道可以在需要时加载文本编码器。然而，我们通常会从第一个 IF 管道中预先计算好文本嵌入。如果是这样，加载管道时不需要文本编码器。
创建管道
pipe = DiffusionPipeline.from_pretrained(
"DeepFloyd/IF-II-L-v1.0",
text_encoder=None, # no use of text encoder => memory savings!
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
我们再次可以检查中间结果。
pil_image = pt_to_pil(image)
pipe.watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)
pil_image[0]
再次，我们删除 Python 指针并释放内存
del pipe
flush()
1.6 阶段 3：超分辨率 256x256 到 1024x1024
IF 的第二个超分辨率模型是之前发布的 Stability AI 的 x4 上采样器。
让我们创建管道并使用 device_map="auto" 直接加载到 GPU 上。
pipe = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-x4-upscaler",
torch_dtype=torch.float16,
device_map="auto"
)
🧨 diffusers 使独立开发的扩散模型可以轻松组合，因为管道可以链接在一起。这里我们可以直接取前一个 PyTorch 张量输出并将其作为 image=image 传递给阶段 3 管道。
💡 注意：x4 上采样器不使用 T5，并且有自己的文本编码器。因此，我们不能使用之前创建的提示嵌入，而必须传递原始提示。
pil_image = pipe(prompt, generator=generator, image=image).images
与 IF 管道不同，Stable Diffusion x4 上采样器管道的输出默认不会添加 IF 水印。
我们可以手动应用水印。
from diffusers.pipelines.deepfloyd_if import IFWatermarker
watermarker = IFWatermarker.from_pretrained("DeepFloyd/IF-I-XL-v1.0", subfolder="watermarker")
watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)
查看输出图像
pil_image[0]
瞧！在免费版的 Google Colab 中生成了一张漂亮的 1024x1024 图像。
我们已经展示了 🧨 diffusers 如何轻松分解和模块化加载资源密集型的扩散模型。
💡 注意：我们不建议在生产环境中使用上述设置。8位量化、手动释放模型权重和磁盘卸载都以牺牲时间（即推理速度）为代价来换取内存。如果扩散管道被重复使用，这一点尤其明显。在生产环境中，我们建议使用一块 40GB 的 A100，并将所有模型组件都保留在 GPU 上。请参阅官方 IF 演示。
2. 图像变体
相同的 IF 检查点也可用于文本引导的图像变体和图像修复。核心扩散过程与文本到图像生成相同，只是初始的噪声图像是根据要变化或修复的图像创建的。
要运行图像变体，请使用 IFImg2ImgPipeline.from_pretrained() 和 IFImg2ImgSuperResolution.from_pretrained() 加载相同的检查点。
用于内存优化的 API 都是相同的！
让我们释放上一节的内存。
del pipe
flush()
对于图像变体，我们从一张我们想要调整的初始图像开始。
在这一节中，我们将改编著名的“拍车顶”表情包。让我们从网上下载它。
import requests
url = "https://i.kym-cdn.com/entries/icons/original/000/026/561/car.jpg"
response = requests.get(url)
并将其加载到 PIL 图像中
from PIL import Image
from io import BytesIO
original_image = Image.open(BytesIO(response.content)).convert("RGB")
original_image = original_image.resize((768, 512))
original_image
图像变体管道既接受 PIL 图像也接受原始张量。有关预期输入的更深入文档，请查看此处的文档字符串。
2.1 文本编码器
图像变体由文本引导，因此我们可以定义一个提示并用 T5 的文本编码器对其进行编码。
我们再次将文本编码器加载为 8 位精度。
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"
)
对于图像变体，我们使用 IFImg2ImgPipeline 加载检查点。当使用 DiffusionPipeline.from_pretrained(...) 时，检查点会加载到它们的默认管道中。IF 的默认管道是文本到图像的 IFPipeline。当使用非默认管道加载检查点时，必须显式指定管道。
from diffusers import IFImg2ImgPipeline
pipe = IFImg2ImgPipeline.from_pretrained(
"DeepFloyd/IF-I-XL-v1.0",
text_encoder=text_encoder,
unet=None,
device_map="auto"
)
让我们把我们的销售员变成一个动漫角色。
prompt = "anime style"
和之前一样，我们用 T5 创建文本嵌入
prompt_embeds, negative_embeds = pipe.encode_prompt(prompt)
并释放 GPU 和 CPU 内存。
首先，移除 Python 指针
del text_encoder
del pipe
然后释放内存
flush()
2.2 阶段 1：主扩散过程
接下来，我们只将阶段 1 的 UNet 权重加载到管道对象中，就像我们在上一节中所做的那样。
pipe = IFImg2ImgPipeline.from_pretrained(
"DeepFloyd/IF-I-XL-v1.0",
text_encoder=None,
variant="fp16",
torch_dtype=torch.float16,
device_map="auto"
)
图像变体管道需要原始图像和提示嵌入。
我们可以选择使用 strength 参数来配置变体的程度。strength 直接控制添加的噪声量。更高的强度意味着更多的噪声，也意味着更多的变体。
generator = torch.Generator().manual_seed(0)
image = pipe(
image=original_image,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_embeds,
output_type="pt",
generator=generator,
).images
让我们再次检查一下中间的 64x64 图像。
pil_image = pt_to_pil(image)
pipe.watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)
pil_image[0]
看起来不错！我们可以释放内存并再次上采样图像了。
del pipe
flush()
2.3 阶段 2：超分辨率
对于超分辨率，使用 IFImg2ImgSuperResolutionPipeline 和与之前相同的检查点加载。
from diffusers import IFImg2ImgSuperResolutionPipeline
pipe = IFImg2ImgSuperResolutionPipeline.from_pretrained(
"DeepFloyd/IF-II-L-v1.0",
text_encoder=None,
variant="fp16",
torch_dtype=torch.float16,
device_map="auto"
)
💡 注意：图像变体超分辨率管道需要生成的图像以及原始图像。
你也可以在这张图片上使用 Stable Diffusion x4 上采样器。欢迎使用 1.6 节中的代码片段进行尝试。
image = pipe(
image=image,
original_image=original_image,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_embeds,
generator=generator,
).images[0]
image
不错！让我们释放内存，看看最后的图像修复管道。
del pipe
flush()
3. 图像修复
IF 修复管道与图像变体相同，只是只对图像的选定区域进行去噪。
我们用图像掩码指定要修复的区域。
让我们展示一下 IF 惊人的“文字生成”能力。我们可以用不同的口号替换这个牌子上的文字。
首先让我们下载图像
import requests
url = "https://i.imgflip.com/5j6x75.jpg"
response = requests.get(url)
并将其转换为 PIL 图像
from PIL import Image
from io import BytesIO
original_image = Image.open(BytesIO(response.content)).convert("RGB")
original_image = original_image.resize((512, 768))
original_image
我们将遮盖标志，以便替换其文本。
为方便起见，我们已预先生成掩码并将其加载到 HF 数据集中。
让我们下载它。
from huggingface_hub import hf_hub_download
mask_image = hf_hub_download("diffusers/docs-images", repo_type="dataset", filename="if/sign_man_mask.png")
mask_image = Image.open(mask_image)
mask_image
💡 注意：您可以通过手动创建灰度图像来自己创建掩码。
from PIL import Image
import numpy as np
height = 64
width = 64
example_mask = np.zeros((height, width), dtype=np.int8)
# Set masked pixels to 255
example_mask[20:30, 30:40] = 255
# Make sure to create the image in mode 'L'
# meaning single channel grayscale
example_mask = Image.fromarray(example_mask, mode='L')
example_mask
现在我们可以开始修复了 🎨🖌
3.1. 文本编码器
再次，我们首先加载文本编码器
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"
)
这一次，我们使用文本编码器权重初始化 IFInpaintingPipeline 修复管道。
from diffusers import IFInpaintingPipeline
pipe = IFInpaintingPipeline.from_pretrained(
"DeepFloyd/IF-I-XL-v1.0",
text_encoder=text_encoder,
unet=None,
device_map="auto"
)
好吧，让我们让这个人转而为更多的层做广告。
prompt = 'the text, "just stack more layers"'
定义了提示后，我们可以创建提示嵌入
prompt_embeds, negative_embeds = pipe.encode_prompt(prompt)
就像之前一样，我们释放内存
del text_encoder
del pipe
flush()
3.2 阶段 1：主扩散过程
和之前一样，我们现在只加载 UNet 来加载阶段 1 管道。
pipe = IFInpaintingPipeline.from_pretrained(
"DeepFloyd/IF-I-XL-v1.0",
text_encoder=None,
variant="fp16",
torch_dtype=torch.float16,
device_map="auto"
)
现在，我们需要传递输入图像、掩码图像和提示嵌入。
image = pipe(
image=original_image,
mask_image=mask_image,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_embeds,
output_type="pt",
generator=generator,
).images
让我们看一下中间输出。
pil_image = pt_to_pil(image)
pipe.watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)
pil_image[0]
看起来不错！文字非常连贯！
让我们释放内存，以便上采样图像
del pipe
flush()
3.3 阶段 2：超分辨率
对于超分辨率，使用 IFInpaintingSuperResolutionPipeline 加载检查点。
from diffusers import IFInpaintingSuperResolutionPipeline
pipe = IFInpaintingSuperResolutionPipeline.from_pretrained(
"DeepFloyd/IF-II-L-v1.0",
text_encoder=None,
variant="fp16",
torch_dtype=torch.float16,
device_map="auto"
)
修复超分辨率管道需要生成的图像、原始图像、掩码图像和提示嵌入。
让我们进行最后一次去噪运行。
image = pipe(
image=image,
original_image=original_image,
mask_image=mask_image,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_embeds,
generator=generator,
).images[0]
image
太棒了，模型生成的文本没有一个拼写错误！
结论
IF 在 32 位浮点精度下总共使用 40 GB 的权重。我们展示了如何仅使用开源模型和库，在免费版的 Google Colab 实例上运行 IF。
ML 生态系统从开放工具和开放模型的共享中获益匪浅。仅此笔记本就使用了来自 DeepFloyd、StabilityAI 和 Google 的模型。所使用的库——Diffusers、Transformers、Accelerate 和 bitsandbytes——都受益于来自不同组织的无数贡献者。
非常感谢 DeepFloyd 团队创建并开源 IF，并为推动优秀机器学习的民主化做出贡献 🤗。
