使用 LCM LoRA 4 步完成 SDXL 推理

2023-11-21 06:00:37

LCM 模型 通過將原始模型蒸餾為另一個需要更少步數 (4 到 8 步,而不是原來的 25 到 50 步) 的版本以減少用 Stable Diffusion (或 SDXL) 生成影象所需的步數。蒸餾是一種訓練過程,其主要思想是嘗試用一個新模型來複制源模型的輸出。蒸餾後的模型要麼尺寸更小 (如 DistilBERT 或最近釋出的 Distil-Whisper),要麼需要執行的步數更少 (本文即是這種情況)。一般來講,蒸餾是一個漫長且成本高昂的過程,需要大量資料、耐心以及一些 GPU 運算。

但以上所述皆為過往,今天我們翻新篇了!

今天,我們很高興地公開一種新方法,其可以從本質上加速 Stable Diffusion 和 SDXL,效果跟用 LCM 蒸餾過一樣!有了它,我們在 3090 上執行 任何 SDXL 模型,不需要 70 秒,也不需要 7 秒,僅需要約 1 秒就行了!在 Mac 上,我們也有 10 倍的加速!聽起來是不是很帶勁?那繼續往下讀吧!

目錄

方法概述

到底用了啥技巧?

在使用原始 LCM 蒸餾時,每個模型都需要單獨蒸餾。而 LCM LoRA 的核心思想是隻對少量介面卡 (即 LoRA 層) 進行訓練,而不用對完整模型進行訓練。推理時,可將生成的 LoRA 用於同一模型的任何微調版本,而無需對每個版本都進行蒸餾。如果你已經迫不及待地想試試這種方法的實際效果了,可以直接跳到 下一節 試一下推理程式碼。如果你想訓練自己的 LoRA,流程如下:

  1. 從 Hub 中選擇一個教師模型。如: 你可以使用 SDXL (base),或其任何微調版或 dreambooth 微調版,隨你喜歡。
  2. 在該模型上 訓練 LCM LoRA 模型。LoRA 是一種引數高效的微調 (PEFT),其實現成本比全模型微調要便宜得多。有關 PEFT 的更詳細資訊,請參閱 此博文diffusers 庫的 LoRA 檔案
  3. 將 LoRA 與任何 SDXL 模型和 LCM 排程器一起組成一個流水線,進行推理。就這樣!用這個流水線,你只需幾步推理即可生成高質量的影象。

欲知更多詳情,請 下載我們的論文

快有啥用?

Stable Diffusion 和 SDXL 的快速推理為新應用和新工作流開啟了大門,僅舉幾例:

  • 更易得: 變快後,生成工具可以被更多人使用,即使他們沒有最新的硬體。
  • 迭代更快: 無論從個人還是商業角度來看,在短時間內生成更多影象或進行更多嘗試對於藝術家和研究人員來說都非常有用。
  • 可以在各種不同的加速器上進行生產化部署,包括 CPU。
  • 影象生成服務會更便宜。

為了衡量我們所說的速度差異,在 M1 Mac 上用 SDXL (base) 生成一張 1024x1024 影象大約需要一分鐘。而用 LCM LoRA,我們只需約 6 秒 (4 步) 即可獲得出色的結果。速度快了一個數量級,我們再也無需等待結果,這帶來了顛覆性的體驗。如果使用 4090,我們幾乎可以得到實時響應 (不到 1 秒)。有了它,SDXL 可以用於需要實時響應的場合。

快速推理 SDXL LCM LoRA 模型

在最新版的 diffusers 中,大家可以非常容易地用上 LCM LoRA:

from diffusers import DiffusionPipeline, LCMScheduler
import torch

model_id = "stabilityai/stable-diffusion-xl-base-1.0"
lcm_lora_id = "latent-consistency/lcm-lora-sdxl"

pipe = DiffusionPipeline.from_pretrained(model_id, variant="fp16")

pipe.load_lora_weights(lcm_lora_id)
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipe.to(device="cuda", dtype=torch.float16)

prompt = "close-up photography of old man standing in the rain at night, in a street lit by lamps, leica 35mm summilux"
images = pipe(
    prompt=prompt,
    num_inference_steps=4,
    guidance_scale=1,
).images[0]

程式碼所做的事情如下:

  • 使用 SDXL 1.0 base 模型去範例化一個標準的 diffusion 流水線。
  • 應用 LCM LoRA。
  • 將排程器改為 LCMScheduler,這是 LCM 模型使用的排程器。
  • 結束!

生成的全解析度影象如下所示:

LCM LORA 微調後的 SDXL 模型用 4 步生成的影象
LCM LORA 微調後的 SDXL 模型用 4 步生成的影象

生成質量

我們看下步數對生成質量的影響。以下程式碼將分別用 1 步到 8 步生成影象:

images = []
for steps in range(8):
    generator = torch.Generator(device=pipe.device).manual_seed(1337)
    image = pipe(
        prompt=prompt,
        num_inference_steps=steps+1,
        guidance_scale=1,
        generator=generator,
    ).images[0]
    images.append(image)

生成的 8 張影象如下所示:

LCM LoRA 1 至 8 步生成的影象
LCM LoRA 1 至 8 步生成的影象

不出所料,僅使用 1 步即可生成細節和紋理欠缺的粗略影象。然而,隨著步數的增加,效果改善迅速,通常只需 4 到 6 步就可以達到滿意的效果。個人經驗是,8 步生成的影象對於我來說有點過飽和及「卡通化」,所以本例中我個人傾向於選擇 5 步和 6 步生成的影象。生成速度非常快,你只需 4 步即可生成一堆影象,並從中選擇你喜歡的,然後根據需要對步數和提示詞進行調整和迭代。

引導比例及反向提示

請注意,在前面的範例中,我們將引導比例 guidance_scale 設為 1 ,實際上就是禁用它。對大多數提示而言,這樣設定就可以了,此時速度最快,但會忽略反向提示。你還可以將其值設為 12 之間,用於探索反向提示的影響——但我們發現再大就不起作用了。

與標準 SDXL 模型的生成質量對比

就生成質量而言,本文的方法與標準 SDXL 流水線相比如何?我們看一個例子!

我們可以通過卸掉 LoRA 權重並切換回預設排程器來將流水線快速恢復為標準 SDXL 流水線:

from diffusers import EulerDiscreteScheduler

pipe.unload_lora_weights()
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)

然後,我們可以像往常一樣對 SDXL 進行推理。我們使用不同的步數並觀察其結果:

images = []
for steps in (1, 4, 8, 15, 20, 25, 30, 50):
    generator = torch.Generator(device=pipe.device).manual_seed(1337)
    image = pipe(
        prompt=prompt,
        num_inference_steps=steps,
        generator=generator,
    ).images[0]
    images.append(image)

不同步數下的 SDXL 結果
SDXL 流水線結果 (相同的提示和隨機種子),步數分別為 1、4、8、15、20、25、30 和 50

如你所見,此範例中的生成的影象在大約 20 步 (第二行) 之前幾乎毫無用處,且隨著步數的增加,質量仍會不斷明顯提高。最終影象中的細節很不錯,但獲得這樣的效果需要 50 步。

其他模型的 LCM LoRA

該技術也適用於任何其他微調後的 SDXL 或 Stable Diffusion 模型。僅舉一例,我們看看如何在 collage-diffusion 上執行推理,該模型是用 Dreambooth 演演算法對 Stable Diffusion v1.5 微調而得。

程式碼與我們在前面範例中看到的程式碼類似。我們先載入微調後的模型,然後載入適合 Stable Diffusion v1.5 的 LCM LoRA 權重。

from diffusers import DiffusionPipeline, LCMScheduler
import torch

model_id = "wavymulder/collage-diffusion"
lcm_lora_id = "latent-consistency/lcm-lora-sdv1-5"

pipe = DiffusionPipeline.from_pretrained(model_id, variant="fp16")
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipe.load_lora_weights(lcm_lora_id)
pipe.to(device="cuda", dtype=torch.float16)

prompt = "collage style kid sits looking at the night sky, full of stars"

generator = torch.Generator(device=pipe.device).manual_seed(1337)
images = pipe(
    prompt=prompt,
    generator=generator,
    negative_prompt=negative_prompt,
    num_inference_steps=4,
    guidance_scale=1,
).images[0]
images

基於 Dreambooth Stable Diffusion v1.5 模型使用 LCM LoRA,4 步推理
基於 Dreambooth Stable Diffusion v1.5 模型使用 LCM LoRA,4 步推理

Diffusers 全整合

LCM 在 diffusers 中的全面整合使得其可以利用 diffusers 工具箱中的許多特性和工作流,如:

  • 對採用 Apple 晶片的 Mac 提供開箱即用的 mps 支援。
  • 記憶體和效能優化,例如 flash 注意力或 torch.compile()
  • 針對低 RAM 場景的其他記憶體節省策略,包括模型解除安裝。
  • ControlNet 或圖生圖等工作流。
  • 訓練和微調指令碼。

測試基準

本節列出了 SDXL LCM LoRA 在各種硬體上的生成速度,給大家一個印象。忍不住再提一句,能如此輕鬆地探索影象生成真是太爽了!

硬體 SDXL LoRA LCM (4 步) 標準 SDXL (25 步)
Mac, M1 Max 6.5s 64s
2080 Ti 4.7s 10.2s
3090 1.4s 7s
4090 0.7s 3.4s
T4 (Google Colab Free Tier) 8.4s 26.5s
A100 (80 GB) 1.2s 3.8s
Intel i9-10980XE CPU (共 36 核,僅用 1 核) 29s 219s

上述所有測試的 batch size 均為 1,使用是 Sayak Paul 開發的 這個指令碼

對於視訊記憶體容量比較大的卡 (例如 A100),一次生成多張影象,效能會有顯著提高,一般來講生產部署時會採取增加 batch size 的方法來增加吞吐。

已公開發布的 LCM LoRA 及 LCM 模型

加分項: 將 LCM LoRA 與常規 SDXL LoRA 結合起來

使用 diffusers + PEFT 整合,你可以將 LCM LoRA 與常規 SDXL LoRA 結合起來,使其也擁有 4 步推理的超能力。

這裡,我們將 CiroN2022/toy_face LoRA 與 LCM LoRA 結合起來:

from diffusers import DiffusionPipeline, LCMScheduler
import torch

model_id = "stabilityai/stable-diffusion-xl-base-1.0"
lcm_lora_id = "latent-consistency/lcm-lora-sdxl"
pipe = DiffusionPipeline.from_pretrained(model_id, variant="fp16")
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)

pipe.load_lora_weights(lcm_lora_id)
pipe.load_lora_weights("CiroN2022/toy-face", weight_name="toy_face_sdxl.safetensors", adapter_name="toy")

pipe.set_adapters(["lora", "toy"], adapter_weights=[1.0, 0.8])
pipe.to(device="cuda", dtype=torch.float16)

prompt = "a toy_face man"
negative_prompt = "blurry, low quality, render, 3D, oversaturated"
images = pipe(
    prompt=prompt,
    negative_prompt=negative_prompt,
    num_inference_steps=4,
    guidance_scale=0.5,
).images[0]
images

結合兩種 LoRA 以實現快速推理
標準 LoRA 和 LCM LoRA 相結合實現 4 步快速推理

想要探索更多有關 LoRA 的新想法嗎?可以試試我們的實驗性 LoRA the Explorer (LCM 版本) 空間,在這裡你可以把玩社群的驚人創作並從中獲取靈感!

如何訓練 LCM 模型及 LCM LoRA

最新的 diffusers 中,我們提供了與 LCM 團隊作者合作開發的訓練和微調指令碼。有了它們,使用者可以:

  • 在 Laion 等大型資料集上執行 Stable Diffusion 或 SDXL 模型的全模型蒸餾。
  • 訓練 LCM LoRA,它比全模型蒸餾更簡單。正如我們在這篇文章中所示,訓練後,可以用它對 Stable Diffusion 實現快速推理,而無需進行蒸餾訓練。

更多詳細資訊,請查閱程式碼庫中的 SDXLStable Diffusion 說明檔案。

我們希望這些指令碼能夠激勵社群探索實現自己的微調。如果你將它們用於自己的專案,請告訴我們!

資源

致謝

LCM 團隊 完成了 LCM 模型的出色工作,請務必查閱他們的程式碼、報告和論文。該專案是 diffusers 團隊、LCM 團隊以及社群貢獻者 Daniel Gu 合作的結果。我們相信,這證明了開源人工智慧的強大力量,它是研究人員、從業者和探客 (tinkerer) 們探索新想法和共同作業的基石。我們還要感謝 @madebyollin 對社群的持續貢獻,其中包括我們在訓練指令碼中使用的 float16 自編碼器。


英文原文: https://hf.co/blog/lcm_lora

原文作者: Pedro Cuenca,Suraj Patil,Simian Luo,Daniel Gu,Yiqin Tan,Sayak Paul,Apolinário

譯者: Matrix Yao (姚偉峰),英特爾深度學習工程師,工作方向為 transformer-family 模型在各模態資料上的應用及大規模模型的訓練推理。

審校/排版: zhongdongy (阿東)