在免費版 Google Colab 上用 🧨 diffusers 執行 IF

釋出於 2023 年 4 月 26 日

在 GitHub 上更新

贊

訪客

訪客

訪客

Apolinário from multimodal AI art

multimodalart

摘要：我們展示瞭如何使用 🧨 diffusers 在免費版的 Google Colab 上執行最強大的開源文字到影像模型之一 IF。

您也可以直接在 Hugging Face Space 中探索該模型的功能。

if-collage
圖片壓縮自官方 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 之前，您需要接受其使用條件。為此：

1. 確保擁有一個 Hugging Face 帳戶並已登入
1. 在 DeepFloyd/IF-I-XL-v1.0 的模型卡上接受許可證。在第一階段模型卡上接受許可證將自動為其他 IF 模型接受。
1. 確保在本地登入。安裝 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 精度下的大小

我們無法在 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，併為推動優秀機器學習的民主化做出貢獻 🤗。

更多部落格文章

使用 Diffusers 和 PEFT 為 Flux 實現快速 LoRA 推理

作者 2025 年 7 月 23 日 • 42

全世界的 LoRA 訓練指令碼，聯合起來！

作者 2024 年 1 月 2 日 • 69

社群

透過拖放到文字輸入框、貼上或點選此處上傳圖片、音訊和影片。

點選或貼上此處以上傳圖片

· 註冊或登入發表評論

贊