TL;DR：我们展示了如何在免费的 Google Colab 上使用🧨扩散器运行最强大的开源文本到图像模型之一IF。
您还可以直接在 Hugging Face Space 中探索该模型的功能。
图片来自官方 IF GitHub 仓库。
简介
IF 是一种基于像素的文本到图像生成模型，由 DeepFloyd 在2023年4月底发布。该模型的架构受到 Google 的闭源 Imagen 的强烈启发。
与 Stable Diffusion 等现有的文本到图像模型相比，IF 具有两个明显优势：
- 该模型直接在“像素空间”（即未压缩图像上）操作，而不是在潜在空间中运行去噪过程，例如 Stable Diffusion。
- 该模型是在比 CLIP 更强大的文本编码器 T5-XXL 的输出上进行训练的，而 Stable Diffusion 使用 CLIP 作为文本编码器。
因此，IF 在生成高频细节的图像（例如人脸和手）方面表现更好，并且是第一个能够可靠生成带有文本的开源图像生成模型。
在像素空间中操作并使用更强大的文本编码器的缺点是 IF 具有数量显着更多的参数。T5、IF 的文本到图像 UNet 和 IF 的放大器 UNet 分别具有 45 亿、43 亿和 12 亿个参数。与 Stable Diffusion 2.1 的文本编码器和 UNet 分别只有 4 亿和 9 亿个参数相比。
尽管如此，如果优化模型以降低内存使用量，也有可能在普通硬件上运行 IF。我们将在本博文中展示如何使用🧨扩散器实现这一点。
在第一部分中，我们将解释如何使用 IF 进行文本到图像生成，在第二和第三部分中，我们将介绍 IF 的图像变化和图像修复能力。
💡注意：我们通过在速度上获得的优势来换取内存上的收益，以便在免费的 Google Colab 上运行 IF。如果您拥有像 A100 这样的高端 GPU，我们建议将所有模型组件保留在 GPU 上以获取最大的速度，就像在官方 IF 演示中所做的那样。
💡注意：为了在博客格式中更快地加载，一些较大的图像已经进行了压缩。在使用官方模型时，它们的质量应该更好！
让我们开始吧！🚀
IF 的文本生成能力
目录
- 接受许可证
- 优化 IF 在内存受限硬件上的运行
- 可用资源
- 安装依赖项
- 文本到图像生成
- 图像变化
- 修复图像
接受许可证
在使用 IF 之前，您需要接受其使用条件。具体步骤如下：
-
- 确保拥有 Hugging Face 帐户并已登录
-
- 接受 DeepFloyd/IF-I-XL-v1.0 模型卡片上的许可证。在阶段 I 模型卡片上接受许可证将自动适用于其他 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 内存受限（13 GB RAM），GPU 内存受限（T4 15 GB RAM），这使得运行整个大于10B IF 模型具有挑战性！
让我们详细说明 IF 模型组件的大小（以完全 float32 精度为准）：
- T5-XXL 文本编码器：20GB
- Stage 1 UNet：17.2 GB
- Stage 2 超分辨率 UNet：2.5 GB
- Stage 3 超分辨率模型：3.4 GB
由于 T5 和 Stage 1 UNet 的权重都大于可用的 CPU 内存，我们无法使用 float32 运行模型。
在 float16 中，T5、Stage1 和 Stage2 UNet 的组件大小分别为 11GB、8.6GB 和 1.25GB，这对于 GPU 来说是可行的，但在加载 T5 时仍然会遇到 CPU 内存溢出错误（一些 CPU 被其他进程占用）。
因此，我们通过使用 bitsandbytes 的 8 位量化来进一步降低 T5 的精度，这样可以将 T5 检查点保存为仅有 8 GB。
现在每个组件都适合于 CPU 和 GPU 内存，我们需要确保在需要时组件具有所有 CPU 和 GPU 内存。
Diffusers 支持模块化加载单个组件，即我们可以加载文本编码器而不加载 UNet。这种模块化加载将确保我们在管道的给定步骤中仅加载所需的组件，以避免耗尽可用的 CPU 内存和 GPU 内存。
让我们试一试 🚀
可用资源
免费版的 Google Colab 具有约 13 GB 的 CPU 内存：
!grep MemTotal /proc/meminfo
MemTotal: 13297192 kB并且具有 NVIDIA T4 显卡，15 GB 的 VRAM：
!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 -q1. 文本到图像生成
我们将逐步介绍使用Diffusers的文本到图像生成。我们将简要介绍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, # 将之前实例化的8位文本编码器传递给它
unet=None,
device_map="auto"
)IF还附带了超分辨率流水线。我们将保存提示嵌入，以便稍后可以直接将它们传递给超分辨率流水线。这将允许加载超分辨率流水线而不需要文本编码器。
不只是一名宇航员骑着马，让我们再给他一个牌子！
让我们定义一个适当的提示：
prompt = "一张宇航员骑马并手持一块写着像素在太空中的牌子的照片"并将其通过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用户界面中验证内存是否可用。
我们将使用一个辅助函数flush()来刷新内存。
import gc
import torch
def flush():
gc.collect()
torch.cuda.empty_cache()并运行它
flush()1.4 阶段1：主要扩散过程
有了我们现在可用的GPU内存，我们可以重新加载DiffusionPipeline，只使用UNet来运行主要的扩散过程。
variant和torch_dtype标志被扩散器用于以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阶段的输出是一个64×64的图像。
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阶段：64×64到256×256的超分辨率
IF还提供了一个单独的扩散过程来进行放大。
我们使用单独的流水线运行每个扩散过程。
如果需要，超分辨率流水线可以加载文本编码器。但是，我们通常会从第一个IF流水线中预先计算文本嵌入。如果是这样，加载流水线时不要加载文本编码器。
创建流水线
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同样，我们可以检查中间结果。
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阶段：256×256到1024×1024的超分辨率
IF的第二个超分辨率模型是之前发布的Stability AI的x4 Upscaler。
让我们创建流水线，并使用device_map="auto"直接加载到GPU上。
pipe = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-x4-upscaler",
torch_dtype=torch.float16,
device_map="auto"
)🧨 diffusers使得独立开发的扩散模型可以轻松地组合在一起，因为流水线可以链接在一起。在这里，我们只需将之前的PyTorch张量输出传递给第3阶段的流水线，如image=image。
💡 注意：x4 Upscaler 不使用 T5，而是使用自己的文本编码器。因此，我们不能使用先前创建的提示嵌入，而必须传递原始提示。
pil_image = pipe(prompt, generator=generator, image=image).images与 IF 管道不同，稳定扩散 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]
Et voila！一个漂亮的 1024×1024 图片在免费的 Google Colab 中。
我们已经展示了 🧨 diffusers 如何轻松分解和模块化加载资源密集型扩散模型。
💡 注意：我们不建议在生产环境中使用上述设置。8 位量化、模型权重的手动释放和磁盘卸载都是以时间（即推理速度）为代价的内存折衷。如果扩散管道被重复使用，这一点尤其明显。在生产环境中，我们建议使用一个带有所有模型组件留在 GPU 上的 40GB A100。请参阅 官方 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 = "动漫风格"与之前一样，我们使用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让我们再次检查中间的 64×64。
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)
# 将遮罩的像素设置为255
example_mask[20:30, 30:40] = 255
# 确保以'L'模式创建图像，即单通道灰度图像
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
<p太好了，这个模型生成的文本没有出现任何拼写错误！
总结
IF 使用 32 位浮点精度总共使用了 40 GB 的权重。我们展示了如何仅使用开源模型和库在免费的 Google Colab 实例上运行 IF。
机器学习生态系统深受开放工具和开源模型的共享所益。这个笔记本单独使用了来自 DeepFloyd、StabilityAI 和 Google 的模型。使用的库——Diffusers、Transformers、Accelerate 和 bitsandbytes——都受益于不同组织的无数贡献者。
非常感谢 DeepFloyd 团队创建和开源 IF，以及为机器学习的民主化做出的贡献 🤗。