Diffusers 简介
Diffusers 是由 Hugging Face 开发的一个开源深度学习工具箱,专门用于构建、训练和部署扩散模型(Diffusion Models)。扩散模型是一种生成式模型,通过逐步添加噪声到数据中(前向过程),然后学习如何从噪声中恢复原始数据(反向过程),从而生成高质量的图像、音频或视频。
Diffusers 的核心用途
- 图像生成与编辑
- 文本到图像生成:根据自然语言描述生成图像(如 Stable Diffusion)。
- 图像修复与编辑:修改现有图像的特定区域(如去除背景、风格迁移)。
- 图像增强:提升图像分辨率或质量。
- 视频生成
- 生成动态视频内容,用于动画、游戏或多媒体项目。
- 音频生成
- 生成音乐、语音或环境音效。
- 研究与开发
- 提供模块化的工具链,支持自定义扩散模型的训练与优化。
Diffusers 的核心组件
- 模型(Models)
- UNet:核心网络结构,用于预测噪声残差(如
UNet2DModel)。 - VAE(变分自编码器):压缩/解压图像数据(如 Stable Diffusion 中的 VAE)。
- CLIP:文本编码器,将自然语言提示转换为嵌入向量(如 Stable Diffusion 中的文本编码器)。
- UNet:核心网络结构,用于预测噪声残差(如
- 调度器(Schedulers)
- DDPM(Denoising Diffusion Probabilistic Models):管理去噪过程的时间步长和噪声添加策略。
- 其他调度器:如
PNDMScheduler(快速去噪)、LMSDiscreteScheduler(低资源消耗)等。
- Pipeline(管道)
- 高级 API:将模型和调度器组合成端到端的推理接口(如
StableDiffusionPipeline)。 - 灵活扩展:允许用户自定义模型组件或调度策略。
- 高级 API:将模型和调度器组合成端到端的推理接口(如
组件数据流
| 组件 | 输入数据类型 | 输出数据类型 | 简单数据样例 | 说明 |
|---|---|---|---|---|
| VAE 编码器 | 原始图像 (RGB, 256x256) | 潜在空间表示 (向量, 64维) | 输入: [256x256x3] 的图像数组 → 输出: [64] 的浮点数向量 | 将高维图像压缩为低维潜在表示,便于后续处理和生成。 |
| CLIP | 文本描述 (“一只猫在沙发上睡觉”) 或 图像 (同上) | 文本或图像嵌入 (向量, 512维) | 输入: “一只猫在沙发上睡觉” → 输出: [512] 的浮点数向量 | 提供文本嵌入,提供图像的语义信息。 |
| 整合给 UNet | 潜在空间表示 + 文本嵌入 | 条件化输入 (向量, 合并后的维度) | 示例: [64] 的潜在向量 + [512] 的文本向量 → 合并后 [576] 的向量 | 结合来自 VAE 编码器的潜在表示和 CLIP 提供的文本嵌入作为条件信息传递给 UNet。 |
| UNet | 条件化输入 + 噪声(可选) | 预测噪声或去噪图像(取决于具体实现) | 输入: [576] 的条件向量 + [256x256x3] 的噪声图像 → 输出: [256x256x3] 的预测噪声 | 根据条件信息逐步去除噪声,恢复清晰图像。 |
| 调度器 | 时间步长 t 和当前噪声水平 | 更新后的噪声水平 | 输入: 当前时间步 t=10, 噪声水平 σ=0.8 → 输出: 更新后的噪声水平 σ’=0.7 | 控制扩散过程中的噪声添加/去除速率。 |
| VAE 解码器 | 去噪后的潜在空间表示 (同 VAE 编码器输出) | 重构图像 (原始图像大小) | 输入: [64] 的去噪潜在向量 → 输出: [256x256x3] 的重构图像 | 将 UNet 处理后的潜在表示转换回原始图像空间。 |
解释:
- VAE 编码器:处理图像,高维变低维,图像嵌入
- CLIP:处理文本,文本嵌入
- 整合:将 VAE 与 CLIP 输出的嵌入向量进行拼接,作为条件信息输入 UNet。
- UNet:根据提供的条件信息和当前的噪声图像,预测每一步应该去除的噪声,逐渐逼近目标图像。
- 调度器:扩散过程中逐步调整噪声水平,确保最终能从纯噪声状态恢复到清晰的目标图像。
- VAE 解码器:最后一步,将 UNet 处理过的潜在表示还原成实际的图像格式,完成整个生成过程。
总结
组件总览
| 组件 | 核心作用 | 是否必须 | 具体任务 | 典型应用场景 |
|---|---|---|---|---|
| VAE | 高维数据 ⇄ 低维潜空间压缩 | 非必需 | 训练:编码图像为潜变量;推理:解码潜变量为图像 | 潜空间扩散模型(如 Stable Diffusion) |
| CLIP | 文本语义编码与跨模态对齐 | 非必需 | 将文本提示转换为语义向量,注入 UNet 引导生成 | 文本到图像生成、语义编辑 |
| UNet | 噪声预测与逐步去噪 | 必需 | 接收含噪输入+时间步+条件,预测噪声并迭代去噪 | 所有扩散模型(如图像/音频/视频生成) |
| 调度器 | 控制噪声添加/移除规则 | 必需 | 定义前向噪声添加公式及反向采样更新策略 | DDPM、DDIM、Euler、DPM-Solver 等方案 |
微调
| 组件 | 是否深度学习模型 | 是否通常微调 | 说明 | 影响 |
|---|---|---|---|---|
| VAE | ✅ 是 | ❌ 否(可选) | 一般冻结,仅在特定需求下微调 | 影响图像画质 |
| CLIP | ✅ 是 | ❌ 否(可选) | 一般冻结,仅在需增强文本理解时微调 | 影响语义理解 |
| UNet | ✅ 是 | ✅ 是 | 核心生成网络,必须微调以适应新任务 | 影响图像内容 |
| 调度器 | ❌ 否 | ❌ 否 | 非神经网络,不可微调,只能更换 | 影响生成速度 |
入门教程
快速教程
扩散模型被训练用于逐步去噪随机高斯噪声,以生成降噪后的样本,例如图像或音频。
Diffusers 库有三个主要组件:
- DiffusionPipeline:使用预训练的扩散模型生成样本以进行推理
- 模型:构建扩散系统的模块化组件
- 调度器:控制如何在训练中添加噪声的算法,以及如何在推理过程中生成去噪图像
安装前置依赖
!pip install --upgrade diffusers accelerate transformers
DiffusionPipeline
| 任务 | 描述 | Pipeline |
|---|---|---|
| 无条件图像生成 | 从高斯噪声生成图像 | unconditional_image_generation |
| 文本引导图像生成 | 根据文本提示生成图像 | conditional_image_generation |
| 文本引导的图像到图像翻译 | 根据文本提示调整图像 | img2img |
| 文本引导图像修复 | 根据图像、蒙版和文本提示填充图像的蒙版部分 | inpaint |
| 文本引导深度到图像翻译 | 通过深度估计保留结构,同时根据文本提示调整图像部分 | depth2img |
加载模型
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", use_safetensors=True)
# DiffusionPipeline 会下载并缓存所有建模、分词和调度组件,print(pipeline)可以查看pipeline中加载的组件
# 在GPU上运行pipeline
pipeline.to("cuda")
# 将文本提示传递给pipeline生成图像
image = pipeline("An image of a squirrel in Picasso style").images[0]
image
# 保存图片
image.save("image_of_squirrel_painting.png")
本地执行
# 下载权重
!git lfs install
!git clone https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5
# 加载权重到管道中
pipeline = DiffusionPipeline.from_pretrained("./stable-diffusion-v1-5", use_safetensors=True)
切换调度器
不同的调度器具有不同的去噪速度和质量权衡。Diffusers 可以轻松切换不同的调度器:
from diffusers import EulerDiscreteScheduler
pipeline = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", use_safetensors=True)
pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config)
模型
大多数模型采用一个含噪声样本,并在每个时间步预测噪声残差,即噪声较小的图像与输出图像之间的差异。可以混合搭配不同的模型来创建其他扩散系统。
# 模型加载,如果不确定模型类型,可以使用AutoModel自动选择合适的模型类型进行加载
from diffusers import UNet2DModel
repo_id = "google/ddpm-cat-256"
model = UNet2DModel.from_pretrained(repo_id, use_safetensors=True)
# 查看模型参数
model.config
模型配置是一个冻结的参数字典,这意味着在模型创建后这些参数不能更改。这是有意为之的,以确保在开始时用于定义模型架构的参数保持不变,同时其他参数仍然可以在推理过程中调整。
模型主要参数:
sample_size:输入样本的高度和宽度维度。in_channels:输入样本的输入通道数。down_block_types和up_block_types:用于创建 UNet 架构的下采样和上采样块的类型。block_out_channels:下采样块的输出通道数;也用于上采样块的输入通道数,顺序相反。layers_per_block:每个 UNet 块中存在的 ResNet 块的数量。
# 构造样例数据
import torch
torch.manual_seed(0)
noisy_sample = torch.randn(1, model.config.in_channels, model.config.sample_size, model.config.sample_size)
noisy_sample.shape
# 进行推理
# timestep 表示输入图像的噪声程度,在开始时噪声较多,在结束时噪声较少。这有助于模型确定其在扩散过程中的位置,是更接近开始还是结束。
with torch.no_grad():
noisy_residual = model(sample=noisy_sample, timestep=2).sample
调度器
调度器负责根据模型输出(针对刚刚生成的 noisy_residual)将一个嘈杂的样本转换为更干净的样本。
使用 from_config() 方法实例化 DDPMScheduler:
from diffusers import DDPMScheduler
scheduler = DDPMScheduler.from_pretrained(repo_id)
scheduler
注意:与模型不同,调度器没有可训练的权重,并且是无参数的。
主要参数是:
num_train_timesteps:去噪过程的长度,换句话说,就是将随机高斯噪声处理成数据样本所需的步数。beta_schedule:推理和训练中使用的噪声调度类型。beta_start和beta_end:噪声调度程序的起始和结束噪声值。
# 要预测一个稍微降噪后的图像,将以下内容传递给调度器的 step() 方法:模型输出、timestep 和当前 sample
less_noisy_sample = scheduler.step(model_output=noisy_residual, timestep=2, sample=noisy_sample).prev_sample
less_noisy_sample.shape
less_noisy_sample 可以传递到下一个 timestep,进一步降噪。现在让我们将所有内容整合起来,并可视化整个去噪过程。
降噪循环与过程可视化
import PIL.Image
import numpy as np
def display_sample(sample, i):
image_processed = sample.cpu().permute(0, 2, 3, 1)
image_processed = (image_processed + 1.0) * 127.5
image_processed = image_processed.numpy().astype(np.uint8)
image_pil = PIL.Image.fromarray(image_processed[0])
display(f"Image at step {i}")
display(image_pil)
# GPU加速
model.to("cuda")
noisy_sample = noisy_sample.to("cuda")
# 创建一个降噪循环,预测低噪声样本的残差,并使用调度器计算低噪声样本:
import tqdm
sample = noisy_sample
for i, t in enumerate(tqdm.tqdm(scheduler.timesteps)):
# 1. predict noise residual
with torch.no_grad():
residual = model(sample, t).sample
# 2. compute less noisy image and set x_t -> x_t-1
sample = scheduler.step(residual, t, sample).prev_sample
# 3. optionally look at image
if (i + 1) % 50 == 0:
display_sample(sample, i + 1)
高效且有效的扩散
以特定风格生成图像或在图像中包含想要的内容,可能需要多次运行 DiffusionPipeline。
加载模型
from diffusers import DiffusionPipeline
model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5"
pipeline = DiffusionPipeline.from_pretrained(model_id, use_safetensors=True)
# 提示词,例如:生成一个印第安老战士首领的肖像
prompt = "portrait photo of a old warrior chief"
加速
可以通过降低精度、调整调度器来加快生成。
1. 通过 GPU 进行加速
将 pipeline 放在 GPU 上执行:
pipeline = pipeline.to("cuda")
使用 Generator 设置种子,确保结果可以复现:
import torch
generator = torch.Generator("cuda").manual_seed(0)
# 生成图像
image = pipeline(prompt, generator=generator).images[0]
image
这个过程在一个 T4 GPU 上大约耗时 30 秒。默认情况下,DiffusionPipeline 使用双精度浮点类型 float32 进行 50 步推理。可以切换到较低的精度(例如 float16)或运行较少的推理步数加速生成。
2. 通过降低精度进行加速
切换到 float16,大约耗时 11 秒,较之前的速度快了 3 倍:
import torch
pipeline = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, use_safetensors=True)
pipeline = pipeline.to("cuda")
generator = torch.Generator("cuda").manual_seed(0)
image = pipeline(prompt, generator=generator).images[0]
image
3. 通过减少推理步数进行加速
切换成性能更好的调度器,可以减少推理步数。例如,DPMSolverMultistepScheduler 调度器只需要 20-25 步。
查看与当前模型兼容的调度器:
pipeline.scheduler.compatibles
以下为输出的兼容的调度器:
[
diffusers.schedulers.scheduling_lms_discrete.LMSDiscreteScheduler,
diffusers.schedulers.scheduling_unipc_multistep.UniPCMultistepScheduler,
diffusers.schedulers.scheduling_k_dpm_2_discrete.KDPM2DiscreteScheduler,
diffusers.schedulers.scheduling_deis_multistep.DEISMultistepScheduler,
diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler,
diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler,
diffusers.schedulers.scheduling_ddpm.DDPMScheduler,
diffusers.schedulers.scheduling_dpmsolver_singlestep.DPMSolverSinglestepScheduler,
diffusers.schedulers.scheduling_k_dpm_2_ancestral_discrete.KDPM2AncestralDiscreteScheduler,
diffusers.utils.dummy_torch_and_torchsde_objects.DPMSolverSDEScheduler,
diffusers.schedulers.scheduling_heun_discrete.HeunDiscreteScheduler,
diffusers.schedulers.scheduling_pndm.PNDMScheduler,
diffusers.schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteScheduler,
diffusers.schedulers.scheduling_ddim.DDIMScheduler,
]
切换为兼容的调度器,PNDMScheduler(50 步推理)→ DPMSolverMultistepScheduler(20-25 步推理):
from diffusers import DPMSolverMultistepScheduler
pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
# 将推理步数设置为20
generator = torch.Generator("cuda").manual_seed(0)
image = pipeline(prompt, generator=generator, num_inference_steps=20).images[0]
image
节省内存
提高 Pipeline 性能的另一个关键在于减少内存消耗,间接意味着更快的速度。可以通过注意力切片来达到这一目的。
要查看一次能生成多少图像,最简单的方式是调整批次大小,直到出现 OutOfMemoryError(OOM)错误:
# 创建函数,根据提示词和Generator生成一批图像。确保为每个Generator分配一个种子,以便结果可以被复现
def get_inputs(batch_size=1):
generator = [torch.Generator("cuda").manual_seed(i) for i in range(batch_size)]
prompts = batch_size * [prompt]
num_inference_steps = 20
return {"prompt": prompts, "generator": generator, "num_inference_steps": num_inference_steps}
# 输出图像
from diffusers.utils import make_image_grid
images = pipeline(**get_inputs(batch_size=4)).images
make_image_grid(images, 2, 2)
如果上面的代码返回了 OOM 错误,说明大部分内存被注意力层占用,可以使用 enable_attention_slicing() 函数,将批量运行改为顺序运行来节省内存:
pipeline.enable_attention_slicing()
此时可以适当增加批次大小:
images = pipeline(**get_inputs(batch_size=8)).images
make_image_grid(images, rows=2, cols=4)
提升质量
可以通过更好的模型、更好的提示词,来生成质量更高的图像。
1. 使用更好的管道组件
使用更新的版本和更佳的模型:
from diffusers import AutoencoderKL
vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16).to("cuda")
pipeline.vae = vae
images = pipeline(**get_inputs(batch_size=8)).images
make_image_grid(images, rows=2, cols=4)
2. 使用更好的提示词
使用具有更多细节的提示词:
prompt = "portrait photo of a old warrior chief"
prompt += ", tribal panther make up, blue on red, side profile, looking away, serious eyes"
prompt += " 50mm portrait photography, hard rim lighting photography--beta --ar 2:3 --beta --upbeta"
# 生成图像
def get_inputs(batch_size=1):
generator = [torch.Generator("cuda").manual_seed(i) for i in range(batch_size)]
prompts = batch_size * [prompt]
num_inference_steps = 20
return {"prompt": prompts, "generator": generator, "num_inference_steps": num_inference_steps}
images = pipeline(**get_inputs(batch_size=8)).images
make_image_grid(images, rows=2, cols=4)
进一步调整提示词,生成图像:
prompts = [
"portrait photo of the oldest warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes 50mm portrait photography, hard rim lighting photography--beta --ar 2:3 --beta --upbeta",
"portrait photo of an old warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes 50mm portrait photography, hard rim lighting photography--beta --ar 2:3 --beta --upbeta",
"portrait photo of a warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes 50mm portrait photography, hard rim lighting photography--beta --ar 2:3 --beta --upbeta",
"portrait photo of a young warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes 50mm portrait photography, hard rim lighting photography--beta --ar 2:3 --beta --upbeta",
]
generator = [torch.Generator("cuda").manual_seed(1) for _ in range(len(prompts))]
images = pipeline(prompt=prompts, generator=generator, num_inference_steps=25).images
make_image_grid(images, 2, 2)
进阶教程
管道、模型、调度器
管道
DDPMPipeline 管道包含了 UNet2DModel 模型和 DDPMScheduler 调度器。管道通过生成指定大小的随机噪声图像,将其多次通过模型对图像进行去噪。每个时间步,模型预测噪声残差,调度器根据噪声残差对图像中的噪声进行去噪,重复这个过程 25 次,生成最终图像。
from diffusers import DDPMPipeline
ddpm = DDPMPipeline.from_pretrained("google/ddpm-cat-256", use_safetensors=False).to("cuda")
image = ddpm(num_inference_steps=25).images[0]
image
不使用 DDPMPipeline 加载模型,改为分别加载模型、调度器的方式。推理流程:加载模型、调度器 → 设置时间步数组 → 模型输出噪声残差 → 调度器去噪 → 根据时间步数据,循环遍历 → 遍历完成,输出图像。
1. 加载模型和调度器
from diffusers import DDPMScheduler, UNet2DModel
scheduler = DDPMScheduler.from_pretrained("google/ddpm-cat-256")
model = UNet2DModel.from_pretrained("google/ddpm-cat-256", use_safetensors=True).to("cuda")
2. 指定去噪过程运行的步数
scheduler.set_timesteps(50)
3. 设置调度器步数
会创建一个包含均匀间隔元素的张量。本例中有 50 个元素,对应了 set_timesteps(50)。每个元素对应模型去噪图像的一个时间步。稍后创建去噪循环时,会遍历这个张量对图像去噪:
scheduler.timesteps
4. 创建一些与期望输出相同形状的随机噪声
import torch
sample_size = model.config.sample_size
noise = torch.randn((1, 3, sample_size, sample_size), device="cuda")
5. 去噪循环
在每个时间步,模型执行 UNet2DModel.forward() 进行前向计算,返回噪声残差。调度器执行 scheduler.step() 接收噪声残差、时间步和输入,预测前一时间步的图像,作为下一次模型的输入。重复多次,直到遍历整个 timesteps 数组:
input = noise
for t in scheduler.timesteps:
with torch.no_grad():
noisy_residual = model(input, t).sample
previous_noisy_sample = scheduler.step(noisy_residual, t, input).prev_sample
input = previous_noisy_sample
6. 最后一步,将去噪后的输出转换为图像
from PIL import Image
import numpy as np
image = (input / 2 + 0.5).clamp(0, 1).squeeze()
image = (image.permute(1, 2, 0) * 255).round().to(torch.uint8).cpu().numpy()
image = Image.fromarray(image)
image
Stable Diffusion 管道
Stable Diffusion 是文本到图像的潜在扩散模型。除模型和调度器以外,VAE 编码器将图像压缩成更小的表示,VAE 解码器将压缩的表示转换回图像。CLIP 分词器和编码器将文本切分并编码为嵌入向量。
Stable Diffusion 模型有三个独立的预训练模型:UNet 残差模型、VAE 编码器解码器、CLIP 文本编码器。
Stable Diffusion 模型具体包含以下组件:
unet:用于生成输入的潜在表示的模型。vae:我们将使用的自动编码器模块,用于将潜在表示解码为真实图像。text_encoder:Stable Diffusion 使用 CLIP 编码器,其他扩散模型可能使用其他编码器,例如 BERT。tokenizer:与text_encoder模型相匹配的词元分析器。scheduler:训练过程中用于逐步向图像添加噪声的调度算法。
1. 加载扩散模型各组件
AutoencoderKL 加载 vae,CLIPTokenizer 加载 CLIP 词元分析器,CLIPTextModel 加载 CLIP 文本编码器,UNet2DConditionModel 加载模型,UniPCMultistepScheduler 加载调度器(默认为 PNDMScheduler):
from PIL import Image
import torch
from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler
vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae", use_safetensors=True)
tokenizer = CLIPTokenizer.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="tokenizer")
text_encoder = CLIPTextModel.from_pretrained(
"CompVis/stable-diffusion-v1-4", subfolder="text_encoder", use_safetensors=True
)
unet = UNet2DConditionModel.from_pretrained(
"CompVis/stable-diffusion-v1-4", subfolder="unet", use_safetensors=True
)
from diffusers import UniPCMultistepScheduler
scheduler = UniPCMultistepScheduler.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="scheduler")
2. GPU 加速
注意,这里将 vae、CLIP 文本编码器、模型加载到 GPU,它们有可训练的权重,而调度器和 CLIP 词元分析器则不需要,整个过程中仅调用:
torch_device = "cuda"
vae.to(torch_device)
text_encoder.to(torch_device)
unet.to(torch_device)
3. 创建文本嵌入
对文本分词,生成嵌入向量:
prompt = ["a photograph of an astronaut riding a horse"]
height = 512 # default height of Stable Diffusion
width = 512 # default width of Stable Diffusion
num_inference_steps = 25 # Number of denoising steps
guidance_scale = 7.5 # Scale for classifier-free guidance
generator = torch.manual_seed(0) # Seed generator to create the initial latent noise
batch_size = len(prompt)
text_input = tokenizer(
prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt"
)
with torch.no_grad():
text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]
生成无条件的文本嵌入,这些嵌入是填充标记的嵌入,需要与 text_embeddings 具有相同的形状 (batch_size, seq_length):
max_length = text_input.input_ids.shape[-1]
uncond_input = tokenizer([""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt")
uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]
将条件嵌入和无条件嵌入连接成一个批次,以避免进行两次前向传递:
text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
4. 创建随机噪声
生成一些初始随机噪声作为扩散过程的起点。这是图像的潜在表示,逐渐去噪后,潜在表示会小于最终图像尺寸,模型稍后将其逐步转换为最终图像尺寸:
latents = torch.randn(
(batch_size, unet.config.in_channels, height // 8, width // 8),
generator=generator,
device=torch_device,
)
5. 图像去噪
首先使用初始噪声分布和噪声尺度,对输入进行缩放,噪声尺度对于某些改进的调度器是必须的,例如 UniPCMultistepScheduler:
latents = latents * scheduler.init_noise_sigma
创建去噪循环,逐步将 latents 中的纯噪声转换为提示描述的图像。去噪循环主要有三项工作:设置去噪的时间步;遍历时间步;每个时间步中调用模型预测噪声残差,将噪声残差传递给调度器计算前一个带噪声的样本:
from tqdm.auto import tqdm
scheduler.set_timesteps(num_inference_steps)
for t in tqdm(scheduler.timesteps):
# expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
latent_model_input = torch.cat([latents] * 2)
latent_model_input = scheduler.scale_model_input(latent_model_input, timestep=t)
# predict the noise residual
with torch.no_grad():
noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
# perform guidance
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
# compute the previous noisy sample x_t -> x_t-1
latents = scheduler.step(noise_pred, t, latents).prev_sample
6. 图像解码
使用 vae 解码器将潜在表示解码为图像,使用 sample 获取解码结果:
# scale and decode the image latents with vae
latents = 1 / 0.18215 * latents
with torch.no_grad():
image = vae.decode(latents).sample
使用 PIL.Image.fromarray 将解码结果转换为图像:
image = (image / 2 + 0.5).clamp(0, 1).squeeze()
image = (image.permute(1, 2, 0) * 255).to(torch.uint8).cpu().numpy()
image = Image.fromarray(image)
image
自动管道
Diffusers 提供了许多基本任务的管道,例如生成图像、视频、音频、图像修复。基于这些管道,还有针对适配器、功能的专用管道。不同的管道甚至可以使用相同的检查点,通过相同的预训练模型实现不同的下游任务。
自动管道有三种:AutoPipelineForText2Image、AutoPipelineForImage2Image、AutoPipelineForInpainting。AutoPipeline 无需知道具体的管道类型,自动检测应使用的具体管道类型。
| AutoPipeline 类型 | 任务类型 | 主要参数 | 调用示例 |
|---|---|---|---|
AutoPipelineForTextToImage | 文本到图像生成 | prompt:输入文本提示height:生成图像的高度width:生成图像的宽度num_inference_steps:去噪步骤数guidance_scale:引导比例 | pipeline = AutoPipelineForTextToImage.from_pretrained("model_name") |
AutoPipelineForImageToImage | 图像到图像转换 | init_image:初始图像strength:强度控制变换程度num_inference_steps:去噪步骤数guidance_scale:引导比例 | pipeline = AutoPipelineForImageToImage.from_pretrained("model_name") |
AutoPipelineForInpainting | 图像修复 | image:需要修复的图像mask_image:对应的掩码图像num_inference_steps:去噪步骤数guidance_scale:引导比例 | pipeline = AutoPipelineForInpainting.from_pretrained("model_name") |
AutoPipelineForDepthToImage | 深度图到图像生成 | depth_map:深度图num_inference_steps:去噪步骤数guidance_scale:引导比例 | pipeline = AutoPipelineForDepthToImage.from_pretrained("model_name") |
以 dreamlike-art/dreamlike-photoreal-2.0 模型为例,使用不同的 AutoPipeline 加载模型,用于不同的下游任务:
AutoPipeline从model_index.json中检测到stable-diffusion类- 加载
StableDiffusionPipeline、StableDiffusionImg2ImgPipeline或StableDiffusionInpaintPipeline,用于适配具体下游任务
示例 1:AutoPipelineForText2Image 文生图推理
from diffusers import AutoPipelineForText2Image
import torch
pipe_txt2img = AutoPipelineForText2Image.from_pretrained(
"dreamlike-art/dreamlike-photoreal-2.0", torch_dtype=torch.float16, use_safetensors=True
).to("cuda")
prompt = "cinematic photo of Godzilla eating sushi with a cat in a izakaya, 35mm photograph, film, professional, 4k, highly detailed"
generator = torch.Generator(device="cpu").manual_seed(37)
image = pipe_txt2img(prompt, generator=generator).images[0]
image
示例 2:AutoPipelineForImage2Image 图生图推理
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import load_image
import torch
pipe_img2img = AutoPipelineForImage2Image.from_pretrained(
"dreamlike-art/dreamlike-photoreal-2.0", torch_dtype=torch.float16, use_safetensors=True
).to("cuda")
init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/autopipeline-text2img.png")
prompt = "cinematic photo of Godzilla eating burgers with a cat in a fast food restaurant, 35mm photograph, film, professional, 4k, highly detailed"
generator = torch.Generator(device="cpu").manual_seed(53)
image = pipe_img2img(prompt, image=init_image, generator=generator).images[0]
image
注意:如果已经使用
AutoPipelineForText2Image文生图管道完成加载,可以使用AutoPipelineForImage2Image.from_pipe()将其转换为图生图 pipeline:
pipe_img2img = AutoPipelineForImage2Image.from_pipe(pipe_txt2img).to("cuda")
image = pipeline(prompt, image=init_image, generator=generator).images[0]
image
示例 3:AutoPipelineForInpainting 图像修复
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image
import torch
pipeline = AutoPipelineForInpainting.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, use_safetensors=True
).to("cuda")
init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/autopipeline-img2img.png")
mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/autopipeline-mask.png")
prompt = "cinematic photo of a owl, 35mm photograph, film, professional, 4k, highly detailed"
generator = torch.Generator(device="cpu").manual_seed(38)
image = pipeline(prompt, image=init_image, mask_image=mask_image, generator=generator, strength=0.4).images[0]
image
训练一个扩散模型
在 Smithsonian Butterflies 蝴蝶图像数据集的子集上训练一个 UNet2DModel,用于生成自己的图像生成模型。
安装依赖
!pip install diffusers[training]
登录到 Hub
用于同步模型:
from huggingface_hub import notebook_login
notebook_login(token=XXX)
训练配置
from dataclasses import dataclass
@dataclass
class TrainingConfig:
image_size = 128 # 指定生成图像的分辨率大小,这里是128x128像素。这个参数影响模型训练和生成图像的尺寸
train_batch_size = 16 # 训练批次大小,表示每次向模型输入多少张图片进行训练。这里设置为16,意味着每次迭代使用16张图片进行训练
eval_batch_size = 16 # 评估时的批次大小,表示在模型评估阶段每次采样多少张图片。此值同样设为16
num_epochs = 50 # 训练周期数,指定整个数据集将被遍历多少次以进行训练。在此配置中设定为50个周期
gradient_accumulation_steps = 1 # 梯度累积步骤,用于模拟较大的批次大小而不需要增加内存消耗。如果设置为N,则相当于每N次更新才执行一次权重更新。默认情况下是1,即每次更新权重
learning_rate = 1e-4 # 学习率,控制权重更新的幅度。这里设置为1e-4(0.0001),是一个常见的较小的学习率值,有助于模型稳定学习
lr_warmup_steps = 500 # 学习率预热步骤数,在开始训练的初期逐步提高学习率直到达到预定值。这有助于防止初始阶段的大梯度破坏权重。此处设置为500步
save_image_epochs = 10 # 每隔多少个epoch保存一次生成的样本图像。这里设定每隔10个epoch保存一次图像,以便于跟踪模型的进展
save_model_epochs = 30 # 每隔多少个epoch保存一次模型。这里设置为30,意味着每经过30个训练周期就保存一次当前的模型状态
mixed_precision = "fp16" # 混合精度类型,可以是"no"(使用float32)或"fp16"(自动混合精度)。使用混合精度可以在一定程度上加速训练过程并减少显存占用
output_dir = "ddpm-butterflies-128" # 输出目录名,既是本地保存模型的位置名称,也是上传到Hugging Face Hub上的仓库名称
push_to_hub = True # 是否在训练结束后将模型推送到Hugging Face模型库
hub_model_id = "<your-username>/<my-awesome-model>" # 如果push_to_hub为True,这是要创建的Hugging Face Hub上的仓库名称。需要替换<your-username>/<my-awesome-model>为实际用户名和希望的模型ID
hub_private_repo = None # 决定Hugging Face Hub上的仓库是否私有。如果设置为None,则遵循Hugging Face的默认行为
overwrite_output_dir = True # 是否覆盖旧模型文件。当重新运行训练脚本时,如果输出目录已存在且该标志为True,则会覆盖原有内容
seed = 0 # 随机种子,用于初始化随机数发生器,确保实验结果的可重复性。这里设置为0
config = TrainingConfig()
加载数据集
from datasets import load_dataset
config.dataset_name = "huggan/smithsonian_butterflies_subset"
dataset = load_dataset(config.dataset_name, split="train")
使用 PIL.Image 加载并可视化:
import matplotlib.pyplot as plt
fig, axs = plt.subplots(1, 4, figsize=(16, 4))
for i, image in enumerate(dataset[:4]["image"]):
axs[i].imshow(image)
axs[i].set_axis_off()
fig.show()
图像预处理
- Resize:将图像大小调整为
config.image_size中定义的大小。 - RandomHorizontalFlip:通过随机镜像图像来增强数据集。
- Normalize:重要的是将像素值重新缩放到
[-1, 1]范围内,这是模型所期望的。
from torchvision import transforms
preprocess = transforms.Compose(
[
transforms.Resize((config.image_size, config.image_size)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]),
]
)
使用 Datasets.set_transform 方法调用 preprocess 函数:
def transform(examples):
images = [preprocess(image.convert("RGB")) for image in examples["image"]]
return {"images": images}
dataset.set_transform(transform)
使用 DataLoader 封装数据集:
import torch
train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=config.train_batch_size, shuffle=True)
创建 UNet 模型
from diffusers import UNet2DModel
model = UNet2DModel(
sample_size=config.image_size, # the target image resolution
in_channels=3, # the number of input channels, 3 for RGB images
out_channels=3, # the number of output channels
layers_per_block=2, # how many ResNet layers to use per UNet block
block_out_channels=(128, 128, 256, 256, 512, 512), # the number of output channels for each UNet block
down_block_types=(
"DownBlock2D", # a regular ResNet downsampling block
"DownBlock2D",
"DownBlock2D",
"DownBlock2D",
"AttnDownBlock2D", # a ResNet downsampling block with spatial self-attention
"DownBlock2D",
),
up_block_types=(
"UpBlock2D", # a regular ResNet upsampling block
"AttnUpBlock2D", # a ResNet upsampling block with spatial self-attention
"UpBlock2D",
"UpBlock2D",
"UpBlock2D",
"UpBlock2D",
),
)
查看数据形状与模型输出形状:
sample_image = dataset[0]["images"].unsqueeze(0)
print("Input shape:", sample_image.shape)
print("Output shape:", model(sample_image, timestep=0).sample.shape)
创建调度器
训练过程中,调度器从特定的扩散过程点获取模型输出,根据规则对图像添加随机噪声;推理过程中,调度器对图像去噪,生成图形:
import torch
from PIL import Image
from diffusers import DDPMScheduler
noise_scheduler = DDPMScheduler(num_train_timesteps=1000)
noise = torch.randn(sample_image.shape)
timesteps = torch.LongTensor([50])
noisy_image = noise_scheduler.add_noise(sample_image, noise, timesteps)
Image.fromarray(((noisy_image.permute(0, 2, 3, 1) + 1.0) * 127.5).type(torch.uint8).numpy()[0])
模型训练目标是预测添加到图像中的噪声,损失值用以下方式计算:
import torch.nn.functional as F
noise_pred = model(noisy_image, timesteps).sample
loss = F.mse_loss(noise_pred, noise)
训练模型
初始化优化器和学习率调度器
from diffusers.optimization import get_cosine_schedule_with_warmup
optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)
lr_scheduler = get_cosine_schedule_with_warmup(
optimizer=optimizer,
num_warmup_steps=config.lr_warmup_steps,
num_training_steps=(len(train_dataloader) * config.num_epochs),
)
评估模型
评估时使用 DDPMPipeline 生成一批样本图像,将其保存为网格:
from diffusers import DDPMPipeline
from diffusers.utils import make_image_grid
import os
def evaluate(config, epoch, pipeline):
# Sample some images from random noise (this is the backward diffusion process).
# The default pipeline output type is `List[PIL.Image]`
images = pipeline(
batch_size=config.eval_batch_size,
generator=torch.Generator(device='cpu').manual_seed(config.seed), # Use a separate torch generator to avoid rewinding the random state of the main training loop
).images
# Make a grid out of the images
image_grid = make_image_grid(images, rows=4, cols=4)
# Save the images
test_dir = os.path.join(config.output_dir, "samples")
os.makedirs(test_dir, exist_ok=True)
image_grid.save(f"{test_dir}/{epoch:04d}.png")
Accelerator 加速与模型上传到 Hub
from accelerate import Accelerator
from huggingface_hub import create_repo, upload_folder
from tqdm.auto import tqdm
from pathlib import Path
import os
def train_loop(config, model, noise_scheduler, optimizer, train_dataloader, lr_scheduler):
# Initialize accelerator and tensorboard logging
accelerator = Accelerator(
mixed_precision=config.mixed_precision,
gradient_accumulation_steps=config.gradient_accumulation_steps,
log_with="tensorboard",
project_dir=os.path.join(config.output_dir, "logs"),
)
if accelerator.is_main_process:
if config.output_dir is not None:
os.makedirs(config.output_dir, exist_ok=True)
if config.push_to_hub:
repo_id = create_repo(
repo_id=config.hub_model_id or Path(config.output_dir).name, exist_ok=True
).repo_id
accelerator.init_trackers("train_example")
# Prepare everything
# There is no specific order to remember, you just need to unpack the
# objects in the same order you gave them to the prepare method.
model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
model, optimizer, train_dataloader, lr_scheduler
)
global_step = 0
# Now you train the model
for epoch in range(config.num_epochs):
progress_bar = tqdm(total=len(train_dataloader), disable=not accelerator.is_local_main_process)
progress_bar.set_description(f"Epoch {epoch}")
for step, batch in enumerate(train_dataloader):
clean_images = batch["images"]
# Sample noise to add to the images
noise = torch.randn(clean_images.shape, device=clean_images.device)
bs = clean_images.shape[0]
# Sample a random timestep for each image
timesteps = torch.randint(
0, noise_scheduler.config.num_train_timesteps, (bs,), device=clean_images.device,
dtype=torch.int64
)
# Add noise to the clean images according to the noise magnitude at each timestep
# (this is the forward diffusion process)
noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)
with accelerator.accumulate(model):
# Predict the noise residual
noise_pred = model(noisy_images, timesteps, return_dict=False)[0]
loss = F.mse_loss(noise_pred, noise)
accelerator.backward(loss)
if accelerator.sync_gradients:
accelerator.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
progress_bar.update(1)
logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0], "step": global_step}
progress_bar.set_postfix(**logs)
accelerator.log(logs, step=global_step)
global_step += 1
# After each epoch you optionally sample some demo images with evaluate() and save the model
if accelerator.is_main_process:
pipeline = DDPMPipeline(unet=accelerator.unwrap_model(model), scheduler=noise_scheduler)
if (epoch + 1) % config.save_image_epochs == 0 or epoch == config.num_epochs - 1:
evaluate(config, epoch, pipeline)
if (epoch + 1) % config.save_model_epochs == 0 or epoch == config.num_epochs - 1:
if config.push_to_hub:
upload_folder(
repo_id=repo_id,
folder_path=config.output_dir,
commit_message=f"Epoch {epoch}",
ignore_patterns=["step_*", "epoch_*"],
)
else:
pipeline.save_pretrained(config.output_dir)
在 Notebook 环境下训练
from accelerate import notebook_launcher
args = (config, model, noise_scheduler, optimizer, train_dataloader, lr_scheduler)
notebook_launcher(train_loop, args, num_processes=1)
查看结果
import glob
sample_images = sorted(glob.glob(f"{config.output_dir}/samples/*.png"))
Image.open(sample_images[-1])
加载管道和适配器
加载管道
通用管道:DiffusionPipeline
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", use_safetensors=True)
init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-init.png")
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipeline("Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", image=init_image).images[0]
特定管道:StableDiffusionImg2ImgPipeline
from diffusers import StableDiffusionImg2ImgPipeline
pipeline = StableDiffusionImg2ImgPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", use_safetensors=True)
指定组件的数据类型
from diffusers import HunyuanVideoPipeline
import torch
pipe = HunyuanVideoPipeline.from_pretrained(
"hunyuanvideo-community/HunyuanVideo",
torch_dtype={"transformer": torch.bfloat16, "default": torch.float16},
)
print(pipe.transformer.dtype, pipe.vae.dtype) # (torch.bfloat16, torch.float16)
本地管道
!git-lfs install
!git clone https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5
from diffusers import DiffusionPipeline
stable_diffusion = DiffusionPipeline.from_pretrained("./stable-diffusion-v1-5", use_safetensors=True)
自定义管道
# 加载调度器和VAE
from diffusers import StableDiffusionXLPipeline, HeunDiscreteScheduler, AutoencoderKL
import torch
scheduler = HeunDiscreteScheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler")
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, use_safetensors=True)
# 将调度器和VAE传递给指定类型的Pipeline
pipeline = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
scheduler=scheduler,
vae=vae,
torch_dtype=torch.float16,
variant="fp16",
use_safetensors=True
).to("cuda")
复用管道(管道类型的转换)
已经使用一个 Pipeline 加载了模型:
from diffusers import DiffusionPipeline, StableDiffusionSAGPipeline
import torch
import gc
from diffusers.utils import load_image
from accelerate.utils import compute_module_sizes
image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/load_neg_embed.png")
pipe_sd = DiffusionPipeline.from_pretrained("SG161222/Realistic_Vision_V6.0_B1_noVAE", torch_dtype=torch.float16)
pipe_sd.load_ip_adapter("h94/IP-Adapter", subfolder="models", weight_name="ip-adapter_sd15.bin")
pipe_sd.set_ip_adapter_scale(0.6)
pipe_sd.to("cuda")
generator = torch.Generator(device="cpu").manual_seed(33)
out_sd = pipe_sd(
prompt="bear eats pizza",
negative_prompt="wrong white balance, dark, sketches,worst quality,low quality",
ip_adapter_image=image,
num_inference_steps=50,
generator=generator,
).images[0]
out_sd
查看此时的内存占用:
def bytes_to_giga_bytes(bytes):
return bytes / 1024 / 1024 / 1024
print(f"Max memory allocated: {bytes_to_giga_bytes(torch.cuda.max_memory_allocated())} GB")
# "Max memory allocated: 4.406213283538818 GB"
使用 from_pipe() 将管道从 StableDiffusionPipeline 类型转换为 StableDiffusionSAGPipeline 类型:
pipe_sag = StableDiffusionSAGPipeline.from_pipe(
pipe_sd
)
generator = torch.Generator(device="cpu").manual_seed(33)
out_sag = pipe_sag(
prompt="bear eats pizza",
negative_prompt="wrong white balance, dark, sketches,worst quality,low quality",
ip_adapter_image=image,
num_inference_steps=50,
generator=generator,
guidance_scale=1.0,
sag_scale=0.75
).images[0]
out_sag
查看此时的内存占用——与之前相同,因为两个不同类型的管道共享相同的管道组件,允许在不使用任何额外的内存开销的前提下进行管道的转换与复用:
print(f"Max memory allocated: {bytes_to_giga_bytes(torch.cuda.max_memory_allocated())} GB")
# "Max memory allocated: 4.406213283538818 GB"
管道转换后,添加适配器模块。例如,加载 IP-Adapter,加载 MotionAdapter,但是仅使用 MotionAdapter 适配器:
from diffusers import AnimateDiffPipeline, MotionAdapter, DDIMScheduler
from diffusers.utils import export_to_gif
pipe_sag.unload_ip_adapter()
adapter = MotionAdapter.from_pretrained("guoyww/animatediff-motion-adapter-v1-5-2", torch_dtype=torch.float16)
pipe_animate = AnimateDiffPipeline.from_pipe(pipe_sd, motion_adapter=adapter)
pipe_animate.scheduler = DDIMScheduler.from_config(pipe_animate.scheduler.config, beta_schedule="linear")
# load IP-Adapter and LoRA weights again
pipe_animate.load_ip_adapter("h94/IP-Adapter", subfolder="models", weight_name="ip-adapter_sd15.bin")
pipe_animate.load_lora_weights("guoyww/animatediff-motion-lora-zoom-out", adapter_name="zoom-out")
pipe_animate.to("cuda")
generator = torch.Generator(device="cpu").manual_seed(33)
pipe_animate.set_adapters("zoom-out", adapter_weights=0.75)
out = pipe_animate(
prompt="bear eats pizza",
num_frames=16,
num_inference_steps=50,
ip_adapter_image=image,
generator=generator,
).frames[0]
export_to_gif(out, "out_animate.gif")
调整管道,卸载 IP-Adapter,仅使用 MotionAdapter 适配器:
pipe.sag_unload_ip_adapter()
generator = torch.Generator(device="cpu").manual_seed(33)
out_sd = pipe_sd(
prompt="bear eats pizza",
negative_prompt="wrong white balance, dark, sketches,worst quality,low quality",
ip_adapter_image=image,
num_inference_steps=50,
generator=generator,
).images[0]
# "AttributeError: 'NoneType' object has no attribute 'image_projection_layers'"
说明:当使用
from_pipe()加载多个管道时,内存使用量由最高的管道决定,与创建的管道数量无关。
| Pipeline | Memory usage (GB) |
|---|---|
| StableDiffusionPipeline | 4.400 |
| StableDiffusionSAGPipeline | 4.400 |
| AnimateDiffPipeline | 15.178 |
例如,当使用 from_pipe() 加载以上 3 个管道后,内存使用量为 AnimateDiffPipeline 的 15.178GB,每个管道可以互换使用,没有额外的内存开销。
安全检查器
Diffusers 为 Stable Diffusion 模型实现了一个安全检查器,该模型可能生成有害内容。安全检查器会筛选生成的输出,以检测已知的硬编码的”不适合工作”(NSFW)内容。
如果你出于任何原因想要禁用安全检查器,可以将 safety_checker=None 传递给 from_pretrained() 方法:
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", safety_checker=None, use_safetensors=True)
加载不同版本的检查点
检查点变体是指 from_pretrained 可以通过参数加载不同类型权重的模型:
| 检查点类型 | 权重名称 | 涉及的参数 | 特点 |
|---|---|---|---|
| fp32 全精度类型 | diffusion_pytorch_model.safetensors | — | 不存在限制 |
| fp16 半精度权重 | diffusion_pytorch_model.fp16.safetensors | variant, torch_dtype | 不支持训练,不支持 CPU |
| 非指数移动平均(non-EMA)权重 | diffusion_pytorch_model.non_ema.safetensors | variant | 不支持推理 |
加载 fp16 类型权重的模型——variant 指定为 fp16,torch_dtype 指定为 torch.float16:
from diffusers import DiffusionPipeline
import torch
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", variant="fp16", torch_dtype=torch.float16, use_safetensors=True
)
加载 non-EMA 权重的模型——variant 指定为 non_ema:
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", variant="non_ema", use_safetensors=True
)
保存 fp16 类型权重的模型:
from diffusers import DiffusionPipeline
pipeline.save_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", variant="fp16")
再次加载保存的 fp16 类型权重的模型:
# 👎 这样不行
pipeline = DiffusionPipeline.from_pretrained(
"./stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
# 👍 这样可以
pipeline = DiffusionPipeline.from_pretrained(
"./stable-diffusion-v1-5", variant="fp16", torch_dtype=torch.float16, use_safetensors=True
)
保存 non-EMA 权重的模型:
pipeline.save_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", variant="non_ema")
再次加载 non-EMA 权重的模型:
# 👎 这样不行
pipeline = DiffusionPipeline.from_pretrained(
"./stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
# 👍 这样可以
pipeline = DiffusionPipeline.from_pretrained(
"./stable-diffusion-v1-5", variant="non_ema", use_safetensors=True
)
管道加载的组件
以 stable-diffusion-v1-5/stable-diffusion-v1-5 为例:
from diffusers import DiffusionPipeline
repo_id = "stable-diffusion-v1-5/stable-diffusion-v1-5"
pipeline = DiffusionPipeline.from_pretrained(repo_id, use_safetensors=True)
print(pipeline)
输出如下:
StableDiffusionPipeline {
"feature_extractor": [
"transformers",
"CLIPImageProcessor"
],
"safety_checker": [
"stable_diffusion",
"StableDiffusionSafetyChecker"
],
"scheduler": [
"diffusers",
"PNDMScheduler"
],
"text_encoder": [
"transformers",
"CLIPTextModel"
],
"tokenizer": [
"transformers",
"CLIPTokenizer"
],
"unet": [
"diffusers",
"UNet2DConditionModel"
],
"vae": [
"diffusers",
"AutoencoderKL"
]
}
管道加载的 7 个组件如下所示:
"feature_extractor":来自 Transformers 的一个CLIPImageProcessor。"safety_checker":用于筛选有害内容的组件。"scheduler":PNDMScheduler的一个实例。"text_encoder":来自 Transformers 的CLIPTextModel。"tokenizer":来自 Transformers 的CLIPTokenizer。"unet":UNet2DConditionModel的一个实例。"vae":AutoencoderKL的一个实例。
管道加载 7 个组件时会根据根目录下的 model_index.json 提供的信息实例化具体的组件类型,以当前模型的 model_index.json 为例:
{
"_class_name": "StableDiffusionPipeline",
"_diffusers_version": "0.6.0",
"feature_extractor": [
"transformers",
"CLIPImageProcessor"
],
"safety_checker": [
"stable_diffusion",
"StableDiffusionSafetyChecker"
],
"scheduler": [
"diffusers",
"PNDMScheduler"
],
"text_encoder": [
"transformers",
"CLIPTextModel"
],
"tokenizer": [
"transformers",
"CLIPTokenizer"
],
"unet": [
"diffusers",
"UNet2DConditionModel"
],
"vae": [
"diffusers",
"AutoencoderKL"
]
}
查看 stable-diffusion-v1-5/stable-diffusion-v1-5 目录结构,每个组件都有对应的文件,目录结构如下所示:
.
├── feature_extractor
│ └── preprocessor_config.json
├── model_index.json
├── safety_checker
│ ├── config.json
| ├── model.fp16.safetensors
│ ├── model.safetensors
│ ├── pytorch_model.bin
| └── pytorch_model.fp16.bin
├── scheduler
│ └── scheduler_config.json
├── text_encoder
│ ├── config.json
| ├── model.fp16.safetensors
│ ├── model.safetensors
│ |── pytorch_model.bin
| └── pytorch_model.fp16.bin
├── tokenizer
│ ├── merges.txt
│ ├── special_tokens_map.json
│ ├── tokenizer_config.json
│ └── vocab.json
├── unet
│ ├── config.json
│ ├── diffusion_pytorch_model.bin
| |── diffusion_pytorch_model.fp16.bin
│ |── diffusion_pytorch_model.f16.safetensors
│ |── diffusion_pytorch_model.non_ema.bin
│ |── diffusion_pytorch_model.non_ema.safetensors
│ └── diffusion_pytorch_model.safetensors
|── vae
. ├── config.json
. ├── diffusion_pytorch_model.bin
├── diffusion_pytorch_model.fp16.bin
├── diffusion_pytorch_model.fp16.safetensors
└── diffusion_pytorch_model.safetensors
加载社区管道和组件
加载 Hub 管道
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"google/ddpm-cifar10-32", custom_pipeline="hf-internal-testing/diffusers-dummy-pipeline", use_safetensors=True
)
加载本地文件
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5",
custom_pipeline="./path/to/pipeline_directory/",
clip_model=clip_model,
feature_extractor=feature_extractor,
use_safetensors=True,
)
加载主分支版本
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5",
custom_pipeline="clip_guided_stable_diffusion",
custom_revision="main",
clip_model=clip_model,
feature_extractor=feature_extractor,
use_safetensors=True,
)
加载历史版本
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5",
custom_pipeline="clip_guided_stable_diffusion",
custom_revision="v0.25.0",
clip_model=clip_model,
feature_extractor=feature_extractor,
use_safetensors=True,
)
加载调度器和模型
扩散模型的管道是一组可以调整的调度器和模型的组合。调度器封装了整个去噪过程,例如去噪步数和寻找去噪样本的算法。调度器没有参数,不需要训练,不会占用太多内存。模型通常只关注从噪声输入到较少噪声样本的前向传递。
1. 管道加载
import torch
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
).to("cuda")
2. 通过 pipeline.scheduler 属性查看管道使用的组件信息
pipeline.scheduler
使用的调度器为 PNDM:
PNDMScheduler {
"_class_name": "PNDMScheduler",
"_diffusers_version": "0.21.4",
"beta_end": 0.012,
"beta_schedule": "scaled_linear",
"beta_start": 0.00085,
"clip_sample": false,
"num_train_timesteps": 1000,
"set_alpha_to_one": false,
"skip_prk_steps": true,
"steps_offset": 1,
"timestep_spacing": "leading",
"trained_betas": null
}
3. 加载另一个调度器,以 DDIM 调度器为例
from diffusers import DDIMScheduler, DiffusionPipeline
ddim = DDIMScheduler.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="scheduler")
4. 使用 DDIM 调度器替换 pipeline 中的 PNDM 调度器
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", scheduler=ddim, torch_dtype=torch.float16, use_safetensors=True
).to("cuda")
5. 查看与管道兼容的调度器
pipeline.scheduler.compatibles
6. 对比多个调度器的效果
默认使用 PNDM 调度器:
import torch
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
).to("cuda")
prompt = "A photograph of an astronaut riding a horse on Mars, high resolution, high definition."
generator = torch.Generator(device="cuda").manual_seed(8)
使用 LMSDiscreteScheduler 调度器生成图片:
from diffusers import LMSDiscreteScheduler
pipeline.scheduler = LMSDiscreteScheduler.from_config(pipeline.scheduler.config)
image = pipeline(prompt, generator=generator).images[0]
image
使用 EulerDiscreteScheduler 调度器生成图片:
from diffusers import EulerDiscreteScheduler
pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config)
image = pipeline(prompt, generator=generator).images[0]
image
使用 EulerAncestralDiscreteScheduler 调度器生成图片:
from diffusers import EulerAncestralDiscreteScheduler
pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(pipeline.scheduler.config)
image = pipeline(prompt, generator=generator).images[0]
image
使用 DPMSolverMultistepScheduler 调度器生成图片:
from diffusers import DPMSolverMultistepScheduler
pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
image = pipeline(prompt, generator=generator).images[0]
image
7. 补充:从子文件夹中加载模型
从 unet 目录中加载模型:
from diffusers import UNet2DConditionModel
unet = UNet2DConditionModel.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="unet", use_safetensors=True)
从 unet 目录中加载模型变体:
from diffusers import UNet2DConditionModel
unet = UNet2DConditionModel.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="unet", variant="non_ema", use_safetensors=True
)
unet.save_pretrained("./local-unet", variant="non_ema")
模型文件和布局
加载 safetensors 文件
# 安装依赖
!pip install safetensors
多个 safetensors 文件
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5",
use_safetensors=True
)
单个 safetensors 文件
from diffusers import StableDiffusionPipeline
pipeline = StableDiffusionPipeline.from_single_file(
"https://huggingface.co/WarriorMama777/OrangeMixs/blob/main/Models/AbyssOrangeMix/AbyssOrangeMix.safetensors"
)
加载 LoRA 文件
通过 load_lora_weights() 方法加载到基础模型中:
from diffusers import StableDiffusionXLPipeline
import torch
# 加载基础模型
pipeline = StableDiffusionXLPipeline.from_pretrained(
"Lykon/dreamshaper-xl-1-0", torch_dtype=torch.float16, variant="fp16"
).to("cuda")
# 下载LoRA权重文件(以safetensors文件为例)
!wget https://civitai.com/api/download/models/168776 -O blueprintify.safetensors
# 加载LoRA权重文件
pipeline.load_lora_weights(".", weight_name="blueprintify.safetensors")
prompt = "bl3uprint, a highly detailed blueprint of the empire state building, explaining how to build all parts, many txt, blueprint grid backdrop"
negative_prompt = "lowres, cropped, worst quality, low quality, normal quality, artifacts, signature, watermark, username, blurry, more than one bridge, bad architecture"
image = pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
generator=torch.manual_seed(0),
).images[0]
image
加载 ckpt 文件
和加载单个 safetensors 文件一样:
from diffusers import StableDiffusionPipeline
pipeline = StableDiffusionPipeline.from_single_file(
"https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/blob/main/v1-5-pruned.ckpt"
)
推送文件至 Hub
远程登录 Hub
from huggingface_hub import notebook_login
notebook_login()
模型推送至 Hub
权重默认保存为 safetensors 格式文件。使用 push_to_hub() 方法:
from diffusers import ControlNetModel
controlnet = ControlNetModel(
block_out_channels=(32, 64),
layers_per_block=2,
in_channels=4,
down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),
cross_attention_dim=32,
conditioning_embedding_out_channels=(16, 32),
)
controlnet.push_to_hub("my-controlnet-model")
将权重变体推送至 Hub,以 fp16 为例:
controlnet.push_to_hub("my-controlnet-model", variant="fp16")
重新加载模型:
model = ControlNetModel.from_pretrained("your-namespace/my-controlnet-model")
调度器推送至 Hub
from diffusers import DDIMScheduler
scheduler = DDIMScheduler(
beta_start=0.00085,
beta_end=0.012,
beta_schedule="scaled_linear",
clip_sample=False,
set_alpha_to_one=False,
)
scheduler.push_to_hub("my-controlnet-scheduler")
重新加载调度器:
scheduler = DDIMScheduler.from_pretrained("your-namepsace/my-controlnet-scheduler")
Pipeline 推送至 Hub
整个管道及其组件都会被推送:
# 加载样例组件
from diffusers import (
UNet2DConditionModel,
AutoencoderKL,
DDIMScheduler,
StableDiffusionPipeline,
)
from transformers import CLIPTextModel, CLIPTextConfig, CLIPTokenizer
unet = UNet2DConditionModel(
block_out_channels=(32, 64),
layers_per_block=2,
sample_size=32,
in_channels=4,
out_channels=4,
down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),
up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"),
cross_attention_dim=32,
)
scheduler = DDIMScheduler(
beta_start=0.00085,
beta_end=0.012,
beta_schedule="scaled_linear",
clip_sample=False,
set_alpha_to_one=False,
)
vae = AutoencoderKL(
block_out_channels=[32, 64],
in_channels=3,
out_channels=3,
down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"],
up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"],
latent_channels=4,
)
text_encoder_config = CLIPTextConfig(
bos_token_id=0,
eos_token_id=2,
hidden_size=32,
intermediate_size=37,
layer_norm_eps=1e-05,
num_attention_heads=4,
num_hidden_layers=5,
pad_token_id=1,
vocab_size=1000,
)
text_encoder = CLIPTextModel(text_encoder_config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
构造 pipeline,推送至 Hub:
components = {
"unet": unet,
"scheduler": scheduler,
"vae": vae,
"text_encoder": text_encoder,
"tokenizer": tokenizer,
"safety_checker": None,
"feature_extractor": None,
}
pipeline = StableDiffusionPipeline(**components)
pipeline.push_to_hub("my-pipeline")
重新加载 Pipeline:
pipeline = StableDiffusionPipeline.from_pretrained("your-namespace/my-pipeline")
在 Hub 保存为 private 项目
指定 private 参数为 True:
controlnet.push_to_hub("my-controlnet-model-private", private=True)
适配器
LoRA
LoRA(低秩适配)是一种快速训练新任务模型的方法。通过冻结原始模型权重并添加少量新的可训练参数,将现有模型适配到新任务(例如以新风格生成图像)。LoRA 权重文件通常只有几百 MB 的大小,使用 load_lora_weights() 可以将这组权重加载到现有的基础模型。
【推荐】load_lora_weights() 方法加载 LoRA 权重
load_lora_weights() 方法会自动根据权重文件中的键(keys)进行模式匹配,分别加载 UNet 和 text_encoder 权重。
示例 1:文生图加载 LoRA 权重
import torch
from diffusers import AutoPipelineForText2Image
pipeline = AutoPipelineForText2Image.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16
).to("cuda")
pipeline.load_lora_weights(
"ostris/super-cereal-sdxl-lora",
weight_name="cereal_box_sdxl_v1.safetensors",
adapter_name="cereal" # 适配器的唯一名称,后续主要通过set_adapters(<adapter_name>)等方法,启用适配器
)
pipeline("bears, pizza bites").images[0]
示例 2:文本到视频加载 LoRA 权重
import torch
from diffusers import LTXConditionPipeline
from diffusers.utils import export_to_video, load_image
pipeline = LTXConditionPipeline.from_pretrained(
"Lightricks/LTX-Video-0.9.5", torch_dtype=torch.bfloat16
)
pipeline.load_lora_weights(
"Lightricks/LTX-Video-Cakeify-LoRA",
weight_name="ltxv_095_cakeify_lora.safetensors",
adapter_name="cakeify" # 适配器的唯一名称,后续主要通过set_adapters(<adapter_name>)等方法,启用适配器
)
pipeline.set_adapters("cakeify")
# use "CAKEIFY" to trigger the LoRA
prompt = "CAKEIFY a person using a knife to cut a cake shaped like a Pikachu plushie"
image = load_image("https://huggingface.co/Lightricks/LTX-Video-Cakeify-LoRA/resolve/main/assets/images/pikachu.png")
video = pipeline(
prompt=prompt,
image=image,
width=576,
height=576,
num_frames=161,
decode_timestep=0.03,
decode_noise_scale=0.025,
num_inference_steps=50,
).frames[0]
export_to_video(video, "output.mp4", fps=26)
说明:
load_lora_weights()方法可以将 LoRA 权重加载到 UNet 模型和文本编码器,它可以自动处理以下情况:
情况 i:LoRA 权重有区分 UNet 和文本编码器标识符
LoRA 权重文件中的所有键(keys)有明确标识符的情况(理想情况),如下所示:
lora_unet_down_blocks_0_resnets_0_conv1.alpha
lora_unet_down_blocks_0_resnets_0_conv1.lora_down.weight
lora_unet_down_blocks_0_resnets_0_conv1.lora_up.weight
lora_unet_down_blocks_0_attentions_0_proj_q.alpha
lora_unet_down_blocks_0_attentions_0_proj_q.lora_down.weight
lora_unet_down_blocks_0_attentions_0_proj_q.lora_up.weight
lora_te_text_model_encoder_layers_0_mlp_fc1.alpha
lora_te_text_model_encoder_layers_0_mlp_fc1.lora_down.weight
lora_te_text_model_encoder_layers_0_mlp_fc1.lora_up.weight
lora_te_text_model_encoder_layers_0_self_attn_q_proj.alpha
lora_te_text_model_encoder_layers_0_self_attn_q_proj.lora_down.weight
lora_te_text_model_encoder_layers_0_self_attn_q_proj.lora_up.weight
lora_unet* 会被识别为 UNet 权重,lora_te* 会被识别为 text_encoder 权重,此时 load_lora_weights() 方法会直接根据前缀将权重分别加载到 UNet 和 text_encoder 组件上。
情况 ii:LoRA 权重没有区分 UNet 和文本编码器的标识符
LoRA 权重文件中的键(keys)没有使用明确标识符的情况(模糊情况),例如,所有的键都是 lora_down* 形式:
lora_down.down_blocks_0_resnets_0_conv1.alpha
lora_down.down_blocks_0_resnets_0_conv1.lora_down.weight
lora_down.down_blocks_0_attentions_0_proj_q.alpha
lora_down.text_model_encoder_layers_0_mlp_fc1.alpha
lora_down.text_model_encoder_layers_0_self_attn_q_proj.alpha
没有 unet 或 te 这样的明确标识符,仅从 lora_down* 前缀无法区分权重属于哪个模型。此时,load_lora_weights() 方法会进一步解析键的剩余部分:
down_blocks_0_resnets_0_conv1这部分结构(包含down_blocks、resnets、conv1等)是 UNet 架构特有的命名方式;text_model_encoder_layers_0_mlp_fc1这部分结构(包含text_model_encoder_layers)是 CLIP Text Encoder 特有的命名方式;
以此为依据,load_lora_weights() 方法将不同的键对应的权重分别加载到 UNet 和 text_encoder 组件上。
load_lora_weights() 方法的智能解析过程
- 加载权重字典:读取
.safetensors或.bin文件,得到一个包含所有键值对的字典 - 模式匹配与分类:
- 查找是否有标准的、明确的前缀,例如
lora_unet.或lora_te.或lora_text_encoder. - 如果找到了明确的前缀,就按前缀分类
- 如果没找到明确的前缀,则遍历字典中的每个键,检查键的字符串是否包含能区分模型组件的子字符串或路径
- 查找是否有标准的、明确的前缀,例如
- 分离权重:
- 将属于 UNet 的权重分离出来,保存为一个字典
- 将属于 text encoder 的权重分离出来,保存为另一个字典
- 调用专用加载函数:
- 将 UNet 的权重字典传递给
unet.load_attn_procs() - 将 Text Encoder 的权重字典传递给
text_encoder.load_attn_procs()
- 将 UNet 的权重字典传递给
总结:
load_lora_weights()方法通过模式识别的方式区分权重:
- 前缀匹配(例如
lora_unet_、lora_te_)- 路径匹配(例如
down_blocks、text_model_encoder_layers)
load_lora_adapter() 加载 LoRA 权重
load_lora_adapter() 方法要求模型是 PeftAdapterMixin 子类,它会构建适配器所需的模型配置,将适配器加载到 UNet 中。该方法会忽略文本编码器相关的键,使用 prefix 指定的关键字 "unet" 过滤并加载符合要求的权重字典:
import torch
from diffusers import AutoPipelineForText2Image
pipeline = AutoPipelineForText2Image.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16
).to("cuda")
pipeline.unet.load_lora_adapter(
"jbilcke-hf/sdxl-cinematic-1",
weight_name="pytorch_lora_weights.safetensors",
adapter_name="cinematic",
prefix="unet"
)
# use cnmt in the prompt to trigger the LoRA
pipeline("A cute cnmt eating a slice of pizza, stunning color scheme, masterpiece, illustration").images[0]
torch.compile 通过编译加速推理
torch.compile 对模型进行编译,加速推理。在编译之前,需要先将 LoRA 权重融合到基础模型中并卸载:
import torch
from diffusers import DiffusionPipeline
# 加载基础模型和LoRA权重
pipeline = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16
).to("cuda")
pipeline.load_lora_weights(
"ostris/ikea-instructions-lora-sdxl",
weight_name="ikea_instructions_xl_v1_5.safetensors",
adapter_name="ikea"
)
# 激活LoRA,调整adapter权重
# adapter_weights控制 LoRA 适配器对原始模型的影响强度,最终的权重是:W_final = W_original + scale * ΔW_lora,LoRA 的更新量(ΔW)乘以 0.7,即只应用 70% 的微调效果
pipeline.set_adapters("ikea", adapter_weights=0.7)
# 融合LoRAs,卸载权重
# fuse_lora是将 LoRA 的低秩矩阵永久地融合(fuse)到原始模型的权重中,lora_scale=1.0 表示在融合时,LoRA 的权重按 1.0 的比例合并
# 融合后,模型的前向传播不在需要额外计算LoRA的增量,因此需要unload_lora_weights卸载LoRA权重
# 融合后,推理速度会更快,卸载后,会释放GPU显存
pipeline.fuse_lora(adapter_names=["ikea"], lora_scale=1.0)
pipeline.unload_lora_weights()
# 通常,编译UNet,因为它是管道中计算量最大的组件
# .to(memory_format=torch.channels_last)表示将张量布局改为channels_last,这种格式有时能被torch.compile更好的优化,进一步提升性能
pipeline.unet.to(memory_format=torch.channels_last)
# torch.compile是Pytorch 2.0以后引入的功能,可以将模型的前向计算图进行优化,生成更高效的代码,推理速度提升20%-50%
# mode="reduce-overhead"表示优化重点是减少内核启动开销,适合推理
# fullgraph=True表示整个unet前向传播编译成一个完整的计算图,避免中途退出编译模式
pipeline.unet = torch.compile(pipeline.unet, mode="reduce-overhead", fullgraph=True)
pipeline("A bowl of ramen shaped like a cute kawaii bear").images[0]
权重放缩
scale 参数用于控制应用多少比例的 LoRA 微调增量。0 表示仅使用基础模型权重,1 表示完全使用 LoRA 微调带来的增量。
示例 1:简单使用
将 cross_attention_kwargs={"scale": 1.0} 传递给管道:
import torch
from diffusers import AutoPipelineForText2Image
pipeline = AutoPipelineForText2Image.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16
).to("cuda")
pipeline.load_lora_weights(
"ostris/super-cereal-sdxl-lora",
weight_name="cereal_box_sdxl_v1.safetensors",
adapter_name="cereal"
)
pipeline("bears, pizza bites", cross_attention_kwargs={"scale": 1.0}).images[0]
示例 2:对 UNet 或文本编码器的每个独立组件进行更精细的控制
UNet 中的 "down" 模块被缩放 0.9,"up" 模块指定 "block_0" 和 "block_1" 中转换器的缩放比例。如果未指定,例如 "mid" 这样的模块,则使用默认值 1.0:
import torch
from diffusers import AutoPipelineForText2Image
pipeline = AutoPipelineForText2Image.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16
).to("cuda")
pipeline.load_lora_weights(
"ostris/super-cereal-sdxl-lora",
weight_name="cereal_box_sdxl_v1.safetensors",
adapter_name="cereal"
)
scales = {
"text_encoder": 0.5,
"text_encoder_2": 0.5,
"unet": {
"down": 0.9,
"up": {
"block_0": 0.6,
"block_1": [0.4, 0.8, 1.0],
}
}
}
pipeline.set_adapters("cereal", scales)
pipeline("bears, pizza bites").images[0]
示例 3:动态放缩
起初 LoRA 以较高的权重开始迭代,中期在最初的 20 步中逐渐衰减,后期仅应用 0.2 的比例,避免向 LoRA 未训练过的图像其他部分添加过多的 LoRA 特征:
import torch
from diffusers import FluxPipeline
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16
).to("cuda")
pipelne.load_lora_weights("alvarobartt/ghibli-characters-flux-lora", "lora")
num_inference_steps = 30
lora_steps = 20
lora_scales = torch.linspace(1.5, 0.7, lora_steps).tolist()
lora_scales += [0.2] * (num_inference_steps - lora_steps + 1)
pipeline.set_adapters("lora", lora_scales[0])
def callback(pipeline: FluxPipeline, step: int, timestep: torch.LongTensor, callback_kwargs: dict):
pipeline.set_adapters("lora", lora_scales[step + 1])
return callback_kwargs
prompt = """
Ghibli style The Grinch, a mischievous green creature with a sly grin, peeking out from behind a snow-covered tree while plotting his antics,
in a quaint snowy village decorated for the holidays, warm light glowing from cozy homes, with playful snowflakes dancing in the air
"""
pipeline(
prompt=prompt,
guidance_scale=3.0,
num_inference_steps=num_inference_steps,
generator=torch.Generator().manual_seed(42),
callback_on_step_end=callback,
).images[0]
示例 4:热插拔
假如第一次加载 LoRA 权重后进行编译,第二次重新加载 LoRA 权重后会自动重新编译:
import torch
from diffusers import DiffusionPipeline
# load base model and LoRAs
pipeline = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16
).to("cuda")
# 1. enable_lora_hotswap
# 在 enable_lora_hotswap()函数中, target_rank 设置为 max_rank 会将其设置为最高值。
# 对于秩不同的 LoRAs,你需要将其设置为多个LoRA中最高的秩值。默认的秩值是 128。
pipeline.enable_lora_hotswap(target_rank=max_rank)
pipeline.load_lora_weights(
"ostris/ikea-instructions-lora-sdxl",
weight_name="ikea_instructions_xl_v1_5.safetensors",
adapter_name="ikea"
)
# 2. torch.compile
pipeline.unet = torch.compile(pipeline.unet, mode="reduce-overhead", fullgraph=True)
# 3. hotswap
# 注意:使用相同的 adapter_name 参数来指定要替换的LoRA权重,此时会自动重新编译
pipeline.load_lora_weights(
"lordjia/by-feng-zikai",
hotswap=True,
adapter_name="ikea"
)
适配器合并
多个 LoRA 的权重可以进行合并,以生成混合风格图像。合并 LoRAs 有几种方法,每种方法在权重合并方式上有所不同(可能影响生成质量)。
set_adapters 简单合并
set_adapters() 方法通过指定的权重合并多个 LoRA:
# 将 LoRA 名称列表传递给 set_adapters(),并使用 adapter_weights 参数来控制每个 LoRA 的缩放。
# 例如,如果 adapter_weights=[0.5, 0.5] ,输出是两个 LoRAs 的平均值。
import torch
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16
).to("cuda")
pipeline.load_lora_weights(
"ostris/ikea-instructions-lora-sdxl",
weight_name="ikea_instructions_xl_v1_5.safetensors",
adapter_name="ikea"
)
pipeline.load_lora_weights(
"lordjia/by-feng-zikai",
weight_name="fengzikai_v1.0_XL.safetensors",
adapter_name="feng"
)
pipeline.set_adapters(["ikea", "feng"], adapter_weights=[0.7, 0.8])
# use by Feng Zikai to activate the lordjia/by-feng-zikai LoRA
pipeline("A bowl of ramen shaped like a cute kawaii bear, by Feng Zikai", cross_attention_kwargs={"scale": 1.0}).images[0]
fuse_lora 融合至基础模型
fuse_lora() 方法直接将 LoRA 权重与底层模型的原始 UNet 和文本编码器权重融合。这减少了每次加载 LoRA 时底层模型的加载开销,因为它只需加载一次模型,从而降低内存使用并提高推理速度:
import torch
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16
).to("cuda")
pipeline.load_lora_weights(
"ostris/ikea-instructions-lora-sdxl",
weight_name="ikea_instructions_xl_v1_5.safetensors",
adapter_name="ikea"
)
pipeline.load_lora_weights(
"lordjia/by-feng-zikai",
weight_name="fengzikai_v1.0_XL.safetensors",
adapter_name="feng"
)
pipeline.set_adapters(["ikea", "feng"], adapter_weights=[0.7, 0.8])
# 调用 fuse_lora() 将它们融合。通过lora_scale 参数控制 LoRA 权重的缩放。
pipeline.fuse_lora(adapter_names=["ikea", "feng"], lora_scale=1.0)
# 融合后需要卸载LoRA权重,因为它们已经和基础模型融合
pipeline.unload_lora_weights()
# 保存到本地
pipeline.save_pretrained("path/to/fused-pipeline")
# 保存到Hub
pipeline.push_to_hub("fused-ikea-feng")
# 重新加载
pipeline = DiffusionPipeline.from_pretrained(
"username/fused-ikea-feng", torch_dtype=torch.float16,
).to("cuda")
pipeline("A bowl of ramen shaped like a cute kawaii bear, by Feng Zikai").images[0]
# 恢复基础模型的权重,如果未生效,则需要重新加载基础模型
pipeline.unfuse_lora()
适配器管理
set_adapters:如果存在多个 LoRA,使用该方法激活指定的 LoRA
import torch
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16
).to("cuda")
pipeline.load_lora_weights(
"ostris/ikea-instructions-lora-sdxl",
weight_name="ikea_instructions_xl_v1_5.safetensors",
adapter_name="ikea"
)
pipeline.load_lora_weights(
"lordjia/by-feng-zikai",
weight_name="fengzikai_v1.0_XL.safetensors",
adapter_name="feng"
)
# activates the feng LoRA instead of the ikea LoRA
pipeline.set_adapters("feng")
save_lora_adapter:保存 LoRA 适配器
import torch
from diffusers import AutoPipelineForText2Image
pipeline = AutoPipelineForText2Image.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16
).to("cuda")
pipeline.unet.load_lora_adapter(
"jbilcke-hf/sdxl-cinematic-1",
weight_name="pytorch_lora_weights.safetensors",
adapter_name="cinematic",
prefix="unet"
)
pipeline.save_lora_adapter("path/to/save", adapter_name="cinematic")
unload_lora_weights:卸载 LoRA 权重,恢复基础模型权重
pipeline.unload_lora_weights()
disable_lora:禁用所有 LoRAs(但仍然保留在管道中),并恢复底层模型权重
pipeline.disable_lora()
get_active_adapters:返回管道上附加的激活 LoRAs 的列表
pipeline.get_active_adapters()
# ["cereal", "ikea"]
get_list_adapters:返回管道中每个组件上附加的激活 LoRAs 的列表
pipeline.get_list_adapters()
# {"unet": ["cereal", "ikea"], "text_encoder_2": ["cereal"]}
delete_adapters:从模型中移除一个 LoRA 及其层
pipeline.delete_adapters("ikea")
IP-Adapter
IP-Adapter(IP,image and prompt)是一款轻量级适配器,它允许模型不仅基于文本提示,还能结合图像特征进行更精确的控制。
IP-Adapter 原理
适配器模块的添加
- 插入点:IP-Adapter 通常在 UNet 模型中的某些层之间插入适配器模块,通常是注意力层或跨注意力层(cross-attention layers),因为这些层负责处理输入的不同部分之间的关系
- 结构:每个适配器模块包含一个小型神经网络,该网络设计用于处理特定类型的外部信息(例如 CLIP 图像编码器)。小型网络一般为几个线形成和激活函数,将外部输入转换为与原模型兼容的形式
外部信息的整合
- 图像特征提取:使用 IP-Adapter 时,需要先从参考图像中提取特征,通常是通过预训练的图像编码器(例如 CLIP 图像编码器)完成。提取参考图像的内容、风格等信息
- 特征映射:提取的图像特征随后被传递给适配器模块,处理并映射到 UNet 模型内部使用的空间维度。特征需要调整大小以匹配 UNet 中相应层的维度
对预训练模型的影响
- 附加路径:IP-Adapter 为模型提供了新的路径,使得除了文本提示以外,还可以提供图像特征,不改变原有模型权重,增加了模型处理的信息维度
- 动态调节:前向传播过程中,UNet 处理每一层时,适配器模块会根据当前层的状态和提供的图像特征生成相应的调整信号,这些信号被加权后加入到原始的中间表示中,从而影响最终的输出
- 参数控制:通过设置不同的超参数(例如适配器权重的比例)来控制图像特征对最终结果的影响程度
IP-Adapter 对交叉注意力层的影响
IP-Adapter 为交叉注意力层提供额外的参数进行加权处理。
- 位置:在 UNet 模型中,存在大量的 Cross-Attention 层,这些层通常位于 UNet 的瓶颈(bottleneck)和上采样(upsampling)阶段
- 功能:标准的 Cross-Attention 层负责将文本提示(text embeddings)注入到图像特征(latent features)中。计算图像特征(作为 Query)与文本嵌入(作为 Key/Value)之间的注意力权重,从而让生成过程”关注”文本描述
- 原始流程:
Query (latent) + Key/Value (text) → Attention Output - IP-Adapter 流程:
Query (latent) + Key/Value (text) → Text AttentionQuery (latent) + Key/Value (image features) → Image AttentionAttention Output = Text Attention + scale * Image Attention
IP-Adapter 加载
使用 load_ip_adapter() 用于加载 IP-Adapter。使用 set_ip_adapter_scale() 参数在生成过程中缩放 IP 适配器的影响。值为 1.0 表示模型仅以图像提示为条件,值为 0.5 通常会在文本和图像提示之间产生平衡的结果:
import torch
from diffusers import AutoPipelineForText2Image
from diffusers.utils import load_image
pipeline = AutoPipelineForText2Image.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16
).to("cuda")
pipeline.load_ip_adapter(
"h94/IP-Adapter",
subfolder="sdxl_models",
weight_name="ip-adapter_sdxl.bin"
)
pipeline.set_ip_adapter_scale(0.8)
image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/ip_adapter_diner.png")
pipeline(
prompt="a polar bear sitting in a chair drinking a milkshake",
ip_adapter_image=image,
negative_prompt="deformed, ugly, wrong proportion, low res, bad anatomy, worst quality, low quality",
).images[0]
IP-Adapter 适配下游任务
示例 1:使用 IP-Adapter 执行图像到图像的推理
import torch
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import load_image
pipeline = AutoPipelineForImage2Image.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16
).to("cuda")
pipeline.load_ip_adapter(
"h94/IP-Adapter",
subfolder="sdxl_models",
weight_name="ip-adapter_sdxl.bin"
)
pipeline.set_ip_adapter_scale(0.8)
image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/ip_adapter_bear_1.png")
ip_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/ip_adapter_gummy.png")
pipeline(
prompt="best quality, high quality",
image=image,
ip_adapter_image=ip_image,
strength=0.5,
).images[0]
示例 2:使用 IP-Adapter 执行图像修复的推理
import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image
pipeline = AutoPipelineForInpainting.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16
).to("cuda")
pipeline.load_ip_adapter(
"h94/IP-Adapter",
subfolder="sdxl_models",
weight_name="ip-adapter_sdxl.bin"
)
pipeline.set_ip_adapter_scale(0.6)
mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/ip_adapter_mask.png")
image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/ip_adapter_bear_1.png")
ip_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/ip_adapter_gummy.png")
pipeline(
prompt="a cute gummy bear waving",
image=image,
mask_image=mask_image,
ip_adapter_image=ip_image,
).images[0]
示例 3:使用 IP-Adapter 执行视频生成的推理
注意:
enable_model_cpu_offload()方法有助于减少内存占用,应在 IP 适配器加载后启用。否则,IP 适配器的图像编码器也会被卸载到 CPU 上并返回错误。
import torch
from diffusers import AnimateDiffPipeline, DDIMScheduler, MotionAdapter
from diffusers.utils import export_to_gif
from diffusers.utils import load_image
adapter = MotionAdapter.from_pretrained(
"guoyww/animatediff-motion-adapter-v1-5-2",
torch_dtype=torch.float16
)
pipeline = AnimateDiffPipeline.from_pretrained(
"emilianJR/epiCRealism",
motion_adapter=adapter,
torch_dtype=torch.float16
)
scheduler = DDIMScheduler.from_pretrained(
"emilianJR/epiCRealism",
subfolder="scheduler",
clip_sample=False,
timestep_spacing="linspace",
beta_schedule="linear",
steps_offset=1,
)
pipeline.scheduler = scheduler
pipeline.enable_vae_slicing()
pipeline.load_ip_adapter("h94/IP-Adapter", subfolder="models", weight_name="ip-adapter_sd15.bin")
pipeline.enable_model_cpu_offload()
ip_adapter_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/ip_adapter_inpaint.png")
pipeline(
prompt="A cute gummy bear waving",
negative_prompt="bad quality, worse quality, low resolution",
ip_adapter_image=ip_adapter_image,
num_frames=16,
guidance_scale=7.5,
num_inference_steps=50,
).frames[0]
模型变体
变体 Plus
Plus 版本使用面片嵌入和 ViT-H 图像编码器:
import torch
from transformers import CLIPVisionModelWithProjection, AutoPipelineForText2Image
image_encoder = CLIPVisionModelWithProjection.from_pretrained(
"h94/IP-Adapter",
subfolder="models/image_encoder",
torch_dtype=torch.float16
)
pipeline = AutoPipelineForText2Image.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
image_encoder=image_encoder,
torch_dtype=torch.float16
).to("cuda")
pipeline.load_ip_adapter(
"h94/IP-Adapter",
subfolder="sdxl_models",
weight_name="ip-adapter-plus_sdxl_vit-h.safetensors"
)
变体 FaceID
FaceID 版本使用 InsightFace 生成的人脸嵌入:
import torch
from transformers import AutoPipelineForText2Image
pipeline = AutoPipelineForText2Image.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16
).to("cuda")
pipeline.load_ip_adapter(
"h94/IP-Adapter-FaceID",
subfolder=None,
weight_name="ip-adapter-faceid_sdxl.bin",
image_encoder_folder=None
)
变体 FaceID Plus
需要加载 CLIP 图像编码器以及 CLIPVisionModelWithProjection:
from transformers import AutoPipelineForText2Image, CLIPVisionModelWithProjection
image_encoder = CLIPVisionModelWithProjection.from_pretrained(
"laion/CLIP-ViT-H-14-laion2B-s32B-b79K",
torch_dtype=torch.float16,
)
pipeline = AutoPipelineForText2Image.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5",
image_encoder=image_encoder,
torch_dtype=torch.float16
).to("cuda")
pipeline.load_ip_adapter(
"h94/IP-Adapter-FaceID",
subfolder=None,
weight_name="ip-adapter-faceid-plus_sd15.bin"
)
创建可复用的图像嵌入
prepare_ip_adapter_image_embeds 创建可复用的图像嵌入:
import torch
from diffusers import AutoPipelineForText2Image
pipeline = AutoPipelineForImage2Image.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16
).to("cuda")
# 通过将图像嵌入传递给 ip_adapter_image_embeds 参数来重新加载图像嵌入。将 image_encoder_folder 设置为 None ,因为您不再需要图像编码器来生成图像嵌入。
image_embeds = pipeline.prepare_ip_adapter_image_embeds(
ip_adapter_image=image,
ip_adapter_image_embeds=None,
device="cuda",
num_images_per_prompt=1,
do_classifier_free_guidance=True,
)
torch.save(image_embeds, "image_embeds.ipadpt")
加载并复用图像嵌入:
pipeline.load_ip_adapter(
"h94/IP-Adapter",
subfolder="sdxl_models",
image_encoder_folder=None,
weight_name="ip-adapter_sdxl.bin"
)
pipeline.set_ip_adapter_scale(0.8)
image_embeds = torch.load("image_embeds.ipadpt")
pipeline(
prompt="a polar bear sitting in a chair drinking a milkshake",
ip_adapter_image_embeds=image_embeds,
negative_prompt="deformed, ugly, wrong proportion, low res, bad anatomy, worst quality, low quality",
num_inference_steps=100,
generator=generator,
).images[0]
Masking 蒙版
二进制掩码可以将 IP 适配器映像分配到输出映像的特定区域,从而方便组合多个 IP 适配器映像。每个 IP 适配器映像都需要一个二进制掩码。
加载 IPAdapterMaskProcessor 来预处理图像蒙版。为了获得最佳效果,请提供输出 height 和 width,以确保不同宽高比的蒙版大小合适。如果输入蒙版已与生成图像的宽高比匹配,则无需设置 height 和 width:
import torch
from diffusers import AutoPipelineForText2Image
from diffusers.image_processor import IPAdapterMaskProcessor
from diffusers.utils import load_image
pipeline = AutoPipelineForImage2Image.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16
).to("cuda")
mask1 = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/ip_mask_mask1.png")
mask2 = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/ip_mask_mask2.png")
# 处理器IPAdapterMaskProcessor用于将图片进行批量处理
# 处理器使用list接收多个图片输入,调整为指定的height和width大小,归一化为[0,1]范围浮点数值,转换为PyTorch张量
# 输出张量masks的形状为[num_masks, 1, height, width],num_masks为蒙版数量(这里是2个),1是通道数(蒙版是单通道),height和width是指定的大小
processor = IPAdapterMaskProcessor()
masks = processor.preprocess([mask1, mask2], height=1024, width=1024)
以列表形式提供 IP 适配器图像及其比例。将预处理后的掩码传递给管道中的 cross_attention_kwargs:
face_image1 = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/ip_mask_girl1.png")
face_image2 = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/ip_mask_girl2.png")
pipeline.load_ip_adapter(
"h94/IP-Adapter",
subfolder="sdxl_models",
# load_ip_adapter支持加载多个适配器权重文件,所以这里的weight_name是一个list类型,这个safetensors文件名表示是IP-Adapter Plus模型,专门针对人脸优化的版本
weight_name=["ip-adapter-plus-face_sdxl_vit-h.safetensors"]
)
# 外层列表[]对应不同的IP-Adapter适配器权重,
# 内层列表[0.7, 0.7]对应每张参考图像的权重
pipeline.set_ip_adapter_scale([[0.7, 0.7]])
ip_images = [[face_image1, face_image2]]
# reshape是为了将masks和ip_images的形状对齐,确保每张参考图都有对应的掩码
# masks原本的形状是[2,1,1024,1024],reshape后调整为[1,2,1,1024,1024],即[batch=1, num_masks=2, channel=1, height, width]
masks = [masks.reshape(1, masks.shape[0], masks.shape[2], masks.shape[3])]
pipeline(
prompt="2 girls",
ip_adapter_image=ip_images,
negative_prompt="monochrome, lowres, bad anatomy, worst quality, low quality",
# 将处理好的掩码masks传递给模型的交叉注意力层
cross_attention_kwargs={"ip_adapter_masks": masks}
).images[0]
人脸模型
import torch
from diffusers import StableDiffusionPipeline, DDIMScheduler
from diffusers.utils import load_image
pipeline = StableDiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5",
torch_dtype=torch.float16,
).to("cuda")
pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config)
pipeline.load_ip_adapter(
"h94/IP-Adapter",
subfolder="models",
weight_name="ip-adapter-full-face_sd15.bin"
)
pipeline.set_ip_adapter_scale(0.5)
image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/ip_adapter_einstein_base.png")
pipeline(
prompt="A photo of Einstein as a chef, wearing an apron, cooking in a French restaurant",
ip_adapter_image=image,
negative_prompt="lowres, bad anatomy, worst quality, low quality",
num_inference_steps=100,
).images[0]
多个 IP-Adapter
使用多个 IP-Adapter 生成多样化的图像。例如,可以使用 IP-Adapter Face 来生成一致的面孔和角色,并使用 IP-Adapter Plus 来将这些面孔生成特定风格:
import torch
from diffusers import AutoPipelineForText2Image, DDIMScheduler
from transformers import CLIPVisionModelWithProjection
from diffusers.utils import load_image
image_encoder = CLIPVisionModelWithProjection.from_pretrained(
"h94/IP-Adapter",
subfolder="models/image_encoder",
torch_dtype=torch.float16,
)
# 加载基础模型、调度器和以下 IP-Adapters。
# 1. ip-adapter-plus_sdxl_vit-h 使用 patch 嵌入和一个 ViT-H 图像编码器
# 2. ip-adapter-plus-face_sdxl_vit-h 使用 patch 嵌入和一个 ViT-H 图像编码器,但它基于裁剪的面部图像进行条件化
pipeline = AutoPipelineForText2Image.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16,
image_encoder=image_encoder,
)
pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config)
pipeline.load_ip_adapter(
"h94/IP-Adapter",
subfolder="sdxl_models",
weight_name=["ip-adapter-plus_sdxl_vit-h.safetensors", "ip-adapter-plus-face_sdxl_vit-h.safetensors"]
)
pipeline.set_ip_adapter_scale([0.7, 0.3])
# enable_model_cpu_offload to reduce memory usage
pipeline.enable_model_cpu_offload()
# 加载一张图片和一个包含特定风格图片的文件夹以应用
face_image = load_image("https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/women_input.png")
style_folder = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/style_ziggy"
style_images = [load_image(f"{style_folder}/img{i}.png") for i in range(10)]
# 将风格和人脸图片作为列表传递给 ip_adapter_image
generator = torch.Generator(device="cpu").manual_seed(0)
pipeline(
prompt="wonderwoman",
ip_adapter_image=[style_images, face_image],
negative_prompt="monochrome, lowres, bad anatomy, worst quality, low quality",
).images[0]
即时生成
潜在一致性模型(LCM)可以在 4 步或更少的时间内生成图像,而其他扩散模型则需要更多步骤,使其感觉”即时”。IP-Adapters 与 LCM 模型兼容,可即时生成图像。
加载 IP-Adapter 权重,并使用 load_lora_weights() 加载 LoRA 权重:
import torch
from diffusers import DiffusionPipeline, LCMScheduler
from diffusers.utils import load_image
pipeline = DiffusionPipeline.from_pretrained(
"sd-dreambooth-library/herge-style",
torch_dtype=torch.float16
)
pipeline.load_ip_adapter(
"h94/IP-Adapter",
subfolder="models",
weight_name="ip-adapter_sd15.bin"
)
pipeline.load_lora_weights("latent-consistency/lcm-lora-sdv1-5")
pipeline.scheduler = LCMScheduler.from_config(pipeline.scheduler.config)
# enable_model_cpu_offload to reduce memory usage
pipeline.enable_model_cpu_offload()
# 尝试使用较低的 IP-Adapter 比例,以便更侧重于你想要应用的风格,并记得在提示中使用特殊标记来触发其生成。
pipeline.set_ip_adapter_scale(0.4)
prompt = "herge_style woman in armor, best quality, high quality"
ip_adapter_image = load_image("https://user-images.githubusercontent.com/24734142/266492875-2d50d223-8475-44f0-a7c6-08b51cb53572.png")
pipeline(
prompt=prompt,
ip_adapter_image=ip_adapter_image,
num_inference_steps=4,
guidance_scale=1,
).images[0]
结构控制
对于结构控制,将 IP-Adapter 与基于深度图、边缘图、姿态估计等的 ControlNet 结合使用。
加载一个基于深度图进行条件化的 ControlNetModel,并将其与 IP-Adapter 结合使用:
import torch
from diffusers.utils import load_image
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
controlnet = ControlNetModel.from_pretrained(
"lllyasviel/control_v11f1p_sd15_depth",
torch_dtype=torch.float16
)
pipeline = StableDiffusionControlNetPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5",
controlnet=controlnet,
torch_dtype=torch.float16
).to("cuda")
pipeline.load_ip_adapter(
"h94/IP-Adapter",
subfolder="models",
weight_name="ip-adapter_sd15.bin"
)
# 将深度图和 IP-Adapter 图像传递给管道
pipeline(
prompt="best quality, high quality",
image=depth_map,
ip_adapter_image=ip_adapter_image,
negative_prompt="monochrome, lowres, bad anatomy, worst quality, low quality",
).images[0]
风格和布局控制
为了进行风格和布局控制,请将 IP-Adapter 与 InstantStyle 结合使用。InstantStyle 将风格(颜色、纹理、整体感觉)与内容分离。它仅在模型的特定风格区块中应用风格,以防止其扭曲图像的其他区域。这能生成具有更强、更一致风格以及更好布局控制效果的图像。
IP-Adapter 仅对模型的特定部分起作用。使用 set_ip_adapter_scale() 方法来调整 IP-Adapter 在不同层中的影响。下面的示例在模型第 2 层(向下 block_2 层和向上 block_0 层)激活 IP-Adapter。向下 block_2 层是 IP-Adapter 注入布局信息的地方,向上 block_0 层是注入风格的地方:
import torch
from diffusers import AutoPipelineForText2Image
from diffusers.utils import load_image
pipeline = AutoPipelineForText2Image.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16
).to("cuda")
pipeline.load_ip_adapter(
"h94/IP-Adapter",
subfolder="sdxl_models",
weight_name="ip-adapter_sdxl.bin"
)
scale = {
"down": {"block_2": [0.0, 1.0]},
"up": {"block_0": [0.0, 1.0, 0.0]},
}
pipeline.set_ip_adapter_scale(scale)
# 加载风格图像并生成图像。
style_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg")
pipeline(
prompt="a cat, masterpiece, best quality, high quality",
ip_adapter_image=style_image,
negative_prompt="text, watermark, lowres, low quality, worst quality, deformed, glitch, low contrast, noisy, saturation, blurry",
guidance_scale=5,
).images[0]
你也可以将 IP-Adapter 插入所有模型层。这往往会生成更专注于图像提示的图像,并可能减少生成图像的多样性。仅在 up block_0 或风格层中激活 IP-Adapter:
scale = {
"up": {"block_0": [0.0, 1.0, 0.0]},
}
pipeline.set_ip_adapter_scale(scale)
pipeline(
prompt="a cat, masterpiece, best quality, high quality",
ip_adapter_image=style_image,
negative_prompt="text, watermark, lowres, low quality, worst quality, deformed, glitch, low contrast, noisy, saturation, blurry",
guidance_scale=5,
).images[0]
ControlNet
ControlNet 是一个能够实现可控生成的适配器,例如生成特定姿势的猫的图像或遵循特定猫的素描线条。它通过添加一个较小的”零卷积”层网络,并逐步训练这些层以避免干扰原始模型来工作。原始模型参数被冻结以避免重新训练它。
ControlNet 基于额外的视觉信息或”结构控制”(如边缘检测、深度图、人体姿态等)进行条件化,这些信息可以与文本提示结合,生成由视觉输入引导的图像。
文生图
使用 opencv-python 生成一个 canny 图像:
import cv2
import numpy as np
from PIL import Image
from diffusers.utils import load_image
original_image = load_image(
"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/non-enhanced-prompt.png"
)
image = np.array(original_image)
low_threshold = 100
high_threshold = 200
image = cv2.Canny(image, low_threshold, high_threshold)
image = image[:, :, None]
image = np.concatenate([image, image, image], axis=2)
canny_image = Image.fromarray(image)
将 canny 图像传递给管道。使用 controlnet_conditioning_scale 参数来确定分配给控制的权重:
import torch
from diffusers.utils import load_image
from diffusers import FluxControlNetPipeline, FluxControlNetModel
controlnet = FluxControlNetModel.from_pretrained(
"InstantX/FLUX.1-dev-Controlnet-Canny", torch_dtype=torch.bfloat16
)
pipeline = FluxControlNetPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev", controlnet=controlnet, torch_dtype=torch.bfloat16
).to("cuda")
prompt = """
A photorealistic overhead image of a cat reclining sideways in a flamingo pool floatie holding a margarita.
The cat is floating leisurely in the pool and completely relaxed and happy.
"""
pipeline(
prompt,
control_image=canny_image,
controlnet_conditioning_scale=0.5,
num_inference_steps=50,
guidance_scale=3.5,
).images[0]
图生图
使用 Transformer 生成深度图,并从深度估计流程中生成:
import torch
import numpy as np
from PIL import Image
from transformers import DPTImageProcessor, DPTForDepthEstimation
from diffusers import ControlNetModel, StableDiffusionXLControlNetImg2ImgPipeline, AutoencoderKL
from diffusers.utils import load_image
depth_estimator = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to("cuda")
feature_extractor = DPTImageProcessor.from_pretrained("Intel/dpt-hybrid-midas")
def get_depth_map(image):
image = feature_extractor(images=image, return_tensors="pt").pixel_values.to("cuda")
with torch.no_grad(), torch.autocast("cuda"):
depth_map = depth_estimator(image).predicted_depth
depth_map = torch.nn.functional.interpolate(
depth_map.unsqueeze(1),
size=(1024, 1024),
mode="bicubic",
align_corners=False,
)
depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True)
depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True)
depth_map = (depth_map - depth_min) / (depth_max - depth_min)
image = torch.cat([depth_map] * 3, dim=1)
image = image.permute(0, 2, 3, 1).cpu().numpy()[0]
image = Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8))
return image
depth_image = get_depth_map(image)
将深度图传递给流程。使用 controlnet_conditioning_scale 参数来确定分配给控制的权重:
controlnet = ControlNetModel.from_pretrained(
"diffusers/controlnet-depth-sdxl-1.0-small",
torch_dtype=torch.float16,
)
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
pipeline = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
controlnet=controlnet,
vae=vae,
torch_dtype=torch.float16,
).to("cuda")
prompt = """
A photorealistic overhead image of a cat reclining sideways in a flamingo pool floatie holding a margarita.
The cat is floating leisurely in the pool and completely relaxed and happy.
"""
image = load_image(
"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/non-enhanced-prompt.png"
).resize((1024, 1024))
controlnet_conditioning_scale = 0.5
pipeline(
prompt,
image=image,
control_image=depth_image,
controlnet_conditioning_scale=controlnet_conditioning_scale,
strength=0.99,
num_inference_steps=100,
).images[0]
图片修复
生成一个掩码图像,并将其转换为张量,以标记原始图像中如果掩码图像的对应像素超过某个阈值,则将这些像素标记为已遮罩:
import cv2
import torch
import numpy as np
from PIL import Image
from diffusers.utils import load_image
from diffusers import StableDiffusionXLControlNetInpaintPipeline, ControlNetModel
init_image = load_image(
"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/non-enhanced-prompt.png"
)
init_image = init_image.resize((1024, 1024))
mask_image = load_image(
"/content/cat_mask.png"
)
mask_image = mask_image.resize((1024, 1024))
def make_canny_condition(image):
image = np.array(image)
image = cv2.Canny(image, 100, 200)
image = image[:, :, None]
image = np.concatenate([image, image, image], axis=2)
image = Image.fromarray(image)
return image
control_image = make_canny_condition(init_image)
将遮罩图像和控制图像传递给管道。使用 controlnet_conditioning_scale 参数来确定分配给控制的权重:
controlnet = ControlNetModel.from_pretrained(
"diffusers/controlnet-canny-sdxl-1.0", torch_dtype=torch.float16
)
pipeline = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", controlnet=controlnet, torch_dtype=torch.float16
)
pipeline(
"a cute and fluffy bunny rabbit",
num_inference_steps=100,
strength=0.99,
controlnet_conditioning_scale=0.5,
image=init_image,
mask_image=mask_image,
control_image=control_image,
).images[0]
多个 ControlNet
组合多个 ControlNet,例如边缘检测图像和深度图,以创建一个 MultiControlNet。将 ControlNets 作为列表传递给管道,并将图像调整到预期的输入大小:
import torch
from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel, AutoencoderKL
controlnets = [
ControlNetModel.from_pretrained(
"diffusers/controlnet-depth-sdxl-1.0-small", torch_dtype=torch.float16
),
ControlNetModel.from_pretrained(
"diffusers/controlnet-canny-sdxl-1.0", torch_dtype=torch.float16,
),
]
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
pipeline = StableDiffusionXLControlNetPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", controlnet=controlnets, vae=vae, torch_dtype=torch.float16
).to("cuda")
prompt = """
a relaxed rabbit sitting on a striped towel next to a pool with a tropical drink nearby,
bright sunny day, vacation scene, 35mm photograph, film, professional, 4k, highly detailed
"""
negative_prompt = "lowres, bad anatomy, worst quality, low quality, deformed, ugly"
images = [canny_image.resize((1024, 1024)), depth_image.resize((1024, 1024))]
pipeline(
prompt,
negative_prompt=negative_prompt,
image=images,
num_inference_steps=100,
controlnet_conditioning_scale=[0.5, 0.5],
strength=0.7,
).images[0]
guess_mode 猜测模式
猜测模式仅根据控制输入(如边缘检测图、深度图、姿态等)生成图像,且不受提示词的指导。它通过固定的比例根据区块深度调整 ControlNet 输出残差的规模。早期的 DownBlock 仅被缩放 0.1,而 MidBlock 则被完全缩放 1.0:
import torch
from diffusers.utils import load_image
from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
controlnet = ControlNetModel.from_pretrained(
"diffusers/controlnet-canny-sdxl-1.0", torch_dtype=torch.float16
)
pipeline = StableDiffusionXLControlNetPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
controlnet=controlnet,
torch_dtype=torch.float16
).to("cuda")
canny_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/canny-cat.png")
pipeline(
"",
image=canny_image,
guess_mode=True
).images[0]
T2I-Adapter
T2I-Adapter 是一个能够实现可控生成的适配器,类似于 ControlNet。T2I-Adapter 通过学习控制信号(例如深度图)与预训练模型内部知识之间的映射来工作。该适配器插入到基础模型中,在生成过程中根据控制信号提供额外的指导。
T2I-Adapter 加载
加载一个针对特定控制(如 canny 边缘)进行条件化的 T2I-Adapter,并通过 from_pretrained() 将其传递给管道:
import torch
from diffusers import T2IAdapter, StableDiffusionXLAdapterPipeline, AutoencoderKL
t2i_adapter = T2IAdapter.from_pretrained(
"TencentARC/t2i-adapter-canny-sdxl-1.0",
torch_dtype=torch.float16,
)
使用 opencv-python 生成一个 canny 图像:
import cv2
import numpy as np
from PIL import Image
from diffusers.utils import load_image
original_image = load_image(
"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/non-enhanced-prompt.png"
)
image = np.array(original_image)
low_threshold = 100
high_threshold = 200
image = cv2.Canny(image, low_threshold, high_threshold)
image = image[:, :, None]
image = np.concatenate([image, image, image], axis=2)
canny_image = Image.fromarray(image)
将 canny 图像传递给管道以生成图像:
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
pipeline = StableDiffusionXLAdapterPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
adapter=t2i_adapter,
vae=vae,
torch_dtype=torch.float16,
).to("cuda")
prompt = """
A photorealistic overhead image of a cat reclining sideways in a flamingo pool floatie holding a margarita.
The cat is floating leisurely in the pool and completely relaxed and happy.
"""
pipeline(
prompt,
image=canny_image,
num_inference_steps=100,
guidance_scale=10,
).images[0]
MultiAdapter
使用 MultiAdapter 类将多个控制组合起来,例如 canny 图像和深度图。将控制图像和 T2I-Adapters 作为列表加载:
import torch
from diffusers.utils import load_image
from diffusers import StableDiffusionXLAdapterPipeline, AutoencoderKL, MultiAdapter, T2IAdapter
canny_image = load_image(
"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/canny-cat.png"
)
depth_image = load_image(
"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/sdxl_depth_image.png"
)
controls = [canny_image, depth_image]
prompt = ["""
a relaxed rabbit sitting on a striped towel next to a pool with a tropical drink nearby,
bright sunny day, vacation scene, 35mm photograph, film, professional, 4k, highly detailed
"""]
adapters = MultiAdapter(
[
T2IAdapter.from_pretrained("TencentARC/t2i-adapter-canny-sdxl-1.0", torch_dtype=torch.float16),
T2IAdapter.from_pretrained("TencentARC/t2i-adapter-depth-midas-sdxl-1.0", torch_dtype=torch.float16),
]
)
将适配器、提示和控制图像传递给 StableDiffusionXLAdapterPipeline。使用 adapter_conditioning_scale 参数来确定为每个控制分配多少权重:
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
pipeline = StableDiffusionXLAdapterPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16,
vae=vae,
adapter=adapters,
).to("cuda")
pipeline(
prompt,
image=controls,
height=1024,
width=1024,
adapter_conditioning_scale=[0.7, 0.7]
).images[0]
DreamBooth
DreamBooth 是一种生成特定实例个性化图像的方法。它通过在 3-5 张与唯一标识符(sks cat)关联的物体图像(例如,一只猫)上微调模型来工作。这使你能够在提示中使用 sks cat 来触发模型生成你猫在不同场景、光照、姿势和风格下的图像。
DreamBooth 检查点通常有几 GB 的大小,因为它包含了完整的模型权重。使用 from_pretrained() 加载 DreamBooth 检查点,并在提示中包含唯一标识符以激活其生成功能:
import torch
from diffusers import AutoPipelineForText2Image
pipeline = AutoPipelineForText2Image.from_pretrained(
"sd-dreambooth-library/herge-style",
torch_dtype=torch.float16
).to("cuda")
prompt = "A cute sks herge_style brown bear eating a slice of pizza, stunning color scheme, masterpiece, illustration"
pipeline(prompt).images[0]
Textual Inversion(文本反转)
文本反转是一种生成概念个性化图像的方法。它通过在 3-5 张与唯一标记(<sks>)相关联的概念图像(例如,像素艺术)上微调模型的词嵌入来工作。这使您能够在提示中使用 <sks> 标记来触发模型生成像素艺术图像。
文本反转权重非常轻量级,通常只有几 KB,因为它们只是词嵌入。然而,这也意味着词嵌入需要在使用 from_pretrained() 加载模型后加载:
import torch
from diffusers import AutoPipelineForText2Image
pipeline = AutoPipelineForText2Image.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5",
torch_dtype=torch.float16
).to("cuda")
# 使用 load_textual_inversion()加载词嵌入,并在提示中包含唯一标记以激活其生成。
pipeline.load_textual_inversion("sd-concepts-library/gta5-artwork")
prompt = "A cute brown bear eating a slice of pizza, stunning color scheme, masterpiece, illustration, <gta5-artwork> style"
pipeline(prompt).images[0]
负向嵌入
文本反转还可以训练以学习负嵌入,以避免生成不希望的特征,例如”模糊”或”丑陋”。它对于提高图像质量很有用。EasyNegative 是一种广泛使用的负向嵌入,包含多个学习到的负向概念。加载负向嵌入,并指定与负向嵌入相关联的文件名和标记。将标记传递给 negative_prompt 以在您的管道中激活它:
import torch
from diffusers import AutoPipelineForText2Image
pipeline = AutoPipelineForText2Image.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5",
torch_dtype=torch.float16
).to("cuda")
pipeline.load_textual_inversion(
"EvilEngine/easynegative",
weight_name="easynegative.safetensors",
token="easynegative"
)
prompt = "A cute brown bear eating a slice of pizza, stunning color scheme, masterpiece, illustration"
negative_prompt = "easynegative"
pipeline(prompt, negative_prompt).images[0]
生成任务
无约束的图片生成
无条件图像生成生成的图像看起来像是模型训练所用训练数据的随机样本,因为去噪过程没有受到任何额外上下文(如文本或图像)的指导。
使用 DiffusionPipeline 加载 anton-l/ddpm-butterflies-128 检查点来生成蝴蝶图像。DiffusionPipeline 会下载并缓存生成图像所需的所有模型组件:
from diffusers import DiffusionPipeline
generator = DiffusionPipeline.from_pretrained("anton-l/ddpm-butterflies-128").to("cuda")
image = generator().images[0]
image
输出图像是一个 PIL.Image 对象,可以保存:
image.save("generated_image.png")
你也可以尝试调整 num_inference_steps 参数,该参数控制去噪步数。更多的去噪步数通常会生成更高质量的图像,但生成时间会更长:
image = generator(num_inference_steps=100).images[0]
image
文本生成图像
当你想到扩散模型时,通常首先想到的是文本到图像。文本到图像根据文本描述(例如,“丛林中的宇航员,冷色调,柔和色彩,细节丰富,8k”)生成图像,这被称为提示。
从非常高的层面上讲,扩散模型接收一个提示和一些随机的初始噪声,并迭代地去除噪声来构建图像。去噪过程由提示引导,一旦在预定的时间步数后去噪过程结束,图像表示就被解码成图像。
您可以在 Diffusers 中通过两步从提示词生成图像:
a. 将检查点加载到 AutoPipelineForText2Image 类中,该类会根据检查点自动检测要使用的适当管道类:
from diffusers import AutoPipelineForText2Image
import torch
pipeline = AutoPipelineForText2Image.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16"
).to("cuda")
b. 将提示词传递给管道以生成图像:
image = pipeline(
"stained glass of darth vader, backlight, centered composition, masterpiece, photorealistic, 8k"
).images[0]
image
流行模型
最常见的文本到图像模型是 Stable Diffusion v1.5、Stable Diffusion XL (SDXL) 和 Kandinsky 2.2。还有一些 ControlNet 模型或适配器可以与文本到图像模型一起使用,以便在生成图像时进行更直接的控制。由于它们的架构和训练过程不同,每个模型的结果略有差异,但无论你选择哪个模型,它们的用法大致相同。让我们使用相同的提示词对每个模型进行测试,并比较它们的结果。
Stable Diffusion v1.5
Stable Diffusion v1.5 是一个从 Stable Diffusion v1-4 初始化的潜在扩散模型,并在 LAION-Aesthetics V2 数据集的 512x512 图像上进行了 595K 步的微调:
from diffusers import AutoPipelineForText2Image
import torch
pipeline = AutoPipelineForText2Image.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16"
).to("cuda")
generator = torch.Generator("cuda").manual_seed(31)
image = pipeline("Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", generator=generator).images[0]
image
Stable Diffusion XL
SDXL 是之前 Stable Diffusion 模型的更大版本,它涉及一个两阶段模型过程,向图像中添加更多细节。它还包括一些额外的微观条件调节,以生成以中心主题的高质量图像:
from diffusers import AutoPipelineForText2Image
import torch
pipeline = AutoPipelineForText2Image.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16"
).to("cuda")
generator = torch.Generator("cuda").manual_seed(31)
image = pipeline("Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", generator=generator).images[0]
image
Kandinsky 2.2
Kandinsky 模型与 Stable Diffusion 模型略有不同,因为它也使用图像先验模型来创建嵌入,这些嵌入用于在扩散模型中更好地对齐文本和图像:
from diffusers import AutoPipelineForText2Image
import torch
pipeline = AutoPipelineForText2Image.from_pretrained(
"kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16
).to("cuda")
generator = torch.Generator("cuda").manual_seed(31)
image = pipeline("Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", generator=generator).images[0]
image
ControlNet
ControlNet 模型是辅助模型或适配器,在文本到图像模型(如 Stable Diffusion v1.5)的基础上进行微调。使用 ControlNet 模型与文本到图像模型结合,为生成图像提供了更多显式的控制选项。通过 ControlNet,你向模型添加一个额外的条件输入图像。例如,如果你将一张人体姿态图像(通常表示为多个连接成骨架的关键点)作为条件输入,模型将生成一个遵循该图像姿态的图像。
在这个例子中,让我们用人体姿态估计图像来对 ControlNet 进行条件化。加载在人体姿态估计上预训练的 ControlNet 模型:
from diffusers import ControlNetModel, AutoPipelineForText2Image
from diffusers.utils import load_image
import torch
controlnet = ControlNetModel.from_pretrained(
"lllyasviel/control_v11p_sd15_openpose", torch_dtype=torch.float16, variant="fp16"
).to("cuda")
pose_image = load_image("https://huggingface.co/lllyasviel/control_v11p_sd15_openpose/resolve/main/images/control.png")
将 controlnet 传递给 AutoPipelineForText2Image,并提供提示和姿态估计图像:
pipeline = AutoPipelineForText2Image.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16, variant="fp16"
).to("cuda")
generator = torch.Generator("cuda").manual_seed(31)
image = pipeline("Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", image=pose_image, generator=generator).images[0]
image
配置管道参数
高度和宽度
height 和 width 参数控制生成图像的高度和宽度(以像素为单位)。默认情况下,Stable Diffusion v1.5 模型输出 512x512 的图像,但您可以将其更改为任何 8 的倍数尺寸。例如,要创建一个矩形图像:
from diffusers import AutoPipelineForText2Image
import torch
pipeline = AutoPipelineForText2Image.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16"
).to("cuda")
image = pipeline(
"Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", height=768, width=512
).images[0]
image
引导尺度
guidance_scale 参数影响提示词对图像生成的影响程度。较低的值给模型”创造力”,以生成与提示词更松散关联的图像。较高的 guidance_scale 值促使模型更紧密地遵循提示词,如果这个值过高,你可能会在生成的图像中观察到一些伪影:
from diffusers import AutoPipelineForText2Image
import torch
pipeline = AutoPipelineForText2Image.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16
).to("cuda")
image = pipeline(
"Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", guidance_scale=3.5
).images[0]
image
负面提示
就像提示引导生成一样,负面提示会引导模型避开你不希望模型生成的内容。这通常用于通过移除”低分辨率”或”糟糕的细节”等糟糕或不好的图像特征来提高整体图像质量。你也可以使用负面提示来移除或修改图像的内容和风格:
from diffusers import AutoPipelineForText2Image
import torch
pipeline = AutoPipelineForText2Image.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16
).to("cuda")
image = pipeline(
prompt="Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
negative_prompt="ugly, deformed, disfigured, poor details, bad anatomy",
).images[0]
image
生成器
一个 torch.Generator 对象通过设置手动种子,使管道中的可重复性得以实现。您可以使用 Generator 生成一批图像,并迭代改进从种子生成的图像。使用 Generator 创建图像时,每次都应返回相同的结果,而不是随机生成新图像:
from diffusers import AutoPipelineForText2Image
import torch
pipeline = AutoPipelineForText2Image.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16
).to("cuda")
generator = torch.Generator(device="cuda").manual_seed(30)
image = pipeline(
"Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
generator=generator,
).images[0]
image
控制图像生成
提示词权重
提示词权重是一种技术,用于增加或减少提示词中概念的重要性,以强调或最小化图像中的某些特征。我们建议使用 Compel 库来帮助您生成加权提示词嵌入。
创建嵌入后,您可以将它们传递到管道中的 prompt_embeds(如果您使用负提示,则传递到 negative_prompt_embeds)参数:
from diffusers import AutoPipelineForText2Image
import torch
pipeline = AutoPipelineForText2Image.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16
).to("cuda")
image = pipeline(
prompt_embeds=prompt_embeds, # generated from Compel
negative_prompt_embeds=negative_prompt_embeds, # generated from Compel
).images[0]
ControlNet
正如您在 ControlNet 部分所见,这些模型通过结合额外的条件图像输入,提供了一种更灵活和准确生成图像的方法。每个 ControlNet 模型都在特定类型的条件图像上进行预训练,以生成与其相似的新图像。例如,如果您使用在深度图上预训练的 ControlNet 模型,可以将深度图作为条件输入提供给模型,模型将生成保留其中空间信息的图像。这比在提示中指定深度信息更快、更容易。您甚至可以使用 MultiControlNet 组合多个条件输入!
有多种类型的条件输入可供使用,Diffusers 支持用于 Stable Diffusion 和 SDXL 模型的 ControlNet。
优化
扩散模型体积庞大,图像去噪的迭代过程计算成本高且资源密集。但这并不意味着你需要强大的——甚至多个——GPU 才能使用它们。有许多优化技术可以在消费级和免费级资源上运行扩散模型。例如,你可以以半精度加载模型权重以节省 GPU 内存并提高速度,或者将整个模型卸载到 GPU 以节省更多内存。
PyTorch 2.0 还支持一种更内存高效的注意力机制,称为缩放点积注意力,如果你使用 PyTorch 2.0,它将自动启用。你可以将其与 torch.compile 结合使用,以进一步加速你的代码:
from diffusers import AutoPipelineForText2Image
import torch
pipeline = AutoPipelineForText2Image.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16").to("cuda")
pipeline.unet = torch.compile(pipeline.unet, mode="reduce-overhead", fullgraph=True)
视频生成
视频生成模型扩展了图像生成(可以视为 1 帧视频),以处理与空间和时间相关的数据。确保所有这些数据——文本、空间、时间——在帧与帧之间保持一致和协调,是生成长且高分辨率视频中的一个重大挑战。
现代视频模型通过扩散 Transformer(DiT)架构应对这一挑战。这降低了计算成本,并允许更高效地扩展到更大和更高质量的图像和视频数据。
Wan2.1
# pip install ftfy
import torch
import numpy as np
from diffusers import AutoModel, WanPipeline
from diffusers.hooks.group_offloading import apply_group_offloading
from diffusers.utils import export_to_video, load_image
from transformers import UMT5EncoderModel
text_encoder = UMT5EncoderModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="text_encoder", torch_dtype=torch.bfloat16)
vae = AutoModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="vae", torch_dtype=torch.float32)
transformer = AutoModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16)
# group-offloading
onload_device = torch.device("cuda")
offload_device = torch.device("cpu")
apply_group_offloading(text_encoder,
onload_device=onload_device,
offload_device=offload_device,
offload_type="block_level",
num_blocks_per_group=4
)
transformer.enable_group_offload(
onload_device=onload_device,
offload_device=offload_device,
offload_type="leaf_level",
use_stream=True
)
pipeline = WanPipeline.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers",
vae=vae,
transformer=transformer,
text_encoder=text_encoder,
torch_dtype=torch.bfloat16
)
pipeline.to("cuda")
prompt = """
The camera rushes from far to near in a low-angle shot,
revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in
for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground.
Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts dynamic
shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
"""
negative_prompt = """
Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality,
low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured,
misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards
"""
output = pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
num_frames=81,
guidance_scale=5.0,
).frames[0]
export_to_video(output, "output.mp4", fps=16)
HunyuanVideo
import torch
from diffusers import AutoModel, HunyuanVideoPipeline
from diffusers.quantizers import PipelineQuantizationConfig
from diffusers.utils import export_to_video
# quantize weights to int4 with bitsandbytes
pipeline_quant_config = PipelineQuantizationConfig(
quant_backend="bitsandbytes_4bit",
quant_kwargs={
"load_in_4bit": True,
"bnb_4bit_quant_type": "nf4",
"bnb_4bit_compute_dtype": torch.bfloat16
},
components_to_quantize=["transformer"]
)
pipeline = HunyuanVideoPipeline.from_pretrained(
"hunyuanvideo-community/HunyuanVideo",
quantization_config=pipeline_quant_config,
torch_dtype=torch.bfloat16,
)
# model-offloading and tiling
pipeline.enable_model_cpu_offload()
pipeline.vae.enable_tiling()
prompt = "A fluffy teddy bear sits on a bed of soft pillows surrounded by children's toys."
video = pipeline(prompt=prompt, num_frames=61, num_inference_steps=30).frames[0]
export_to_video(video, "output.mp4", fps=15)
LTX-Video
import torch
from diffusers import LTXPipeline, AutoModel
from diffusers.hooks import apply_group_offloading
from diffusers.utils import export_to_video
# fp8 layerwise weight-casting
transformer = AutoModel.from_pretrained(
"Lightricks/LTX-Video",
subfolder="transformer",
torch_dtype=torch.bfloat16
)
transformer.enable_layerwise_casting(
storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16
)
pipeline = LTXPipeline.from_pretrained("Lightricks/LTX-Video", transformer=transformer, torch_dtype=torch.bfloat16)
# group-offloading
onload_device = torch.device("cuda")
offload_device = torch.device("cpu")
pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", use_stream=True)
apply_group_offloading(pipeline.text_encoder, onload_device=onload_device, offload_type="block_level", num_blocks_per_group=2)
apply_group_offloading(pipeline.vae, onload_device=onload_device, offload_type="leaf_level")
prompt = """
A woman with long brown hair and light skin smiles at another woman with long blonde hair. The woman with brown hair wears a black jacket and has a small, barely noticeable mole on her right cheek. The camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm and natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be real-life footage
"""
negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted"
video = pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
width=768,
height=512,
num_frames=161,
decode_timestep=0.03,
decode_noise_scale=0.025,
num_inference_steps=50,
).frames[0]
export_to_video(video, "output.mp4", fps=24)
管道参数
管道中有几个参数需要配置,这些参数会影响视频生成的质量或速度。尝试不同的参数值对于发现适当的质量和速度权衡非常重要。
num_frames 帧数
一帧是静止图像,它在其他帧的序列中播放,以创建运动或视频。使用 num_frames 控制每秒生成的帧数。增加 num_frames 会提高感知的运动平滑度和视觉连贯性,这对具有动态内容视频尤为重要。较高的 num_frames 值也会增加视频时长。
某些视频模型需要更具体的 num_frames 值进行推理。例如,HunyuanVideoPipeline 推荐使用 (4 * num_frames) + 1 计算 num_frames。始终检查管道的 API 模型卡,以查看是否有推荐值。
import torch
from diffusers import LTXPipeline
from diffusers.utils import export_to_video
pipeline = LTXPipeline.from_pretrained(
"Lightricks/LTX-Video", torch_dtype=torch.bfloat16
).to("cuda")
prompt = """
A woman with long brown hair and light skin smiles at another woman with long blonde hair. The woman
with brown hair wears a black jacket and has a small, barely noticeable mole on her right cheek. The
camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm and
natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be
real-life footage
"""
video = pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
width=768,
height=512,
num_frames=161,
decode_timestep=0.03,
decode_noise_scale=0.025,
num_inference_steps=50,
).frames[0]
export_to_video(video, "output.mp4", fps=24)
guidance_scale 引导尺度
引导尺度或”cfg”控制生成帧与输入条件(文本、图像或两者)的贴合程度。增加 guidance_scale 会使生成的帧更贴近输入条件,包含更精细的细节,但可能引入伪影并减少输出多样性。较低 guidance_scale 的值会鼓励更松散的提示遵循和增加输出多样性,但细节可能不够精细。如果它太低,可能会完全忽略你的提示并生成随机噪声。
import torch
from diffusers import CogVideoXPipeline, CogVideoXTransformer3DModel
from diffusers.utils import export_to_video
pipeline = CogVideoXPipeline.from_pretrained(
"THUDM/CogVideoX-2b",
torch_dtype=torch.float16
).to("cuda")
prompt = """
A detailed wooden toy ship with intricately carved masts and sails is seen gliding smoothly over
a plush, blue carpet that mimics the waves of the sea. The ship's hull is painted a rich brown,
with tiny windows. The carpet, soft and textured, provides a perfect backdrop, resembling an
oceanic expanse. Surrounding the ship are various other toys and children's items, hinting at
a playful environment. The scene captures the innocence and imagination of childhood,
with the toy ship's journey symbolizing endless adventures in a whimsical, indoor setting.
"""
video = pipeline(
prompt=prompt,
guidance_scale=6,
num_inference_steps=50
).frames[0]
export_to_video(video, "output.mp4", fps=8)
文本引导深度到图像生成
StableDiffusionDepth2ImgPipeline 允许您传递文本提示和初始图像来调节新图像的生成。此外,您还可以传递 depth_map 来保留图像结构。如果没有提供 depth_map,该管道会自动通过集成的深度估计模型预测深度。
首先创建一个 StableDiffusionDepth2ImgPipeline 的实例:
import torch
from diffusers import StableDiffusionDepth2ImgPipeline
from diffusers.utils import load_image, make_image_grid
pipeline = StableDiffusionDepth2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-depth",
torch_dtype=torch.float16,
use_safetensors=True,
).to("cuda")
现在将您的提示传递给该管道。您还可以传递 negative_prompt 来防止某些词语指导图像的生成:
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
init_image = load_image(url)
prompt = "two tigers"
negative_prompt = "bad, deformed, ugly, bad anatomy"
image = pipeline(prompt=prompt, image=init_image, negative_prompt=negative_prompt, strength=0.7).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
推理技术
总览
推理管道支持两类技术:
- 管道功能:这些技术修改管道或将其扩展用于其他应用。例如,管道回调为管道添加新功能,管道也可以扩展用于分布式推理。
- 提高推理质量:这些技术增加了生成图像的视觉质量。例如,您可以使用 GPT2 增强提示,以较低的努力创建更好的图像。
创建服务
扩散器的管道可以用作服务器的推理引擎。它支持并发和多线程请求,以同时为多个用户生成图像。
示例中使用 StableDiffusion3Pipeline:
- 导航到
examples/server文件夹并安装所有依赖项。
pip install .
pip install -f requirements.txt
- 使用以下命令启动服务器。
python server.py
- 服务器可通过
http://localhost:8000访问。您可以使用以下命令curl此模型。
curl -X POST -H "Content-Type: application/json" --data '{"model": "something", "prompt": "a kitten in front of a fireplace"}' http://localhost:8000/v1/images/generations
分布式推理
展示如何使用 Accelerate 和 PyTorch Distributed 进行分布式推理。
Accelerate
Accelerate 是一个旨在简化分布式设置中训练或运行推理的库。它简化了设置分布式环境的流程,让您能够专注于您的 PyTorch 代码。
- 首先,创建一个 Python 文件并初始化
accelerate.PartialState来创建分布式环境;设置将自动检测,无需显式定义rank或world_size。将DiffusionPipeline移动到distributed_state.device以分配给每个进程一个 GPU。
使用 split_between_processes 实用程序作为上下文管理器,自动将提示分配给进程数。
import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)
with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
result = pipeline(prompt).images[0]
result.save(f"result_{distributed_state.process_index}.png")
- 使用
--num_processes参数指定要使用的 GPU 数量,并调用accelerate launch来运行脚本:
accelerate launch run_distributed.py --num_processes=2
PyTorch 分布式
PyTorch 支持 DistributedDataParallel,它能够实现数据并行。
- 首先,创建一个 Python 文件并导入
torch.distributed和torch.multiprocessing,以设置分布式进程组并在每个 GPU 上启动推理进程。你还应该初始化一个DiffusionPipeline:
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from diffusers import DiffusionPipeline
sd = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
- 创建一个用于运行推理的函数;
init_process_group负责创建分布式环境,包括使用的后端类型、当前进程的rank以及参与进程的world_size或进程数量。如果你在 2 个 GPU 上并行运行推理,那么world_size的值就是 2。
将 DiffusionPipeline 移至 rank,并使用 get_rank 为每个进程分配 GPU,其中每个进程处理不同的提示:
def run_inference(rank, world_size):
dist.init_process_group("nccl", rank=rank, world_size=world_size)
sd.to(rank)
if torch.distributed.get_rank() == 0:
prompt = "a dog"
elif torch.distributed.get_rank() == 1:
prompt = "a cat"
image = sd(prompt).images[0]
image.save(f"./{'_'.join(prompt)}.png")
- 要运行分布式推理,调用
mp.spawn在world_size定义的 GPU 数量上运行run_inference函数:
def main():
world_size = 2
mp.spawn(run_inference, args=(world_size,), nprocs=world_size, join=True)
if __name__ == "__main__":
main()
- 完成推理脚本后,使用
--nproc_per_node参数指定要使用的 GPU 数量,并调用torchrun运行脚本:
torchrun run_distributed.py --nproc_per_node=2
Model sharding 模型分片
现代扩散系统如 Flux 规模非常大,包含多个模型。例如,Flux.1-Dev 由两个文本编码器——T5-XXL 和 CLIP-L、一个扩散变换器以及一个 VAE 组成。对于如此规模的模型,在消费级 GPU 上运行推理会面临挑战。
模型分片是一种将模型分布到多个 GPU 上的技术,当模型无法在单个 GPU 上运行时使用。以下示例假设有两块 16GB 的 GPU 用于推理。
- 首先使用文本编码器计算文本嵌入。通过设置
device_map="balanced"将文本编码器保留在两个 GPU 上。balanced策略将模型均匀分布在所有可用的 GPU 上。使用max_memory参数为每个 GPU 上的文本编码器分配最大内存量。
# 仅加载此步骤的文本编码器!扩散 Transformer 和 VAE 将在后续步骤中加载以节省内存。
from diffusers import FluxPipeline
import torch
prompt = "a photo of a dog with cat-like look"
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
transformer=None,
vae=None,
device_map="balanced",
max_memory={0: "16GB", 1: "16GB"},
torch_dtype=torch.bfloat16
)
with torch.no_grad():
print("Encoding prompts.")
prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
prompt=prompt, prompt_2=None, max_sequence_length=512
)
- 一旦文本嵌入计算完成,将其从 GPU 中移除,为扩散 Transformer 腾出空间。
import gc
def flush():
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
torch.cuda.reset_peak_memory_stats()
del pipeline.text_encoder
del pipeline.text_encoder_2
del pipeline.tokenizer
del pipeline.tokenizer_2
del pipeline
flush()
- 加载扩散 Transformer,它有 125 亿参数。这次将
device_map="auto"设置为自动将模型分布在两个 16GB 的 GPU 上。auto策略由 Accelerate 支持,作为 Big Model Inference 功能的一部分提供。它首先将模型分布到最快的设备(GPU)上,如果需要,再移动到较慢的设备如 CPU 和硬盘。在较慢的设备上存储模型参数的权衡是推理延迟变慢。
from diffusers import AutoModel
import torch
transformer = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
device_map="auto",
torch_dtype=torch.bfloat16
)
# 任何时候,你可以尝试 print(pipeline.hf_device_map) 来查看各种模型如何在设备间分布。这对于跟踪模型的设备放置很有用。
# 你也可以尝试 print(transformer.hf_device_map) 来查看 Transformer 模型如何在设备间分片。
- 将转换器模型添加到管道中进行去噪,但将其他模型级组件(如文本编码器和 VAE)设置为
None,因为你现在不需要它们。
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
text_encoder=None,
text_encoder_2=None,
tokenizer=None,
tokenizer_2=None,
vae=None,
transformer=transformer,
torch_dtype=torch.bfloat16
)
print("Running denoising.")
height, width = 768, 1360
latents = pipeline(
prompt_embeds=prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
num_inference_steps=50,
guidance_scale=3.5,
height=height,
width=width,
output_type="latent",
).images
# 由于不再需要,从内存中移除管道和转换器。
del pipeline.transformer
del pipeline
flush()
- 最后,使用 VAE 将潜在变量解码为图像。VAE 通常足够小,可以加载到单个 GPU 上。
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor
import torch
vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", torch_dtype=torch.bfloat16).to("cuda")
vae_scale_factor = 2 ** (len(vae.config.block_out_channels))
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
with torch.no_grad():
print("Running decoding.")
latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
image = vae.decode(latents, return_dict=False)[0]
image = image_processor.postprocess(image, output_type="pil")
image[0].save("split_transformer.png")
- 通过在特定阶段选择性地加载和卸载所需的模型,并将最大的模型跨多个 GPU 分片,可以在消费级 GPU 上运行大型模型的推理。
文本引导深度到图像生成
StableDiffusionDepth2ImgPipeline 允许您传递文本提示和初始图像来调节新图像的生成。此外,您还可以传递 depth_map 来保留图像结构。如果没有提供 depth_map,该管道会自动通过集成的深度估计模型预测深度。
首先创建一个 StableDiffusionDepth2ImgPipeline 的实例:
import torch
from diffusers import StableDiffusionDepth2ImgPipeline
from diffusers.utils import load_image, make_image_grid
pipeline = StableDiffusionDepth2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-depth",
torch_dtype=torch.float16,
use_safetensors=True,
).to("cuda")
现在将您的提示传递给该管道。您还可以传递 negative_prompt 来防止某些词语指导图像的生成:
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
init_image = load_image(url)
prompt = "two tigers"
negative_prompt = "bad, deformed, ugly, bad anatomy"
image = pipeline(prompt=prompt, image=init_image, negative_prompt=negative_prompt, strength=0.7).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
推理技术
总览
推理管道支持两类技术:
- 管道功能:这些技术修改管道或将其扩展用于其他应用。例如,管道回调为管道添加新功能,管道也可以扩展用于分布式推理。
- 提高推理质量:这些技术增加了生成图像的视觉质量。例如,您可以使用 GPT2 增强提示,以较低的努力创建更好的图像。
创建服务
扩散器的管道可以用作服务器的推理引擎。它支持并发和多线程请求,以同时为多个用户生成图像。
示例中使用 StableDiffusion3Pipeline:
- 导航到
examples/server文件夹并安装所有依赖项。
pip install .
pip install -f requirements.txt
- 使用以下命令启动服务器。
python server.py
- 服务器可通过
http://localhost:8000访问。您可以使用以下命令curl此模型。
curl -X POST -H "Content-Type: application/json" --data '{"model": "something", "prompt": "a kitten in front of a fireplace"}' http://localhost:8000/v1/images/generations
分布式推理
展示如何使用 Accelerate 和 PyTorch Distributed 进行分布式推理。
Accelerate
Accelerate 是一个旨在简化分布式设置中训练或运行推理的库。它简化了设置分布式环境的流程,让您能够专注于您的 PyTorch 代码。
- 首先,创建一个 Python 文件并初始化
accelerate.PartialState来创建分布式环境;设置将自动检测,无需显式定义rank或world_size。将DiffusionPipeline移动到distributed_state.device以分配给每个进程一个 GPU。
使用 split_between_processes 实用程序作为上下文管理器,自动将提示分配给进程数。
import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)
with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
result = pipeline(prompt).images[0]
result.save(f"result_{distributed_state.process_index}.png")
- 使用
--num_processes参数指定要使用的 GPU 数量,并调用accelerate launch来运行脚本:
accelerate launch run_distributed.py --num_processes=2
PyTorch 分布式
PyTorch 支持 DistributedDataParallel,它能够实现数据并行。
- 首先,创建一个 Python 文件并导入
torch.distributed和torch.multiprocessing,以设置分布式进程组并在每个 GPU 上启动推理进程。你还应该初始化一个DiffusionPipeline:
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from diffusers import DiffusionPipeline
sd = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
- 创建一个用于运行推理的函数;
init_process_group负责创建分布式环境,包括使用的后端类型、当前进程的rank以及参与进程的world_size或进程数量。如果你在 2 个 GPU 上并行运行推理,那么world_size的值就是 2。
将 DiffusionPipeline 移至 rank,并使用 get_rank 为每个进程分配 GPU,其中每个进程处理不同的提示:
def run_inference(rank, world_size):
dist.init_process_group("nccl", rank=rank, world_size=world_size)
sd.to(rank)
if torch.distributed.get_rank() == 0:
prompt = "a dog"
elif torch.distributed.get_rank() == 1:
prompt = "a cat"
image = sd(prompt).images[0]
image.save(f"./{'_'.join(prompt)}.png")
- 要运行分布式推理,调用
mp.spawn在world_size定义的 GPU 数量上运行run_inference函数:
def main():
world_size = 2
mp.spawn(run_inference, args=(world_size,), nprocs=world_size, join=True)
if __name__ == "__main__":
main()
- 完成推理脚本后,使用
--nproc_per_node参数指定要使用的 GPU 数量,并调用torchrun运行脚本:
torchrun run_distributed.py --nproc_per_node=2
Model sharding 模型分片
现代扩散系统如 Flux 规模非常大,包含多个模型。例如,Flux.1-Dev 由两个文本编码器——T5-XXL 和 CLIP-L、一个扩散变换器以及一个 VAE 组成。对于如此规模的模型,在消费级 GPU 上运行推理会面临挑战。
模型分片是一种将模型分布到多个 GPU 上的技术,当模型无法在单个 GPU 上运行时使用。以下示例假设有两块 16GB 的 GPU 用于推理。
- 首先使用文本编码器计算文本嵌入。通过设置
device_map="balanced"将文本编码器保留在两个 GPU 上。balanced策略将模型均匀分布在所有可用的 GPU 上。使用max_memory参数为每个 GPU 上的文本编码器分配最大内存量。
# 仅加载此步骤的文本编码器!扩散 Transformer 和 VAE 将在后续步骤中加载以节省内存。
from diffusers import FluxPipeline
import torch
prompt = "a photo of a dog with cat-like look"
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
transformer=None,
vae=None,
device_map="balanced",
max_memory={0: "16GB", 1: "16GB"},
torch_dtype=torch.bfloat16
)
with torch.no_grad():
print("Encoding prompts.")
prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
prompt=prompt, prompt_2=None, max_sequence_length=512
)
- 一旦文本嵌入计算完成,将其从 GPU 中移除,为扩散 Transformer 腾出空间。
import gc
def flush():
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
torch.cuda.reset_peak_memory_stats()
del pipeline.text_encoder
del pipeline.text_encoder_2
del pipeline.tokenizer
del pipeline.tokenizer_2
del pipeline
flush()
- 加载扩散 Transformer,它有 125 亿参数。这次将
device_map="auto"设置为自动将模型分布在两个 16GB 的 GPU 上。auto策略由 Accelerate 支持,作为 Big Model Inference 功能的一部分提供。它首先将模型分布到最快的设备(GPU)上,如果需要,再移动到较慢的设备如 CPU 和硬盘。在较慢的设备上存储模型参数的权衡是推理延迟变慢。
from diffusers import AutoModel
import torch
transformer = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
device_map="auto",
torch_dtype=torch.bfloat16
)
# 任何时候,你可以尝试 print(pipeline.hf_device_map) 来查看各种模型如何在设备间分布。这对于跟踪模型的设备放置很有用。
# 你也可以尝试 print(transformer.hf_device_map) 来查看 Transformer 模型如何在设备间分片。
- 将转换器模型添加到管道中进行去噪,但将其他模型级组件(如文本编码器和 VAE)设置为
None,因为你现在不需要它们。
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
text_encoder=None,
text_encoder_2=None,
tokenizer=None,
tokenizer_2=None,
vae=None,
transformer=transformer,
torch_dtype=torch.bfloat16
)
print("Running denoising.")
height, width = 768, 1360
latents = pipeline(
prompt_embeds=prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
num_inference_steps=50,
guidance_scale=3.5,
height=height,
width=width,
output_type="latent",
).images
# 由于不再需要,从内存中移除管道和转换器。
del pipeline.transformer
del pipeline
flush()
- 最后,使用 VAE 将潜在变量解码为图像。VAE 通常足够小,可以加载到单个 GPU 上。
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor
import torch
vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", torch_dtype=torch.bfloat16).to("cuda")
vae_scale_factor = 2 ** (len(vae.config.block_out_channels))
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
with torch.no_grad():
print("Running decoding.")
latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
image = vae.decode(latents, return_dict=False)[0]
image = image_processor.postprocess(image, output_type="pil")
image[0].save("split_transformer.png")
- 通过在特定阶段选择性地加载和卸载所需的模型,并将最大的模型跨多个 GPU 分片,可以在消费级 GPU 上运行大型模型的推理。
negative_prompt 负面提示
负面提示有助于排除您不希望在生成视频中看到的内容。它通常用于通过将模型从”模糊、扭曲、丑陋”等不希望出现的元素中推开,来提高生成视频的质量和一致性。这可以创建更干净、更专注的视频。
# pip install ftfy
import torch
from diffusers import WanPipeline
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
from diffusers.utils import export_to_video
vae = AutoencoderKLWan.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="vae", torch_dtype=torch.float32
)
pipeline = WanPipeline.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers", vae=vae, torch_dtype=torch.bfloat16
)
pipeline.scheduler = UniPCMultistepScheduler.from_config(
pipeline.scheduler.config, flow_shift=5.0
)
pipeline.to("cuda")
pipeline.load_lora_weights("benjamin-paine/steamboat-willie-14b", adapter_name="steamboat-willie")
pipeline.set_adapters("steamboat-willie")
pipeline.enable_model_cpu_offload()
# use "steamboat willie style" to trigger the LoRA
prompt = """
steamboat willie style, golden era animation, The camera rushes from far to near in a low-angle shot,
revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in
for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground.
Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts
dynamic shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
"""
output = pipeline(
prompt=prompt,
num_frames=81,
guidance_scale=5.0,
).frames[0]
export_to_video(output, "output.mp4", fps=16)
减少内存使用
近期具有百亿以上参数的视频模型,如 HunyuanVideoPipeline 和 WanPipeline,需要大量内存,且常常超出消费级硬件的可用内存。Diffusers 提供了多种技术来降低这些大型模型的内存需求。
组卸载
其中一种技术是组卸载,它在不使用时将内部模型层组(例如 torch.nn.Sequential)卸载到 CPU。这些层仅在需要计算时才会被加载,以避免将所有模型组件存储在 GPU 上。对于像 WanPipeline 这样的百四十亿参数模型,组卸载可将所需内存降低至约 13GB 显存。
# pip install ftfy
import torch
import numpy as np
from diffusers import AutoModel, WanPipeline
from diffusers.hooks.group_offloading import apply_group_offloading
from diffusers.utils import export_to_video, load_image
from transformers import UMT5EncoderModel
text_encoder = UMT5EncoderModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="text_encoder", torch_dtype=torch.bfloat16)
vae = AutoModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="vae", torch_dtype=torch.float32)
transformer = AutoModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16)
# group-offloading
onload_device = torch.device("cuda")
offload_device = torch.device("cpu")
apply_group_offloading(text_encoder,
onload_device=onload_device,
offload_device=offload_device,
offload_type="block_level",
num_blocks_per_group=4
)
transformer.enable_group_offload(
onload_device=onload_device,
offload_device=offload_device,
offload_type="leaf_level",
use_stream=True
)
pipeline = WanPipeline.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers",
vae=vae,
transformer=transformer,
text_encoder=text_encoder,
torch_dtype=torch.bfloat16
)
pipeline.to("cuda")
prompt = """
The camera rushes from far to near in a low-angle shot,
revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in
for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground.
Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts dynamic
shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
"""
negative_prompt = """
Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality,
low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured,
misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards
"""
output = pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
num_frames=81,
guidance_scale=5.0,
).frames[0]
export_to_video(output, "output.mp4", fps=16)
模型量化
减少内存的另一种选择是考虑对模型进行量化,即将模型权重存储在低精度的数据类型中。然而,量化可能会根据具体的视频模型影响视频质量。请参考量化概述,了解更多关于不同支持量化的后端信息。
下面的示例使用 bitsandbytes 对模型进行量化。
# pip install ftfy
import torch
from diffusers import AutoModel, WanPipeline
from diffusers.quantizers import PipelineQuantizationConfig
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
from transformers import UMT5EncoderModel
from diffusers.utils import export_to_video
# quantize transformer and text encoder weights with bitsandbytes
pipeline_quant_config = PipelineQuantizationConfig(
quant_backend="bitsandbytes_4bit",
quant_kwargs={"load_in_4bit": True},
components_to_quantize=["transformer", "text_encoder"]
)
vae = AutoModel.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="vae", torch_dtype=torch.float32
)
pipeline = WanPipeline.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers", vae=vae, quantization_config=pipeline_quant_config, torch_dtype=torch.bfloat16
)
pipeline.scheduler = UniPCMultistepScheduler.from_config(
pipeline.scheduler.config, flow_shift=5.0
)
pipeline.to("cuda")
pipeline.load_lora_weights("benjamin-paine/steamboat-willie-14b", adapter_name="steamboat-willie")
pipeline.set_adapters("steamboat-willie")
pipeline.enable_model_cpu_offload()
# use "steamboat willie style" to trigger the LoRA
prompt = """
steamboat willie style, golden era animation, The camera rushes from far to near in a low-angle shot,
revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in
for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground.
Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts
dynamic shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
"""
output = pipeline(
prompt=prompt,
num_frames=81,
guidance_scale=5.0,
).frames[0]
export_to_video(output, "output.mp4", fps=16)
推理速度
torch.compile 可以通过使用优化过的内核来加速推理。第一次编译需要更长时间,但一旦编译完成,就会快得多。最好是一次编译整个管道,然后多次使用该管道而不做任何更改。任何更改,例如图像大小,都会触发重新编译。
下面的示例编译了管道中的转换器,并使用 "max-autotune" 模式来最大化性能。
import torch
from diffusers import CogVideoXPipeline, CogVideoXTransformer3DModel
from diffusers.utils import export_to_video
pipeline = CogVideoXPipeline.from_pretrained(
"THUDM/CogVideoX-2b",
torch_dtype=torch.float16
).to("cuda")
# torch.compile
pipeline.transformer.to(memory_format=torch.channels_last)
pipeline.transformer = torch.compile(
pipeline.transformer, mode="max-autotune", fullgraph=True
)
prompt = """
A detailed wooden toy ship with intricately carved masts and sails is seen gliding smoothly over a plush, blue carpet that mimics the waves of the sea.
The ship's hull is painted a rich brown, with tiny windows. The carpet, soft and textured, provides a perfect backdrop, resembling an oceanic expanse.
Surrounding the ship are various other toys and children's items, hinting at a playful environment. The scene captures the innocence and imagination of childhood,
with the toy ship's journey symbolizing endless adventures in a whimsical, indoor setting.
"""
video = pipeline(
prompt=prompt,
guidance_scale=6,
num_inference_steps=50
).frames[0]
export_to_video(video, "output.mp4", fps=8)
文本引导深度到图像生成
StableDiffusionDepth2ImgPipeline 允许您传递文本提示和初始图像来调节新图像的生成。此外,您还可以传递 depth_map 来保留图像结构。如果没有提供 depth_map,该管道会自动通过集成的深度估计模型预测深度。
首先创建一个 StableDiffusionDepth2ImgPipeline 的实例:
import torch
from diffusers import StableDiffusionDepth2ImgPipeline
from diffusers.utils import load_image, make_image_grid
pipeline = StableDiffusionDepth2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-depth",
torch_dtype=torch.float16,
use_safetensors=True,
).to("cuda")
现在将您的提示传递给该管道。您还可以传递 negative_prompt 来防止某些词语指导图像的生成:
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
init_image = load_image(url)
prompt = "two tigers"
negative_prompt = "bad, deformed, ugly, bad anatomy"
image = pipeline(prompt=prompt, image=init_image, negative_prompt=negative_prompt, strength=0.7).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
推理技术
总览
推理管道支持两类技术:
- 管道功能:这些技术修改管道或将其扩展用于其他应用。例如,管道回调为管道添加新功能,管道也可以扩展用于分布式推理。
- 提高推理质量:这些技术增加了生成图像的视觉质量。例如,您可以使用 GPT2 增强提示,以较低的努力创建更好的图像。
创建服务
扩散器的管道可以用作服务器的推理引擎。它支持并发和多线程请求,以同时为多个用户生成图像。
示例中使用 StableDiffusion3Pipeline:
- 导航到
examples/server文件夹并安装所有依赖项。
pip install .
pip install -f requirements.txt
- 使用以下命令启动服务器。
python server.py
- 服务器可通过
http://localhost:8000访问。您可以使用以下命令curl此模型。
curl -X POST -H "Content-Type: application/json" --data '{"model": "something", "prompt": "a kitten in front of a fireplace"}' http://localhost:8000/v1/images/generations
分布式推理
展示如何使用 Accelerate 和 PyTorch Distributed 进行分布式推理。
Accelerate
Accelerate 是一个旨在简化分布式设置中训练或运行推理的库。它简化了设置分布式环境的流程,让您能够专注于您的 PyTorch 代码。
- 首先,创建一个 Python 文件并初始化
accelerate.PartialState来创建分布式环境;设置将自动检测,无需显式定义rank或world_size。将DiffusionPipeline移动到distributed_state.device以分配给每个进程一个 GPU。
使用 split_between_processes 实用程序作为上下文管理器,自动将提示分配给进程数。
import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)
with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
result = pipeline(prompt).images[0]
result.save(f"result_{distributed_state.process_index}.png")
- 使用
--num_processes参数指定要使用的 GPU 数量,并调用accelerate launch来运行脚本:
accelerate launch run_distributed.py --num_processes=2
PyTorch 分布式
PyTorch 支持 DistributedDataParallel,它能够实现数据并行。
- 首先,创建一个 Python 文件并导入
torch.distributed和torch.multiprocessing,以设置分布式进程组并在每个 GPU 上启动推理进程。你还应该初始化一个DiffusionPipeline:
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from diffusers import DiffusionPipeline
sd = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
- 创建一个用于运行推理的函数;
init_process_group负责创建分布式环境,包括使用的后端类型、当前进程的rank以及参与进程的world_size或进程数量。如果你在 2 个 GPU 上并行运行推理,那么world_size的值就是 2。
将 DiffusionPipeline 移至 rank,并使用 get_rank 为每个进程分配 GPU,其中每个进程处理不同的提示:
def run_inference(rank, world_size):
dist.init_process_group("nccl", rank=rank, world_size=world_size)
sd.to(rank)
if torch.distributed.get_rank() == 0:
prompt = "a dog"
elif torch.distributed.get_rank() == 1:
prompt = "a cat"
image = sd(prompt).images[0]
image.save(f"./{'_'.join(prompt)}.png")
- 要运行分布式推理,调用
mp.spawn在world_size定义的 GPU 数量上运行run_inference函数:
def main():
world_size = 2
mp.spawn(run_inference, args=(world_size,), nprocs=world_size, join=True)
if __name__ == "__main__":
main()
- 完成推理脚本后,使用
--nproc_per_node参数指定要使用的 GPU 数量,并调用torchrun运行脚本:
torchrun run_distributed.py --nproc_per_node=2
Model sharding 模型分片
现代扩散系统如 Flux 规模非常大,包含多个模型。例如,Flux.1-Dev 由两个文本编码器——T5-XXL 和 CLIP-L、一个扩散变换器以及一个 VAE 组成。对于如此规模的模型,在消费级 GPU 上运行推理会面临挑战。
模型分片是一种将模型分布到多个 GPU 上的技术,当模型无法在单个 GPU 上运行时使用。以下示例假设有两块 16GB 的 GPU 用于推理。
- 首先使用文本编码器计算文本嵌入。通过设置
device_map="balanced"将文本编码器保留在两个 GPU 上。balanced策略将模型均匀分布在所有可用的 GPU 上。使用max_memory参数为每个 GPU 上的文本编码器分配最大内存量。
# 仅加载此步骤的文本编码器!扩散 Transformer 和 VAE 将在后续步骤中加载以节省内存。
from diffusers import FluxPipeline
import torch
prompt = "a photo of a dog with cat-like look"
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
transformer=None,
vae=None,
device_map="balanced",
max_memory={0: "16GB", 1: "16GB"},
torch_dtype=torch.bfloat16
)
with torch.no_grad():
print("Encoding prompts.")
prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
prompt=prompt, prompt_2=None, max_sequence_length=512
)
- 一旦文本嵌入计算完成,将其从 GPU 中移除,为扩散 Transformer 腾出空间。
import gc
def flush():
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
torch.cuda.reset_peak_memory_stats()
del pipeline.text_encoder
del pipeline.text_encoder_2
del pipeline.tokenizer
del pipeline.tokenizer_2
del pipeline
flush()
- 加载扩散 Transformer,它有 125 亿参数。这次将
device_map="auto"设置为自动将模型分布在两个 16GB 的 GPU 上。auto策略由 Accelerate 支持,作为 Big Model Inference 功能的一部分提供。它首先将模型分布到最快的设备(GPU)上,如果需要,再移动到较慢的设备如 CPU 和硬盘。在较慢的设备上存储模型参数的权衡是推理延迟变慢。
from diffusers import AutoModel
import torch
transformer = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
device_map="auto",
torch_dtype=torch.bfloat16
)
# 任何时候,你可以尝试 print(pipeline.hf_device_map) 来查看各种模型如何在设备间分布。这对于跟踪模型的设备放置很有用。
# 你也可以尝试 print(transformer.hf_device_map) 来查看 Transformer 模型如何在设备间分片。
- 将转换器模型添加到管道中进行去噪,但将其他模型级组件(如文本编码器和 VAE)设置为
None,因为你现在不需要它们。
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
text_encoder=None,
text_encoder_2=None,
tokenizer=None,
tokenizer_2=None,
vae=None,
transformer=transformer,
torch_dtype=torch.bfloat16
)
print("Running denoising.")
height, width = 768, 1360
latents = pipeline(
prompt_embeds=prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
num_inference_steps=50,
guidance_scale=3.5,
height=height,
width=width,
output_type="latent",
).images
# 由于不再需要,从内存中移除管道和转换器。
del pipeline.transformer
del pipeline
flush()
- 最后,使用 VAE 将潜在变量解码为图像。VAE 通常足够小,可以加载到单个 GPU 上。
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor
import torch
vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", torch_dtype=torch.bfloat16).to("cuda")
vae_scale_factor = 2 ** (len(vae.config.block_out_channels))
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
with torch.no_grad():
print("Running decoding.")
latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
image = vae.decode(latents, return_dict=False)[0]
image = image_processor.postprocess(image, output_type="pil")
image[0].save("split_transformer.png")
- 通过在特定阶段选择性地加载和卸载所需的模型,并将最大的模型跨多个 GPU 分片,可以在消费级 GPU 上运行大型模型的推理。
文本引导深度到图像生成
StableDiffusionDepth2ImgPipeline 允许您传递文本提示和初始图像来调节新图像的生成。此外,您还可以传递 depth_map 来保留图像结构。如果没有提供 depth_map,该管道会自动通过集成的深度估计模型预测深度。
首先创建一个 StableDiffusionDepth2ImgPipeline 的实例:
import torch
from diffusers import StableDiffusionDepth2ImgPipeline
from diffusers.utils import load_image, make_image_grid
pipeline = StableDiffusionDepth2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-depth",
torch_dtype=torch.float16,
use_safetensors=True,
).to("cuda")
现在将您的提示传递给该管道。您还可以传递 negative_prompt 来防止某些词语指导图像的生成:
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
init_image = load_image(url)
prompt = "two tigers"
negative_prompt = "bad, deformed, ugly, bad anatomy"
image = pipeline(prompt=prompt, image=init_image, negative_prompt=negative_prompt, strength=0.7).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
推理技术
总览
推理管道支持两类技术:
- 管道功能:这些技术修改管道或将其扩展用于其他应用。例如,管道回调为管道添加新功能,管道也可以扩展用于分布式推理。
- 提高推理质量:这些技术增加了生成图像的视觉质量。例如,您可以使用 GPT2 增强提示,以较低的努力创建更好的图像。
创建服务
扩散器的管道可以用作服务器的推理引擎。它支持并发和多线程请求,以同时为多个用户生成图像。
示例中使用 StableDiffusion3Pipeline:
- 导航到
examples/server文件夹并安装所有依赖项。
pip install .
pip install -f requirements.txt
- 使用以下命令启动服务器。
python server.py
- 服务器可通过
http://localhost:8000访问。您可以使用以下命令curl此模型。
curl -X POST -H "Content-Type: application/json" --data '{"model": "something", "prompt": "a kitten in front of a fireplace"}' http://localhost:8000/v1/images/generations
分布式推理
展示如何使用 Accelerate 和 PyTorch Distributed 进行分布式推理。
Accelerate
Accelerate 是一个旨在简化分布式设置中训练或运行推理的库。它简化了设置分布式环境的流程,让您能够专注于您的 PyTorch 代码。
- 首先,创建一个 Python 文件并初始化
accelerate.PartialState来创建分布式环境;设置将自动检测,无需显式定义rank或world_size。将DiffusionPipeline移动到distributed_state.device以分配给每个进程一个 GPU。
使用 split_between_processes 实用程序作为上下文管理器,自动将提示分配给进程数。
import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)
with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
result = pipeline(prompt).images[0]
result.save(f"result_{distributed_state.process_index}.png")
- 使用
--num_processes参数指定要使用的 GPU 数量,并调用accelerate launch来运行脚本:
accelerate launch run_distributed.py --num_processes=2
PyTorch 分布式
PyTorch 支持 DistributedDataParallel,它能够实现数据并行。
- 首先,创建一个 Python 文件并导入
torch.distributed和torch.multiprocessing,以设置分布式进程组并在每个 GPU 上启动推理进程。你还应该初始化一个DiffusionPipeline:
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from diffusers import DiffusionPipeline
sd = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
- 创建一个用于运行推理的函数;
init_process_group负责创建分布式环境,包括使用的后端类型、当前进程的rank以及参与进程的world_size或进程数量。如果你在 2 个 GPU 上并行运行推理,那么world_size的值就是 2。
将 DiffusionPipeline 移至 rank,并使用 get_rank 为每个进程分配 GPU,其中每个进程处理不同的提示:
def run_inference(rank, world_size):
dist.init_process_group("nccl", rank=rank, world_size=world_size)
sd.to(rank)
if torch.distributed.get_rank() == 0:
prompt = "a dog"
elif torch.distributed.get_rank() == 1:
prompt = "a cat"
image = sd(prompt).images[0]
image.save(f"./{'_'.join(prompt)}.png")
- 要运行分布式推理,调用
mp.spawn在world_size定义的 GPU 数量上运行run_inference函数:
def main():
world_size = 2
mp.spawn(run_inference, args=(world_size,), nprocs=world_size, join=True)
if __name__ == "__main__":
main()
- 完成推理脚本后,使用
--nproc_per_node参数指定要使用的 GPU 数量,并调用torchrun运行脚本:
torchrun run_distributed.py --nproc_per_node=2
Model sharding 模型分片
现代扩散系统如 Flux 规模非常大,包含多个模型。例如,Flux.1-Dev 由两个文本编码器——T5-XXL 和 CLIP-L、一个扩散变换器以及一个 VAE 组成。对于如此规模的模型,在消费级 GPU 上运行推理会面临挑战。
模型分片是一种将模型分布到多个 GPU 上的技术,当模型无法在单个 GPU 上运行时使用。以下示例假设有两块 16GB 的 GPU 用于推理。
- 首先使用文本编码器计算文本嵌入。通过设置
device_map="balanced"将文本编码器保留在两个 GPU 上。balanced策略将模型均匀分布在所有可用的 GPU 上。使用max_memory参数为每个 GPU 上的文本编码器分配最大内存量。
# 仅加载此步骤的文本编码器!扩散 Transformer 和 VAE 将在后续步骤中加载以节省内存。
from diffusers import FluxPipeline
import torch
prompt = "a photo of a dog with cat-like look"
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
transformer=None,
vae=None,
device_map="balanced",
max_memory={0: "16GB", 1: "16GB"},
torch_dtype=torch.bfloat16
)
with torch.no_grad():
print("Encoding prompts.")
prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
prompt=prompt, prompt_2=None, max_sequence_length=512
)
- 一旦文本嵌入计算完成,将其从 GPU 中移除,为扩散 Transformer 腾出空间。
import gc
def flush():
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
torch.cuda.reset_peak_memory_stats()
del pipeline.text_encoder
del pipeline.text_encoder_2
del pipeline.tokenizer
del pipeline.tokenizer_2
del pipeline
flush()
- 加载扩散 Transformer,它有 125 亿参数。这次将
device_map="auto"设置为自动将模型分布在两个 16GB 的 GPU 上。auto策略由 Accelerate 支持,作为 Big Model Inference 功能的一部分提供。它首先将模型分布到最快的设备(GPU)上,如果需要,再移动到较慢的设备如 CPU 和硬盘。在较慢的设备上存储模型参数的权衡是推理延迟变慢。
from diffusers import AutoModel
import torch
transformer = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
device_map="auto",
torch_dtype=torch.bfloat16
)
# 任何时候,你可以尝试 print(pipeline.hf_device_map) 来查看各种模型如何在设备间分布。这对于跟踪模型的设备放置很有用。
# 你也可以尝试 print(transformer.hf_device_map) 来查看 Transformer 模型如何在设备间分片。
- 将转换器模型添加到管道中进行去噪,但将其他模型级组件(如文本编码器和 VAE)设置为
None,因为你现在不需要它们。
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
text_encoder=None,
text_encoder_2=None,
tokenizer=None,
tokenizer_2=None,
vae=None,
transformer=transformer,
torch_dtype=torch.bfloat16
)
print("Running denoising.")
height, width = 768, 1360
latents = pipeline(
prompt_embeds=prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
num_inference_steps=50,
guidance_scale=3.5,
height=height,
width=width,
output_type="latent",
).images
# 由于不再需要,从内存中移除管道和转换器。
del pipeline.transformer
del pipeline
flush()
- 最后,使用 VAE 将潜在变量解码为图像。VAE 通常足够小,可以加载到单个 GPU 上。
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor
import torch
vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", torch_dtype=torch.bfloat16).to("cuda")
vae_scale_factor = 2 ** (len(vae.config.block_out_channels))
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
with torch.no_grad():
print("Running decoding.")
latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
image = vae.decode(latents, return_dict=False)[0]
image = image_processor.postprocess(image, output_type="pil")
image[0].save("split_transformer.png")
- 通过在特定阶段选择性地加载和卸载所需的模型,并将最大的模型跨多个 GPU 分片,可以在消费级 GPU 上运行大型模型的推理。
图像生成图像
图像到图像与文本到图像类似,但除了提示词外,您还可以传递一张初始图像作为扩散过程的起点。
初始图像被编码到潜在空间,并添加噪声。然后,潜在扩散模型接收提示词和带噪声的潜在图像,预测添加的噪声,并从初始潜在图像中移除预测的噪声,以获得新的潜在图像。最后,解码器将新的潜在图像解码回图像。
使用 Diffusers 执行图像生成图像任务:
a. 将检查点加载到 AutoPipelineForImage2Image 类中;这个管道会根据检查点自动处理加载正确的管道类:
import torch
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import load_image, make_image_grid
pipeline = AutoPipelineForImage2Image.from_pretrained(
"kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16, use_safetensors=True
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
b. 加载一张图片传递给流程:
init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png")
c. 传递提示和图片给流程以生成图片:
prompt = "cat wizard, gandalf, lord of the rings, detailed, fantasy, cute, adorable, Pixar, Disney, 8k"
image = pipeline(prompt, image=init_image).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
流行模型
最受欢迎的图像到图像模型是 Stable Diffusion v1.5、Stable Diffusion XL (SDXL) 和 Kandinsky 2.2。由于架构差异和训练过程的不同,Stable Diffusion 和 Kandinsky 模型的结果会有所不同;通常情况下,SDXL 生成的图像质量会高于 Stable Diffusion v1.5。
Stable Diffusion v1.5
Stable Diffusion v1.5 是一个基于早期检查点初始化的潜在扩散模型,并在 512x512 图像上进行了 595K 步的微调:
import torch
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import make_image_grid, load_image
pipeline = AutoPipelineForImage2Image.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# prepare image
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-init.png"
init_image = load_image(url)
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
# pass prompt and image to pipeline
image = pipeline(prompt, image=init_image).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
Stable Diffusion XL (SDXL)
SDXL 是 Stable Diffusion 模型的更强大版本。它使用更大的基础模型,并附加一个精炼模型来提高基础模型输出的质量:
import torch
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import make_image_grid, load_image
pipeline = AutoPipelineForImage2Image.from_pretrained(
"stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# prepare image
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-sdxl-init.png"
init_image = load_image(url)
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
# pass prompt and image to pipeline
image = pipeline(prompt, image=init_image, strength=0.5).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
Kandinsky 2.2
Kandinsky 模型与 Stable Diffusion 模型不同,因为它使用图像先验模型来创建图像嵌入。这些嵌入有助于在文本和图像之间建立更好的对齐,使潜在扩散模型能够生成更好的图像:
import torch
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import make_image_grid, load_image
pipeline = AutoPipelineForImage2Image.from_pretrained(
"kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16, use_safetensors=True
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# prepare image
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-init.png"
init_image = load_image(url)
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
# pass prompt and image to pipeline
image = pipeline(prompt, image=init_image).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
配置管道参数
强度
strength 是需要考虑的最重要参数之一,它将对生成的图像产生巨大影响。它决定了生成的图像与初始图像的相似程度:
- 较高的
strength值会给模型更多”创造力”,以生成与初始图像不同的图像;strength值为 1.0 表示初始图像几乎被忽略 - 较低的
strength值意味着生成的图像与初始图像更相似
strength 和 num_inference_steps 参数是相关的,因为 strength 决定了要添加的噪声步数。例如,如果 num_inference_steps 是 50 且 strength 是 0.8,那么这意味着要对初始图像添加 40(50 * 0.8)步的噪声,然后进行 40 步去噪以获得新生成的图像:
import torch
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import make_image_grid, load_image
pipeline = AutoPipelineForImage2Image.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# prepare image
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-init.png"
init_image = load_image(url)
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
# pass prompt and image to pipeline
image = pipeline(prompt, image=init_image, strength=0.8).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
引导尺度
guidance_scale 参数用于控制生成图像与文本提示的匹配程度。guidance_scale 值越高,生成图像与提示的匹配度越高;guidance_scale 值越低,生成图像偏离提示的空间越大。
你可以结合 guidance_scale 和 strength 实现更精确的控制,以调整模型的表达程度。例如,结合高 strength + guidance_scale 以实现最大创意,或使用低 strength 和低 guidance_scale 的组合生成与初始图像相似但受提示约束较松的图像:
import torch
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import make_image_grid, load_image
pipeline = AutoPipelineForImage2Image.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# prepare image
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-init.png"
init_image = load_image(url)
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
# pass prompt and image to pipeline
image = pipeline(prompt, image=init_image, guidance_scale=8.0).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
负面提示
负面提示用于指导模型在图像中不包含某些内容,它可用于提升图像质量或修改图像。例如,通过添加”细节差”或”模糊”等负面提示,可以鼓励模型生成更高质量的图像。或者,您可以通过指定要排除的元素来修改图像:
import torch
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import make_image_grid, load_image
pipeline = AutoPipelineForImage2Image.from_pretrained(
"stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# prepare image
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-init.png"
init_image = load_image(url)
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
negative_prompt = "ugly, deformed, disfigured, poor details, bad anatomy"
# pass prompt and image to pipeline
image = pipeline(prompt, negative_prompt=negative_prompt, image=init_image).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
串联的图像到图像流程
文本到图像到图像
将文本到图像和图像到图像的流程串联起来,可以让你从文本生成图像,并使用生成的图像作为图像到图像流程的初始图像。如果你想要完全从零开始生成图像,这会很有用。例如,让我们串联一个 Stable Diffusion 和一个 Kandinsky 模型。
先使用文本到图像流程生成一张图像:
from diffusers import AutoPipelineForText2Image, AutoPipelineForImage2Image
import torch
from diffusers.utils import make_image_grid
pipeline = AutoPipelineForText2Image.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
text2image = pipeline("Astronaut in a jungle, cold color palette, muted colors, detailed, 8k").images[0]
text2image
现在你可以将这张生成的图像传递给图像到图像的流程:
pipeline = AutoPipelineForImage2Image.from_pretrained(
"kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16, use_safetensors=True
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
image2image = pipeline("Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", image=text2image).images[0]
make_image_grid([text2image, image2image], rows=1, cols=2)
图像到图像到图像
你也可以将多个图像到图像的流程链在一起,以创建更有趣的图像。这适用于对图像进行迭代式风格迁移、生成短 GIF、为图像恢复颜色或恢复图像缺失区域。
生成一张图片:
import torch
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import make_image_grid, load_image
pipeline = AutoPipelineForImage2Image.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# prepare image
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-init.png"
init_image = load_image(url)
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
# pass prompt and image to pipeline
image = pipeline(prompt, image=init_image, output_type="latent").images[0]
将此流程的潜在输出传递到下一个流程,以生成漫画艺术风格的图像:
pipeline = AutoPipelineForImage2Image.from_pretrained(
"ogkalu/Comic-Diffusion", torch_dtype=torch.float16
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# need to include the token "charliebo artstyle" in the prompt to use this checkpoint
image = pipeline("Astronaut in a jungle, charliebo artstyle", image=image, output_type="latent").images[0]
再重复一次,以像素艺术风格生成最终图像:
pipeline = AutoPipelineForImage2Image.from_pretrained(
"kohbanye/pixel-art-style", torch_dtype=torch.float16
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# need to include the token "pixelartstyle" in the prompt to use this checkpoint
image = pipeline("Astronaut in a jungle, pixelartstyle", image=image).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
图像到放大器到超分辨率
另一种你可以串联你的图像到图像流程的方式是使用一个上采样器和超分辨率流程,以真正增加图像中的细节水平。
从一个图像到图像流程开始:
import torch
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import make_image_grid, load_image
pipeline = AutoPipelineForImage2Image.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# prepare image
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-init.png"
init_image = load_image(url)
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
# pass prompt and image to pipeline
image_1 = pipeline(prompt, image=init_image, output_type="latent").images[0]
将其连接到上采样管道以增加图像分辨率:
from diffusers import StableDiffusionLatentUpscalePipeline
upscaler = StableDiffusionLatentUpscalePipeline.from_pretrained(
"stabilityai/sd-x2-latent-upscaler", torch_dtype=torch.float16, use_safetensors=True
)
upscaler.enable_model_cpu_offload()
upscaler.enable_xformers_memory_efficient_attention()
image_2 = upscaler(prompt, image=image_1).images[0]
最后,将其连接到超分辨率流程中,以进一步提升分辨率:
from diffusers import StableDiffusionUpscalePipeline
super_res = StableDiffusionUpscalePipeline.from_pretrained(
"stabilityai/stable-diffusion-x4-upscaler", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
)
super_res.enable_model_cpu_offload()
super_res.enable_xformers_memory_efficient_attention()
image_3 = super_res(prompt, image=image_2).images[0]
make_image_grid([init_image, image_3.resize((512, 512))], rows=1, cols=2)
控制图像生成
想要生成一个完全符合你期望的图像可能很困难,这也是可控生成技术和模型如此有用的原因。虽然你可以使用 negative_prompt 来部分控制图像生成,但还有更稳健的方法,如提示权重调整和 ControlNets。
提示权重
提示权重允许你调整提示中每个概念的表现。例如,在类似”宇航员在丛林中,冷色调,柔和色彩,精细,8k”的提示中,你可以选择增加或减少”宇航员”和”丛林”的嵌入。Compel 库提供了一个简单的语法来调整提示权重并生成嵌入。
AutoPipelineForImage2Image 有一个 prompt_embeds(如果你使用负提示,还有 negative_prompt_embeds)参数,你可以通过这个参数传递替换 prompt 参数的嵌入:
from diffusers import AutoPipelineForImage2Image
import torch
pipeline = AutoPipelineForImage2Image.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
image = pipeline(prompt_embeds=prompt_embeds, # generated from Compel
negative_prompt_embeds=negative_prompt_embeds, # generated from Compel
image=init_image,
).images[0]
ControlNet
ControlNets 提供了一种更灵活和准确的方式来控制图像生成,因为你可以使用额外的条件图像。条件图像可以是一张边缘检测图像、深度图、图像分割结果,甚至是涂鸦!无论你选择哪种类型的条件图像,ControlNet 都会生成一张保留其中信息的图像。
使用深度图来调节图像,以保留图像中的空间信息:
from diffusers.utils import load_image, make_image_grid
# prepare image
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-init.png"
init_image = load_image(url)
init_image = init_image.resize((958, 960)) # resize to depth image dimensions
depth_image = load_image("https://huggingface.co/lllyasviel/control_v11f1p_sd15_depth/resolve/main/images/control.png")
make_image_grid([init_image, depth_image], rows=1, cols=2)
加载一个基于深度图调节的 ControlNet 模型和 AutoPipelineForImage2Image:
from diffusers import ControlNetModel, AutoPipelineForImage2Image
import torch
controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11f1p_sd15_depth", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
pipeline = AutoPipelineForImage2Image.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16, variant="fp16", use_safetensors=True
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
现在根据深度图、初始图像和提示生成一张新图像:
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image_control_net = pipeline(prompt, image=init_image, control_image=depth_image).images[0]
make_image_grid([init_image, depth_image, image_control_net], rows=1, cols=3)
让我们将一种新风格应用于由 ControlNet 生成的图像,通过将其与图像到图像的流程链式连接:
pipeline = AutoPipelineForImage2Image.from_pretrained(
"nitrosocke/elden-ring-diffusion", torch_dtype=torch.float16,
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
prompt = "elden ring style astronaut in a jungle" # include the token "elden ring style" in the prompt
negative_prompt = "ugly, deformed, disfigured, poor details, bad anatomy"
image_elden_ring = pipeline(prompt, negative_prompt=negative_prompt, image=image_control_net, strength=0.45, guidance_scale=10.5).images[0]
make_image_grid([init_image, depth_image, image_control_net, image_elden_ring], rows=2, cols=2)
优化
运行扩散模型计算成本高且资源密集,但通过一些优化技巧,完全可以在消费级和免费级 GPU 上运行。例如,你可以使用更节省内存的注意力机制,如 PyTorch 2.0 的缩放点积注意力或 xFormers(你可以使用其中一种,但无需同时使用两种)。你也可以在 GPU 上卸载模型,同时让其他管道组件在 CPU 上等待:
pipeline.enable_model_cpu_offload()
pipeline.enable_xformers_memory_efficient_attention()
使用 torch.compile,你可以通过将其包裹在 UNet 中进一步提升推理速度:
pipeline.unet = torch.compile(pipeline.unet, mode="reduce-overhead", fullgraph=True)
视频生成
视频生成模型扩展了图像生成(可以视为 1 帧视频),以处理与空间和时间相关的数据。确保所有这些数据——文本、空间、时间——在帧与帧之间保持一致和协调,是生成长且高分辨率视频中的一个重大挑战。
现代视频模型通过扩散 Transformer(DiT)架构应对这一挑战。这降低了计算成本,并允许更高效地扩展到更大和更高质量的图像和视频数据。
Wan2.1
# pip install ftfy
import torch
import numpy as np
from diffusers import AutoModel, WanPipeline
from diffusers.hooks.group_offloading import apply_group_offloading
from diffusers.utils import export_to_video, load_image
from transformers import UMT5EncoderModel
text_encoder = UMT5EncoderModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="text_encoder", torch_dtype=torch.bfloat16)
vae = AutoModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="vae", torch_dtype=torch.float32)
transformer = AutoModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16)
# group-offloading
onload_device = torch.device("cuda")
offload_device = torch.device("cpu")
apply_group_offloading(text_encoder,
onload_device=onload_device,
offload_device=offload_device,
offload_type="block_level",
num_blocks_per_group=4
)
transformer.enable_group_offload(
onload_device=onload_device,
offload_device=offload_device,
offload_type="leaf_level",
use_stream=True
)
pipeline = WanPipeline.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers",
vae=vae,
transformer=transformer,
text_encoder=text_encoder,
torch_dtype=torch.bfloat16
)
pipeline.to("cuda")
prompt = """
The camera rushes from far to near in a low-angle shot,
revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in
for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground.
Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts dynamic
shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
"""
negative_prompt = """
Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality,
low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured,
misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards
"""
output = pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
num_frames=81,
guidance_scale=5.0,
).frames[0]
export_to_video(output, "output.mp4", fps=16)
HunyuanVideo
import torch
from diffusers import AutoModel, HunyuanVideoPipeline
from diffusers.quantizers import PipelineQuantizationConfig
from diffusers.utils import export_to_video
# quantize weights to int4 with bitsandbytes
pipeline_quant_config = PipelineQuantizationConfig(
quant_backend="bitsandbytes_4bit",
quant_kwargs={
"load_in_4bit": True,
"bnb_4bit_quant_type": "nf4",
"bnb_4bit_compute_dtype": torch.bfloat16
},
components_to_quantize=["transformer"]
)
pipeline = HunyuanVideoPipeline.from_pretrained(
"hunyuanvideo-community/HunyuanVideo",
quantization_config=pipeline_quant_config,
torch_dtype=torch.bfloat16,
)
# model-offloading and tiling
pipeline.enable_model_cpu_offload()
pipeline.vae.enable_tiling()
prompt = "A fluffy teddy bear sits on a bed of soft pillows surrounded by children's toys."
video = pipeline(prompt=prompt, num_frames=61, num_inference_steps=30).frames[0]
export_to_video(video, "output.mp4", fps=15)
LTX-Video
import torch
from diffusers import LTXPipeline, AutoModel
from diffusers.hooks import apply_group_offloading
from diffusers.utils import export_to_video
# fp8 layerwise weight-casting
transformer = AutoModel.from_pretrained(
"Lightricks/LTX-Video",
subfolder="transformer",
torch_dtype=torch.bfloat16
)
transformer.enable_layerwise_casting(
storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16
)
pipeline = LTXPipeline.from_pretrained("Lightricks/LTX-Video", transformer=transformer, torch_dtype=torch.bfloat16)
# group-offloading
onload_device = torch.device("cuda")
offload_device = torch.device("cpu")
pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", use_stream=True)
apply_group_offloading(pipeline.text_encoder, onload_device=onload_device, offload_type="block_level", num_blocks_per_group=2)
apply_group_offloading(pipeline.vae, onload_device=onload_device, offload_type="leaf_level")
prompt = """
A woman with long brown hair and light skin smiles at another woman with long blonde hair. The woman with brown hair wears a black jacket and has a small, barely noticeable mole on her right cheek. The camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm and natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be real-life footage
"""
negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted"
video = pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
width=768,
height=512,
num_frames=161,
decode_timestep=0.03,
decode_noise_scale=0.025,
num_inference_steps=50,
).frames[0]
export_to_video(video, "output.mp4", fps=24)
管道参数
管道中有几个参数需要配置,这些参数会影响视频生成的质量或速度。尝试不同的参数值对于发现适当的质量和速度权衡非常重要。
num_frames 帧数
一帧是静止图像,它在其他帧的序列中播放,以创建运动或视频。使用 num_frames 控制每秒生成的帧数。增加 num_frames 会提高感知的运动平滑度和视觉连贯性,这对具有动态内容视频尤为重要。较高的 num_frames 值也会增加视频时长。
某些视频模型需要更具体的 num_frames 值进行推理。例如,HunyuanVideoPipeline 推荐使用 (4 * num_frames) + 1 计算 num_frames。始终检查管道的 API 模型卡,以查看是否有推荐值。
import torch
from diffusers import LTXPipeline
from diffusers.utils import export_to_video
pipeline = LTXPipeline.from_pretrained(
"Lightricks/LTX-Video", torch_dtype=torch.bfloat16
).to("cuda")
prompt = """
A woman with long brown hair and light skin smiles at another woman with long blonde hair. The woman
with brown hair wears a black jacket and has a small, barely noticeable mole on her right cheek. The
camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm and
natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be
real-life footage
"""
video = pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
width=768,
height=512,
num_frames=161,
decode_timestep=0.03,
decode_noise_scale=0.025,
num_inference_steps=50,
).frames[0]
export_to_video(video, "output.mp4", fps=24)
guidance_scale 引导尺度
引导尺度或”cfg”控制生成帧与输入条件(文本、图像或两者)的贴合程度。增加 guidance_scale 会使生成的帧更贴近输入条件,包含更精细的细节,但可能引入伪影并减少输出多样性。较低 guidance_scale 的值会鼓励更松散的提示遵循和增加输出多样性,但细节可能不够精细。如果它太低,可能会完全忽略你的提示并生成随机噪声。
import torch
from diffusers import CogVideoXPipeline, CogVideoXTransformer3DModel
from diffusers.utils import export_to_video
pipeline = CogVideoXPipeline.from_pretrained(
"THUDM/CogVideoX-2b",
torch_dtype=torch.float16
).to("cuda")
prompt = """
A detailed wooden toy ship with intricately carved masts and sails is seen gliding smoothly over
a plush, blue carpet that mimics the waves of the sea. The ship's hull is painted a rich brown,
with tiny windows. The carpet, soft and textured, provides a perfect backdrop, resembling an
oceanic expanse. Surrounding the ship are various other toys and children's items, hinting at
a playful environment. The scene captures the innocence and imagination of childhood,
with the toy ship's journey symbolizing endless adventures in a whimsical, indoor setting.
"""
video = pipeline(
prompt=prompt,
guidance_scale=6,
num_inference_steps=50
).frames[0]
export_to_video(video, "output.mp4", fps=8)
文本引导深度到图像生成
StableDiffusionDepth2ImgPipeline 允许您传递文本提示和初始图像来调节新图像的生成。此外,您还可以传递 depth_map 来保留图像结构。如果没有提供 depth_map,该管道会自动通过集成的深度估计模型预测深度。
首先创建一个 StableDiffusionDepth2ImgPipeline 的实例:
import torch
from diffusers import StableDiffusionDepth2ImgPipeline
from diffusers.utils import load_image, make_image_grid
pipeline = StableDiffusionDepth2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-depth",
torch_dtype=torch.float16,
use_safetensors=True,
).to("cuda")
现在将您的提示传递给该管道。您还可以传递 negative_prompt 来防止某些词语指导图像的生成:
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
init_image = load_image(url)
prompt = "two tigers"
negative_prompt = "bad, deformed, ugly, bad anatomy"
image = pipeline(prompt=prompt, image=init_image, negative_prompt=negative_prompt, strength=0.7).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
推理技术
总览
推理管道支持两类技术:
- 管道功能:这些技术修改管道或将其扩展用于其他应用。例如,管道回调为管道添加新功能,管道也可以扩展用于分布式推理。
- 提高推理质量:这些技术增加了生成图像的视觉质量。例如,您可以使用 GPT2 增强提示,以较低的努力创建更好的图像。
创建服务
扩散器的管道可以用作服务器的推理引擎。它支持并发和多线程请求,以同时为多个用户生成图像。
示例中使用 StableDiffusion3Pipeline:
- 导航到
examples/server文件夹并安装所有依赖项。
pip install .
pip install -f requirements.txt
- 使用以下命令启动服务器。
python server.py
- 服务器可通过
http://localhost:8000访问。您可以使用以下命令curl此模型。
curl -X POST -H "Content-Type: application/json" --data '{"model": "something", "prompt": "a kitten in front of a fireplace"}' http://localhost:8000/v1/images/generations
分布式推理
展示如何使用 Accelerate 和 PyTorch Distributed 进行分布式推理。
Accelerate
Accelerate 是一个旨在简化分布式设置中训练或运行推理的库。它简化了设置分布式环境的流程,让您能够专注于您的 PyTorch 代码。
- 首先,创建一个 Python 文件并初始化
accelerate.PartialState来创建分布式环境;设置将自动检测,无需显式定义rank或world_size。将DiffusionPipeline移动到distributed_state.device以分配给每个进程一个 GPU。
使用 split_between_processes 实用程序作为上下文管理器,自动将提示分配给进程数。
import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)
with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
result = pipeline(prompt).images[0]
result.save(f"result_{distributed_state.process_index}.png")
- 使用
--num_processes参数指定要使用的 GPU 数量,并调用accelerate launch来运行脚本:
accelerate launch run_distributed.py --num_processes=2
PyTorch 分布式
PyTorch 支持 DistributedDataParallel,它能够实现数据并行。
- 首先,创建一个 Python 文件并导入
torch.distributed和torch.multiprocessing,以设置分布式进程组并在每个 GPU 上启动推理进程。你还应该初始化一个DiffusionPipeline:
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from diffusers import DiffusionPipeline
sd = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
- 创建一个用于运行推理的函数;
init_process_group负责创建分布式环境,包括使用的后端类型、当前进程的rank以及参与进程的world_size或进程数量。如果你在 2 个 GPU 上并行运行推理,那么world_size的值就是 2。
将 DiffusionPipeline 移至 rank,并使用 get_rank 为每个进程分配 GPU,其中每个进程处理不同的提示:
def run_inference(rank, world_size):
dist.init_process_group("nccl", rank=rank, world_size=world_size)
sd.to(rank)
if torch.distributed.get_rank() == 0:
prompt = "a dog"
elif torch.distributed.get_rank() == 1:
prompt = "a cat"
image = sd(prompt).images[0]
image.save(f"./{'_'.join(prompt)}.png")
- 要运行分布式推理,调用
mp.spawn在world_size定义的 GPU 数量上运行run_inference函数:
def main():
world_size = 2
mp.spawn(run_inference, args=(world_size,), nprocs=world_size, join=True)
if __name__ == "__main__":
main()
- 完成推理脚本后,使用
--nproc_per_node参数指定要使用的 GPU 数量,并调用torchrun运行脚本:
torchrun run_distributed.py --nproc_per_node=2
Model sharding 模型分片
现代扩散系统如 Flux 规模非常大,包含多个模型。例如,Flux.1-Dev 由两个文本编码器——T5-XXL 和 CLIP-L、一个扩散变换器以及一个 VAE 组成。对于如此规模的模型,在消费级 GPU 上运行推理会面临挑战。
模型分片是一种将模型分布到多个 GPU 上的技术,当模型无法在单个 GPU 上运行时使用。以下示例假设有两块 16GB 的 GPU 用于推理。
- 首先使用文本编码器计算文本嵌入。通过设置
device_map="balanced"将文本编码器保留在两个 GPU 上。balanced策略将模型均匀分布在所有可用的 GPU 上。使用max_memory参数为每个 GPU 上的文本编码器分配最大内存量。
# 仅加载此步骤的文本编码器!扩散 Transformer 和 VAE 将在后续步骤中加载以节省内存。
from diffusers import FluxPipeline
import torch
prompt = "a photo of a dog with cat-like look"
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
transformer=None,
vae=None,
device_map="balanced",
max_memory={0: "16GB", 1: "16GB"},
torch_dtype=torch.bfloat16
)
with torch.no_grad():
print("Encoding prompts.")
prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
prompt=prompt, prompt_2=None, max_sequence_length=512
)
- 一旦文本嵌入计算完成,将其从 GPU 中移除,为扩散 Transformer 腾出空间。
import gc
def flush():
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
torch.cuda.reset_peak_memory_stats()
del pipeline.text_encoder
del pipeline.text_encoder_2
del pipeline.tokenizer
del pipeline.tokenizer_2
del pipeline
flush()
- 加载扩散 Transformer,它有 125 亿参数。这次将
device_map="auto"设置为自动将模型分布在两个 16GB 的 GPU 上。auto策略由 Accelerate 支持,作为 Big Model Inference 功能的一部分提供。它首先将模型分布到最快的设备(GPU)上,如果需要,再移动到较慢的设备如 CPU 和硬盘。在较慢的设备上存储模型参数的权衡是推理延迟变慢。
from diffusers import AutoModel
import torch
transformer = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
device_map="auto",
torch_dtype=torch.bfloat16
)
# 任何时候,你可以尝试 print(pipeline.hf_device_map) 来查看各种模型如何在设备间分布。这对于跟踪模型的设备放置很有用。
# 你也可以尝试 print(transformer.hf_device_map) 来查看 Transformer 模型如何在设备间分片。
- 将转换器模型添加到管道中进行去噪,但将其他模型级组件(如文本编码器和 VAE)设置为
None,因为你现在不需要它们。
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
text_encoder=None,
text_encoder_2=None,
tokenizer=None,
tokenizer_2=None,
vae=None,
transformer=transformer,
torch_dtype=torch.bfloat16
)
print("Running denoising.")
height, width = 768, 1360
latents = pipeline(
prompt_embeds=prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
num_inference_steps=50,
guidance_scale=3.5,
height=height,
width=width,
output_type="latent",
).images
# 由于不再需要,从内存中移除管道和转换器。
del pipeline.transformer
del pipeline
flush()
- 最后,使用 VAE 将潜在变量解码为图像。VAE 通常足够小,可以加载到单个 GPU 上。
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor
import torch
vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", torch_dtype=torch.bfloat16).to("cuda")
vae_scale_factor = 2 ** (len(vae.config.block_out_channels))
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
with torch.no_grad():
print("Running decoding.")
latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
image = vae.decode(latents, return_dict=False)[0]
image = image_processor.postprocess(image, output_type="pil")
image[0].save("split_transformer.png")
- 通过在特定阶段选择性地加载和卸载所需的模型,并将最大的模型跨多个 GPU 分片,可以在消费级 GPU 上运行大型模型的推理。
文本引导深度到图像生成
StableDiffusionDepth2ImgPipeline 允许您传递文本提示和初始图像来调节新图像的生成。此外,您还可以传递 depth_map 来保留图像结构。如果没有提供 depth_map,该管道会自动通过集成的深度估计模型预测深度。
首先创建一个 StableDiffusionDepth2ImgPipeline 的实例:
import torch
from diffusers import StableDiffusionDepth2ImgPipeline
from diffusers.utils import load_image, make_image_grid
pipeline = StableDiffusionDepth2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-depth",
torch_dtype=torch.float16,
use_safetensors=True,
).to("cuda")
现在将您的提示传递给该管道。您还可以传递 negative_prompt 来防止某些词语指导图像的生成:
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
init_image = load_image(url)
prompt = "two tigers"
negative_prompt = "bad, deformed, ugly, bad anatomy"
image = pipeline(prompt=prompt, image=init_image, negative_prompt=negative_prompt, strength=0.7).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
推理技术
总览
推理管道支持两类技术:
- 管道功能:这些技术修改管道或将其扩展用于其他应用。例如,管道回调为管道添加新功能,管道也可以扩展用于分布式推理。
- 提高推理质量:这些技术增加了生成图像的视觉质量。例如,您可以使用 GPT2 增强提示,以较低的努力创建更好的图像。
创建服务
扩散器的管道可以用作服务器的推理引擎。它支持并发和多线程请求,以同时为多个用户生成图像。
示例中使用 StableDiffusion3Pipeline:
- 导航到
examples/server文件夹并安装所有依赖项。
pip install .
pip install -f requirements.txt
- 使用以下命令启动服务器。
python server.py
- 服务器可通过
http://localhost:8000访问。您可以使用以下命令curl此模型。
curl -X POST -H "Content-Type: application/json" --data '{"model": "something", "prompt": "a kitten in front of a fireplace"}' http://localhost:8000/v1/images/generations
分布式推理
展示如何使用 Accelerate 和 PyTorch Distributed 进行分布式推理。
Accelerate
Accelerate 是一个旨在简化分布式设置中训练或运行推理的库。它简化了设置分布式环境的流程,让您能够专注于您的 PyTorch 代码。
- 首先,创建一个 Python 文件并初始化
accelerate.PartialState来创建分布式环境;设置将自动检测,无需显式定义rank或world_size。将DiffusionPipeline移动到distributed_state.device以分配给每个进程一个 GPU。
使用 split_between_processes 实用程序作为上下文管理器,自动将提示分配给进程数。
import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)
with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
result = pipeline(prompt).images[0]
result.save(f"result_{distributed_state.process_index}.png")
- 使用
--num_processes参数指定要使用的 GPU 数量,并调用accelerate launch来运行脚本:
accelerate launch run_distributed.py --num_processes=2
PyTorch 分布式
PyTorch 支持 DistributedDataParallel,它能够实现数据并行。
- 首先,创建一个 Python 文件并导入
torch.distributed和torch.multiprocessing,以设置分布式进程组并在每个 GPU 上启动推理进程。你还应该初始化一个DiffusionPipeline:
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from diffusers import DiffusionPipeline
sd = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
- 创建一个用于运行推理的函数;
init_process_group负责创建分布式环境,包括使用的后端类型、当前进程的rank以及参与进程的world_size或进程数量。如果你在 2 个 GPU 上并行运行推理,那么world_size的值就是 2。
将 DiffusionPipeline 移至 rank,并使用 get_rank 为每个进程分配 GPU,其中每个进程处理不同的提示:
def run_inference(rank, world_size):
dist.init_process_group("nccl", rank=rank, world_size=world_size)
sd.to(rank)
if torch.distributed.get_rank() == 0:
prompt = "a dog"
elif torch.distributed.get_rank() == 1:
prompt = "a cat"
image = sd(prompt).images[0]
image.save(f"./{'_'.join(prompt)}.png")
- 要运行分布式推理,调用
mp.spawn在world_size定义的 GPU 数量上运行run_inference函数:
def main():
world_size = 2
mp.spawn(run_inference, args=(world_size,), nprocs=world_size, join=True)
if __name__ == "__main__":
main()
- 完成推理脚本后,使用
--nproc_per_node参数指定要使用的 GPU 数量,并调用torchrun运行脚本:
torchrun run_distributed.py --nproc_per_node=2
Model sharding 模型分片
现代扩散系统如 Flux 规模非常大,包含多个模型。例如,Flux.1-Dev 由两个文本编码器——T5-XXL 和 CLIP-L、一个扩散变换器以及一个 VAE 组成。对于如此规模的模型,在消费级 GPU 上运行推理会面临挑战。
模型分片是一种将模型分布到多个 GPU 上的技术,当模型无法在单个 GPU 上运行时使用。以下示例假设有两块 16GB 的 GPU 用于推理。
- 首先使用文本编码器计算文本嵌入。通过设置
device_map="balanced"将文本编码器保留在两个 GPU 上。balanced策略将模型均匀分布在所有可用的 GPU 上。使用max_memory参数为每个 GPU 上的文本编码器分配最大内存量。
# 仅加载此步骤的文本编码器!扩散 Transformer 和 VAE 将在后续步骤中加载以节省内存。
from diffusers import FluxPipeline
import torch
prompt = "a photo of a dog with cat-like look"
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
transformer=None,
vae=None,
device_map="balanced",
max_memory={0: "16GB", 1: "16GB"},
torch_dtype=torch.bfloat16
)
with torch.no_grad():
print("Encoding prompts.")
prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
prompt=prompt, prompt_2=None, max_sequence_length=512
)
- 一旦文本嵌入计算完成,将其从 GPU 中移除,为扩散 Transformer 腾出空间。
import gc
def flush():
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
torch.cuda.reset_peak_memory_stats()
del pipeline.text_encoder
del pipeline.text_encoder_2
del pipeline.tokenizer
del pipeline.tokenizer_2
del pipeline
flush()
- 加载扩散 Transformer,它有 125 亿参数。这次将
device_map="auto"设置为自动将模型分布在两个 16GB 的 GPU 上。auto策略由 Accelerate 支持,作为 Big Model Inference 功能的一部分提供。它首先将模型分布到最快的设备(GPU)上,如果需要,再移动到较慢的设备如 CPU 和硬盘。在较慢的设备上存储模型参数的权衡是推理延迟变慢。
from diffusers import AutoModel
import torch
transformer = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
device_map="auto",
torch_dtype=torch.bfloat16
)
# 任何时候,你可以尝试 print(pipeline.hf_device_map) 来查看各种模型如何在设备间分布。这对于跟踪模型的设备放置很有用。
# 你也可以尝试 print(transformer.hf_device_map) 来查看 Transformer 模型如何在设备间分片。
- 将转换器模型添加到管道中进行去噪,但将其他模型级组件(如文本编码器和 VAE)设置为
None,因为你现在不需要它们。
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
text_encoder=None,
text_encoder_2=None,
tokenizer=None,
tokenizer_2=None,
vae=None,
transformer=transformer,
torch_dtype=torch.bfloat16
)
print("Running denoising.")
height, width = 768, 1360
latents = pipeline(
prompt_embeds=prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
num_inference_steps=50,
guidance_scale=3.5,
height=height,
width=width,
output_type="latent",
).images
# 由于不再需要,从内存中移除管道和转换器。
del pipeline.transformer
del pipeline
flush()
- 最后,使用 VAE 将潜在变量解码为图像。VAE 通常足够小,可以加载到单个 GPU 上。
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor
import torch
vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", torch_dtype=torch.bfloat16).to("cuda")
vae_scale_factor = 2 ** (len(vae.config.block_out_channels))
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
with torch.no_grad():
print("Running decoding.")
latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
image = vae.decode(latents, return_dict=False)[0]
image = image_processor.postprocess(image, output_type="pil")
image[0].save("split_transformer.png")
- 通过在特定阶段选择性地加载和卸载所需的模型,并将最大的模型跨多个 GPU 分片,可以在消费级 GPU 上运行大型模型的推理。
negative_prompt 负面提示
负面提示有助于排除您不希望在生成视频中看到的内容。它通常用于通过将模型从”模糊、扭曲、丑陋”等不希望出现的元素中推开,来提高生成视频的质量和一致性。这可以创建更干净、更专注的视频。
# pip install ftfy
import torch
from diffusers import WanPipeline
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
from diffusers.utils import export_to_video
vae = AutoencoderKLWan.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="vae", torch_dtype=torch.float32
)
pipeline = WanPipeline.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers", vae=vae, torch_dtype=torch.bfloat16
)
pipeline.scheduler = UniPCMultistepScheduler.from_config(
pipeline.scheduler.config, flow_shift=5.0
)
pipeline.to("cuda")
pipeline.load_lora_weights("benjamin-paine/steamboat-willie-14b", adapter_name="steamboat-willie")
pipeline.set_adapters("steamboat-willie")
pipeline.enable_model_cpu_offload()
# use "steamboat willie style" to trigger the LoRA
prompt = """
steamboat willie style, golden era animation, The camera rushes from far to near in a low-angle shot,
revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in
for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground.
Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts
dynamic shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
"""
output = pipeline(
prompt=prompt,
num_frames=81,
guidance_scale=5.0,
).frames[0]
export_to_video(output, "output.mp4", fps=16)
减少内存使用
近期具有百亿以上参数的视频模型,如 HunyuanVideoPipeline 和 WanPipeline,需要大量内存,且常常超出消费级硬件的可用内存。Diffusers 提供了多种技术来降低这些大型模型的内存需求。
组卸载
其中一种技术是组卸载,它在不使用时将内部模型层组(例如 torch.nn.Sequential)卸载到 CPU。这些层仅在需要计算时才会被加载,以避免将所有模型组件存储在 GPU 上。对于像 WanPipeline 这样的百四十亿参数模型,组卸载可将所需内存降低至约 13GB 显存。
# pip install ftfy
import torch
import numpy as np
from diffusers import AutoModel, WanPipeline
from diffusers.hooks.group_offloading import apply_group_offloading
from diffusers.utils import export_to_video, load_image
from transformers import UMT5EncoderModel
text_encoder = UMT5EncoderModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="text_encoder", torch_dtype=torch.bfloat16)
vae = AutoModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="vae", torch_dtype=torch.float32)
transformer = AutoModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16)
# group-offloading
onload_device = torch.device("cuda")
offload_device = torch.device("cpu")
apply_group_offloading(text_encoder,
onload_device=onload_device,
offload_device=offload_device,
offload_type="block_level",
num_blocks_per_group=4
)
transformer.enable_group_offload(
onload_device=onload_device,
offload_device=offload_device,
offload_type="leaf_level",
use_stream=True
)
pipeline = WanPipeline.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers",
vae=vae,
transformer=transformer,
text_encoder=text_encoder,
torch_dtype=torch.bfloat16
)
pipeline.to("cuda")
prompt = """
The camera rushes from far to near in a low-angle shot,
revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in
for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground.
Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts dynamic
shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
"""
negative_prompt = """
Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality,
low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured,
misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards
"""
output = pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
num_frames=81,
guidance_scale=5.0,
).frames[0]
export_to_video(output, "output.mp4", fps=16)
模型量化
减少内存的另一种选择是考虑对模型进行量化,即将模型权重存储在低精度的数据类型中。然而,量化可能会根据具体的视频模型影响视频质量。请参考量化概述,了解更多关于不同支持量化的后端信息。
下面的示例使用 bitsandbytes 对模型进行量化。
# pip install ftfy
import torch
from diffusers import AutoModel, WanPipeline
from diffusers.quantizers import PipelineQuantizationConfig
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
from transformers import UMT5EncoderModel
from diffusers.utils import export_to_video
# quantize transformer and text encoder weights with bitsandbytes
pipeline_quant_config = PipelineQuantizationConfig(
quant_backend="bitsandbytes_4bit",
quant_kwargs={"load_in_4bit": True},
components_to_quantize=["transformer", "text_encoder"]
)
vae = AutoModel.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="vae", torch_dtype=torch.float32
)
pipeline = WanPipeline.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers", vae=vae, quantization_config=pipeline_quant_config, torch_dtype=torch.bfloat16
)
pipeline.scheduler = UniPCMultistepScheduler.from_config(
pipeline.scheduler.config, flow_shift=5.0
)
pipeline.to("cuda")
pipeline.load_lora_weights("benjamin-paine/steamboat-willie-14b", adapter_name="steamboat-willie")
pipeline.set_adapters("steamboat-willie")
pipeline.enable_model_cpu_offload()
# use "steamboat willie style" to trigger the LoRA
prompt = """
steamboat willie style, golden era animation, The camera rushes from far to near in a low-angle shot,
revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in
for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground.
Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts
dynamic shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
"""
output = pipeline(
prompt=prompt,
num_frames=81,
guidance_scale=5.0,
).frames[0]
export_to_video(output, "output.mp4", fps=16)
推理速度
torch.compile 可以通过使用优化过的内核来加速推理。第一次编译需要更长时间,但一旦编译完成,就会快得多。最好是一次编译整个管道,然后多次使用该管道而不做任何更改。任何更改,例如图像大小,都会触发重新编译。
下面的示例编译了管道中的转换器,并使用 "max-autotune" 模式来最大化性能。
import torch
from diffusers import CogVideoXPipeline, CogVideoXTransformer3DModel
from diffusers.utils import export_to_video
pipeline = CogVideoXPipeline.from_pretrained(
"THUDM/CogVideoX-2b",
torch_dtype=torch.float16
).to("cuda")
# torch.compile
pipeline.transformer.to(memory_format=torch.channels_last)
pipeline.transformer = torch.compile(
pipeline.transformer, mode="max-autotune", fullgraph=True
)
prompt = """
A detailed wooden toy ship with intricately carved masts and sails is seen gliding smoothly over a plush, blue carpet that mimics the waves of the sea.
The ship's hull is painted a rich brown, with tiny windows. The carpet, soft and textured, provides a perfect backdrop, resembling an oceanic expanse.
Surrounding the ship are various other toys and children's items, hinting at a playful environment. The scene captures the innocence and imagination of childhood,
with the toy ship's journey symbolizing endless adventures in a whimsical, indoor setting.
"""
video = pipeline(
prompt=prompt,
guidance_scale=6,
num_inference_steps=50
).frames[0]
export_to_video(video, "output.mp4", fps=8)
文本引导深度到图像生成
StableDiffusionDepth2ImgPipeline 允许您传递文本提示和初始图像来调节新图像的生成。此外,您还可以传递 depth_map 来保留图像结构。如果没有提供 depth_map,该管道会自动通过集成的深度估计模型预测深度。
首先创建一个 StableDiffusionDepth2ImgPipeline 的实例:
import torch
from diffusers import StableDiffusionDepth2ImgPipeline
from diffusers.utils import load_image, make_image_grid
pipeline = StableDiffusionDepth2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-depth",
torch_dtype=torch.float16,
use_safetensors=True,
).to("cuda")
现在将您的提示传递给该管道。您还可以传递 negative_prompt 来防止某些词语指导图像的生成:
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
init_image = load_image(url)
prompt = "two tigers"
negative_prompt = "bad, deformed, ugly, bad anatomy"
image = pipeline(prompt=prompt, image=init_image, negative_prompt=negative_prompt, strength=0.7).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
推理技术
总览
推理管道支持两类技术:
- 管道功能:这些技术修改管道或将其扩展用于其他应用。例如,管道回调为管道添加新功能,管道也可以扩展用于分布式推理。
- 提高推理质量:这些技术增加了生成图像的视觉质量。例如,您可以使用 GPT2 增强提示,以较低的努力创建更好的图像。
创建服务
扩散器的管道可以用作服务器的推理引擎。它支持并发和多线程请求,以同时为多个用户生成图像。
示例中使用 StableDiffusion3Pipeline:
- 导航到
examples/server文件夹并安装所有依赖项。
pip install .
pip install -f requirements.txt
- 使用以下命令启动服务器。
python server.py
- 服务器可通过
http://localhost:8000访问。您可以使用以下命令curl此模型。
curl -X POST -H "Content-Type: application/json" --data '{"model": "something", "prompt": "a kitten in front of a fireplace"}' http://localhost:8000/v1/images/generations
分布式推理
展示如何使用 Accelerate 和 PyTorch Distributed 进行分布式推理。
Accelerate
Accelerate 是一个旨在简化分布式设置中训练或运行推理的库。它简化了设置分布式环境的流程,让您能够专注于您的 PyTorch 代码。
- 首先,创建一个 Python 文件并初始化
accelerate.PartialState来创建分布式环境;设置将自动检测,无需显式定义rank或world_size。将DiffusionPipeline移动到distributed_state.device以分配给每个进程一个 GPU。
使用 split_between_processes 实用程序作为上下文管理器,自动将提示分配给进程数。
import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)
with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
result = pipeline(prompt).images[0]
result.save(f"result_{distributed_state.process_index}.png")
- 使用
--num_processes参数指定要使用的 GPU 数量,并调用accelerate launch来运行脚本:
accelerate launch run_distributed.py --num_processes=2
PyTorch 分布式
PyTorch 支持 DistributedDataParallel,它能够实现数据并行。
- 首先,创建一个 Python 文件并导入
torch.distributed和torch.multiprocessing,以设置分布式进程组并在每个 GPU 上启动推理进程。你还应该初始化一个DiffusionPipeline:
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from diffusers import DiffusionPipeline
sd = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
- 创建一个用于运行推理的函数;
init_process_group负责创建分布式环境,包括使用的后端类型、当前进程的rank以及参与进程的world_size或进程数量。如果你在 2 个 GPU 上并行运行推理,那么world_size的值就是 2。
将 DiffusionPipeline 移至 rank,并使用 get_rank 为每个进程分配 GPU,其中每个进程处理不同的提示:
def run_inference(rank, world_size):
dist.init_process_group("nccl", rank=rank, world_size=world_size)
sd.to(rank)
if torch.distributed.get_rank() == 0:
prompt = "a dog"
elif torch.distributed.get_rank() == 1:
prompt = "a cat"
image = sd(prompt).images[0]
image.save(f"./{'_'.join(prompt)}.png")
- 要运行分布式推理,调用
mp.spawn在world_size定义的 GPU 数量上运行run_inference函数:
def main():
world_size = 2
mp.spawn(run_inference, args=(world_size,), nprocs=world_size, join=True)
if __name__ == "__main__":
main()
- 完成推理脚本后,使用
--nproc_per_node参数指定要使用的 GPU 数量,并调用torchrun运行脚本:
torchrun run_distributed.py --nproc_per_node=2
Model sharding 模型分片
现代扩散系统如 Flux 规模非常大,包含多个模型。例如,Flux.1-Dev 由两个文本编码器——T5-XXL 和 CLIP-L、一个扩散变换器以及一个 VAE 组成。对于如此规模的模型,在消费级 GPU 上运行推理会面临挑战。
模型分片是一种将模型分布到多个 GPU 上的技术,当模型无法在单个 GPU 上运行时使用。以下示例假设有两块 16GB 的 GPU 用于推理。
- 首先使用文本编码器计算文本嵌入。通过设置
device_map="balanced"将文本编码器保留在两个 GPU 上。balanced策略将模型均匀分布在所有可用的 GPU 上。使用max_memory参数为每个 GPU 上的文本编码器分配最大内存量。
# 仅加载此步骤的文本编码器!扩散 Transformer 和 VAE 将在后续步骤中加载以节省内存。
from diffusers import FluxPipeline
import torch
prompt = "a photo of a dog with cat-like look"
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
transformer=None,
vae=None,
device_map="balanced",
max_memory={0: "16GB", 1: "16GB"},
torch_dtype=torch.bfloat16
)
with torch.no_grad():
print("Encoding prompts.")
prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
prompt=prompt, prompt_2=None, max_sequence_length=512
)
- 一旦文本嵌入计算完成,将其从 GPU 中移除,为扩散 Transformer 腾出空间。
import gc
def flush():
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
torch.cuda.reset_peak_memory_stats()
del pipeline.text_encoder
del pipeline.text_encoder_2
del pipeline.tokenizer
del pipeline.tokenizer_2
del pipeline
flush()
- 加载扩散 Transformer,它有 125 亿参数。这次将
device_map="auto"设置为自动将模型分布在两个 16GB 的 GPU 上。auto策略由 Accelerate 支持,作为 Big Model Inference 功能的一部分提供。它首先将模型分布到最快的设备(GPU)上,如果需要,再移动到较慢的设备如 CPU 和硬盘。在较慢的设备上存储模型参数的权衡是推理延迟变慢。
from diffusers import AutoModel
import torch
transformer = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
device_map="auto",
torch_dtype=torch.bfloat16
)
# 任何时候,你可以尝试 print(pipeline.hf_device_map) 来查看各种模型如何在设备间分布。这对于跟踪模型的设备放置很有用。
# 你也可以尝试 print(transformer.hf_device_map) 来查看 Transformer 模型如何在设备间分片。
- 将转换器模型添加到管道中进行去噪,但将其他模型级组件(如文本编码器和 VAE)设置为
None,因为你现在不需要它们。
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
text_encoder=None,
text_encoder_2=None,
tokenizer=None,
tokenizer_2=None,
vae=None,
transformer=transformer,
torch_dtype=torch.bfloat16
)
print("Running denoising.")
height, width = 768, 1360
latents = pipeline(
prompt_embeds=prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
num_inference_steps=50,
guidance_scale=3.5,
height=height,
width=width,
output_type="latent",
).images
# 由于不再需要,从内存中移除管道和转换器。
del pipeline.transformer
del pipeline
flush()
- 最后,使用 VAE 将潜在变量解码为图像。VAE 通常足够小,可以加载到单个 GPU 上。
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor
import torch
vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", torch_dtype=torch.bfloat16).to("cuda")
vae_scale_factor = 2 ** (len(vae.config.block_out_channels))
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
with torch.no_grad():
print("Running decoding.")
latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
image = vae.decode(latents, return_dict=False)[0]
image = image_processor.postprocess(image, output_type="pil")
image[0].save("split_transformer.png")
- 通过在特定阶段选择性地加载和卸载所需的模型,并将最大的模型跨多个 GPU 分片,可以在消费级 GPU 上运行大型模型的推理。
文本引导深度到图像生成
StableDiffusionDepth2ImgPipeline 允许您传递文本提示和初始图像来调节新图像的生成。此外,您还可以传递 depth_map 来保留图像结构。如果没有提供 depth_map,该管道会自动通过集成的深度估计模型预测深度。
首先创建一个 StableDiffusionDepth2ImgPipeline 的实例:
import torch
from diffusers import StableDiffusionDepth2ImgPipeline
from diffusers.utils import load_image, make_image_grid
pipeline = StableDiffusionDepth2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-depth",
torch_dtype=torch.float16,
use_safetensors=True,
).to("cuda")
现在将您的提示传递给该管道。您还可以传递 negative_prompt 来防止某些词语指导图像的生成:
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
init_image = load_image(url)
prompt = "two tigers"
negative_prompt = "bad, deformed, ugly, bad anatomy"
image = pipeline(prompt=prompt, image=init_image, negative_prompt=negative_prompt, strength=0.7).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
推理技术
总览
推理管道支持两类技术:
- 管道功能:这些技术修改管道或将其扩展用于其他应用。例如,管道回调为管道添加新功能,管道也可以扩展用于分布式推理。
- 提高推理质量:这些技术增加了生成图像的视觉质量。例如,您可以使用 GPT2 增强提示,以较低的努力创建更好的图像。
创建服务
扩散器的管道可以用作服务器的推理引擎。它支持并发和多线程请求,以同时为多个用户生成图像。
示例中使用 StableDiffusion3Pipeline:
- 导航到
examples/server文件夹并安装所有依赖项。
pip install .
pip install -f requirements.txt
- 使用以下命令启动服务器。
python server.py
- 服务器可通过
http://localhost:8000访问。您可以使用以下命令curl此模型。
curl -X POST -H "Content-Type: application/json" --data '{"model": "something", "prompt": "a kitten in front of a fireplace"}' http://localhost:8000/v1/images/generations
分布式推理
展示如何使用 Accelerate 和 PyTorch Distributed 进行分布式推理。
Accelerate
Accelerate 是一个旨在简化分布式设置中训练或运行推理的库。它简化了设置分布式环境的流程,让您能够专注于您的 PyTorch 代码。
- 首先,创建一个 Python 文件并初始化
accelerate.PartialState来创建分布式环境;设置将自动检测,无需显式定义rank或world_size。将DiffusionPipeline移动到distributed_state.device以分配给每个进程一个 GPU。
使用 split_between_processes 实用程序作为上下文管理器,自动将提示分配给进程数。
import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)
with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
result = pipeline(prompt).images[0]
result.save(f"result_{distributed_state.process_index}.png")
- 使用
--num_processes参数指定要使用的 GPU 数量,并调用accelerate launch来运行脚本:
accelerate launch run_distributed.py --num_processes=2
PyTorch 分布式
PyTorch 支持 DistributedDataParallel,它能够实现数据并行。
- 首先,创建一个 Python 文件并导入
torch.distributed和torch.multiprocessing,以设置分布式进程组并在每个 GPU 上启动推理进程。你还应该初始化一个DiffusionPipeline:
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from diffusers import DiffusionPipeline
sd = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
- 创建一个用于运行推理的函数;
init_process_group负责创建分布式环境,包括使用的后端类型、当前进程的rank以及参与进程的world_size或进程数量。如果你在 2 个 GPU 上并行运行推理,那么world_size的值就是 2。
将 DiffusionPipeline 移至 rank,并使用 get_rank 为每个进程分配 GPU,其中每个进程处理不同的提示:
def run_inference(rank, world_size):
dist.init_process_group("nccl", rank=rank, world_size=world_size)
sd.to(rank)
if torch.distributed.get_rank() == 0:
prompt = "a dog"
elif torch.distributed.get_rank() == 1:
prompt = "a cat"
image = sd(prompt).images[0]
image.save(f"./{'_'.join(prompt)}.png")
- 要运行分布式推理,调用
mp.spawn在world_size定义的 GPU 数量上运行run_inference函数:
def main():
world_size = 2
mp.spawn(run_inference, args=(world_size,), nprocs=world_size, join=True)
if __name__ == "__main__":
main()
- 完成推理脚本后,使用
--nproc_per_node参数指定要使用的 GPU 数量,并调用torchrun运行脚本:
torchrun run_distributed.py --nproc_per_node=2
Model sharding 模型分片
现代扩散系统如 Flux 规模非常大,包含多个模型。例如,Flux.1-Dev 由两个文本编码器——T5-XXL 和 CLIP-L、一个扩散变换器以及一个 VAE 组成。对于如此规模的模型,在消费级 GPU 上运行推理会面临挑战。
模型分片是一种将模型分布到多个 GPU 上的技术,当模型无法在单个 GPU 上运行时使用。以下示例假设有两块 16GB 的 GPU 用于推理。
- 首先使用文本编码器计算文本嵌入。通过设置
device_map="balanced"将文本编码器保留在两个 GPU 上。balanced策略将模型均匀分布在所有可用的 GPU 上。使用max_memory参数为每个 GPU 上的文本编码器分配最大内存量。
# 仅加载此步骤的文本编码器!扩散 Transformer 和 VAE 将在后续步骤中加载以节省内存。
from diffusers import FluxPipeline
import torch
prompt = "a photo of a dog with cat-like look"
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
transformer=None,
vae=None,
device_map="balanced",
max_memory={0: "16GB", 1: "16GB"},
torch_dtype=torch.bfloat16
)
with torch.no_grad():
print("Encoding prompts.")
prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
prompt=prompt, prompt_2=None, max_sequence_length=512
)
- 一旦文本嵌入计算完成,将其从 GPU 中移除,为扩散 Transformer 腾出空间。
import gc
def flush():
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
torch.cuda.reset_peak_memory_stats()
del pipeline.text_encoder
del pipeline.text_encoder_2
del pipeline.tokenizer
del pipeline.tokenizer_2
del pipeline
flush()
- 加载扩散 Transformer,它有 125 亿参数。这次将
device_map="auto"设置为自动将模型分布在两个 16GB 的 GPU 上。auto策略由 Accelerate 支持,作为 Big Model Inference 功能的一部分提供。它首先将模型分布到最快的设备(GPU)上,如果需要,再移动到较慢的设备如 CPU 和硬盘。在较慢的设备上存储模型参数的权衡是推理延迟变慢。
from diffusers import AutoModel
import torch
transformer = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
device_map="auto",
torch_dtype=torch.bfloat16
)
# 任何时候,你可以尝试 print(pipeline.hf_device_map) 来查看各种模型如何在设备间分布。这对于跟踪模型的设备放置很有用。
# 你也可以尝试 print(transformer.hf_device_map) 来查看 Transformer 模型如何在设备间分片。
- 将转换器模型添加到管道中进行去噪,但将其他模型级组件(如文本编码器和 VAE)设置为
None,因为你现在不需要它们。
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
text_encoder=None,
text_encoder_2=None,
tokenizer=None,
tokenizer_2=None,
vae=None,
transformer=transformer,
torch_dtype=torch.bfloat16
)
print("Running denoising.")
height, width = 768, 1360
latents = pipeline(
prompt_embeds=prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
num_inference_steps=50,
guidance_scale=3.5,
height=height,
width=width,
output_type="latent",
).images
# 由于不再需要,从内存中移除管道和转换器。
del pipeline.transformer
del pipeline
flush()
- 最后,使用 VAE 将潜在变量解码为图像。VAE 通常足够小,可以加载到单个 GPU 上。
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor
import torch
vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", torch_dtype=torch.bfloat16).to("cuda")
vae_scale_factor = 2 ** (len(vae.config.block_out_channels))
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
with torch.no_grad():
print("Running decoding.")
latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
image = vae.decode(latents, return_dict=False)[0]
image = image_processor.postprocess(image, output_type="pil")
image[0].save("split_transformer.png")
- 通过在特定阶段选择性地加载和卸载所需的模型,并将最大的模型跨多个 GPU 分片,可以在消费级 GPU 上运行大型模型的推理。
图像修复
修复图替换或编辑图像的特定区域。这使得它成为图像修复的有用工具,例如去除缺陷和伪影,甚至用全新的内容替换图像区域。修复图依赖于掩码来确定要填充的图像区域;要修复的区域由白色像素表示,要保留的区域由黑色像素表示。白色像素由提示填充。
使用 Diffusers 进行修复图:
a. 使用 AutoPipelineForInpainting 类加载修复图检查点。这将自动根据检查点检测要加载的适当管道类:
import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image, make_image_grid
pipeline = AutoPipelineForInpainting.from_pretrained(
"kandinsky-community/kandinsky-2-2-decoder-inpaint", torch_dtype=torch.float16
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
b. 加载基础图像和蒙版图像:
init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png")
mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png")
c. 创建一个用于填充图像的提示,并将其与基础图像和蒙版图像一起传递给管道:
prompt = "a black cat with glowing eyes, cute, adorable, disney, pixar, highly detailed, 8k"
negative_prompt = "bad anatomy, deformed, ugly, disfigured"
image = pipeline(prompt=prompt, negative_prompt=negative_prompt, image=init_image, mask_image=mask_image).images[0]
make_image_grid([init_image, mask_image, image], rows=1, cols=3)
创建蒙版图像
您需要为需要修复的图像创建蒙版图像。VaeImageProcessor.blur 方法提供了如何混合原始图像和修复区域的选项。模糊的程度由 blur_factor 参数决定。增加 blur_factor 会增加应用于蒙版边缘的模糊量,使原始图像和修复区域之间的过渡更加柔和。低值或零 blur_factor 保留了蒙版更清晰的边缘。
要使用此功能,请使用图像处理器创建一个模糊蒙版:
import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image
from PIL import Image
pipeline = AutoPipelineForInpainting.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16).to('cuda')
mask = load_image("https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/seashore_mask.png")
blurred_mask = pipeline.mask_processor.blur(mask, blur_factor=33)
blurred_mask
流行模型
Stable Diffusion 修复、Stable Diffusion XL (SDXL) 修复以及 Kandinsky 2.2 修复是最受欢迎的修复模型之一。SDXL 通常比 Stable Diffusion v1.5 生成更高分辨率的图像,而 Kandinsky 2.2 也能够生成高质量的图像。
Stable Diffusion 修复
Stable Diffusion 修复是一个在 512x512 图像上针对修复任务微调的潜在扩散模型。它是一个不错的起点,因为它的速度相对较快,并且能生成高质量的图像:
import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image, make_image_grid
pipeline = AutoPipelineForInpainting.from_pretrained(
"runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16"
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# load base and mask image
init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png")
mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png")
generator = torch.Generator("cuda").manual_seed(92)
prompt = "concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k"
image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, generator=generator).images[0]
make_image_grid([init_image, mask_image, image], rows=1, cols=3)
Stable Diffusion XL (SDXL) 图像修复
SDXL 是 Stable Diffusion v1.5 的一个更大、更强大的版本。该模型可以遵循一个两阶段模型流程(尽管每个模型也可以单独使用);基础模型生成图像,然后精炼模型将图像进一步增强细节和质量:
import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image, make_image_grid
pipeline = AutoPipelineForInpainting.from_pretrained(
"diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype=torch.float16, variant="fp16"
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# load base and mask image
init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png")
mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png")
generator = torch.Generator("cuda").manual_seed(92)
prompt = "concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k"
image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, generator=generator).images[0]
make_image_grid([init_image, mask_image, image], rows=1, cols=3)
Kandinsky 2.2 图像修复
Kandinsky 模型家族与 SDXL 类似,因为它也使用两个模型;图像先验模型创建图像嵌入,扩散模型从这些嵌入中生成图像。你可以分别加载图像先验模型和扩散模型,但使用 Kandinsky 2.2 最简单的方法是将其加载到 AutoPipelineForInpainting 类中,该类在底层使用 KandinskyV22InpaintCombinedPipeline:
import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image, make_image_grid
pipeline = AutoPipelineForInpainting.from_pretrained(
"kandinsky-community/kandinsky-2-2-decoder-inpaint", torch_dtype=torch.float16
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# load base and mask image
init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png")
mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png")
generator = torch.Generator("cuda").manual_seed(92)
prompt = "concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k"
image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, generator=generator).images[0]
make_image_grid([init_image, mask_image, image], rows=1, cols=3)
非图像修复特定检查点
常规检查点
import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image, make_image_grid
pipeline = AutoPipelineForInpainting.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16"
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# load base and mask image
init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png")
mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png")
generator = torch.Generator("cuda").manual_seed(92)
prompt = "concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k"
image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, generator=generator).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
修复检查点
import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image, make_image_grid
pipeline = AutoPipelineForInpainting.from_pretrained(
"runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16"
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# load base and mask image
init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png")
mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png")
generator = torch.Generator("cuda").manual_seed(92)
prompt = "concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k"
image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, generator=generator).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
生成图片对比:左侧图像是从常规检查点生成的,右侧图像是从修复检查点生成的。你会立刻注意到左侧图像不够干净,仍然能看到模型应该修复区域的轮廓。右侧图像则干净得多,修复区域看起来更自然。
对于更基本的任务,比如从图像中擦除物体(例如路面上的石头),常规的检查点就能得到相当不错的结果。常规检查点和 inpaint 检查点之间的差异并不那么明显。
常规模型擦除物体
import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image, make_image_grid
pipeline = AutoPipelineForInpainting.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16"
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# load base and mask image
init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png")
mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/road-mask.png")
image = pipeline(prompt="road", image=init_image, mask_image=mask_image).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
修复模型擦除物体
import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image, make_image_grid
pipeline = AutoPipelineForInpainting.from_pretrained(
"runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16"
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# load base and mask image
init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png")
mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/road-mask.png")
image = pipeline(prompt="road", image=init_image, mask_image=mask_image).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
擦除效果对比:使用非特定修复的检查点的权衡在于整体图像质量可能较低,但它通常倾向于保留蒙版区域(这就是为什么你能看到蒙版轮廓)。特定修复的检查点有意训练以生成更高质量的修复图像,这包括在蒙版和非蒙版区域之间创建更自然的过渡。因此,这些检查点更有可能改变你的非蒙版区域。
如果保留未遮盖区域对你的任务很重要,你可以使用 VaeImageProcessor.apply_overlay 方法来强制图像的未遮盖区域保持不变,但这会以在遮盖区域和未遮盖区域之间产生一些更不自然的过渡为代价:
import PIL
import numpy as np
import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image, make_image_grid
device = "cuda"
pipeline = AutoPipelineForInpainting.from_pretrained(
"runwayml/stable-diffusion-inpainting",
torch_dtype=torch.float16,
variant="fp16"
)
pipeline = pipeline.to(device)
img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"
init_image = load_image(img_url).resize((512, 512))
mask_image = load_image(mask_url).resize((512, 512))
prompt = "Face of a yellow cat, high resolution, sitting on a park bench"
repainted_image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image).images[0]
repainted_image.save("repainted_image.png")
unmasked_unchanged_image = pipeline.image_processor.apply_overlay(mask_image, init_image, repainted_image)
unmasked_unchanged_image.save("force_unmasked_unchanged.png")
make_image_grid([init_image, mask_image, repainted_image, unmasked_unchanged_image], rows=2, cols=2)
配置管道参数
强度
strength 是衡量向基础图像添加多少噪声的指标,这会影响输出与基础图像的相似程度:
- 较高的
strength值意味着向图像中添加了更多噪声,去噪过程需要更长时间,但您将获得更高质量、与原始图像差异更大的图像 - 较低的
strength值意味着添加到图像中的噪声较少,去噪过程更快,但图像质量可能不会那么好,生成的图像与原始图像更相似
import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image, make_image_grid
pipeline = AutoPipelineForInpainting.from_pretrained(
"runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16"
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# load base and mask image
init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png")
mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png")
prompt = "concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k"
image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, strength=0.6).images[0]
make_image_grid([init_image, mask_image, image], rows=1, cols=3)
引导尺度
guidance_scale 影响文本提示和生成图像的对齐程度:
- 高
guidance_scale值意味着提示和生成图像高度一致,因此输出是对提示更严格的解释 - 低
guidance_scale值意味着提示和生成的图像之间的关联较为松散,因此输出可能与提示的偏差更大
你可以将 strength 和 guidance_scale 结合使用,以更好地控制模型的表达程度。例如,高 strength 和 guidance_scale 值的组合会给模型提供最大的创作自由度:
import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image, make_image_grid
pipeline = AutoPipelineForInpainting.from_pretrained(
"runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16"
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# load base and mask image
init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png")
mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png")
prompt = "concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k"
image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, guidance_scale=2.5).images[0]
make_image_grid([init_image, mask_image, image], rows=1, cols=3)
负面提示
负面提示与提示相反,它引导模型避免在图像中生成某些内容。这有助于快速提高图像质量,并防止模型生成你不希望出现的内容:
import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image, make_image_grid
pipeline = AutoPipelineForInpainting.from_pretrained(
"runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16"
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# load base and mask image
init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png")
mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png")
prompt = "concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k"
negative_prompt = "bad architecture, unstable, poor details, blurry"
image = pipeline(prompt=prompt, negative_prompt=negative_prompt, image=init_image, mask_image=mask_image).images[0]
make_image_grid([init_image, mask_image, image], rows=1, cols=3)
填充掩码裁剪
提高填充图像质量的一种方法是使用 padding_mask_crop 参数。启用此选项时,它会使用用户指定的填充裁剪掩码区域,并从原始图像中裁剪相同区域。图像和掩码都会被提升到更高的分辨率进行填充,然后覆盖在原始图像上。这是一种快速简便的方法,可以在不使用单独的管道(如 StableDiffusionUpscalePipeline)的情况下提高图像质量:
import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image
from PIL import Image
generator = torch.Generator(device='cuda').manual_seed(0)
pipeline = AutoPipelineForInpainting.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16).to('cuda')
base = load_image("https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/seashore.png")
mask = load_image("https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/seashore_mask.png")
image = pipeline("boat", image=base, mask_image=mask, strength=0.75, generator=generator, padding_mask_crop=32).images[0]
image
串联式图像修复流程
AutoPipelineForInpainting 可以与其他 Diffusers 管道串联,以编辑它们的输出。这通常有助于提高其他扩散管道的输出质量,如果你使用多个管道,将它们串联起来以保持输出在潜在空间中,并重用相同的管道组件,可以更节省内存。
文本到图像到修复
将文本到图像和修复管道串联起来,可以让你对生成的图像进行修复,并且你不需要从一开始就提供基础图像。这使得编辑你最喜欢的文本到图像输出变得方便,无需生成全新的图像。
从文本到图像的流程开始创建城堡:
import torch
from diffusers import AutoPipelineForText2Image, AutoPipelineForInpainting
from diffusers.utils import load_image, make_image_grid
pipeline = AutoPipelineForText2Image.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
text2image = pipeline("concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k").images[0]
加载上面输出的掩码图像,然后用瀑布效果来修复遮盖区域:
mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_text-chain-mask.png")
pipeline = AutoPipelineForInpainting.from_pretrained(
"kandinsky-community/kandinsky-2-2-decoder-inpaint", torch_dtype=torch.float16
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
prompt = "digital painting of a fantasy waterfall, cloudy"
image = pipeline(prompt=prompt, image=text2image, mask_image=mask_image).images[0]
make_image_grid([text2image, mask_image, image], rows=1, cols=3)
修复到图像到图像
你也可以在图像到图像转换或其他流程之前串联一个修复流程,以提升质量。
首先对图像进行修复:
import torch
from diffusers import AutoPipelineForInpainting, AutoPipelineForImage2Image
from diffusers.utils import load_image, make_image_grid
pipeline = AutoPipelineForInpainting.from_pretrained(
"runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16"
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# load base and mask image
init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png")
mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png")
prompt = "concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k"
image_inpainting = pipeline(prompt=prompt, image=init_image, mask_image=mask_image).images[0]
# resize image to 1024x1024 for SDXL
image_inpainting = image_inpainting.resize((1024, 1024))
现在让我们将图像传递给另一个使用 SDXL 精炼模型的图像修复流程,以增强图像的细节和质量:
pipeline = AutoPipelineForInpainting.from_pretrained(
"stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, variant="fp16"
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
image = pipeline(prompt=prompt, image=image_inpainting, mask_image=mask_image, output_type="latent").images[0]
最后,你可以将这张图像传递给一个图像到图像的流程,为它完成最后的修饰。使用 from_pipe() 方法来重用现有的流程组件更有效率,避免再次将所有流程组件不必要地加载到内存中:
pipeline = AutoPipelineForImage2Image.from_pipe(pipeline)
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
image = pipeline(prompt=prompt, image=image).images[0]
make_image_grid([init_image, mask_image, image_inpainting, image], rows=2, cols=2)
图像到图像和图像修复实际上是非常相似的任务。图像到图像生成一个与提供的现有图像相似的新图像。图像修复做同样的事情,但它只变换由蒙版定义的图像区域,其余的图像保持不变。你可以将图像修复视为一个更精确的工具来做出特定改变,而图像到图像则有更广泛的范围来做出更广泛的改变。
控制图像生成
让图像完全符合你的要求是具有挑战性的,因为去噪过程是随机的。虽然你可以通过配置像 negative_prompt 这样的参数来控制生成过程的某些方面,但控制图像生成有更好、更有效的方法。
提示权重
提示权重提供了一种可量化的方法来调整提示中概念的表现。你可以用它来增加或减少提示中每个概念的文本嵌入向量的幅度,这进而决定了每个概念生成多少。Compel 库提供了直观的语法来调整提示权重并生成嵌入。
一旦生成了嵌入,将它们传递给 AutoPipelineForInpainting 中的 prompt_embeds(如果你使用负提示,则传递 negative_prompt_embeds)参数。嵌入将替换 prompt 参数:
import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import make_image_grid
pipeline = AutoPipelineForInpainting.from_pretrained(
"runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16,
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
image = pipeline(prompt_embeds=prompt_embeds, # generated from Compel
negative_prompt_embeds=negative_prompt_embeds, # generated from Compel
image=init_image,
mask_image=mask_image
).images[0]
make_image_grid([init_image, mask_image, image], rows=1, cols=3)
ControlNet
ControlNet 模型与 Stable Diffusion 等其他扩散模型一起使用,它们提供了一种更灵活和准确的方式来控制图像的生成方式。ControlNet 接受一个额外的条件图像输入,指导扩散模型保留其中的特征。
例如,让我们使用在 inpaint 图像上预训练的 ControlNet 来对图像进行条件化:
import torch
import numpy as np
from diffusers import ControlNetModel, StableDiffusionControlNetInpaintPipeline
from diffusers.utils import load_image, make_image_grid
# load ControlNet
controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16, variant="fp16")
# pass ControlNet to the pipeline
pipeline = StableDiffusionControlNetInpaintPipeline.from_pretrained(
"runwayml/stable-diffusion-inpainting", controlnet=controlnet, torch_dtype=torch.float16, variant="fp16"
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
# load base and mask image
init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png")
mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png")
# prepare control image
def make_inpaint_condition(init_image, mask_image):
init_image = np.array(init_image.convert("RGB")).astype(np.float32) / 255.0
mask_image = np.array(mask_image.convert("L")).astype(np.float32) / 255.0
assert init_image.shape[0:1] == mask_image.shape[0:1], "image and image_mask must have the same image size"
init_image[mask_image > 0.5] = -1.0 # set as masked pixel
init_image = np.expand_dims(init_image, 0).transpose(0, 3, 1, 2)
init_image = torch.from_numpy(init_image)
return init_image
control_image = make_inpaint_condition(init_image, mask_image)
现在根据基础图像、蒙版和控制图像生成一张图片。你会注意到生成的图像中强烈保留了基础图像的特征:
prompt = "concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k"
image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, control_image=control_image).images[0]
make_image_grid([init_image, mask_image, PIL.Image.fromarray(np.uint8(control_image[0][0])).convert('RGB'), image], rows=2, cols=2)
你可以更进一步,将其与图像到图像的流程链式结合,以应用新的风格:
from diffusers import AutoPipelineForImage2Image
pipeline = AutoPipelineForImage2Image.from_pretrained(
"nitrosocke/elden-ring-diffusion", torch_dtype=torch.float16,
)
pipeline.enable_model_cpu_offload()
# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()
prompt = "elden ring style castle" # include the token "elden ring style" in the prompt
negative_prompt = "bad architecture, deformed, disfigured, poor details"
image_elden_ring = pipeline(prompt, negative_prompt=negative_prompt, image=image).images[0]
make_image_grid([init_image, mask_image, image, image_elden_ring], rows=2, cols=2)
优化
如果你资源有限,运行扩散模型可能会很困难且缓慢,但通过一些优化技巧就不必这样了。你可以启用的一种最大(且最简单)的优化是切换到内存高效的注意力机制。如果你使用的是 PyTorch 2.0,则自动启用了缩放点积注意力,你不需要做任何其他操作。对于非 PyTorch 2.0 用户,你可以安装并使用 xFormers 的内存高效注意力实现。这两种选项都能减少内存使用并加速推理。
你也可以将模型卸载到 CPU 上以节省更多内存:
pipeline.enable_xformers_memory_efficient_attention()
pipeline.enable_model_cpu_offload()
要进一步加速你的推理代码,使用 torch.compile。你应该将 torch.compile 包装在管道中最密集的组件周围,这通常是 UNet:
pipeline.unet = torch.compile(pipeline.unet, mode="reduce-overhead", fullgraph=True)
视频生成
视频生成模型扩展了图像生成(可以视为 1 帧视频),以处理与空间和时间相关的数据。确保所有这些数据——文本、空间、时间——在帧与帧之间保持一致和协调,是生成长且高分辨率视频中的一个重大挑战。
现代视频模型通过扩散 Transformer(DiT)架构应对这一挑战。这降低了计算成本,并允许更高效地扩展到更大和更高质量的图像和视频数据。
Wan2.1
# pip install ftfy
import torch
import numpy as np
from diffusers import AutoModel, WanPipeline
from diffusers.hooks.group_offloading import apply_group_offloading
from diffusers.utils import export_to_video, load_image
from transformers import UMT5EncoderModel
text_encoder = UMT5EncoderModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="text_encoder", torch_dtype=torch.bfloat16)
vae = AutoModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="vae", torch_dtype=torch.float32)
transformer = AutoModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16)
# group-offloading
onload_device = torch.device("cuda")
offload_device = torch.device("cpu")
apply_group_offloading(text_encoder,
onload_device=onload_device,
offload_device=offload_device,
offload_type="block_level",
num_blocks_per_group=4
)
transformer.enable_group_offload(
onload_device=onload_device,
offload_device=offload_device,
offload_type="leaf_level",
use_stream=True
)
pipeline = WanPipeline.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers",
vae=vae,
transformer=transformer,
text_encoder=text_encoder,
torch_dtype=torch.bfloat16
)
pipeline.to("cuda")
prompt = """
The camera rushes from far to near in a low-angle shot,
revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in
for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground.
Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts dynamic
shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
"""
negative_prompt = """
Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality,
low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured,
misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards
"""
output = pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
num_frames=81,
guidance_scale=5.0,
).frames[0]
export_to_video(output, "output.mp4", fps=16)
HunyuanVideo
import torch
from diffusers import AutoModel, HunyuanVideoPipeline
from diffusers.quantizers import PipelineQuantizationConfig
from diffusers.utils import export_to_video
# quantize weights to int4 with bitsandbytes
pipeline_quant_config = PipelineQuantizationConfig(
quant_backend="bitsandbytes_4bit",
quant_kwargs={
"load_in_4bit": True,
"bnb_4bit_quant_type": "nf4",
"bnb_4bit_compute_dtype": torch.bfloat16
},
components_to_quantize=["transformer"]
)
pipeline = HunyuanVideoPipeline.from_pretrained(
"hunyuanvideo-community/HunyuanVideo",
quantization_config=pipeline_quant_config,
torch_dtype=torch.bfloat16,
)
# model-offloading and tiling
pipeline.enable_model_cpu_offload()
pipeline.vae.enable_tiling()
prompt = "A fluffy teddy bear sits on a bed of soft pillows surrounded by children's toys."
video = pipeline(prompt=prompt, num_frames=61, num_inference_steps=30).frames[0]
export_to_video(video, "output.mp4", fps=15)
LTX-Video
import torch
from diffusers import LTXPipeline, AutoModel
from diffusers.hooks import apply_group_offloading
from diffusers.utils import export_to_video
# fp8 layerwise weight-casting
transformer = AutoModel.from_pretrained(
"Lightricks/LTX-Video",
subfolder="transformer",
torch_dtype=torch.bfloat16
)
transformer.enable_layerwise_casting(
storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16
)
pipeline = LTXPipeline.from_pretrained("Lightricks/LTX-Video", transformer=transformer, torch_dtype=torch.bfloat16)
# group-offloading
onload_device = torch.device("cuda")
offload_device = torch.device("cpu")
pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", use_stream=True)
apply_group_offloading(pipeline.text_encoder, onload_device=onload_device, offload_type="block_level", num_blocks_per_group=2)
apply_group_offloading(pipeline.vae, onload_device=onload_device, offload_type="leaf_level")
prompt = """
A woman with long brown hair and light skin smiles at another woman with long blonde hair. The woman with brown hair wears a black jacket and has a small, barely noticeable mole on her right cheek. The camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm and natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be real-life footage
"""
negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted"
video = pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
width=768,
height=512,
num_frames=161,
decode_timestep=0.03,
decode_noise_scale=0.025,
num_inference_steps=50,
).frames[0]
export_to_video(video, "output.mp4", fps=24)
管道参数
管道中有几个参数需要配置,这些参数会影响视频生成的质量或速度。尝试不同的参数值对于发现适当的质量和速度权衡非常重要。
num_frames 帧数
一帧是静止图像,它在其他帧的序列中播放,以创建运动或视频。使用 num_frames 控制每秒生成的帧数。增加 num_frames 会提高感知的运动平滑度和视觉连贯性,这对具有动态内容视频尤为重要。较高的 num_frames 值也会增加视频时长。
某些视频模型需要更具体的 num_frames 值进行推理。例如,HunyuanVideoPipeline 推荐使用 (4 * num_frames) + 1 计算 num_frames。始终检查管道的 API 模型卡,以查看是否有推荐值。
import torch
from diffusers import LTXPipeline
from diffusers.utils import export_to_video
pipeline = LTXPipeline.from_pretrained(
"Lightricks/LTX-Video", torch_dtype=torch.bfloat16
).to("cuda")
prompt = """
A woman with long brown hair and light skin smiles at another woman with long blonde hair. The woman
with brown hair wears a black jacket and has a small, barely noticeable mole on her right cheek. The
camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm and
natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be
real-life footage
"""
video = pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
width=768,
height=512,
num_frames=161,
decode_timestep=0.03,
decode_noise_scale=0.025,
num_inference_steps=50,
).frames[0]
export_to_video(video, "output.mp4", fps=24)
guidance_scale 引导尺度
引导尺度或”cfg”控制生成帧与输入条件(文本、图像或两者)的贴合程度。增加 guidance_scale 会使生成的帧更贴近输入条件,包含更精细的细节,但可能引入伪影并减少输出多样性。较低 guidance_scale 的值会鼓励更松散的提示遵循和增加输出多样性,但细节可能不够精细。如果它太低,可能会完全忽略你的提示并生成随机噪声。
import torch
from diffusers import CogVideoXPipeline, CogVideoXTransformer3DModel
from diffusers.utils import export_to_video
pipeline = CogVideoXPipeline.from_pretrained(
"THUDM/CogVideoX-2b",
torch_dtype=torch.float16
).to("cuda")
prompt = """
A detailed wooden toy ship with intricately carved masts and sails is seen gliding smoothly over
a plush, blue carpet that mimics the waves of the sea. The ship's hull is painted a rich brown,
with tiny windows. The carpet, soft and textured, provides a perfect backdrop, resembling an
oceanic expanse. Surrounding the ship are various other toys and children's items, hinting at
a playful environment. The scene captures the innocence and imagination of childhood,
with the toy ship's journey symbolizing endless adventures in a whimsical, indoor setting.
"""
video = pipeline(
prompt=prompt,
guidance_scale=6,
num_inference_steps=50
).frames[0]
export_to_video(video, "output.mp4", fps=8)
文本引导深度到图像生成
StableDiffusionDepth2ImgPipeline 允许您传递文本提示和初始图像来调节新图像的生成。此外,您还可以传递 depth_map 来保留图像结构。如果没有提供 depth_map,该管道会自动通过集成的深度估计模型预测深度。
首先创建一个 StableDiffusionDepth2ImgPipeline 的实例:
import torch
from diffusers import StableDiffusionDepth2ImgPipeline
from diffusers.utils import load_image, make_image_grid
pipeline = StableDiffusionDepth2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-depth",
torch_dtype=torch.float16,
use_safetensors=True,
).to("cuda")
现在将您的提示传递给该管道。您还可以传递 negative_prompt 来防止某些词语指导图像的生成:
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
init_image = load_image(url)
prompt = "two tigers"
negative_prompt = "bad, deformed, ugly, bad anatomy"
image = pipeline(prompt=prompt, image=init_image, negative_prompt=negative_prompt, strength=0.7).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
推理技术
总览
推理管道支持两类技术:
- 管道功能:这些技术修改管道或将其扩展用于其他应用。例如,管道回调为管道添加新功能,管道也可以扩展用于分布式推理。
- 提高推理质量:这些技术增加了生成图像的视觉质量。例如,您可以使用 GPT2 增强提示,以较低的努力创建更好的图像。
创建服务
扩散器的管道可以用作服务器的推理引擎。它支持并发和多线程请求,以同时为多个用户生成图像。
示例中使用 StableDiffusion3Pipeline:
- 导航到
examples/server文件夹并安装所有依赖项。
pip install .
pip install -f requirements.txt
- 使用以下命令启动服务器。
python server.py
- 服务器可通过
http://localhost:8000访问。您可以使用以下命令curl此模型。
curl -X POST -H "Content-Type: application/json" --data '{"model": "something", "prompt": "a kitten in front of a fireplace"}' http://localhost:8000/v1/images/generations
分布式推理
展示如何使用 Accelerate 和 PyTorch Distributed 进行分布式推理。
Accelerate
Accelerate 是一个旨在简化分布式设置中训练或运行推理的库。它简化了设置分布式环境的流程,让您能够专注于您的 PyTorch 代码。
- 首先,创建一个 Python 文件并初始化
accelerate.PartialState来创建分布式环境;设置将自动检测,无需显式定义rank或world_size。将DiffusionPipeline移动到distributed_state.device以分配给每个进程一个 GPU。
使用 split_between_processes 实用程序作为上下文管理器,自动将提示分配给进程数。
import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)
with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
result = pipeline(prompt).images[0]
result.save(f"result_{distributed_state.process_index}.png")
- 使用
--num_processes参数指定要使用的 GPU 数量,并调用accelerate launch来运行脚本:
accelerate launch run_distributed.py --num_processes=2
PyTorch 分布式
PyTorch 支持 DistributedDataParallel,它能够实现数据并行。
- 首先,创建一个 Python 文件并导入
torch.distributed和torch.multiprocessing,以设置分布式进程组并在每个 GPU 上启动推理进程。你还应该初始化一个DiffusionPipeline:
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from diffusers import DiffusionPipeline
sd = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
- 创建一个用于运行推理的函数;
init_process_group负责创建分布式环境,包括使用的后端类型、当前进程的rank以及参与进程的world_size或进程数量。如果你在 2 个 GPU 上并行运行推理,那么world_size的值就是 2。
将 DiffusionPipeline 移至 rank,并使用 get_rank 为每个进程分配 GPU,其中每个进程处理不同的提示:
def run_inference(rank, world_size):
dist.init_process_group("nccl", rank=rank, world_size=world_size)
sd.to(rank)
if torch.distributed.get_rank() == 0:
prompt = "a dog"
elif torch.distributed.get_rank() == 1:
prompt = "a cat"
image = sd(prompt).images[0]
image.save(f"./{'_'.join(prompt)}.png")
- 要运行分布式推理,调用
mp.spawn在world_size定义的 GPU 数量上运行run_inference函数:
def main():
world_size = 2
mp.spawn(run_inference, args=(world_size,), nprocs=world_size, join=True)
if __name__ == "__main__":
main()
- 完成推理脚本后,使用
--nproc_per_node参数指定要使用的 GPU 数量,并调用torchrun运行脚本:
torchrun run_distributed.py --nproc_per_node=2
Model sharding 模型分片
现代扩散系统如 Flux 规模非常大,包含多个模型。例如,Flux.1-Dev 由两个文本编码器——T5-XXL 和 CLIP-L、一个扩散变换器以及一个 VAE 组成。对于如此规模的模型,在消费级 GPU 上运行推理会面临挑战。
模型分片是一种将模型分布到多个 GPU 上的技术,当模型无法在单个 GPU 上运行时使用。以下示例假设有两块 16GB 的 GPU 用于推理。
- 首先使用文本编码器计算文本嵌入。通过设置
device_map="balanced"将文本编码器保留在两个 GPU 上。balanced策略将模型均匀分布在所有可用的 GPU 上。使用max_memory参数为每个 GPU 上的文本编码器分配最大内存量。
# 仅加载此步骤的文本编码器!扩散 Transformer 和 VAE 将在后续步骤中加载以节省内存。
from diffusers import FluxPipeline
import torch
prompt = "a photo of a dog with cat-like look"
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
transformer=None,
vae=None,
device_map="balanced",
max_memory={0: "16GB", 1: "16GB"},
torch_dtype=torch.bfloat16
)
with torch.no_grad():
print("Encoding prompts.")
prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
prompt=prompt, prompt_2=None, max_sequence_length=512
)
- 一旦文本嵌入计算完成,将其从 GPU 中移除,为扩散 Transformer 腾出空间。
import gc
def flush():
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
torch.cuda.reset_peak_memory_stats()
del pipeline.text_encoder
del pipeline.text_encoder_2
del pipeline.tokenizer
del pipeline.tokenizer_2
del pipeline
flush()
- 加载扩散 Transformer,它有 125 亿参数。这次将
device_map="auto"设置为自动将模型分布在两个 16GB 的 GPU 上。auto策略由 Accelerate 支持,作为 Big Model Inference 功能的一部分提供。它首先将模型分布到最快的设备(GPU)上,如果需要,再移动到较慢的设备如 CPU 和硬盘。在较慢的设备上存储模型参数的权衡是推理延迟变慢。
from diffusers import AutoModel
import torch
transformer = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
device_map="auto",
torch_dtype=torch.bfloat16
)
# 任何时候,你可以尝试 print(pipeline.hf_device_map) 来查看各种模型如何在设备间分布。这对于跟踪模型的设备放置很有用。
# 你也可以尝试 print(transformer.hf_device_map) 来查看 Transformer 模型如何在设备间分片。
- 将转换器模型添加到管道中进行去噪,但将其他模型级组件(如文本编码器和 VAE)设置为
None,因为你现在不需要它们。
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
text_encoder=None,
text_encoder_2=None,
tokenizer=None,
tokenizer_2=None,
vae=None,
transformer=transformer,
torch_dtype=torch.bfloat16
)
print("Running denoising.")
height, width = 768, 1360
latents = pipeline(
prompt_embeds=prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
num_inference_steps=50,
guidance_scale=3.5,
height=height,
width=width,
output_type="latent",
).images
# 由于不再需要,从内存中移除管道和转换器。
del pipeline.transformer
del pipeline
flush()
- 最后,使用 VAE 将潜在变量解码为图像。VAE 通常足够小,可以加载到单个 GPU 上。
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor
import torch
vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", torch_dtype=torch.bfloat16).to("cuda")
vae_scale_factor = 2 ** (len(vae.config.block_out_channels))
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
with torch.no_grad():
print("Running decoding.")
latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
image = vae.decode(latents, return_dict=False)[0]
image = image_processor.postprocess(image, output_type="pil")
image[0].save("split_transformer.png")
- 通过在特定阶段选择性地加载和卸载所需的模型,并将最大的模型跨多个 GPU 分片,可以在消费级 GPU 上运行大型模型的推理。
文本引导深度到图像生成
StableDiffusionDepth2ImgPipeline 允许您传递文本提示和初始图像来调节新图像的生成。此外,您还可以传递 depth_map 来保留图像结构。如果没有提供 depth_map,该管道会自动通过集成的深度估计模型预测深度。
首先创建一个 StableDiffusionDepth2ImgPipeline 的实例:
import torch
from diffusers import StableDiffusionDepth2ImgPipeline
from diffusers.utils import load_image, make_image_grid
pipeline = StableDiffusionDepth2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-depth",
torch_dtype=torch.float16,
use_safetensors=True,
).to("cuda")
现在将您的提示传递给该管道。您还可以传递 negative_prompt 来防止某些词语指导图像的生成:
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
init_image = load_image(url)
prompt = "two tigers"
negative_prompt = "bad, deformed, ugly, bad anatomy"
image = pipeline(prompt=prompt, image=init_image, negative_prompt=negative_prompt, strength=0.7).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
推理技术
总览
推理管道支持两类技术:
- 管道功能:这些技术修改管道或将其扩展用于其他应用。例如,管道回调为管道添加新功能,管道也可以扩展用于分布式推理。
- 提高推理质量:这些技术增加了生成图像的视觉质量。例如,您可以使用 GPT2 增强提示,以较低的努力创建更好的图像。
创建服务
扩散器的管道可以用作服务器的推理引擎。它支持并发和多线程请求,以同时为多个用户生成图像。
示例中使用 StableDiffusion3Pipeline:
- 导航到
examples/server文件夹并安装所有依赖项。
pip install .
pip install -f requirements.txt
- 使用以下命令启动服务器。
python server.py
- 服务器可通过
http://localhost:8000访问。您可以使用以下命令curl此模型。
curl -X POST -H "Content-Type: application/json" --data '{"model": "something", "prompt": "a kitten in front of a fireplace"}' http://localhost:8000/v1/images/generations
分布式推理
展示如何使用 Accelerate 和 PyTorch Distributed 进行分布式推理。
Accelerate
Accelerate 是一个旨在简化分布式设置中训练或运行推理的库。它简化了设置分布式环境的流程,让您能够专注于您的 PyTorch 代码。
- 首先,创建一个 Python 文件并初始化
accelerate.PartialState来创建分布式环境;设置将自动检测,无需显式定义rank或world_size。将DiffusionPipeline移动到distributed_state.device以分配给每个进程一个 GPU。
使用 split_between_processes 实用程序作为上下文管理器,自动将提示分配给进程数。
import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)
with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
result = pipeline(prompt).images[0]
result.save(f"result_{distributed_state.process_index}.png")
- 使用
--num_processes参数指定要使用的 GPU 数量,并调用accelerate launch来运行脚本:
accelerate launch run_distributed.py --num_processes=2
PyTorch 分布式
PyTorch 支持 DistributedDataParallel,它能够实现数据并行。
- 首先,创建一个 Python 文件并导入
torch.distributed和torch.multiprocessing,以设置分布式进程组并在每个 GPU 上启动推理进程。你还应该初始化一个DiffusionPipeline:
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from diffusers import DiffusionPipeline
sd = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
- 创建一个用于运行推理的函数;
init_process_group负责创建分布式环境,包括使用的后端类型、当前进程的rank以及参与进程的world_size或进程数量。如果你在 2 个 GPU 上并行运行推理,那么world_size的值就是 2。
将 DiffusionPipeline 移至 rank,并使用 get_rank 为每个进程分配 GPU,其中每个进程处理不同的提示:
def run_inference(rank, world_size):
dist.init_process_group("nccl", rank=rank, world_size=world_size)
sd.to(rank)
if torch.distributed.get_rank() == 0:
prompt = "a dog"
elif torch.distributed.get_rank() == 1:
prompt = "a cat"
image = sd(prompt).images[0]
image.save(f"./{'_'.join(prompt)}.png")
- 要运行分布式推理,调用
mp.spawn在world_size定义的 GPU 数量上运行run_inference函数:
def main():
world_size = 2
mp.spawn(run_inference, args=(world_size,), nprocs=world_size, join=True)
if __name__ == "__main__":
main()
- 完成推理脚本后,使用
--nproc_per_node参数指定要使用的 GPU 数量,并调用torchrun运行脚本:
torchrun run_distributed.py --nproc_per_node=2
Model sharding 模型分片
现代扩散系统如 Flux 规模非常大,包含多个模型。例如,Flux.1-Dev 由两个文本编码器——T5-XXL 和 CLIP-L、一个扩散变换器以及一个 VAE 组成。对于如此规模的模型,在消费级 GPU 上运行推理会面临挑战。
模型分片是一种将模型分布到多个 GPU 上的技术,当模型无法在单个 GPU 上运行时使用。以下示例假设有两块 16GB 的 GPU 用于推理。
- 首先使用文本编码器计算文本嵌入。通过设置
device_map="balanced"将文本编码器保留在两个 GPU 上。balanced策略将模型均匀分布在所有可用的 GPU 上。使用max_memory参数为每个 GPU 上的文本编码器分配最大内存量。
# 仅加载此步骤的文本编码器!扩散 Transformer 和 VAE 将在后续步骤中加载以节省内存。
from diffusers import FluxPipeline
import torch
prompt = "a photo of a dog with cat-like look"
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
transformer=None,
vae=None,
device_map="balanced",
max_memory={0: "16GB", 1: "16GB"},
torch_dtype=torch.bfloat16
)
with torch.no_grad():
print("Encoding prompts.")
prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
prompt=prompt, prompt_2=None, max_sequence_length=512
)
- 一旦文本嵌入计算完成,将其从 GPU 中移除,为扩散 Transformer 腾出空间。
import gc
def flush():
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
torch.cuda.reset_peak_memory_stats()
del pipeline.text_encoder
del pipeline.text_encoder_2
del pipeline.tokenizer
del pipeline.tokenizer_2
del pipeline
flush()
- 加载扩散 Transformer,它有 125 亿参数。这次将
device_map="auto"设置为自动将模型分布在两个 16GB 的 GPU 上。auto策略由 Accelerate 支持,作为 Big Model Inference 功能的一部分提供。它首先将模型分布到最快的设备(GPU)上,如果需要,再移动到较慢的设备如 CPU 和硬盘。在较慢的设备上存储模型参数的权衡是推理延迟变慢。
from diffusers import AutoModel
import torch
transformer = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
device_map="auto",
torch_dtype=torch.bfloat16
)
# 任何时候,你可以尝试 print(pipeline.hf_device_map) 来查看各种模型如何在设备间分布。这对于跟踪模型的设备放置很有用。
# 你也可以尝试 print(transformer.hf_device_map) 来查看 Transformer 模型如何在设备间分片。
- 将转换器模型添加到管道中进行去噪,但将其他模型级组件(如文本编码器和 VAE)设置为
None,因为你现在不需要它们。
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
text_encoder=None,
text_encoder_2=None,
tokenizer=None,
tokenizer_2=None,
vae=None,
transformer=transformer,
torch_dtype=torch.bfloat16
)
print("Running denoising.")
height, width = 768, 1360
latents = pipeline(
prompt_embeds=prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
num_inference_steps=50,
guidance_scale=3.5,
height=height,
width=width,
output_type="latent",
).images
# 由于不再需要,从内存中移除管道和转换器。
del pipeline.transformer
del pipeline
flush()
- 最后,使用 VAE 将潜在变量解码为图像。VAE 通常足够小,可以加载到单个 GPU 上。
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor
import torch
vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", torch_dtype=torch.bfloat16).to("cuda")
vae_scale_factor = 2 ** (len(vae.config.block_out_channels))
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
with torch.no_grad():
print("Running decoding.")
latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
image = vae.decode(latents, return_dict=False)[0]
image = image_processor.postprocess(image, output_type="pil")
image[0].save("split_transformer.png")
- 通过在特定阶段选择性地加载和卸载所需的模型,并将最大的模型跨多个 GPU 分片,可以在消费级 GPU 上运行大型模型的推理。
negative_prompt 负面提示
负面提示有助于排除您不希望在生成视频中看到的内容。它通常用于通过将模型从”模糊、扭曲、丑陋”等不希望出现的元素中推开,来提高生成视频的质量和一致性。这可以创建更干净、更专注的视频。
# pip install ftfy
import torch
from diffusers import WanPipeline
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
from diffusers.utils import export_to_video
vae = AutoencoderKLWan.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="vae", torch_dtype=torch.float32
)
pipeline = WanPipeline.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers", vae=vae, torch_dtype=torch.bfloat16
)
pipeline.scheduler = UniPCMultistepScheduler.from_config(
pipeline.scheduler.config, flow_shift=5.0
)
pipeline.to("cuda")
pipeline.load_lora_weights("benjamin-paine/steamboat-willie-14b", adapter_name="steamboat-willie")
pipeline.set_adapters("steamboat-willie")
pipeline.enable_model_cpu_offload()
# use "steamboat willie style" to trigger the LoRA
prompt = """
steamboat willie style, golden era animation, The camera rushes from far to near in a low-angle shot,
revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in
for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground.
Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts
dynamic shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
"""
output = pipeline(
prompt=prompt,
num_frames=81,
guidance_scale=5.0,
).frames[0]
export_to_video(output, "output.mp4", fps=16)
减少内存使用
近期具有百亿以上参数的视频模型,如 HunyuanVideoPipeline 和 WanPipeline,需要大量内存,且常常超出消费级硬件的可用内存。Diffusers 提供了多种技术来降低这些大型模型的内存需求。
组卸载
其中一种技术是组卸载,它在不使用时将内部模型层组(例如 torch.nn.Sequential)卸载到 CPU。这些层仅在需要计算时才会被加载,以避免将所有模型组件存储在 GPU 上。对于像 WanPipeline 这样的百四十亿参数模型,组卸载可将所需内存降低至约 13GB 显存。
# pip install ftfy
import torch
import numpy as np
from diffusers import AutoModel, WanPipeline
from diffusers.hooks.group_offloading import apply_group_offloading
from diffusers.utils import export_to_video, load_image
from transformers import UMT5EncoderModel
text_encoder = UMT5EncoderModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="text_encoder", torch_dtype=torch.bfloat16)
vae = AutoModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="vae", torch_dtype=torch.float32)
transformer = AutoModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16)
# group-offloading
onload_device = torch.device("cuda")
offload_device = torch.device("cpu")
apply_group_offloading(text_encoder,
onload_device=onload_device,
offload_device=offload_device,
offload_type="block_level",
num_blocks_per_group=4
)
transformer.enable_group_offload(
onload_device=onload_device,
offload_device=offload_device,
offload_type="leaf_level",
use_stream=True
)
pipeline = WanPipeline.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers",
vae=vae,
transformer=transformer,
text_encoder=text_encoder,
torch_dtype=torch.bfloat16
)
pipeline.to("cuda")
prompt = """
The camera rushes from far to near in a low-angle shot,
revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in
for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground.
Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts dynamic
shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
"""
negative_prompt = """
Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality,
low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured,
misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards
"""
output = pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
num_frames=81,
guidance_scale=5.0,
).frames[0]
export_to_video(output, "output.mp4", fps=16)
模型量化
减少内存的另一种选择是考虑对模型进行量化,即将模型权重存储在低精度的数据类型中。然而,量化可能会根据具体的视频模型影响视频质量。请参考量化概述,了解更多关于不同支持量化的后端信息。
下面的示例使用 bitsandbytes 对模型进行量化。
# pip install ftfy
import torch
from diffusers import AutoModel, WanPipeline
from diffusers.quantizers import PipelineQuantizationConfig
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
from transformers import UMT5EncoderModel
from diffusers.utils import export_to_video
# quantize transformer and text encoder weights with bitsandbytes
pipeline_quant_config = PipelineQuantizationConfig(
quant_backend="bitsandbytes_4bit",
quant_kwargs={"load_in_4bit": True},
components_to_quantize=["transformer", "text_encoder"]
)
vae = AutoModel.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="vae", torch_dtype=torch.float32
)
pipeline = WanPipeline.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers", vae=vae, quantization_config=pipeline_quant_config, torch_dtype=torch.bfloat16
)
pipeline.scheduler = UniPCMultistepScheduler.from_config(
pipeline.scheduler.config, flow_shift=5.0
)
pipeline.to("cuda")
pipeline.load_lora_weights("benjamin-paine/steamboat-willie-14b", adapter_name="steamboat-willie")
pipeline.set_adapters("steamboat-willie")
pipeline.enable_model_cpu_offload()
# use "steamboat willie style" to trigger the LoRA
prompt = """
steamboat willie style, golden era animation, The camera rushes from far to near in a low-angle shot,
revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in
for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground.
Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts
dynamic shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
"""
output = pipeline(
prompt=prompt,
num_frames=81,
guidance_scale=5.0,
).frames[0]
export_to_video(output, "output.mp4", fps=16)
推理速度
torch.compile 可以通过使用优化过的内核来加速推理。第一次编译需要更长时间,但一旦编译完成,就会快得多。最好是一次编译整个管道,然后多次使用该管道而不做任何更改。任何更改,例如图像大小,都会触发重新编译。
下面的示例编译了管道中的转换器,并使用 "max-autotune" 模式来最大化性能。
import torch
from diffusers import CogVideoXPipeline, CogVideoXTransformer3DModel
from diffusers.utils import export_to_video
pipeline = CogVideoXPipeline.from_pretrained(
"THUDM/CogVideoX-2b",
torch_dtype=torch.float16
).to("cuda")
# torch.compile
pipeline.transformer.to(memory_format=torch.channels_last)
pipeline.transformer = torch.compile(
pipeline.transformer, mode="max-autotune", fullgraph=True
)
prompt = """
A detailed wooden toy ship with intricately carved masts and sails is seen gliding smoothly over a plush, blue carpet that mimics the waves of the sea.
The ship's hull is painted a rich brown, with tiny windows. The carpet, soft and textured, provides a perfect backdrop, resembling an oceanic expanse.
Surrounding the ship are various other toys and children's items, hinting at a playful environment. The scene captures the innocence and imagination of childhood,
with the toy ship's journey symbolizing endless adventures in a whimsical, indoor setting.
"""
video = pipeline(
prompt=prompt,
guidance_scale=6,
num_inference_steps=50
).frames[0]
export_to_video(video, "output.mp4", fps=8)
文本引导深度到图像生成
StableDiffusionDepth2ImgPipeline 允许您传递文本提示和初始图像来调节新图像的生成。此外,您还可以传递 depth_map 来保留图像结构。如果没有提供 depth_map,该管道会自动通过集成的深度估计模型预测深度。
首先创建一个 StableDiffusionDepth2ImgPipeline 的实例:
import torch
from diffusers import StableDiffusionDepth2ImgPipeline
from diffusers.utils import load_image, make_image_grid
pipeline = StableDiffusionDepth2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-depth",
torch_dtype=torch.float16,
use_safetensors=True,
).to("cuda")
现在将您的提示传递给该管道。您还可以传递 negative_prompt 来防止某些词语指导图像的生成:
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
init_image = load_image(url)
prompt = "two tigers"
negative_prompt = "bad, deformed, ugly, bad anatomy"
image = pipeline(prompt=prompt, image=init_image, negative_prompt=negative_prompt, strength=0.7).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
推理技术
总览
推理管道支持两类技术:
- 管道功能:这些技术修改管道或将其扩展用于其他应用。例如,管道回调为管道添加新功能,管道也可以扩展用于分布式推理。
- 提高推理质量:这些技术增加了生成图像的视觉质量。例如,您可以使用 GPT2 增强提示,以较低的努力创建更好的图像。
创建服务
扩散器的管道可以用作服务器的推理引擎。它支持并发和多线程请求,以同时为多个用户生成图像。
示例中使用 StableDiffusion3Pipeline:
- 导航到
examples/server文件夹并安装所有依赖项。
pip install .
pip install -f requirements.txt
- 使用以下命令启动服务器。
python server.py
- 服务器可通过
http://localhost:8000访问。您可以使用以下命令curl此模型。
curl -X POST -H "Content-Type: application/json" --data '{"model": "something", "prompt": "a kitten in front of a fireplace"}' http://localhost:8000/v1/images/generations
分布式推理
展示如何使用 Accelerate 和 PyTorch Distributed 进行分布式推理。
Accelerate
Accelerate 是一个旨在简化分布式设置中训练或运行推理的库。它简化了设置分布式环境的流程,让您能够专注于您的 PyTorch 代码。
- 首先,创建一个 Python 文件并初始化
accelerate.PartialState来创建分布式环境;设置将自动检测,无需显式定义rank或world_size。将DiffusionPipeline移动到distributed_state.device以分配给每个进程一个 GPU。
使用 split_between_processes 实用程序作为上下文管理器,自动将提示分配给进程数。
import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)
with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
result = pipeline(prompt).images[0]
result.save(f"result_{distributed_state.process_index}.png")
- 使用
--num_processes参数指定要使用的 GPU 数量,并调用accelerate launch来运行脚本:
accelerate launch run_distributed.py --num_processes=2
PyTorch 分布式
PyTorch 支持 DistributedDataParallel,它能够实现数据并行。
- 首先,创建一个 Python 文件并导入
torch.distributed和torch.multiprocessing,以设置分布式进程组并在每个 GPU 上启动推理进程。你还应该初始化一个DiffusionPipeline:
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from diffusers import DiffusionPipeline
sd = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
- 创建一个用于运行推理的函数;
init_process_group负责创建分布式环境,包括使用的后端类型、当前进程的rank以及参与进程的world_size或进程数量。如果你在 2 个 GPU 上并行运行推理,那么world_size的值就是 2。
将 DiffusionPipeline 移至 rank,并使用 get_rank 为每个进程分配 GPU,其中每个进程处理不同的提示:
def run_inference(rank, world_size):
dist.init_process_group("nccl", rank=rank, world_size=world_size)
sd.to(rank)
if torch.distributed.get_rank() == 0:
prompt = "a dog"
elif torch.distributed.get_rank() == 1:
prompt = "a cat"
image = sd(prompt).images[0]
image.save(f"./{'_'.join(prompt)}.png")
- 要运行分布式推理,调用
mp.spawn在world_size定义的 GPU 数量上运行run_inference函数:
def main():
world_size = 2
mp.spawn(run_inference, args=(world_size,), nprocs=world_size, join=True)
if __name__ == "__main__":
main()
- 完成推理脚本后,使用
--nproc_per_node参数指定要使用的 GPU 数量,并调用torchrun运行脚本:
torchrun run_distributed.py --nproc_per_node=2
Model sharding 模型分片
现代扩散系统如 Flux 规模非常大,包含多个模型。例如,Flux.1-Dev 由两个文本编码器——T5-XXL 和 CLIP-L、一个扩散变换器以及一个 VAE 组成。对于如此规模的模型,在消费级 GPU 上运行推理会面临挑战。
模型分片是一种将模型分布到多个 GPU 上的技术,当模型无法在单个 GPU 上运行时使用。以下示例假设有两块 16GB 的 GPU 用于推理。
- 首先使用文本编码器计算文本嵌入。通过设置
device_map="balanced"将文本编码器保留在两个 GPU 上。balanced策略将模型均匀分布在所有可用的 GPU 上。使用max_memory参数为每个 GPU 上的文本编码器分配最大内存量。
# 仅加载此步骤的文本编码器!扩散 Transformer 和 VAE 将在后续步骤中加载以节省内存。
from diffusers import FluxPipeline
import torch
prompt = "a photo of a dog with cat-like look"
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
transformer=None,
vae=None,
device_map="balanced",
max_memory={0: "16GB", 1: "16GB"},
torch_dtype=torch.bfloat16
)
with torch.no_grad():
print("Encoding prompts.")
prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
prompt=prompt, prompt_2=None, max_sequence_length=512
)
- 一旦文本嵌入计算完成,将其从 GPU 中移除,为扩散 Transformer 腾出空间。
import gc
def flush():
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
torch.cuda.reset_peak_memory_stats()
del pipeline.text_encoder
del pipeline.text_encoder_2
del pipeline.tokenizer
del pipeline.tokenizer_2
del pipeline
flush()
- 加载扩散 Transformer,它有 125 亿参数。这次将
device_map="auto"设置为自动将模型分布在两个 16GB 的 GPU 上。auto策略由 Accelerate 支持,作为 Big Model Inference 功能的一部分提供。它首先将模型分布到最快的设备(GPU)上,如果需要,再移动到较慢的设备如 CPU 和硬盘。在较慢的设备上存储模型参数的权衡是推理延迟变慢。
from diffusers import AutoModel
import torch
transformer = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
device_map="auto",
torch_dtype=torch.bfloat16
)
# 任何时候,你可以尝试 print(pipeline.hf_device_map) 来查看各种模型如何在设备间分布。这对于跟踪模型的设备放置很有用。
# 你也可以尝试 print(transformer.hf_device_map) 来查看 Transformer 模型如何在设备间分片。
- 将转换器模型添加到管道中进行去噪,但将其他模型级组件(如文本编码器和 VAE)设置为
None,因为你现在不需要它们。
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
text_encoder=None,
text_encoder_2=None,
tokenizer=None,
tokenizer_2=None,
vae=None,
transformer=transformer,
torch_dtype=torch.bfloat16
)
print("Running denoising.")
height, width = 768, 1360
latents = pipeline(
prompt_embeds=prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
num_inference_steps=50,
guidance_scale=3.5,
height=height,
width=width,
output_type="latent",
).images
# 由于不再需要,从内存中移除管道和转换器。
del pipeline.transformer
del pipeline
flush()
- 最后,使用 VAE 将潜在变量解码为图像。VAE 通常足够小,可以加载到单个 GPU 上。
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor
import torch
vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", torch_dtype=torch.bfloat16).to("cuda")
vae_scale_factor = 2 ** (len(vae.config.block_out_channels))
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
with torch.no_grad():
print("Running decoding.")
latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
image = vae.decode(latents, return_dict=False)[0]
image = image_processor.postprocess(image, output_type="pil")
image[0].save("split_transformer.png")
- 通过在特定阶段选择性地加载和卸载所需的模型,并将最大的模型跨多个 GPU 分片,可以在消费级 GPU 上运行大型模型的推理。
文本引导深度到图像生成
StableDiffusionDepth2ImgPipeline 允许您传递文本提示和初始图像来调节新图像的生成。此外,您还可以传递 depth_map 来保留图像结构。如果没有提供 depth_map,该管道会自动通过集成的深度估计模型预测深度。
首先创建一个 StableDiffusionDepth2ImgPipeline 的实例:
import torch
from diffusers import StableDiffusionDepth2ImgPipeline
from diffusers.utils import load_image, make_image_grid
pipeline = StableDiffusionDepth2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-depth",
torch_dtype=torch.float16,
use_safetensors=True,
).to("cuda")
现在将您的提示传递给该管道。您还可以传递 negative_prompt 来防止某些词语指导图像的生成:
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
init_image = load_image(url)
prompt = "two tigers"
negative_prompt = "bad, deformed, ugly, bad anatomy"
image = pipeline(prompt=prompt, image=init_image, negative_prompt=negative_prompt, strength=0.7).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
推理技术
总览
推理管道支持两类技术:
- 管道功能:这些技术修改管道或将其扩展用于其他应用。例如,管道回调为管道添加新功能,管道也可以扩展用于分布式推理。
- 提高推理质量:这些技术增加了生成图像的视觉质量。例如,您可以使用 GPT2 增强提示,以较低的努力创建更好的图像。
创建服务
扩散器的管道可以用作服务器的推理引擎。它支持并发和多线程请求,以同时为多个用户生成图像。
示例中使用 StableDiffusion3Pipeline:
- 导航到
examples/server文件夹并安装所有依赖项。
pip install .
pip install -f requirements.txt
- 使用以下命令启动服务器。
python server.py
- 服务器可通过
http://localhost:8000访问。您可以使用以下命令curl此模型。
curl -X POST -H "Content-Type: application/json" --data '{"model": "something", "prompt": "a kitten in front of a fireplace"}' http://localhost:8000/v1/images/generations
分布式推理
展示如何使用 Accelerate 和 PyTorch Distributed 进行分布式推理。
Accelerate
Accelerate 是一个旨在简化分布式设置中训练或运行推理的库。它简化了设置分布式环境的流程,让您能够专注于您的 PyTorch 代码。
- 首先,创建一个 Python 文件并初始化
accelerate.PartialState来创建分布式环境;设置将自动检测,无需显式定义rank或world_size。将DiffusionPipeline移动到distributed_state.device以分配给每个进程一个 GPU。
使用 split_between_processes 实用程序作为上下文管理器,自动将提示分配给进程数。
import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)
with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
result = pipeline(prompt).images[0]
result.save(f"result_{distributed_state.process_index}.png")
- 使用
--num_processes参数指定要使用的 GPU 数量,并调用accelerate launch来运行脚本:
accelerate launch run_distributed.py --num_processes=2
PyTorch 分布式
PyTorch 支持 DistributedDataParallel,它能够实现数据并行。
- 首先,创建一个 Python 文件并导入
torch.distributed和torch.multiprocessing,以设置分布式进程组并在每个 GPU 上启动推理进程。你还应该初始化一个DiffusionPipeline:
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from diffusers import DiffusionPipeline
sd = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
- 创建一个用于运行推理的函数;
init_process_group负责创建分布式环境,包括使用的后端类型、当前进程的rank以及参与进程的world_size或进程数量。如果你在 2 个 GPU 上并行运行推理,那么world_size的值就是 2。
将 DiffusionPipeline 移至 rank,并使用 get_rank 为每个进程分配 GPU,其中每个进程处理不同的提示:
def run_inference(rank, world_size):
dist.init_process_group("nccl", rank=rank, world_size=world_size)
sd.to(rank)
if torch.distributed.get_rank() == 0:
prompt = "a dog"
elif torch.distributed.get_rank() == 1:
prompt = "a cat"
image = sd(prompt).images[0]
image.save(f"./{'_'.join(prompt)}.png")
- 要运行分布式推理,调用
mp.spawn在world_size定义的 GPU 数量上运行run_inference函数:
def main():
world_size = 2
mp.spawn(run_inference, args=(world_size,), nprocs=world_size, join=True)
if __name__ == "__main__":
main()
- 完成推理脚本后,使用
--nproc_per_node参数指定要使用的 GPU 数量,并调用torchrun运行脚本:
torchrun run_distributed.py --nproc_per_node=2
Model sharding 模型分片
现代扩散系统如 Flux 规模非常大,包含多个模型。例如,Flux.1-Dev 由两个文本编码器——T5-XXL 和 CLIP-L、一个扩散变换器以及一个 VAE 组成。对于如此规模的模型,在消费级 GPU 上运行推理会面临挑战。
模型分片是一种将模型分布到多个 GPU 上的技术,当模型无法在单个 GPU 上运行时使用。以下示例假设有两块 16GB 的 GPU 用于推理。
- 首先使用文本编码器计算文本嵌入。通过设置
device_map="balanced"将文本编码器保留在两个 GPU 上。balanced策略将模型均匀分布在所有可用的 GPU 上。使用max_memory参数为每个 GPU 上的文本编码器分配最大内存量。
# 仅加载此步骤的文本编码器!扩散 Transformer 和 VAE 将在后续步骤中加载以节省内存。
from diffusers import FluxPipeline
import torch
prompt = "a photo of a dog with cat-like look"
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
transformer=None,
vae=None,
device_map="balanced",
max_memory={0: "16GB", 1: "16GB"},
torch_dtype=torch.bfloat16
)
with torch.no_grad():
print("Encoding prompts.")
prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
prompt=prompt, prompt_2=None, max_sequence_length=512
)
- 一旦文本嵌入计算完成,将其从 GPU 中移除,为扩散 Transformer 腾出空间。
import gc
def flush():
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
torch.cuda.reset_peak_memory_stats()
del pipeline.text_encoder
del pipeline.text_encoder_2
del pipeline.tokenizer
del pipeline.tokenizer_2
del pipeline
flush()
- 加载扩散 Transformer,它有 125 亿参数。这次将
device_map="auto"设置为自动将模型分布在两个 16GB 的 GPU 上。auto策略由 Accelerate 支持,作为 Big Model Inference 功能的一部分提供。它首先将模型分布到最快的设备(GPU)上,如果需要,再移动到较慢的设备如 CPU 和硬盘。在较慢的设备上存储模型参数的权衡是推理延迟变慢。
from diffusers import AutoModel
import torch
transformer = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
device_map="auto",
torch_dtype=torch.bfloat16
)
# 任何时候,你可以尝试 print(pipeline.hf_device_map) 来查看各种模型如何在设备间分布。这对于跟踪模型的设备放置很有用。
# 你也可以尝试 print(transformer.hf_device_map) 来查看 Transformer 模型如何在设备间分片。
- 将转换器模型添加到管道中进行去噪,但将其他模型级组件(如文本编码器和 VAE)设置为
None,因为你现在不需要它们。
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
text_encoder=None,
text_encoder_2=None,
tokenizer=None,
tokenizer_2=None,
vae=None,
transformer=transformer,
torch_dtype=torch.bfloat16
)
print("Running denoising.")
height, width = 768, 1360
latents = pipeline(
prompt_embeds=prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
num_inference_steps=50,
guidance_scale=3.5,
height=height,
width=width,
output_type="latent",
).images
# 由于不再需要,从内存中移除管道和转换器。
del pipeline.transformer
del pipeline
flush()
- 最后,使用 VAE 将潜在变量解码为图像。VAE 通常足够小,可以加载到单个 GPU 上。
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor
import torch
vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", torch_dtype=torch.bfloat16).to("cuda")
vae_scale_factor = 2 ** (len(vae.config.block_out_channels))
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
with torch.no_grad():
print("Running decoding.")
latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
image = vae.decode(latents, return_dict=False)[0]
image = image_processor.postprocess(image, output_type="pil")
image[0].save("split_transformer.png")
- 通过在特定阶段选择性地加载和卸载所需的模型,并将最大的模型跨多个 GPU 分片,可以在消费级 GPU 上运行大型模型的推理。
文本引导深度到图像生成
StableDiffusionDepth2ImgPipeline 允许您传递文本提示和初始图像来调节新图像的生成。此外,您还可以传递 depth_map 来保留图像结构。如果没有提供 depth_map,该管道会自动通过集成的深度估计模型预测深度。
首先创建一个 StableDiffusionDepth2ImgPipeline 的实例:
import torch
from diffusers import StableDiffusionDepth2ImgPipeline
from diffusers.utils import load_image, make_image_grid
pipeline = StableDiffusionDepth2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-depth",
torch_dtype=torch.float16,
use_safetensors=True,
).to("cuda")
现在将您的提示传递给该管道。您还可以传递 negative_prompt 来防止某些词语指导图像的生成:
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
init_image = load_image(url)
prompt = "two tigers"
negative_prompt = "bad, deformed, ugly, bad anatomy"
image = pipeline(prompt=prompt, image=init_image, negative_prompt=negative_prompt, strength=0.7).images[0]
make_image_grid([init_image, image], rows=1, cols=2)
推理技术
总览
推理管道支持两类技术:
- 管道功能:这些技术修改管道或将其扩展用于其他应用。例如,管道回调为管道添加新功能,管道也可以扩展用于分布式推理。
- 提高推理质量:这些技术增加了生成图像的视觉质量。例如,您可以使用 GPT2 增强提示,以较低的努力创建更好的图像。
创建服务
扩散器的管道可以用作服务器的推理引擎。它支持并发和多线程请求,以同时为多个用户生成图像。
示例中使用 StableDiffusion3Pipeline:
- 导航到
examples/server文件夹并安装所有依赖项。
pip install .
pip install -f requirements.txt
- 使用以下命令启动服务器。
python server.py
- 服务器可通过
http://localhost:8000访问。您可以使用以下命令curl此模型。
curl -X POST -H "Content-Type: application/json" --data '{"model": "something", "prompt": "a kitten in front of a fireplace"}' http://localhost:8000/v1/images/generations
分布式推理
展示如何使用 Accelerate 和 PyTorch Distributed 进行分布式推理。
Accelerate
Accelerate 是一个旨在简化分布式设置中训练或运行推理的库。它简化了设置分布式环境的流程,让您能够专注于您的 PyTorch 代码。
- 首先,创建一个 Python 文件并初始化
accelerate.PartialState来创建分布式环境;设置将自动检测,无需显式定义rank或world_size。将DiffusionPipeline移动到distributed_state.device以分配给每个进程一个 GPU。
使用 split_between_processes 实用程序作为上下文管理器,自动将提示分配给进程数。
import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)
with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
result = pipeline(prompt).images[0]
result.save(f"result_{distributed_state.process_index}.png")
- 使用
--num_processes参数指定要使用的 GPU 数量,并调用accelerate launch来运行脚本:
accelerate launch run_distributed.py --num_processes=2
PyTorch 分布式
PyTorch 支持 DistributedDataParallel,它能够实现数据并行。
- 首先,创建一个 Python 文件并导入
torch.distributed和torch.multiprocessing,以设置分布式进程组并在每个 GPU 上启动推理进程。你还应该初始化一个DiffusionPipeline:
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from diffusers import DiffusionPipeline
sd = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
- 创建一个用于运行推理的函数;
init_process_group负责创建分布式环境,包括使用的后端类型、当前进程的rank以及参与进程的world_size或进程数量。如果你在 2 个 GPU 上并行运行推理,那么world_size的值就是 2。
将 DiffusionPipeline 移至 rank,并使用 get_rank 为每个进程分配 GPU,其中每个进程处理不同的提示:
def run_inference(rank, world_size):
dist.init_process_group("nccl", rank=rank, world_size=world_size)
sd.to(rank)
if torch.distributed.get_rank() == 0:
prompt = "a dog"
elif torch.distributed.get_rank() == 1:
prompt = "a cat"
image = sd(prompt).images[0]
image.save(f"./{'_'.join(prompt)}.png")
- 要运行分布式推理,调用
mp.spawn在world_size定义的 GPU 数量上运行run_inference函数:
def main():
world_size = 2
mp.spawn(run_inference, args=(world_size,), nprocs=world_size, join=True)
if __name__ == "__main__":
main()
- 完成推理脚本后,使用
--nproc_per_node参数指定要使用的 GPU 数量,并调用torchrun运行脚本:
torchrun run_distributed.py --nproc_per_node=2
Model sharding 模型分片
现代扩散系统如 Flux 规模非常大,包含多个模型。例如,Flux.1-Dev 由两个文本编码器——T5-XXL 和 CLIP-L、一个扩散变换器以及一个 VAE 组成。对于如此规模的模型,在消费级 GPU 上运行推理会面临挑战。
模型分片是一种将模型分布到多个 GPU 上的技术,当模型无法在单个 GPU 上运行时使用。以下示例假设有两块 16GB 的 GPU 用于推理。
- 首先使用文本编码器计算文本嵌入。通过设置
device_map="balanced"将文本编码器保留在两个 GPU 上。balanced策略将模型均匀分布在所有可用的 GPU 上。使用max_memory参数为每个 GPU 上的文本编码器分配最大内存量。
# 仅加载此步骤的文本编码器!扩散 Transformer 和 VAE 将在后续步骤中加载以节省内存。
from diffusers import FluxPipeline
import torch
prompt = "a photo of a dog with cat-like look"
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
transformer=None,
vae=None,
device_map="balanced",
max_memory={0: "16GB", 1: "16GB"},
torch_dtype=torch.bfloat16
)
with torch.no_grad():
print("Encoding prompts.")
prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
prompt=prompt, prompt_2=None, max_sequence_length=512
)
- 一旦文本嵌入计算完成,将其从 GPU 中移除,为扩散 Transformer 腾出空间。
import gc
def flush():
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
torch.cuda.reset_peak_memory_stats()
del pipeline.text_encoder
del pipeline.text_encoder_2
del pipeline.tokenizer
del pipeline.tokenizer_2
del pipeline
flush()
- 加载扩散 Transformer,它有 125 亿参数。这次将
device_map="auto"设置为自动将模型分布在两个 16GB 的 GPU 上。auto策略由 Accelerate 支持,作为 Big Model Inference 功能的一部分提供。它首先将模型分布到最快的设备(GPU)上,如果需要,再移动到较慢的设备如 CPU 和硬盘。在较慢的设备上存储模型参数的权衡是推理延迟变慢。
from diffusers import AutoModel
import torch
transformer = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
device_map="auto",
torch_dtype=torch.bfloat16
)
# 任何时候,你可以尝试 print(pipeline.hf_device_map) 来查看各种模型如何在设备间分布。这对于跟踪模型的设备放置很有用。
# 你也可以尝试 print(transformer.hf_device_map) 来查看 Transformer 模型如何在设备间分片。
- 将转换器模型添加到管道中进行去噪,但将其他模型级组件(如文本编码器和 VAE)设置为
None,因为你现在不需要它们。
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
text_encoder=None,
text_encoder_2=None,
tokenizer=None,
tokenizer_2=None,
vae=None,
transformer=transformer,
torch_dtype=torch.bfloat16
)
print("Running denoising.")
height, width = 768, 1360
latents = pipeline(
prompt_embeds=prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
num_inference_steps=50,
guidance_scale=3.5,
height=height,
width=width,
output_type="latent",
).images
# 由于不再需要,从内存中移除管道和转换器。
del pipeline.transformer
del pipeline
flush()
- 最后,使用 VAE 将潜在变量解码为图像。VAE 通常足够小,可以加载到单个 GPU 上。
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor
import torch
vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", torch_dtype=torch.bfloat16).to("cuda")
vae_scale_factor = 2 ** (len(vae.config.block_out_channels))
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
with torch.no_grad():
print("Running decoding.")
latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
image = vae.decode(latents, return_dict=False)[0]
image = image_processor.postprocess(image, output_type="pil")
image[0].save("split_transformer.png")
- 通过在特定阶段选择性地加载和卸载所需的模型,并将最大的模型跨多个 GPU 分片,可以在消费级 GPU 上运行大型模型的推理。
调度器特性
调度器是任何扩散模型的重要组件,因为它控制了整个去噪(或采样)过程。调度器有多种类型,有些针对速度进行了优化,有些则针对质量进行了优化。使用 Diffusers 时,您可以修改调度器配置以使用自定义噪声调度、sigma,并重新调整噪声调度。更改这些参数会对推理质量和速度产生深远影响。
时间步
时间步长或噪声调度决定了每一步采样时的噪声量。调度器使用这个信息来生成每一步具有相应噪声量的图像。时间步长调度由调度器的默认配置生成,但你可以自定义调度器以使用 Diffusers 中尚未包含的新和优化的采样调度。
例如,对齐你的步骤(AYS)是一种优化采样调度以在尽可能少的 10 步内生成高质量图像的方法。Stable Diffusion XL 的最佳 10 步调度为:
from diffusers.schedulers import AysSchedules
sampling_schedule = AysSchedules["StableDiffusionXLTimesteps"]
print(sampling_schedule)
"[999, 845, 730, 587, 443, 310, 193, 116, 53, 13]"
你可以通过将其传递给 timesteps 参数在管道中使用 AYS 采样调度。
pipeline = StableDiffusionXLPipeline.from_pretrained(
"SG161222/RealVisXL_V4.0",
torch_dtype=torch.float16,
variant="fp16",
).to("cuda")
pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config, algorithm_type="sde-dpmsolver++")
prompt = "A cinematic shot of a cute little rabbit wearing a jacket and doing a thumbs up"
generator = torch.Generator(device="cpu").manual_seed(2487854446)
image = pipeline(
prompt=prompt,
negative_prompt="",
generator=generator,
timesteps=sampling_schedule,
).images[0]
时间步长间隔
样本步骤在计划中的选择方式可能会影响生成图像的质量,特别是在重新调整噪声计划方面,这可以使模型生成更明亮或更暗的图像。Diffusers 提供三种时间步长间隔方法:
- leading:创建均匀间隔的步骤
- linspace:包括第一个和最后一个步骤,并均匀选择其余中间步骤
- trailing:仅包括最后一个步骤,并从末尾开始均匀选择其余中间步骤
建议使用 trailing 间隔方法,因为它在样本步骤较少时能生成更高质量、细节更丰富的图像。但对于更多标准的样本步骤值,质量差异并不明显。
import torch
from diffusers import StableDiffusionXLPipeline, DPMSolverMultistepScheduler
pipeline = StableDiffusionXLPipeline.from_pretrained(
"SG161222/RealVisXL_V4.0",
torch_dtype=torch.float16,
variant="fp16",
).to("cuda")
pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config, timestep_spacing="trailing")
prompt = "A cinematic shot of a cute little black cat sitting on a pumpkin at night"
generator = torch.Generator(device="cpu").manual_seed(2487854446)
image = pipeline(
prompt=prompt,
negative_prompt="",
generator=generator,
num_inference_steps=5,
).images[0]
image
Sigmas
sigmas 参数是根据时间步长计划在每个时间步长添加的噪声量。与 timesteps 参数类似,你可以自定义 sigmas 参数来控制每步添加的噪声量。当你使用自定义的 sigmas 值时,timesteps 是根据自定义的 sigmas 值计算的,默认的调度器配置将被忽略。
例如,你可以手动将之前 10 步 AYS 计划的 sigmas 传递给管道。
import torch
from diffusers import DiffusionPipeline, EulerDiscreteScheduler
model_id = "stabilityai/stable-diffusion-xl-base-1.0"
pipeline = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16,
variant="fp16",
).to("cuda")
pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config)
sigmas = [14.615, 6.315, 3.771, 2.181, 1.342, 0.862, 0.555, 0.380, 0.234, 0.113, 0.0]
prompt = "anthropomorphic capybara wearing a suit and working with a computer"
generator = torch.Generator(device='cuda').manual_seed(123)
image = pipeline(
prompt=prompt,
num_inference_steps=10,
sigmas=sigmas,
generator=generator
).images[0]
当你查看调度器的 timesteps 参数时,你会发现它与 AYS 时间步长调度相同,因为 timestep 调度是根据 sigmas 计算得出的。
print(f" timesteps: {pipe.scheduler.timesteps}")
"timesteps: tensor([999., 845., 730., 587., 443., 310., 193., 116., 53., 13.], device='cuda:0')"
Karras sigmas
Karras sigmas 不应用于未使用它们进行训练的模型。例如,基础 Stable Diffusion XL 模型不应使用 Karras sigmas,而 DreamShaperXL 模型可以使用,因为它们是使用 Karras sigmas 进行训练的。
Karras 调度器使用来自《阐明扩散式生成模型的设计空间》论文中的时间步长调度和 sigma 值。与其他调度器相比,这种调度器变体在采样过程接近结束时每步添加的噪声更少,并且可以提高生成图像的细节水平。
通过在调度器中设置 use_karras_sigmas=True 来启用 Karras sigmas。
import torch
from diffusers import StableDiffusionXLPipeline, DPMSolverMultistepScheduler
pipeline = StableDiffusionXLPipeline.from_pretrained(
"SG161222/RealVisXL_V4.0",
torch_dtype=torch.float16,
variant="fp16",
).to("cuda")
pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config, algorithm_type="sde-dpmsolver++", use_karras_sigmas=True)
prompt = "A cinematic shot of a cute little rabbit wearing a jacket and doing a thumbs up"
generator = torch.Generator(device="cpu").manual_seed(2487854446)
image = pipeline(
prompt=prompt,
negative_prompt="",
generator=generator,
).images[0]
重置噪声调度
在《Common Diffusion Noise Schedules and Sample Steps are Flawed》论文中,作者发现常见的噪声调度允许部分信号泄漏到最后一个时间步。这种在推理过程中的信号泄漏会导致模型只能生成中等亮度的图像。通过强制时间步调度为零信噪比(SNR)并从最后一个时间步进行采样,可以改进模型以生成非常明亮或黑暗的图像。
对于推理,你需要一个使用 v_prediction 训练过的模型。要使用 v_prediction 训练自己的模型,请在 train_text_to_image.py 或 train_text_to_image_lora.py 脚本中添加以下标志。
--prediction_type="v_prediction"
例如,加载使用 v_prediction 和 DDIMScheduler 训练的 ptx0/pseudo-journey-v2 检查点。在 DDIMScheduler 中配置以下参数:
rescale_betas_zero_snr=True:将噪声调度重新缩放到零信噪比timestep_spacing="trailing":从最后一个时间步开始采样
在管道中将 guidance_rescale 设置为防止过度曝光。较低值会增加亮度,但部分细节可能会显得失真。
from diffusers import DiffusionPipeline, DDIMScheduler
pipeline = DiffusionPipeline.from_pretrained("ptx0/pseudo-journey-v2", use_safetensors=True)
pipeline.scheduler = DDIMScheduler.from_config(
pipeline.scheduler.config, rescale_betas_zero_snr=True, timestep_spacing="trailing"
)
pipeline.to("cuda")
prompt = "cinematic photo of a snowy mountain at night with the northern lights aurora borealis overhead, 35mm photograph, film, professional, 4k, highly detailed"
generator = torch.Generator(device="cpu").manual_seed(23)
image = pipeline(prompt, guidance_rescale=0.7, generator=generator).images[0]
image
管道回调
管道的去噪循环可以通过使用 callback_on_step_end 参数和自定义定义的函数进行修改。回调函数在每个步骤结束时执行,并修改管道属性和变量以供下一步使用。这对于动态调整某些管道属性或修改张量变量非常有用。这种多功能性允许实现有趣的用例,例如在每个时间步更改提示嵌入、为提示嵌入分配不同的权重以及编辑指导比例。使用回调,您可以在不修改底层代码的情况下实现新功能!
官方回调
- SDCFGCutoffCallback:对所有 SD 1.5 流程(包括文本到图像、图像到图像、修复和 ControlNet)在特定步数后禁用 CFG。
- SDXLCFGCutoffCallback:对所有 SDXL 流程(包括文本到图像、图像到图像、修复和 ControlNet)在特定步数后禁用 CFG。
- IPAdapterScaleCutoffCallback:对所有支持 IP-Adapter 的流水线,在特定步骤数后禁用 IP 适配器。
要设置一个回调,您需要指定在多少去噪步骤后回调生效。您可以通过使用这两个参数中的任意一个来实现:
cutoff_step_ratio:步骤数的浮点数比例。cutoff_step_index:整数,表示步骤的确切编号。
import torch
from diffusers import DPMSolverMultistepScheduler, StableDiffusionXLPipeline
from diffusers.callbacks import SDXLCFGCutoffCallback
callback = SDXLCFGCutoffCallback(cutoff_step_ratio=0.4)
# can also be used with cutoff_step_index
# callback = SDXLCFGCutoffCallback(cutoff_step_ratio=None, cutoff_step_index=10)
pipeline = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16,
variant="fp16",
).to("cuda")
pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config, use_karras_sigmas=True)
prompt = "a sports car at the road, best quality, high quality, high detail, 8k resolution"
generator = torch.Generator(device="cpu").manual_seed(2628670641)
out = pipeline(
prompt=prompt,
negative_prompt="",
guidance_scale=6.5,
num_inference_steps=25,
generator=generator,
callback_on_step_end=callback,
)
out.images[0].save("official_callback.png")
动态自由分类器引导
动态自由分类器引导(CFG)是一项功能,允许你在一定数量的推理步骤后禁用 CFG,这可以帮助你以最小的性能损失来节省计算资源。用于此功能的回调函数应具有以下参数:
pipeline(或管道实例)提供对重要属性如num_timesteps和guidance_scale的访问。你可以通过更新底层属性来修改这些属性。在这个例子中,你将通过设置pipeline._guidance_scale=0.0来禁用 CFG。step_index和timestep会告诉你当前在去噪循环中的位置。使用step_index在达到num_timesteps的 40% 后关闭 CFG。callback_kwargs是一个包含在去噪循环中可以修改的张量变量的字典。它仅包括在callback_on_step_end_tensor_inputs参数中指定的变量,该参数传递给管道的__call__方法。不同的管道可能使用不同的变量集,因此请检查管道的_callback_tensor_inputs属性以获取可修改的变量列表。一些常见的变量包括latents和prompt_embeds。对于这个函数,在设置guidance_scale=0.0后更改prompt_embeds的批处理大小,以便它能够正常工作。
def callback_dynamic_cfg(pipe, step_index, timestep, callback_kwargs):
# adjust the batch_size of prompt_embeds according to guidance_scale
if step_index == int(pipeline.num_timesteps * 0.4):
prompt_embeds = callback_kwargs["prompt_embeds"]
prompt_embeds = prompt_embeds.chunk(2)[-1]
# update guidance_scale and prompt_embeds
pipeline._guidance_scale = 0.0
callback_kwargs["prompt_embeds"] = prompt_embeds
return callback_kwargs
现在,你可以将回调函数传递给 callback_on_step_end 参数,以及 prompt_embeds 到 callback_on_step_end_tensor_inputs。
import torch
from diffusers import StableDiffusionPipeline
pipeline = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16)
pipeline = pipeline.to("cuda")
prompt = "a photo of an astronaut riding a horse on mars"
generator = torch.Generator(device="cuda").manual_seed(1)
out = pipeline(
prompt,
generator=generator,
callback_on_step_end=callback_dynamic_cfg,
callback_on_step_end_tensor_inputs=['prompt_embeds']
)
out.images[0].save("out_custom_cfg.png")
中断扩散过程
在构建与 Diffusers 合作的 UI 时,提前停止扩散过程很有用,因为它允许用户在他们对中间结果不满意时停止生成过程。你可以通过回调将此功能集成到你的管道中。
此回调函数应接受以下参数:pipeline、i、t 和 callback_kwargs(必须返回)。将管道的 _interrupt 属性设置为 True,以在特定步数后停止扩散过程。你也可以在回调中实现自己的自定义停止逻辑。
在这个例子中,尽管 num_inference_steps 设置为 50,但扩散过程在 10 步后停止。
from diffusers import StableDiffusionPipeline
pipeline = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5")
pipeline.enable_model_cpu_offload()
num_inference_steps = 50
def interrupt_callback(pipeline, i, t, callback_kwargs):
stop_idx = 10
if i == stop_idx:
pipeline._interrupt = True
return callback_kwargs
pipeline(
"A photo of a cat",
num_inference_steps=num_inference_steps,
callback_on_step_end=interrupt_callback,
)
适配器截止
IP Adapter 是一种图像提示适配器,可用于无需对底层模型进行任何修改的扩散模型。我们可以使用 IP Adapter 截断回调在特定步数后禁用 IP Adapter。要设置该回调,您需要指定在回调生效后的去噪步数。您可以通过使用以下两个参数中的任意一个来实现:
cutoff_step_ratio:步骤数的浮点数比例。cutoff_step_index:整数,表示步骤的确切编号。
# 下载扩散模型并为其加载 ip_adapter,如下所示:
from diffusers import AutoPipelineForText2Image
from diffusers.utils import load_image
import torch
pipeline = AutoPipelineForText2Image.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16).to("cuda")
pipeline.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name="ip-adapter_sdxl.bin")
pipeline.set_ip_adapter_scale(0.6)
# 回调设置
from diffusers import AutoPipelineForText2Image
from diffusers.callbacks import IPAdapterScaleCutoffCallback
from diffusers.utils import load_image
import torch
pipeline = AutoPipelineForText2Image.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16
).to("cuda")
pipeline.load_ip_adapter(
"h94/IP-Adapter",
subfolder="sdxl_models",
weight_name="ip-adapter_sdxl.bin"
)
pipeline.set_ip_adapter_scale(0.6)
callback = IPAdapterScaleCutoffCallback(
cutoff_step_ratio=None,
cutoff_step_index=5
)
image = load_image(
"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/ip_adapter_diner.png"
)
generator = torch.Generator(device="cuda").manual_seed(2628670641)
images = pipeline(
prompt="a tiger sitting in a chair drinking orange juice",
ip_adapter_image=image,
negative_prompt="deformed, ugly, wrong proportion, low res, bad anatomy, worst quality, low quality",
generator=generator,
num_inference_steps=50,
callback_on_step_end=callback,
).images
images[0].save("custom_callback_img.png")
在每一步生成后显示图像
通过在每一步后访问并将潜在值转换为图像来在每一步生成后显示图像。潜在空间被压缩到 128x128,因此图像也是 128x128,这对于快速预览很有用。
- 使用以下函数将 SDXL 潜在值(4 个通道)转换为 RGB 张量(3 个通道),如 SDXL 潜在空间解释博客中所述。
def latents_to_rgb(latents):
weights = (
(60, -60, 25, -70),
(60, -5, 15, -50),
(60, 10, -5, -35),
)
weights_tensor = torch.t(torch.tensor(weights, dtype=latents.dtype).to(latents.device))
biases_tensor = torch.tensor((150, 140, 130), dtype=latents.dtype).to(latents.device)
rgb_tensor = torch.einsum("...lxy,lr -> ...rxy", latents, weights_tensor) + biases_tensor.unsqueeze(-1).unsqueeze(-1)
image_array = rgb_tensor.clamp(0, 255).byte().cpu().numpy().transpose(1, 2, 0)
return Image.fromarray(image_array)
- 创建一个函数来解码并将潜空间保存为图像。
def decode_tensors(pipe, step, timestep, callback_kwargs):
latents = callback_kwargs["latents"]
image = latents_to_rgb(latents[0])
image.save(f"{step}.png")
return callback_kwargs
- 将
decode_tensors函数传递给callback_on_step_end参数,以便在每个步骤后解码张量。您还需要在callback_on_step_end_tensor_inputs参数中指定您想要修改的内容,在本例中是潜在变量。
from diffusers import AutoPipelineForText2Image
import torch
from PIL import Image
pipeline = AutoPipelineForText2Image.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16,
variant="fp16",
use_safetensors=True
).to("cuda")
image = pipeline(
prompt="A croissant shaped like a cute bear.",
negative_prompt="Deformed, ugly, bad anatomy",
callback_on_step_end=decode_tensors,
callback_on_step_end_tensor_inputs=["latents"],
).images[0]
管道复现
扩散模型本质上具有随机性,这也是它每次运行时能生成不同输出的原因。但在某些情况下,你希望每次都能生成相同的输出,比如在测试、复制结果,甚至提升图像质量时。虽然你无法期望跨平台获得完全一致的结果,但可以期望在一定容差范围内跨版本和平台获得可复现的结果(尽管即使这一点也可能存在差异)。
本指南将向你展示如何在 CPU 和 GPU 上控制随机性以实现确定性生成。
控制随机性
在推理过程中,管道严重依赖随机采样操作,包括创建高斯噪声张量以去噪以及在调度步骤中添加噪声。
查看 DDIMPipeline 在两次推理步骤后的张量值。
from diffusers import DDIMPipeline
import numpy as np
ddim = DDIMPipeline.from_pretrained("google/ddpm-cifar10-32", use_safetensors=True)
image = ddim(num_inference_steps=2, output_type="np").images
print(np.abs(image).sum())
运行上述代码会打印一个值,但如果你再次运行它,你会得到不同的值。
每次运行管道时,torch.randn 都会使用不同的随机种子来创建高斯噪声张量。这导致每次运行的结果都不同,并使扩散管道能够每次生成不同的随机图像。
但如果您需要可靠地生成相同的图像,那取决于您是在 CPU 还是 GPU 上运行管道。
CPU 上可重复结果:你需要使用 PyTorch Generator 并设置一个种子。现在当你运行代码时,它总是打印 1491.1711 的值,因为带有种子的 Generator 对象被传递到管道中的所有随机函数。无论你在什么硬件和 PyTorch 版本上运行,都应该得到相似(如果不是完全相同)的结果。
import torch
import numpy as np
from diffusers import DDIMPipeline
ddim = DDIMPipeline.from_pretrained("google/ddpm-cifar10-32", use_safetensors=True)
generator = torch.Generator(device="cpu").manual_seed(0)
image = ddim(num_inference_steps=2, output_type="np", generator=generator).images
print(np.abs(image).sum())
GPU 上可重复结果:在 GPU 上编写可复现的流程稍微有些棘手,并且无法保证在不同硬件上完全可复现,因为矩阵乘法——扩散流程需要大量使用——在 GPU 上的确定性不如 CPU。例如,如果你运行 CPU 示例中相同的代码示例,即使种子相同,你也会得到不同的结果。这是因为 GPU 使用的是与 CPU 不同的随机数生成器。
import torch
import numpy as np
from diffusers import DDIMPipeline
ddim = DDIMPipeline.from_pretrained("google/ddpm-cifar10-32", use_safetensors=True)
ddim.to("cuda")
generator = torch.Generator(device="cuda").manual_seed(0)
image = ddim(num_inference_steps=2, output_type="np", generator=generator).images
print(np.abs(image).sum())
为了避免这个问题,Diffusers 提供了 randn_tensor() 函数用于在 CPU 上创建随机噪声,并在必要时将张量移动到 GPU。randn_tensor() 函数在管道内部被广泛使用。现在你可以调用 torch.manual_seed,它会自动创建一个 CPU Generator,即使管道在 GPU 上运行也可以将其传递给管道。
import torch
import numpy as np
from diffusers import DDIMPipeline
ddim = DDIMPipeline.from_pretrained("google/ddpm-cifar10-32", use_safetensors=True)
ddim.to("cuda")
generator = torch.manual_seed(0)
image = ddim(num_inference_steps=2, output_type="np", generator=generator).images
print(np.abs(image).sum())
确定性算法
你也可以配置 PyTorch 使用确定性算法来创建可重复的管道。缺点是确定性算法可能比非确定性算法慢,并且你可能会观察到性能下降。
当操作在多个 CUDA 流中启动时会发生非确定性行为。为了避免这种情况,将环境变量 CUBLAS_WORKSPACE_CONFIG 设置为 :16:8,以在运行时只使用一个缓冲区大小。
PyTorch 通常会测试多种算法来选择最快的那个,但如果您需要可重复性,应该禁用这个功能,因为测试可能会每次选择不同的算法。设置 Diffusers enable_full_determinism 为 True 可以启用确定性算法。
import torch
from diffusers import DDIMScheduler, StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", use_safetensors=True).to("cuda")
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
g = torch.Generator(device="cuda")
prompt = "A bear is playing a guitar on Times Square"
g.manual_seed(0)
result1 = pipe(prompt=prompt, num_inference_steps=50, generator=g, output_type="latent").images
g.manual_seed(0)
result2 = pipe(prompt=prompt, num_inference_steps=50, generator=g, output_type="latent").images
print("L_inf dist =", abs(result1 - result2).max())
"L_inf dist = tensor(0., device='cuda:0')"
确定性批量生成
创建可重复流程的一个实际应用是确定性批量生成。您生成一批图像,并选择其中一张使用更详细的提示来改进。主要思路是将一个 Generator 列表传递给流程,并将每个 Generator 与一个种子绑定,以便可以重复使用。
让我们使用 stable-diffusion-v1-5/stable-diffusion-v1-5 检查点来生成一批图像。
import torch
from diffusers import DiffusionPipeline
from diffusers.utils import make_image_grid
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
pipeline = pipeline.to("cuda")
定义四个不同的 Generator,并为每个 Generator 分配一个种子(0 到 3)。然后生成一批图像,并选择其中一个进行迭代。
generator = [torch.Generator(device="cuda").manual_seed(i) for i in range(4)]
prompt = "Labrador in the style of Vermeer"
images = pipeline(prompt, generator=generator, num_images_per_prompt=4).images[0]
make_image_grid(images, rows=2, cols=2)
让我们改进与 Generator 对应的第一张图片(你可以选择任何你想要的图片),使用种子 0。在你的提示中添加一些额外文字,然后确保你使用相同的 Generator 和种子 0 进行重用。所有生成的图片都应该与第一张图片相似。
prompt = [prompt + t for t in [", highly realistic", ", artsy", ", trending", ", colorful"]]
generator = [torch.Generator(device="cuda").manual_seed(0) for i in range(4)]
images = pipeline(prompt, generator=generator).images
make_image_grid(images, rows=2, cols=2)
控制图片质量
扩散模型的组件,如 UNet 和调度器,可以优化以提升生成图像的质量,从而获得更好的细节。这些技术在你没有资源直接使用更大模型进行推理时尤其有用。你可以在推理过程中启用这些技术,而无需任何额外的训练。
FreeU
FreeU 通过重新平衡 UNet 的主干和跳跃连接的权重来提升图像细节。跳跃连接可能导致模型忽略主干的一些语义信息,从而在生成的图像中出现不自然的细节。这种技术无需任何额外的训练,可以在推理过程中即时应用,用于图像到图像和文本到视频等任务。
在你的管道上使用 enable_freeu() 方法,并配置主干网络(b1 和 b2)和跳跃连接(s1 和 s2)的缩放因子。每个缩放因子后面的数字对应于 UNet 中应用该因子的阶段。参考不同模型的超参数,请查看 FreeU 仓库。
禁用 FreeU:调用 pipelines.StableDiffusionMixin.disable_freeu() 方法来禁用 FreeU。
pipeline.disable_freeu()
Stable Diffusion v1-5
import torch
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, safety_checker=None
).to("cuda")
pipeline.enable_freeu(s1=0.9, s2=0.2, b1=1.5, b2=1.6)
generator = torch.Generator(device="cpu").manual_seed(33)
prompt = ""
image = pipeline(prompt, generator=generator).images[0]
image
Stable Diffusion v2-1
import torch
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16, safety_checker=None
).to("cuda")
pipeline.enable_freeu(s1=0.9, s2=0.2, b1=1.4, b2=1.6)
generator = torch.Generator(device="cpu").manual_seed(80)
prompt = "A squirrel eating a burger"
image = pipeline(prompt, generator=generator).images[0]
image
Stable Diffusion XL
import torch
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16,
).to("cuda")
pipeline.enable_freeu(s1=0.9, s2=0.2, b1=1.3, b2=1.4)
generator = torch.Generator(device="cpu").manual_seed(13)
prompt = "A squirrel eating a burger"
image = pipeline(prompt, generator=generator).images[0]
image
Zeroscope
import torch
from diffusers import DiffusionPipeline
from diffusers.utils import export_to_video
pipeline = DiffusionPipeline.from_pretrained(
"damo-vilab/text-to-video-ms-1.7b", torch_dtype=torch.float16
).to("cuda")
# values come from https://github.com/lyn-rgb/FreeU_Diffusers#video-pipelines
pipeline.enable_freeu(b1=1.2, b2=1.4, s1=0.9, s2=0.2)
prompt = "Confident teddy bear surfer rides the wave in the tropics"
generator = torch.Generator(device="cpu").manual_seed(47)
video_frames = pipeline(prompt, generator=generator).frames[0]
export_to_video(video_frames, "teddy_bear.mp4", fps=10)
提示词技术
提示词很重要,因为它们描述了你希望扩散模型生成的内容。最好的提示词应该是详细、具体且结构良好的,以帮助模型实现你的愿景。但制作一个优秀的提示词需要时间和精力,有时可能还不够,因为语言和词汇可能不够精确。这就是你需要借助其他技巧来增强提示词,例如提示词增强和提示词加权,以获得你想要的结果的地方。
提示工程
新的扩散模型在从基本提示生成高质量图像方面做得相当不错,但创建一个精心撰写的提示仍然非常重要,以获得最佳结果。以下是一些撰写良好提示的建议:
- 图像的媒介是什么?是照片、绘画、3D 插图,还是其他什么?
- 图像的主题是什么?是人、动物、物体,还是场景?
- 你想在图像中看到哪些细节?这是你可以真正发挥创意并有很多乐趣尝试不同词语来让图像生动起来的地方。例如,光线怎么样?氛围和美学是什么?你在寻找什么样的艺术或插画风格?你使用的词语越具体和精确,模型就越能理解你想生成的内容。
使用 GPT2 增强提示
提示增强是一种快速提升提示质量的技术,而无需花费过多精力构建提示。它使用在 Stable Diffusion 文本提示上预训练的 GPT2 模型,自动用额外的关键词丰富提示,以生成高质量图像。
该技术通过筛选特定关键词并强制模型生成这些词语来增强原始提示。这样,你的提示可以是”一只猫”,而 GPT2 可以将其增强为”土耳其屋顶上沐浴阳光的猫的电影剧照,高度细节,高预算好莱坞电影,宽银幕,情绪化,史诗般,华丽,胶片颗粒质量,锐利焦点,精美细节,复杂,惊人,史诗般”。
首先,定义某些风格和一组词语(你可以查看 Fooocus 使用的更全面词语和风格列表),以增强提示。
import torch
from transformers import GenerationConfig, GPT2LMHeadModel, GPT2Tokenizer, LogitsProcessor, LogitsProcessorList
from diffusers import StableDiffusionXLPipeline
styles = {
"cinematic": "cinematic film still of {prompt}, highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain",
"anime": "anime artwork of {prompt}, anime style, key visual, vibrant, studio anime, highly detailed",
"photographic": "cinematic photo of {prompt}, 35mm photograph, film, professional, 4k, highly detailed",
"comic": "comic of {prompt}, graphic illustration, comic art, graphic novel art, vibrant, highly detailed",
"lineart": "line art drawing {prompt}, professional, sleek, modern, minimalist, graphic, line art, vector graphics",
"pixelart": " pixel-art {prompt}, low-res, blocky, pixel art style, 8-bit graphics",
}
words = [
"aesthetic", "astonishing", "beautiful", "breathtaking", "composition", "contrasted", "epic", "moody", "enhanced",
"exceptional", "fascinating", "flawless", "glamorous", "glorious", "illumination", "impressive", "improved",
"inspirational", "magnificent", "majestic", "hyperrealistic", "smooth", "sharp", "focus", "stunning", "detailed",
"intricate", "dramatic", "high", "quality", "perfect", "light", "ultra", "highly", "radiant", "satisfying",
"soothing", "sophisticated", "stylish", "sublime", "terrific", "touching", "timeless", "wonderful", "unbelievable",
"elegant", "awesome", "amazing", "dynamic", "trendy",
]
你可能已经注意到在 words 列表中,有些词可以组合在一起形成更有意义的内容。例如,“high” 和 “quality” 可以组合成 “high quality”。让我们将这些词组合起来,并移除那些无法组合的词。
word_pairs = ["highly detailed", "high quality", "enhanced quality", "perfect composition", "dynamic light"]
def find_and_order_pairs(s, pairs):
words = s.split()
found_pairs = []
for pair in pairs:
pair_words = pair.split()
if pair_words[0] in words and pair_words[1] in words:
found_pairs.append(pair)
words.remove(pair_words[0])
words.remove(pair_words[1])
for word in words[:]:
for pair in pairs:
if word in pair.split():
words.remove(word)
break
ordered_pairs = ", ".join(found_pairs)
remaining_s = ", ".join(words)
return ordered_pairs, remaining_s
接下来,实现一个自定义的 LogitsProcessor 类,将 words 列表中的词赋予 0 值,将 words 列表之外的词赋予负值,这样在生成过程中就不会选择这些词。这样,生成过程就会倾向于 words 列表中的词。当列表中的词被使用后,它也会被赋予负值,这样就不会再次被选择。
class CustomLogitsProcessor(LogitsProcessor):
def __init__(self, bias):
super().__init__()
self.bias = bias
def __call__(self, input_ids, scores):
if len(input_ids.shape) == 2:
last_token_id = input_ids[0, -1]
self.bias[last_token_id] = -1e10
return scores + self.bias
word_ids = [tokenizer.encode(word, add_prefix_space=True)[0] for word in words]
bias = torch.full((tokenizer.vocab_size,), -float("Inf")).to("cuda")
bias[word_ids] = 0
processor = CustomLogitsProcessor(bias)
processor_list = LogitsProcessorList([processor])
将提示与之前在 styles 字典中定义的 cinematic 样式提示相结合。
prompt = "a cat basking in the sun on a roof in Turkey"
style = "cinematic"
prompt = styles[style].format(prompt=prompt)
prompt
# "cinematic film still of a cat basking in the sun on a roof in Turkey, highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain"
从 Gustavosta/MagicPrompt-Stable-Diffusion 检查点加载 GPT2 分词器和模型(此特定检查点用于生成提示),以增强提示。
tokenizer = GPT2Tokenizer.from_pretrained("Gustavosta/MagicPrompt-Stable-Diffusion")
model = GPT2LMHeadModel.from_pretrained("Gustavosta/MagicPrompt-Stable-Diffusion", torch_dtype=torch.float16).to(
"cuda"
)
model.eval()
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
token_count = inputs["input_ids"].shape[1]
max_new_tokens = 50 - token_count
generation_config = GenerationConfig(
penalty_alpha=0.7,
top_k=50,
eos_token_id=model.config.eos_token_id,
pad_token_id=model.config.eos_token_id,
pad_token=model.config.pad_token_id,
do_sample=True,
)
with torch.no_grad():
generated_ids = model.generate(
input_ids=inputs["input_ids"],
attention_mask=inputs["attention_mask"],
max_new_tokens=max_new_tokens,
generation_config=generation_config,
logits_processor=proccesor_list,
)
然后你可以将输入提示和生成的提示结合起来。你可以随意查看生成的提示(generated_part)、找到的词对(pairs)以及剩余的词(words)。所有这些都被打包在 enhanced_prompt 中。
output_tokens = [tokenizer.decode(generated_id, skip_special_tokens=True) for generated_id in generated_ids]
input_part, generated_part = output_tokens[0][: len(prompt)], output_tokens[0][len(prompt) :]
pairs, words = find_and_order_pairs(generated_part, word_pairs)
formatted_generated_part = pairs + ", " + words
enhanced_prompt = input_part + ", " + formatted_generated_part
enhanced_prompt
# ["cinematic film still of a cat basking in the sun on a roof in Turkey, highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain quality sharp focus beautiful detailed intricate stunning amazing epic"]
最后,加载一个管道和低权重的偏移噪声 LoRA,以生成具有增强提示的图像。
pipeline = StableDiffusionXLPipeline.from_pretrained(
"RunDiffusion/Juggernaut-XL-v9", torch_dtype=torch.float16, variant="fp16"
).to("cuda")
pipeline.load_lora_weights(
"stabilityai/stable-diffusion-xl-base-1.0",
weight_name="sd_xl_offset_example-lora_1.0.safetensors",
adapter_name="offset",
)
pipeline.set_adapters(["offset"], adapter_weights=[0.2])
image = pipeline(
enhanced_prompt,
width=1152,
height=896,
guidance_scale=7.5,
num_inference_steps=25,
).images[0]
image
提示权重
提示权重提供了一种强调或弱化提示某些部分的方法,从而对生成的图像有更多的控制。提示可以包含多个概念,这些概念被转换为上下文化的文本嵌入。这些嵌入被模型用来调节其交叉注意力层以生成图像(阅读 Stable Diffusion 博客文章了解更多关于其工作原理)。
提示权重通过增加或减少与提示中概念对应的文本嵌入向量的规模来工作,因为你可能并不希望模型平等地关注所有概念。准备提示嵌入最容易的方法是使用 Stable Diffusion 长提示加权嵌入(sd_embed)。一旦你获得了提示加权嵌入,你就可以将它们传递给任何具有 prompt_embeds(以及可选的 negative_prompt_embeds)参数的流程,例如 StableDiffusionPipeline、StableDiffusionControlNetPipeline 和 StableDiffusionXLPipeline。
开始之前,请确保您已安装最新版本的 sd_embed:
pip install git+https://github.com/xhinker/sd_embed.git@main
使用 StableDiffusionXLPipeline:
from diffusers import StableDiffusionXLPipeline, UniPCMultistepScheduler
import torch
pipe = StableDiffusionXLPipeline.from_pretrained("Lykon/dreamshaper-xl-1-0", torch_dtype=torch.float16)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe.to("cuda")
要增加或减少某个概念的权重,请用括号将文本括起来。括号越多,对文本的权重就越大。您也可以在文本后附加一个数值乘数,以指示您希望增加或减少其权重的程度。
| 格式 | 乘数效果 |
|---|---|
(hippo) | 增加 1.1 倍 |
((hippo)) | 增加 1.21 倍 |
(hippo:1.5) | 增加 1.5 倍 |
(hippo:0.25) | 减少 4 倍 |
创建一个提示,并使用括号和数字乘数来提高不同文本的权重。
from sd_embed.embedding_funcs import get_weighted_text_embeddings_sdxl
prompt = """A whimsical and creative image depicting a hybrid creature that is a mix of a waffle and a hippopotamus.
This imaginative creature features the distinctive, bulky body of a hippo,
but with a texture and appearance resembling a golden-brown, crispy waffle.
The creature might have elements like waffle squares across its skin and a syrup-like sheen.
It's set in a surreal environment that playfully combines a natural water habitat of a hippo with elements of a breakfast table setting,
possibly including oversized utensils or plates in the background.
The image should evoke a sense of playful absurdity and culinary fantasy.
"""
neg_prompt = """\
skin spots,acnes,skin blemishes,age spot,(ugly:1.2),(duplicate:1.2),(morbid:1.21),(mutilated:1.2),\
(tranny:1.2),mutated hands,(poorly drawn hands:1.5),blurry,(bad anatomy:1.2),(bad proportions:1.3),\
extra limbs,(disfigured:1.2),(missing arms:1.2),(extra legs:1.2),(fused fingers:1.5),\
(too many fingers:1.5),(unclear eyes:1.2),lowers,bad hands,missing fingers,extra digit,\
bad hands,missing fingers,(extra arms and legs),(worst quality:2),(low quality:2),\
(normal quality:2),lowres,((monochrome)),((grayscale))
"""
使用 get_weighted_text_embeddings_sdxl 函数来生成提示嵌入和负提示嵌入。由于你正在使用 SDXL 模型,它还会生成池化提示嵌入和负池化提示嵌入。
(
prompt_embeds,
prompt_neg_embeds,
pooled_prompt_embeds,
negative_pooled_prompt_embeds
) = get_weighted_text_embeddings_sdxl(
pipe,
prompt=prompt,
neg_prompt=neg_prompt
)
image = pipe(
prompt_embeds=prompt_embeds,
negative_prompt_embeds=prompt_neg_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
num_inference_steps=30,
height=1024,
width=1024 + 512,
guidance_scale=4.0,
generator=torch.Generator("cuda").manual_seed(2)
).images[0]
image
文本反转
文本反转是一种从一些图像中学习特定概念的技术,你可以使用它来根据该概念生成新的图像。
创建一个流程,并使用 load_textual_inversion() 函数加载文本反转嵌入:
import torch
from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5",
torch_dtype=torch.float16,
).to("cuda")
pipe.load_textual_inversion("sd-concepts-library/midjourney-style")
在提示中添加 <midjourney-style> 文本以触发文本反转。
from sd_embed.embedding_funcs import get_weighted_text_embeddings_sd15
prompt = """<midjourney-style> A whimsical and creative image depicting a hybrid creature that is a mix of a waffle and a hippopotamus.
This imaginative creature features the distinctive, bulky body of a hippo,
but with a texture and appearance resembling a golden-brown, crispy waffle.
The creature might have elements like waffle squares across its skin and a syrup-like sheen.
It's set in a surreal environment that playfully combines a natural water habitat of a hippo with elements of a breakfast table setting,
possibly including oversized utensils or plates in the background.
The image should evoke a sense of playful absurdity and culinary fantasy.
"""
neg_prompt = """\
skin spots,acnes,skin blemishes,age spot,(ugly:1.2),(duplicate:1.2),(morbid:1.21),(mutilated:1.2),\
(tranny:1.2),mutated hands,(poorly drawn hands:1.5),blurry,(bad anatomy:1.2),(bad proportions:1.3),\
extra limbs,(disfigured:1.2),(missing arms:1.2),(extra legs:1.2),(fused fingers:1.5),\
(too many fingers:1.5),(unclear eyes:1.2),lowers,bad hands,missing fingers,extra digit,\
bad hands,missing fingers,(extra arms and legs),(worst quality:2),(low quality:2),\
(normal quality:2),lowres,((monochrome)),((grayscale))
"""
使用 get_weighted_text_embeddings_sd15 函数来生成提示嵌入和负提示嵌入。
(
prompt_embeds,
prompt_neg_embeds,
) = get_weighted_text_embeddings_sd15(
pipe,
prompt=prompt,
neg_prompt=neg_prompt
)
image = pipe(
prompt_embeds=prompt_embeds,
negative_prompt_embeds=prompt_neg_embeds,
height=768,
width=896,
guidance_scale=4.0,
generator=torch.Generator("cuda").manual_seed(2)
).images[0]
image
DreamBooth
DreamBooth 是一种在仅提供少量训练图像的情况下,生成与特定主题相关的图像的技术。它与文本反转类似,但 DreamBooth 会训练整个模型,而文本反转仅微调文本嵌入。这意味着你应该使用 from_pretrained() 来加载 DreamBooth 模型:
import torch
from diffusers import DiffusionPipeline, UniPCMultistepScheduler
pipe = DiffusionPipeline.from_pretrained("sd-dreambooth-library/dndcoverart-v1", torch_dtype=torch.float16).to("cuda")
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
根据你使用的模型,你需要将模型的唯一标识符整合到你的提示中。例如,dndcoverart-v1 模型使用标识符 dndcoverart:
from sd_embed.embedding_funcs import get_weighted_text_embeddings_sd15
prompt = """dndcoverart of A whimsical and creative image depicting a hybrid creature that is a mix of a waffle and a hippopotamus.
This imaginative creature features the distinctive, bulky body of a hippo,
but with a texture and appearance resembling a golden-brown, crispy waffle.
The creature might have elements like waffle squares across its skin and a syrup-like sheen.
It's set in a surreal environment that playfully combines a natural water habitat of a hippo with elements of a breakfast table setting,
possibly including oversized utensils or plates in the background.
The image should evoke a sense of playful absurdity and culinary fantasy.
"""
neg_prompt = """\
skin spots,acnes,skin blemishes,age spot,(ugly:1.2),(duplicate:1.2),(morbid:1.21),(mutilated:1.2),\
(tranny:1.2),mutated hands,(poorly drawn hands:1.5),blurry,(bad anatomy:1.2),(bad proportions:1.3),\
extra limbs,(disfigured:1.2),(missing arms:1.2),(extra legs:1.2),(fused fingers:1.5),\
(too many fingers:1.5),(unclear eyes:1.2),lowers,bad hands,missing fingers,extra digit,\
bad hands,missing fingers,(extra arms and legs),(worst quality:2),(low quality:2),\
(normal quality:2),lowres,((monochrome)),((grayscale))
"""
(
prompt_embeds
, prompt_neg_embeds
) = get_weighted_text_embeddings_sd15(
pipe
, prompt = prompt
, neg_prompt = neg_prompt
)
推理进阶
扩图
Outpainting 将图像扩展到其原始边界之外,允许您添加、替换或修改图像中的视觉元素,同时保留原始图像。与 inpainting 类似,您需要用新的视觉元素填充白色区域(在这种情况下,是原始图像之外的区域),同时保持原始图像(由黑色像素的掩码表示)。Outpainting 有几种方法,例如使用 ControlNet 或使用 Differential Diffusion。
下面展示如何使用 inpainting 模型、ControlNet 和 ZoeDepth 估计器进行 Outpainting。
开始之前,请确保已安装 controlnet_aux 库,以便使用 ZoeDepth 估计器。
pip install -q controlnet_aux
图像准备
首先选择一张要扩展的图片,并使用类似 BRIA-RMBG-1.4 的 Space 工具去除背景。
Stable Diffusion XL (SDXL) 模型最适合使用 1024x1024 的图片,但只要你的硬件有足够的内存支持,你可以将图片调整到任何尺寸。图片中的透明背景也应替换为白色背景。创建一个函数(如下所示),将图片缩放并粘贴到白色背景上。
import random
import requests
import torch
from controlnet_aux import ZoeDetector
from PIL import Image, ImageOps
from diffusers import (
AutoencoderKL,
ControlNetModel,
StableDiffusionXLControlNetPipeline,
StableDiffusionXLInpaintPipeline,
)
def scale_and_paste(original_image):
aspect_ratio = original_image.width / original_image.height
if original_image.width > original_image.height:
new_width = 1024
new_height = round(new_width / aspect_ratio)
else:
new_height = 1024
new_width = round(new_height * aspect_ratio)
resized_original = original_image.resize((new_width, new_height), Image.LANCZOS)
white_background = Image.new("RGBA", (1024, 1024), "white")
x = (1024 - new_width) // 2
y = (1024 - new_height) // 2
white_background.paste(resized_original, (x, y), resized_original)
return resized_original, white_background
original_image = Image.open(
requests.get(
"https://huggingface.co/datasets/stevhliu/testing-images/resolve/main/no-background-jordan.png",
stream=True,
).raw
).convert("RGBA")
resized_img, white_bg_image = scale_and_paste(original_image)
为了避免添加不想要的额外细节,使用 ZoeDepth 估计器在生成过程中提供额外指导,并确保鞋子与原始图像保持一致。
zoe = ZoeDetector.from_pretrained("lllyasviel/Annotators")
image_zoe = zoe(white_bg_image, detect_resolution=512, image_resolution=1024)
image_zoe
扩展绘制
当您的图像准备就绪后,您可以使用 controlnet-inpaint-dreamer-sdxl 生成鞋子周围白色区域的内容,这是一种为修复训练的 SDXL ControlNet。
加载修复 ControlNet、ZoeDepth 模型、VAE,并将它们传递给 StableDiffusionXLControlNetPipeline。然后您可以创建一个可选的 generate_image 函数(为了方便)来修复初始图像。
controlnets = [
ControlNetModel.from_pretrained(
"destitech/controlnet-inpaint-dreamer-sdxl", torch_dtype=torch.float16, variant="fp16"
),
ControlNetModel.from_pretrained(
"diffusers/controlnet-zoe-depth-sdxl-1.0", torch_dtype=torch.float16
),
]
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16).to("cuda")
pipeline = StableDiffusionXLControlNetPipeline.from_pretrained(
"SG161222/RealVisXL_V4.0", torch_dtype=torch.float16, variant="fp16", controlnet=controlnets, vae=vae
).to("cuda")
def generate_image(prompt, negative_prompt, inpaint_image, zoe_image, seed: int = None):
if seed is None:
seed = random.randint(0, 2**32 - 1)
generator = torch.Generator(device="cpu").manual_seed(seed)
image = pipeline(
prompt,
negative_prompt=negative_prompt,
image=[inpaint_image, zoe_image],
guidance_scale=6.5,
num_inference_steps=25,
generator=generator,
controlnet_conditioning_scale=[0.5, 0.8],
control_guidance_end=[0.9, 0.6],
).images[0]
return image
prompt = "nike air jordans on a basketball court"
negative_prompt = ""
temp_image = generate_image(prompt, negative_prompt, white_bg_image, image_zoe, 908097)
将原始图像粘贴到初始的扩展图像上。你将在稍后的步骤中改进扩展的背景。
x = (1024 - resized_img.width) // 2
y = (1024 - resized_img.height) // 2
temp_image.paste(resized_img, (x, y), resized_img)
temp_image
如果你内存不足,释放一些内存:
pipeline = None
torch.cuda.empty_cache()
现在你已经有了初始的扩展图像,使用 RealVisXL 模型加载 StableDiffusionXLInpaintPipeline,以生成质量更好的最终扩展图像。
pipeline = StableDiffusionXLInpaintPipeline.from_pretrained(
"OzzyGT/RealVisXL_V4.0_inpainting",
torch_dtype=torch.float16,
variant="fp16",
vae=vae,
).to("cuda")
为最终扩展图像准备一个蒙版。为了在原始图像和扩展背景之间创建更自然的过渡,将蒙版模糊处理,以帮助它更好地融合。
mask = Image.new("L", temp_image.size)
mask.paste(resized_img.split()[3], (x, y))
mask = ImageOps.invert(mask)
final_mask = mask.point(lambda p: p > 128 and 255)
mask_blurred = pipeline.mask_processor.blur(final_mask, blur_factor=20)
mask_blurred
创建一个更好的提示,并将其传递给 generate_outpaint 函数以生成最终的抠图背景。同样,将原始图像粘贴到最终的抠图背景上。
def generate_outpaint(prompt, negative_prompt, image, mask, seed: int = None):
if seed is None:
seed = random.randint(0, 2**32 - 1)
generator = torch.Generator(device="cpu").manual_seed(seed)
image = pipeline(
prompt,
negative_prompt=negative_prompt,
image=image,
mask_image=mask,
guidance_scale=10.0,
strength=0.8,
num_inference_steps=30,
generator=generator,
).images[0]
return image
prompt = "high quality photo of nike air jordans on a basketball court, highly detailed"
negative_prompt = ""
final_image = generate_outpaint(prompt, negative_prompt, temp_image, mask_blurred, 7688778)
x = (1024 - resized_img.width) // 2
y = (1024 - resized_img.height) // 2
final_image.paste(resized_img, (x, y), resized_img)
final_image
混合推理
纵览
为何使用混合推理?
- 🚀 降低要求:无需昂贵的硬件即可访问强大的模型。
- 💎 无需妥协:在不牺牲性能的情况下实现最高质量。
- 💰 成本效益:免费!
- 🎯 多样化应用场景:完全兼容 Diffusers 🧨 及更广泛的社区。
- 🔧 开发者友好:简单请求,快速响应。
可用模型:
- VAE 解码 🖼️:快速将潜在表示解码为高质量图像,同时不牺牲性能或工作流程速度。
- VAE 编码 🔢:高效地将图像编码为潜在表示,用于生成和训练。
- 文本编码器 📃(即将推出):快速准确地计算提示文本的嵌入,确保流畅且高质量的流程。
VAE 解码
VAE 解码是扩散模型的关键组件——将潜在表示转换为图像或视频。
显存
这些表格展示了在不同 GPU 上使用 SD v1 和 SD XL 进行 VAE 解码所需的显存。
对于这些大多数 GPU,过高的内存使用率百分比会导致其他模型(文本编码器、UNet/Transformer)需要卸载,或者必须使用分块解码,这会增加所需时间并影响质量。
SD v1.5:
| GPU | Resolution | Time (seconds) | Memory (%) | Tiled Time (secs) | Tiled Memory (%) |
|---|---|---|---|---|---|
| NVIDIA GeForce RTX 4090 | 512x512 | 0.031 | 5.60% | 0.031 (0%) | 5.60% |
| NVIDIA GeForce RTX 4090 | 1024x1024 | 0.148 | 20.00% | 0.301 (+103%) | 5.60% |
| NVIDIA GeForce RTX 4080 | 512x512 | 0.05 | 8.40% | 0.050 (0%) | 8.40% |
| NVIDIA GeForce RTX 4080 | 1024x1024 | 0.224 | 30.00% | 0.356 (+59%) | 8.40% |
| NVIDIA GeForce RTX 4070 Ti | 512x512 | 0.066 | 11.30% | 0.066 (0%) | 11.30% |
| NVIDIA GeForce RTX 4070 Ti | 1024x1024 | 0.284 | 40.50% | 0.454 (+60%) | 11.40% |
| NVIDIA GeForce RTX 3090 | 512x512 | 0.062 | 5.20% | 0.062 (0%) | 5.20% |
| NVIDIA GeForce RTX 3090 | 1024x1024 | 0.253 | 18.50% | 0.464 (+83%) | 5.20% |
| NVIDIA GeForce RTX 3080 | 512x512 | 0.07 | 12.80% | 0.070 (0%) | 12.80% |
| NVIDIA GeForce RTX 3080 | 1024x1024 | 0.286 | 45.30% | 0.466 (+63%) | 12.90% |
| NVIDIA GeForce RTX 3070 | 512x512 | 0.102 | 15.90% | 0.102 (0%) | 15.90% |
| NVIDIA GeForce RTX 3070 | 1024x1024 | 0.421 | 56.30% | 0.746 (+77%) | 16.00% |
SDXL:
| GPU | Resolution | Time (seconds) | Memory Consumed (%) | Tiled Time (seconds) | Tiled Memory (%) |
|---|---|---|---|---|---|
| NVIDIA GeForce RTX 4090 | 512x512 | 0.057 | 10.00% | 0.057 (0%) | 10.00% |
| NVIDIA GeForce RTX 4090 | 1024x1024 | 0.256 | 35.50% | 0.257 (+0.4%) | 35.50% |
| NVIDIA GeForce RTX 4080 | 512x512 | 0.092 | 15.00% | 0.092 (0%) | 15.00% |
| NVIDIA GeForce RTX 4080 | 1024x1024 | 0.406 | 53.30% | 0.406 (0%) | 53.30% |
| NVIDIA GeForce RTX 4070 Ti | 512x512 | 0.121 | 20.20% | 0.120 (-0.8%) | 20.20% |
| NVIDIA GeForce RTX 4070 Ti | 1024x1024 | 0.519 | 72.00% | 0.519 (0%) | 72.00% |
| NVIDIA GeForce RTX 3090 | 512x512 | 0.107 | 10.50% | 0.107 (0%) | 10.50% |
| NVIDIA GeForce RTX 3090 | 1024x1024 | 0.459 | 38.00% | 0.460 (+0.2%) | 38.00% |
| NVIDIA GeForce RTX 3080 | 512x512 | 0.121 | 25.60% | 0.121 (0%) | 25.60% |
| NVIDIA GeForce RTX 3080 | 1024x1024 | 0.524 | 93.00% | 0.524 (0%) | 93.00% |
| NVIDIA GeForce RTX 3070 | 512x512 | 0.183 | 31.80% | 0.183 (0%) | 31.80% |
| NVIDIA GeForce RTX 3070 | 1024x1024 | 0.794 | 96.40% | 0.794 (0%) | 96.40% |
可用的 VAE:
| Endpoint | Model |
|---|---|
| Stable Diffusion v1 | https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud |
| Stable Diffusion XL | https://x2dmsqunjd6k9prw.us-east-1.aws.endpoints.huggingface.cloud |
| Flux | https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud |
| HunyuanVideo | https://o7ywnmrahorts457.us-east-1.aws.endpoints.huggingface.cloud |
代码
从 main 安装 diffusers 以运行代码:
pip install git+https://github.com/huggingface/diffusers@main
一个辅助方法简化了与 Hybrid Inference 的交互:
from diffusers.utils.remote_utils import remote_decode
基本示例:
image = remote_decode(
endpoint="https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud/",
tensor=torch.randn([1, 4, 64, 64], dtype=torch.float16),
scaling_factor=0.18215,
)
Flux 的使用略有不同。Flux 的潜空间是打包的,因此我们需要发送 height 和 width:
image = remote_decode(
endpoint="https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud/",
tensor=torch.randn([1, 4096, 64], dtype=torch.float16),
height=1024,
width=1024,
scaling_factor=0.3611,
shift_factor=0.1159,
)
HunyuanVideo 示例:
video = remote_decode(
endpoint="https://o7ywnmrahorts457.us-east-1.aws.endpoints.huggingface.cloud/",
tensor=torch.randn([1, 16, 3, 40, 64], dtype=torch.float16),
output_type="mp4",
)
with open("video.mp4", "wb") as f:
f.write(video)
生成
我们希望将 VAE 应用于实际流程以获取真实图像,而不是随机噪声。下面的示例展示了如何使用 SD v1.5 来实现这一点。
from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5",
torch_dtype=torch.float16,
variant="fp16",
vae=None,
).to("cuda")
prompt = "Strawberry ice cream, in a stylish modern glass, coconut, splashing milk cream and honey, in a gradient purple background, fluid motion, dynamic movement, cinematic lighting, Mysterious"
latent = pipe(
prompt=prompt,
output_type="latent",
).images
image = remote_decode(
endpoint="https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud/",
tensor=latent,
scaling_factor=0.18215,
)
image.save("test.jpg")
Flux 示例:
from diffusers import FluxPipeline
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-schnell",
torch_dtype=torch.bfloat16,
vae=None,
).to("cuda")
prompt = "Strawberry ice cream, in a stylish modern glass, coconut, splashing milk cream and honey, in a gradient purple background, fluid motion, dynamic movement, cinematic lighting, Mysterious"
latent = pipe(
prompt=prompt,
guidance_scale=0.0,
num_inference_steps=4,
output_type="latent",
).images
image = remote_decode(
endpoint="https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud/",
tensor=latent,
height=1024,
width=1024,
scaling_factor=0.3611,
shift_factor=0.1159,
)
image.save("test.jpg")
HunyuanVideo 示例:
from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
model_id = "hunyuanvideo-community/HunyuanVideo"
transformer = HunyuanVideoTransformer3DModel.from_pretrained(
model_id, subfolder="transformer", torch_dtype=torch.bfloat16
)
pipe = HunyuanVideoPipeline.from_pretrained(
model_id, transformer=transformer, vae=None, torch_dtype=torch.float16
).to("cuda")
latent = pipe(
prompt="A cat walks on the grass, realistic",
height=320,
width=512,
num_frames=61,
num_inference_steps=30,
output_type="latent",
).frames
video = remote_decode(
endpoint="https://o7ywnmrahorts457.us-east-1.aws.endpoints.huggingface.cloud/",
tensor=latent,
output_type="mp4",
)
if isinstance(video, bytes):
with open("video.mp4", "wb") as f:
f.write(video)
队列
使用远程 VAE 的一个巨大好处是我们可以排队多个生成请求。当当前潜在变量正在被处理以进行解码时,我们可以排队另一个。这有助于提高并发性。
import queue
import threading
from IPython.display import display
from diffusers import StableDiffusionPipeline
def decode_worker(q: queue.Queue):
while True:
item = q.get()
if item is None:
break
image = remote_decode(
endpoint="https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud/",
tensor=item,
scaling_factor=0.18215,
)
display(image)
q.task_done()
q = queue.Queue()
thread = threading.Thread(target=decode_worker, args=(q,), daemon=True)
thread.start()
def decode(latent: torch.Tensor):
q.put(latent)
prompts = [
"Blueberry ice cream, in a stylish modern glass , ice cubes, nuts, mint leaves, splashing milk cream, in a gradient purple background, fluid motion, dynamic movement, cinematic lighting, Mysterious",
"Lemonade in a glass, mint leaves, in an aqua and white background, flowers, ice cubes, halo, fluid motion, dynamic movement, soft lighting, digital painting, rule of thirds composition, Art by Greg rutkowski, Coby whitmore",
"Comic book art, beautiful, vintage, pastel neon colors, extremely detailed pupils, delicate features, light on face, slight smile, Artgerm, Mary Blair, Edmund Dulac, long dark locks, bangs, glowing, fashionable style, fairytale ambience, hot pink.",
"Masterpiece, vanilla cone ice cream garnished with chocolate syrup, crushed nuts, choco flakes, in a brown background, gold, cinematic lighting, Art by WLOP",
"A bowl of milk, falling cornflakes, berries, blueberries, in a white background, soft lighting, intricate details, rule of thirds, octane render, volumetric lighting",
"Cold Coffee with cream, crushed almonds, in a glass, choco flakes, ice cubes, wet, in a wooden background, cinematic lighting, hyper realistic painting, art by Carne Griffiths, octane render, volumetric lighting, fluid motion, dynamic movement, muted colors,",
]
pipe = StableDiffusionPipeline.from_pretrained(
"Lykon/dreamshaper-8",
torch_dtype=torch.float16,
vae=None,
).to("cuda")
pipe.unet = pipe.unet.to(memory_format=torch.channels_last)
pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
_ = pipe(
prompt=prompts[0],
output_type="latent",
)
for prompt in prompts:
latent = pipe(
prompt=prompt,
output_type="latent",
).images
decode(latent)
q.put(None)
thread.join()
VAE 编码
VAE 编码用于训练、图像到图像和图像到视频转换,将图像或视频转换为潜在表示。
显存
这些表格展示了在不同 GPU 上使用 SD v1 和 SD XL 进行 VAE 编码所需的 VRAM。
对于这些大多数 GPU,内存使用率%会导致其他模型(文本编码器、UNet/Transformer)必须卸载,或者需要使用分块编码,这会增加所需时间并影响质量。
SD v1.5:
| GPU | Resolution | Time (seconds) | Memory (%) | Tiled Time (secs) | Tiled Memory (%) |
|---|---|---|---|---|---|
| NVIDIA GeForce RTX 4090 | 512x512 | 0.015 | 3.51901 | 0.015 | 3.51901 |
| NVIDIA GeForce RTX 4090 | 256x256 | 0.004 | 1.3154 | 0.005 | 1.3154 |
| NVIDIA GeForce RTX 4090 | 2048x2048 | 0.402 | 47.1852 | 0.496 | 3.51901 |
| NVIDIA GeForce RTX 4090 | 1024x1024 | 0.078 | 12.2658 | 0.094 | 3.51901 |
| NVIDIA GeForce RTX 4080 SUPER | 512x512 | 0.023 | 5.30105 | 0.023 | 5.30105 |
| NVIDIA GeForce RTX 4080 SUPER | 256x256 | 0.006 | 1.98152 | 0.006 | 1.98152 |
| NVIDIA GeForce RTX 4080 SUPER | 2048x2048 | 0.574 | 71.08 | 0.656 | 5.30105 |
| NVIDIA GeForce RTX 4080 SUPER | 1024x1024 | 0.111 | 18.4772 | 0.14 | 5.30105 |
| NVIDIA GeForce RTX 3090 | 512x512 | 0.032 | 3.52782 | 0.032 | 3.52782 |
| NVIDIA GeForce RTX 3090 | 256x256 | 0.01 | 1.31869 | 0.009 | 1.31869 |
| NVIDIA GeForce RTX 3090 | 2048x2048 | 0.742 | 47.3033 | 0.954 | 3.52782 |
| NVIDIA GeForce RTX 3090 | 1024x1024 | 0.136 | 12.2965 | 0.207 | 3.52782 |
| NVIDIA GeForce RTX 3080 | 512x512 | 0.036 | 8.51761 | 0.036 | 8.51761 |
| NVIDIA GeForce RTX 3080 | 256x256 | 0.01 | 3.18387 | 0.01 | 3.18387 |
| NVIDIA GeForce RTX 3080 | 2048x2048 | 0.863 | 86.7424 | 1.191 | 8.51761 |
| NVIDIA GeForce RTX 3080 | 1024x1024 | 0.157 | 29.6888 | 0.227 | 8.51761 |
| NVIDIA GeForce RTX 3070 | 512x512 | 0.051 | 10.6941 | 0.051 | 10.6941 |
| NVIDIA GeForce RTX 3070 | 256x256 | 0.015 | 3.99743 | 0.015 | 3.99743 |
| NVIDIA GeForce RTX 3070 | 2048x2048 | 1.217 | 96.054 | 1.482 | 10.6941 |
| NVIDIA GeForce RTX 3070 | 1024x1024 | 0.223 | 37.2751 | 0.327 | 10.6941 |
SDXL:
| GPU | Resolution | Time (seconds) | Memory Consumed (%) | Tiled Time (seconds) | Tiled Memory (%) |
|---|---|---|---|---|---|
| NVIDIA GeForce RTX 4090 | 512x512 | 0.029 | 4.95707 | 0.029 | 4.95707 |
| NVIDIA GeForce RTX 4090 | 256x256 | 0.007 | 2.29666 | 0.007 | 2.29666 |
| NVIDIA GeForce RTX 4090 | 2048x2048 | 0.873 | 66.3452 | 0.863 | 15.5649 |
| NVIDIA GeForce RTX 4090 | 1024x1024 | 0.142 | 15.5479 | 0.143 | 15.5479 |
| NVIDIA GeForce RTX 4080 SUPER | 512x512 | 0.044 | 7.46735 | 0.044 | 7.46735 |
| NVIDIA GeForce RTX 4080 SUPER | 256x256 | 0.01 | 3.4597 | 0.01 | 3.4597 |
| NVIDIA GeForce RTX 4080 SUPER | 2048x2048 | 1.317 | 87.1615 | 1.291 | 23.447 |
| NVIDIA GeForce RTX 4080 SUPER | 1024x1024 | 0.213 | 23.4215 | 0.214 | 23.4215 |
| NVIDIA GeForce RTX 3090 | 512x512 | 0.058 | 5.65638 | 0.058 | 5.65638 |
| NVIDIA GeForce RTX 3090 | 256x256 | 0.016 | 2.45081 | 0.016 | 2.45081 |
| NVIDIA GeForce RTX 3090 | 2048x2048 | 1.755 | 77.8239 | 1.614 | 18.4193 |
| NVIDIA GeForce RTX 3090 | 1024x1024 | 0.265 | 18.4023 | 0.265 | 18.4023 |
| NVIDIA GeForce RTX 3080 | 512x512 | 0.064 | 13.6568 | 0.064 | 13.6568 |
| NVIDIA GeForce RTX 3080 | 256x256 | 0.018 | 5.91728 | 0.018 | 5.91728 |
| NVIDIA GeForce RTX 3080 | 2048x2048 | OOM | OOM | 1.866 | 44.4717 |
| NVIDIA GeForce RTX 3080 | 1024x1024 | 0.302 | 44.4308 | 0.302 | 44.4308 |
| NVIDIA GeForce RTX 3070 | 512x512 | 0.093 | 17.1465 | 0.093 | 17.1465 |
| NVIDIA GeForce RTX 3070 | 256x256 | 0.025 | 7.42931 | 0.026 | 7.42931 |
| NVIDIA GeForce RTX 3070 | 2048x2048 | OOM | OOM | 2.674 | 55.8355 |
| NVIDIA GeForce RTX 3070 | 1024x1024 | 0.443 | 55.7841 | 0.443 | 55.7841 |
可用的 VAE:
| 模型 | Endpoint |
|---|---|
| Stable Diffusion v1 | https://qc6479g0aac6qwy9.us-east-1.aws.endpoints.huggingface.cloud |
| Stable Diffusion XL | https://xjqqhmyn62rog84g.us-east-1.aws.endpoints.huggingface.cloud |
| Flux | https://ptccx55jz97f9zgo.us-east-1.aws.endpoints.huggingface.cloud |
代码
一个辅助方法简化了与 Hybrid Inference 的交互:
from diffusers.utils.remote_utils import remote_encode
基本示例:
from diffusers.utils import load_image
from diffusers.utils.remote_utils import remote_decode
image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg?download=true")
latent = remote_encode(
endpoint="https://ptccx55jz97f9zgo.us-east-1.aws.endpoints.huggingface.cloud/",
scaling_factor=0.3611,
shift_factor=0.1159,
)
decoded = remote_decode(
endpoint="https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud/",
tensor=latent,
scaling_factor=0.3611,
shift_factor=0.1159,
)
生成
现在让我们来看一个生成示例,我们将编码图像,生成然后远程解码!
import torch
from diffusers import StableDiffusionImg2ImgPipeline
from diffusers.utils import load_image
from diffusers.utils.remote_utils import remote_decode, remote_encode
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5",
torch_dtype=torch.float16,
variant="fp16",
vae=None,
).to("cuda")
init_image = load_image(
"https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
)
init_image = init_image.resize((768, 512))
init_latent = remote_encode(
endpoint="https://qc6479g0aac6qwy9.us-east-1.aws.endpoints.huggingface.cloud/",
image=init_image,
scaling_factor=0.18215,
)
prompt = "A fantasy landscape, trending on artstation"
latent = pipe(
prompt=prompt,
image=init_latent,
strength=0.75,
output_type="latent",
).images
image = remote_decode(
endpoint="https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud/",
tensor=latent,
scaling_factor=0.18215,
)
image.save("fantasy_landscape.jpg")
特殊管道示例
ConsisID
ConsisID 是一种保持身份信息的文本到视频生成模型,通过频率分解使生成视频中的面部保持一致。ConsisID 的主要特点包括:
- 频率分解:从频域角度分析 DiT 架构的特征,并基于这些特征设计合理的控制信息注入方法。
- 一致性训练策略:采用由粗到细的训练策略、动态掩码损失和动态跨脸损失,进一步提升模型的泛化能力和身份保持性能。
- 无需微调的推理:以往方法在推理前需要对输入 ID 进行逐个微调,导致显著的时间和计算成本。相比之下,ConsisID 无需微调。
参考文档:https://huggingface.co/docs/diffusers/v0.34.0/en/using-diffusers/consisid
Stable Diffusion XL
Stable Diffusion XL (SDXL) 是一种强大的文本到图像生成模型,它在三个方面迭代了之前的 Stable Diffusion 模型:
- UNet 的尺寸是原来的 3 倍,SDXL 结合了第二个文本编码器(OpenCLIP ViT-bigG/14)与原始文本编码器,显著增加了参数数量
- 引入了尺寸和裁剪条件,以保留训练数据不被丢弃,并更多地控制生成图像的裁剪方式
- 引入了双阶段模型流程;基础模型(也可以作为独立模型运行)生成图像作为精炼模型(refiner model)的输入,精炼模型添加额外的细节
参考文档:https://huggingface.co/docs/diffusers/v0.34.0/en/using-diffusers/sdxl
SDXL Turbo
SDXL Turbo 是一种通过对抗时间蒸馏技术生成的 Stable Diffusion XL (SDXL) 模型,能够在仅 1 步推理中运行。
参考文档:https://huggingface.co/docs/diffusers/v0.34.0/en/using-diffusers/sdxl_turbo
Kandinsky
卡文斯基模型是一系列多语言文本到图像生成模型。卡文斯基 2.0 模型使用两个多语言文本编码器,并将这些结果连接起来用于 UNet。
卡文斯基 2.1 改变了架构,加入了一个图像先验模型(CLIP),用于生成文本和图像嵌入之间的映射。该映射提供了更好的文本-图像对齐,并在训练过程中与文本嵌入一起使用,从而获得更高质量的结果。最后,卡文斯基 2.1 使用了一个调制量化向量(MoVQ)解码器——该解码器添加了一个空间条件归一化层以增加照片真实性——将潜在信息解码为图像。
卡文斯基 2.2 通过用更大的 CLIP-ViT-G 模型替换图像先验模型中的图像编码器来改进之前的模型,以提高质量。图像先验模型也在不同分辨率和宽高比的图像上进行重新训练,以生成更高分辨率的图像和不同尺寸的图像。
Kandinsky 3 简化了架构,并摒弃了涉及先验模型和扩散模型的二阶段生成过程。相反,Kandinsky 3 使用 Flan-UL2 对文本进行编码,采用带有 BigGan 深度块的 UNet,以及 Sber-MoVQGAN 将潜在变量解码为图像。文本理解和生成图像质量主要通过使用更大的文本编码器和 UNet 实现。
参考文档:https://huggingface.co/docs/diffusers/v0.34.0/en/using-diffusers/kandinsky
OmniGen
OmniGen 是一个图像生成模型。与现有的文本到图像模型不同,OmniGen 是一个单一模型,设计用于处理多种任务(例如,文本到图像、图像编辑、可控生成)。它具有以下特点:
- 极简模型架构:仅由一个 VAE 和一个 Transformer 模块组成,用于联合建模文本和图像。
- 支持多模态输入:它可以处理任何文本-图像混合数据作为图像生成的指令,而不仅仅依赖文本。
参考文档:https://huggingface.co/docs/diffusers/v0.34.0/en/using-diffusers/omnigen
PAG
扰动注意力引导(PAG)是一种新的扩散采样引导方法,它在不需要额外训练或集成外部模块的情况下,提高了无条件及条件设置下的样本质量。PAG 通过考虑自注意力机制捕获结构信息的能力,设计为在去噪过程中逐步增强合成样本的结构。它通过用单位矩阵替换扩散 U-Net 中选定的自注意力图来生成结构退化的中间样本,并引导去噪过程远离这些退化的样本。
参考文档:https://huggingface.co/docs/diffusers/v0.34.0/en/using-diffusers/pag
Latent Consistency Model
潜在一致性模型(LCMs)通过直接在潜在空间而非像素空间中预测逆向扩散过程,实现快速高质量的图像生成。换句话说,LCMs 试图从含噪图像预测无噪图像,而典型的扩散模型则通过迭代从含噪图像中去除噪声。通过避免迭代采样过程,LCMs 能够在 2-4 步内生成高质量图像,而不是 20-30 步。
LCMs 是从预训练模型中蒸馏出来的,这需要约 32 小时的 A100 计算资源。为了加快这一过程,LCM-LoRAs 训练一个 LoRA 适配器,其参数数量远少于完整模型。训练完成后,LCM-LoRA 可以接入扩散模型。
参考文档:https://huggingface.co/docs/diffusers/v0.34.0/en/using-diffusers/inference_with_lcm
Shap-E
Shap-E 是一个用于生成 3D 资产的生成模型,可用于游戏开发、室内设计和建筑领域。它基于大量的 3D 资产数据集进行训练,并通过后处理渲染每个对象的更多视角,生成 16K 点云而非 4K 点云。Shap-E 模型通过两个步骤进行训练:
- 编码器接收 3D 资产的点云和渲染视图,并输出表示该资产的隐式函数参数
- 扩散模型在编码器生成的潜在空间上进行训练,以生成神经辐射场(NeRFs)或带纹理的 3D 网格,从而更方便地在下游应用中渲染和使用 3D 资产
参考文档:https://huggingface.co/docs/diffusers/v0.34.0/en/using-diffusers/shap-e
DiffEdit
图像编辑通常需要提供要编辑区域的蒙版。DiffEdit 根据文本查询自动生成蒙版,从而无需图像编辑软件即可更轻松地创建蒙版。DiffEdit 算法分为三个步骤:
- 扩散模型根据某些查询文本和参考文本对图像进行去噪,为图像的不同区域产生不同的噪声估计;利用这些差异来推断蒙版,以识别图像中需要更改的区域以匹配查询文本
- 输入图像使用 DDIM 编码到潜在空间
- 潜在空间使用扩散模型根据文本查询解码,以蒙版为引导,使蒙版外的像素保持与输入图像相同
参考文档:https://huggingface.co/docs/diffusers/v0.34.0/en/using-diffusers/diffedit
Trajectory Consistency Distillation-LoRA
轨迹一致性蒸馏(TCD)使模型能够在更少的步骤中生成更高质量、更详细的图像。此外,由于蒸馏过程中有效的误差缓解,TCD 即使在较大推理步数的情况下也表现出优越的性能。
TCD 的主要优势包括:
- 优于教师模型:TCD 在小型和大型推理步数下均表现出更优越的生成质量,并超越了使用稳定扩散 XL(SDXL)的 DPM-Solver++(2S)的性能。在 TCD 训练过程中没有包含额外的判别器或 LPIPS 监督。
- 灵活的推理步骤:TCD 采样推理步骤可以自由调整,而不会影响图像质量。
- 自由调整细节级别:在推理过程中,图像的细节级别可以通过单个超参数
gamma进行调整。
参考文档:https://huggingface.co/docs/diffusers/v0.34.0/en/using-diffusers/inference_with_tcd_lora
Stable Video Diffusion
Stable Video Diffusion (SVD) 是一种强大的图像转视频生成模型,可以在输入图像的条件下生成 2-4 秒的高分辨率(576x1024)视频。
参考文档:https://huggingface.co/docs/diffusers/v0.34.0/en/using-diffusers/svd
Marigold Computer Vision
Marigold 基于扩散的方法,以及一系列为密集型计算机视觉任务设计的流程,包括单目深度预测、表面法线估计和固有图像分解。
参考文档:https://huggingface.co/docs/diffusers/v0.34.0/en/using-diffusers/marigold_usage
训练
总览
Diffusers 提供了一系列训练脚本,供您训练自己的扩散模型。您可以在 diffusers/examples 中找到所有我们的训练脚本。
每个训练脚本都是:
- 自包含的:训练脚本不依赖于任何本地文件,运行脚本所需的全部包都是从
requirements.txt文件中安装的。 - 易于调整的:训练脚本是一个示例,展示了如何为特定任务训练扩散模型,并不能直接适用于所有训练场景。您可能需要根据具体用例调整训练脚本。为了帮助您,我们已经完全暴露了数据预处理代码和训练循环,以便您可以根据自己的需求进行修改。
- 入门友好:训练脚本设计为对初学者友好且易于理解,而不是包含最新的最先进方法以获得最佳和最具竞争力的结果。我们考虑为过于复杂的任何训练方法都故意排除在外。
- 单一目的:每个训练脚本专门设计用于仅一个任务,以保持其可读性和可理解性。
训练脚本集合:
| Training | SDXL-support | LoRA-support | Flax-support |
|---|---|---|---|
| unconditional image generation | |||
| text-to-image | 👍 | 👍 | 👍 |
| textual inversion | 👍 | ||
| DreamBooth | 👍 | 👍 | 👍 |
| ControlNet | 👍 | 👍 | |
| InstructPix2Pix | 👍 | ||
| Custom Diffusion | |||
| T2I-Adapters | 👍 | ||
| Kandinsky 2.2 | 👍 | ||
| Wuerstchen | 👍 |
安装:
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .
安装依赖:
cd 到训练脚本的文件夹(例如,DreamBooth)并安装 requirements.txt 文件。
一些训练脚本对 SDXL、LoRA 或 Flax 有特定的要求文件。如果你使用这些脚本之一,请确保安装其对应的要求文件。
cd examples/dreambooth
pip install -r requirements.txt
# to train SDXL with DreamBooth
pip install -r requirements_sdxl.txt
创建训练数据集
Hub 上有许多数据集可以用来训练模型,但如果您找不到感兴趣的数据集或想使用自己的数据集,您可以使用🤗 Datasets 库创建一个数据集。数据集的结构取决于您想训练模型的任务。最基本的数据集结构是用于无条件图像生成等任务的图像目录。另一种数据集结构可能是用于文本到图像生成等任务的图像目录和一个包含其对应文本说明的文本文件。
本指南将向您展示两种创建数据集以进行微调的方法:
- 向
--train_data_dir参数提供图像文件夹 - 上传数据集到 Hub,并将数据集存储库 ID 传递给
--dataset_name参数
将数据集作为文件夹提供
对于无条件生成,您可以提供自己的数据集作为图像文件夹。训练脚本使用 Datasets 的 ImageFolder 构建器自动从文件夹构建数据集。您的目录结构应如下所示:
data_dir/xxx.png
data_dir/xxy.png
data_dir/[...]/xxz.png
将数据集目录的路径传递给 --train_data_dir 参数,然后就可以开始训练:
accelerate launch train_unconditional.py \
--train_data_dir <path-to-train-directory> \
<other-arguments>
将您的数据上传到 Hub
首先使用 ImageFolder 功能创建数据集,该功能会创建一个包含 PIL 编码图像的 image 列。
您可以使用 data_dir 或 data_files 参数来指定数据集的位置。data_files 参数支持将特定文件映射到数据集分割,如 train 或 test:
from datasets import load_dataset
# example 1: local folder
dataset = load_dataset("imagefolder", data_dir="path_to_your_folder")
# example 2: local files (supported formats are tar, gzip, zip, xz, rar, zstd)
dataset = load_dataset("imagefolder", data_files="path_to_zip_file")
# example 3: remote files (supported formats are tar, gzip, zip, xz, rar, zstd)
dataset = load_dataset(
"imagefolder",
data_files="https://download.microsoft.com/download/3/E/1/3E1C3F21-ECDB-4869-8368-6DEBA77B919F/kagglecatsanddogs_3367a.zip",
)
# example 4: providing several splits
dataset = load_dataset(
"imagefolder", data_files={"train": ["path/to/file1", "path/to/file2"], "test": ["path/to/file3", "path/to/file4"]}
)
然后使用 push_to_hub 方法将数据集上传到 Hub:
# assuming you have ran the huggingface-cli login command in a terminal
dataset.push_to_hub("name_of_your_dataset")
# if you want to push to a private repo, simply pass private=True:
dataset.push_to_hub("name_of_your_dataset", private=True)
现在,通过将数据集名称传递给 --dataset_name 参数,数据集即可用于训练:
accelerate launch --mixed_precision="fp16" train_text_to_image.py \
--pretrained_model_name_or_path="stable-diffusion-v1-5/stable-diffusion-v1-5" \
--dataset_name="name_of_your_dataset" \
<other-arguments>
模型适配下游任务
许多扩散系统共享相同组件,允许您将一个任务上的预训练模型适配到完全不同的任务上。
本指南将展示您如何通过初始化和修改预训练的 UNet2DConditionModel 的架构,来将预训练的文本到图像模型适配用于图像修复。
配置 UNet2DConditionModel 参数
UNet2DConditionModel 默认接受输入样本中的 4 个通道。例如,加载预训练的文本到图像模型 stable-diffusion-v1-5/stable-diffusion-v1-5,并查看 in_channels 的数量:
from diffusers import StableDiffusionPipeline
pipeline = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", use_safetensors=True)
pipeline.unet.config["in_channels"]
# 4
修复图像需要输入样本中有 9 个通道。你可以在像 runwayml/stable-diffusion-inpainting 这样的预训练修复模型中检查这个值:
from diffusers import StableDiffusionPipeline
pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-inpainting", use_safetensors=True)
pipeline.unet.config["in_channels"]
# 9
要使你的文本到图像模型适用于修复图像,你需要将 in_channels 的数量从 4 改为 9。
使用预训练的文本到图像模型权重初始化一个 UNet2DConditionModel,并将 in_channels 改为 9。改变 in_channels 的数量意味着你需要设置 ignore_mismatched_sizes=True 和 low_cpu_mem_usage=False 以避免大小不匹配的错误,因为形状现在不同了。
from diffusers import AutoModel
model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5"
unet = AutoModel.from_pretrained(
model_id,
subfolder="unet",
in_channels=9,
low_cpu_mem_usage=False,
ignore_mismatched_sizes=True,
use_safetensors=True,
)
文本到图像模型的其他组件的预训练权重是从它们的检查点初始化的,但 unet 的输入通道权重(conv_in.weight)是随机初始化的。为了修复图像,重要的是要对模型进行微调,否则模型会返回噪声。
模型训练
无条件图像生成
无条件图像生成模型在训练过程中不受文本或图像的约束。它仅生成与其训练数据分布相似的图像。
使用 train_unconditional.py 脚本,训练自己的模型:
1. 安装依赖
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .
导航到训练脚本所在的文件夹,安装所需的依赖:
cd examples/unconditional_image_generation
pip install -r requirements.txt
2. 初始化 Accelerate 环境
accelerate config
如果设置为默认环境,则无需任何配置:
accelerate config default
或者在代码中配置为默认环境:
from accelerate.utils import write_basic_config
write_basic_config()
3. 脚本参数
训练脚本提供了许多参数来帮助您自定义训练过程。所有参数及其描述都在 parse_args() 函数中找到。
示例:使用 bf16 格式进行混合精度加速训练:
accelerate launch train_unconditional.py --mixed_precision="bf16"
重要参数:
| 参数 | 说明 |
|---|---|
--dataset_name | Hub 上的数据集名称或用于训练的本地数据集路径 |
--output_dir | 训练模型保存位置 |
--push_to_hub | 是否将训练模型推送到 Hub |
--checkpointing_steps | 模型训练过程中保存检查点的频率 |
--resume_from_checkpoint | 如果训练中断,可以通过在训练命令中添加该参数,从该检查点继续训练 |
4. 训练脚本
预处理数据集和训练循环的代码位于 main() 函数。
初始化 UNet2DModel,代码如下:
model = UNet2DModel(
sample_size=args.resolution,
in_channels=3,
out_channels=3,
layers_per_block=2,
block_out_channels=(128, 128, 256, 256, 512, 512),
down_block_types=(
"DownBlock2D",
"DownBlock2D",
"DownBlock2D",
"DownBlock2D",
"AttnDownBlock2D",
"DownBlock2D",
),
up_block_types=(
"UpBlock2D",
"AttnUpBlock2D",
"UpBlock2D",
"UpBlock2D",
"UpBlock2D",
"UpBlock2D",
),
)
初始化调度器和优化器,代码如下:
# Initialize the scheduler
accepts_prediction_type = "prediction_type" in set(inspect.signature(DDPMScheduler.__init__).parameters.keys())
if accepts_prediction_type:
noise_scheduler = DDPMScheduler(
num_train_timesteps=args.ddpm_num_steps,
beta_schedule=args.ddpm_beta_schedule,
prediction_type=args.prediction_type,
)
else:
noise_scheduler = DDPMScheduler(num_train_timesteps=args.ddpm_num_steps, beta_schedule=args.ddpm_beta_schedule)
# Initialize the optimizer
optimizer = torch.optim.AdamW(
model.parameters(),
lr=args.learning_rate,
betas=(args.adam_beta1, args.adam_beta2),
weight_decay=args.adam_weight_decay,
eps=args.adam_epsilon,
)
加载数据集,对数据集进行预处理,代码如下:
dataset = load_dataset("imagefolder", data_dir=args.train_data_dir, cache_dir=args.cache_dir, split="train")
augmentations = transforms.Compose(
[
transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR),
transforms.CenterCrop(args.resolution) if args.center_crop else transforms.RandomCrop(args.resolution),
transforms.RandomHorizontalFlip() if args.random_flip else transforms.Lambda(lambda x: x),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]),
]
)
最后,训练循环处理所有其他事务,例如向图像添加噪声、预测噪声残差、计算损失、在指定步骤保存检查点,以及保存并推送到 Hub 模型。
5. 启动脚本
一个完整的训练运行在 4xV100 GPU 上需要 2 小时。
单 GPU 上启动训练:
accelerate launch train_unconditional.py \
--dataset_name="huggan/flowers-102-categories" \
--output_dir="ddpm-ema-flowers-64" \
--mixed_precision="fp16" \
--push_to_hub
多 GPU 上启动训练:
accelerate launch --multi_gpu train_unconditional.py \
--dataset_name="huggan/flowers-102-categories" \
--output_dir="ddpm-ema-flowers-64" \
--mixed_precision="fp16" \
--push_to_hub
6. 保存模型后重新加载
from diffusers import DiffusionPipeline
import torch
pipeline = DiffusionPipeline.from_pretrained("anton-l/ddpm-butterflies-128").to("cuda")
image = pipeline().images[0]
文本生成图像
文本到图像脚本处于实验阶段,容易过拟合并遇到灾难性遗忘等问题。
训练模型可能会对硬件造成压力,启用 gradient_checkpointing 和 mixed_precision,可以在单个 24GB GPU 上训练模型。建议使用至少 30GB 内存的 GPU 或 TPU v3 进行 Flax 训练。
使用 train_text_to_image.py 脚本,训练自己的模型:
1. 安装依赖
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .
导航到训练脚本所在文件夹,安装依赖:
cd examples/text_to_image
pip install -r requirements.txt
2. 初始化 Accelerate 环境
accelerate config
如果需要设置默认的 Accelerate 环境:
accelerate config default
如果不支持交互式 shell,可以在代码中配置默认环境:
from accelerate.utils import write_basic_config
write_basic_config()
3. 脚本参数
训练脚本提供了许多参数来帮助你自定义训练过程。所有参数及其描述可以在 parse_args() 函数中找到。
示例:使用 fp16 格式进行混合精度训练:
accelerate launch train_text_to_image.py --mixed_precision="fp16"
重要参数:
| 参数 | 说明 |
|---|---|
--pretrained_model_name_or_path | Hub 上的模型名称或预训练模型的本地路径 |
--dataset_name | Hub 上的数据集名称或用于训练的本地数据集路径 |
--image_column | 数据集中用于训练的图像列名称 |
--caption_column | 数据集中用于训练的文本列名称 |
--output_dir | 训练模型保存位置 |
--push_to_hub | 是否将训练模型推送到 Hub |
--checkpointing_steps | 模型训练时保存检查点的频率 |
--resume_from_checkpoint | 如果训练因某种原因中断,可以通过在训练命令中添加参数,来从该检查点继续训练 |
最小信噪比加权:
最小信噪比加权策略可以通过重新平衡损失来帮助训练,实现更快的收敛。训练脚本支持预测 epsilon(噪声)或 v_prediction,但最小信噪比与这两种预测类型都兼容。
示例:添加 --snr_gamma 参数并将其设置为推荐的值 5.0:
accelerate launch train_text_to_image.py --snr_gamma=5.0
对于较小的数据集,最小信噪比的效果可能不如较大的数据集明显。
4. 训练脚本
数据集预处理代码和训练循环位于 main() 函数中。
train_text_to_image 脚本加载调度器和分词器:
noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
tokenizer = CLIPTokenizer.from_pretrained(
args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision
)
加载 UNet 模型:
load_model = UNet2DConditionModel.from_pretrained(input_dir, subfolder="unet")
model.register_to_config(**load_model.config)
model.load_state_dict(load_model.state_dict())
预处理数据集中的文本和图像列:tokenize_captions 函数处理输入的标记化,train_transforms 函数指定应用于图像的变换类型:
def preprocess_train(examples):
images = [image.convert("RGB") for image in examples[image_column]]
examples["pixel_values"] = [train_transforms(image) for image in images]
examples["input_ids"] = tokenize_captions(examples)
return examples
最后,训练循环处理所有其他事务。它将图像编码到潜在空间中,向潜在值添加噪声,计算用于条件化的文本嵌入,更新模型参数,并将模型保存并推送到 Hub。
5. 启动脚本
在 Naruto BLIP 标题数据集上训练,以生成你自己的 Naruto 角色:
export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export dataset_name="lambdalabs/naruto-blip-captions"
accelerate launch --mixed_precision="fp16" train_text_to_image.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--dataset_name=$dataset_name \
--use_ema \
--resolution=512 --center_crop --random_flip \
--train_batch_size=1 \
--gradient_accumulation_steps=4 \
--gradient_checkpointing \
--max_train_steps=15000 \
--learning_rate=1e-05 \
--max_grad_norm=1 \
--enable_xformers_memory_efficient_attention \
--lr_scheduler="constant" --lr_warmup_steps=0 \
--output_dir="sd-naruto-model" \
--push_to_hub
6. 保存模型后重新加载
from diffusers import StableDiffusionPipeline
import torch
pipeline = StableDiffusionPipeline.from_pretrained("path/to/saved_model", torch_dtype=torch.float16, use_safetensors=True).to("cuda")
image = pipeline(prompt="yoda").images[0]
image.save("yoda-naruto.png")
Stable Diffusion XL
Stable Diffusion XL (SDXL) 是 Stable Diffusion 模型的一个更大、更强大的迭代版本,能够生成更高分辨率的图像。
SDXL 的 UNet 大小是原来的 3 倍,并且模型在架构中添加了第二个文本编码器。可能无法在像 Tesla T4 这样的消费级 GPU 上运行。
进行训练需要启用 gradient_checkpointing、mixed_precision 和 gradient_accumulation_steps。您还可以通过启用 xFormers 的内存高效注意力机制、使用 bitsandbytes 的 8 位优化器来进一步减少内存使用。
使用 train_text_to_image_sdxl.py 训练脚本,训练自己的模型:
1. 安装依赖
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .
导航到训练脚本所在文件夹,安装依赖:
cd examples/text_to_image
pip install -r requirements_sdxl.txt
2. 初始化 Accelerate 环境
accelerate config
如果需要设置默认环境:
accelerate config default
如果需要在代码中设置默认环境:
from accelerate.utils import write_basic_config
write_basic_config()
3. 脚本参数
训练脚本提供了许多参数来帮助你自定义训练过程。所有参数及其描述可以在 parse_args() 函数中找到。
示例:使用 bf16 格式进行混合精度加速训练:
accelerate launch train_text_to_image_sdxl.py --mixed_precision="bf16"
重要参数:
| 参数 | 说明 |
|---|---|
--pretrained_vae_model_name_or_path | 预训练 VAE 的路径;SDXL VAE 已知存在数值不稳定性,因此此参数允许你指定一个更好的 VAE |
--proportion_empty_prompts | 将多少比例的图像提示替换为空字符串 |
--timestep_bias_strategy | 在时间步的哪个阶段(早期或晚期)应用偏差,这可以鼓励模型学习低频或高频细节 |
--timestep_bias_multiplier | 应用偏差的权重 |
--timestep_bias_begin | 开始应用偏差的时间步 |
--timestep_bias_end | 结束应用偏差的时间步 |
--timestep_bias_portion | 应用偏差的时间步比例 |
最小信噪比加权:
Min-SNR 加权策略可以通过重新平衡损失来帮助训练,实现更快收敛。训练脚本支持预测 epsilon(噪声)或 v_prediction,但 Min-SNR 与这两种预测类型都兼容。
添加 --snr_gamma 参数并将其设置为推荐的值 5.0:
accelerate launch train_text_to_image_sdxl.py --snr_gamma=5.0
4. 训练脚本
训练脚本也与文本到图像训练指南类似,但已修改以支持 SDXL 训练。
它首先创建函数来对提示进行分词,计算提示嵌入,并使用 VAE 计算图像嵌入。接下来,你需要一个函数来根据时间步数和要应用的时间步偏差策略生成时间步权重。
在 main() 函数中,除了加载分词器,脚本还加载了第二个分词器和文本编码器,因为 SDXL 架构使用两个分词器:
tokenizer_one = AutoTokenizer.from_pretrained(
args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision, use_fast=False
)
tokenizer_two = AutoTokenizer.from_pretrained(
args.pretrained_model_name_or_path, subfolder="tokenizer_2", revision=args.revision, use_fast=False
)
text_encoder_cls_one = import_model_class_from_model_name_or_path(
args.pretrained_model_name_or_path, args.revision
)
text_encoder_cls_two = import_model_class_from_model_name_or_path(
args.pretrained_model_name_or_path, args.revision, subfolder="text_encoder_2"
)
提示词和图像嵌入首先被计算并保存在内存中,这对较小的数据集通常不是问题,但对于较大的数据集可能会导致内存问题。如果出现这种情况,你应该将预计算的嵌入单独保存到磁盘上,并在训练过程中将它们加载到内存中。
text_encoders = [text_encoder_one, text_encoder_two]
tokenizers = [tokenizer_one, tokenizer_two]
compute_embeddings_fn = functools.partial(
encode_prompt,
text_encoders=text_encoders,
tokenizers=tokenizers,
proportion_empty_prompts=args.proportion_empty_prompts,
caption_column=args.caption_column,
)
train_dataset = train_dataset.map(compute_embeddings_fn, batched=True, new_fingerprint=new_fingerprint)
train_dataset = train_dataset.map(
compute_vae_encodings_fn,
batched=True,
batch_size=args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps,
new_fingerprint=new_fingerprint_for_vae,
)
计算嵌入后,删除文本编码器、VAE 和分词器以释放一些内存:
del text_encoders, tokenizers, vae
gc.collect()
torch.cuda.empty_cache()
最后,训练循环负责处理其余部分。如果你选择应用时间步长偏差策略,你会看到时间步长权重被计算并作为噪声添加:
weights = generate_timestep_weights(args, noise_scheduler.config.num_train_timesteps).to(
model_input.device
)
timesteps = torch.multinomial(weights, bsz, replacement=True).long()
noisy_model_input = noise_scheduler.add_noise(model_input, noise, timesteps)
5. 启动脚本
在 Naruto BLIP 标题数据集上训练,以生成你自己的 Naruto 角色:
export MODEL_NAME="stabilityai/stable-diffusion-xl-base-1.0"
export VAE_NAME="madebyollin/sdxl-vae-fp16-fix"
export DATASET_NAME="lambdalabs/naruto-blip-captions"
accelerate launch train_text_to_image_sdxl.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--pretrained_vae_model_name_or_path=$VAE_NAME \
--dataset_name=$DATASET_NAME \
--enable_xformers_memory_efficient_attention \
--resolution=512 \
--center_crop \
--random_flip \
--proportion_empty_prompts=0.2 \
--train_batch_size=1 \
--gradient_accumulation_steps=4 \
--gradient_checkpointing \
--max_train_steps=10000 \
--use_8bit_adam \
--learning_rate=1e-06 \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--mixed_precision="fp16" \
--report_to="wandb" \
--validation_prompt="a cute Sundar Pichai creature" \
--validation_epochs 5 \
--checkpointing_steps=5000 \
--output_dir="sdxl-naruto-model" \
--push_to_hub
6. 保存模型重新加载
from diffusers import DiffusionPipeline
import torch
pipeline = DiffusionPipeline.from_pretrained("path/to/your/model", torch_dtype=torch.float16).to("cuda")
prompt = "A naruto with green eyes and red legs."
image = pipeline(prompt, num_inference_steps=30, guidance_scale=7.5).images[0]
image.save("naruto.png")
Kandinsky 2.2
Kandinsky 2.2 是一个多语言文本到图像模型,能够生成更逼真的图像。该模型包含一个用于从文本提示中创建图像嵌入的图像先验模型,以及一个基于先验模型的嵌入生成图像的解码器模型。
Diffusers 中为 Kandinsky 2.2 提供两个独立的脚本,一个用于训练先验模型,一个用于训练解码器模型。可以分别训练这两个模型,但要获得最佳结果,你应该同时训练先验和解码器模型。
根据 GPU,启用 gradient_checkpointing(⚠️ 先验模型不支持!)、mixed_precision 和 gradient_accumulation_steps 来帮助将模型适配到内存中并加速训练。通过启用 xFormers 的内存高效注意力机制,进一步减少内存使用。
使用 train_text_to_image_prior.py 和 train_text_to_image_decoder.py 脚本,训练自己的模型:
1. 安装依赖
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .
导航到包含训练脚本的文件夹,安装依赖:
cd examples/kandinsky2_2/text_to_image
pip install -r requirements.txt
2. 初始化 Accelerate 环境
accelerate config
生成默认配置:
accelerate config default
使用代码的方式生成默认配置:
from accelerate.utils import write_basic_config
write_basic_config()
3. 脚本参数
训练脚本提供了许多参数来帮助您自定义训练过程。所有参数及其描述都可以在 parse_args() 函数中找到。
示例:使用 bf16 格式进行混合精度训练:
accelerate launch train_text_to_image_prior.py --mixed_precision="fp16"
主要参数配置与文生图模型训练参数相同。
最小信噪比加权:
最小信噪比加权策略可以通过重新平衡损失来帮助训练,实现更快的收敛。训练脚本支持预测 epsilon(噪声)或 v_prediction,但最小信噪比与这两种预测类型都兼容。
添加 --snr_gamma 参数并将其设置为推荐的值 5.0:
accelerate launch train_text_to_image_prior.py --snr_gamma=5.0
4. 训练脚本
4.1 训练先验模型
main() 函数包含准备数据集和训练模型的代码。
除了调度器和分词器,训练脚本还加载 CLIPImageProcessor 用于预处理图像,以及 CLIPVisionModelWithProjection 模型用于编码图像:
noise_scheduler = DDPMScheduler(beta_schedule="squaredcos_cap_v2", prediction_type="sample")
image_processor = CLIPImageProcessor.from_pretrained(
args.pretrained_prior_model_name_or_path, subfolder="image_processor"
)
tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_prior_model_name_or_path, subfolder="tokenizer")
with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
image_encoder = CLIPVisionModelWithProjection.from_pretrained(
args.pretrained_prior_model_name_or_path, subfolder="image_encoder", torch_dtype=weight_dtype
).eval()
text_encoder = CLIPTextModelWithProjection.from_pretrained(
args.pretrained_prior_model_name_or_path, subfolder="text_encoder", torch_dtype=weight_dtype
).eval()
Kandinsky 使用 PriorTransformer 生成图像嵌入,因此你需要设置优化器来学习先验模式的参数:
prior = PriorTransformer.from_pretrained(args.pretrained_prior_model_name_or_path, subfolder="prior")
prior.train()
optimizer = optimizer_cls(
prior.parameters(),
lr=args.learning_rate,
betas=(args.adam_beta1, args.adam_beta2),
weight_decay=args.adam_weight_decay,
eps=args.adam_epsilon,
)
接下来,输入的标题被分词,图像由 CLIPImageProcessor 进行预处理:
def preprocess_train(examples):
images = [image.convert("RGB") for image in examples[image_column]]
examples["clip_pixel_values"] = image_processor(images, return_tensors="pt").pixel_values
examples["text_input_ids"], examples["text_mask"] = tokenize_captions(examples)
return examples
最后,训练循环将输入图像转换为潜在表示,对图像嵌入添加噪声,并做出预测:
model_pred = prior(
noisy_latents,
timestep=timesteps,
proj_embedding=prompt_embeds,
encoder_hidden_states=text_encoder_hidden_states,
attention_mask=text_mask,
).predicted_image_embedding
4.2 训练解码器模型
main() 函数包含准备数据集和训练模型的代码。
与之前的模型不同,解码器初始化一个 VQModel 来将潜在值解码为图像,并创建一个 UNet2DConditionModel:
with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
vae = VQModel.from_pretrained(
args.pretrained_decoder_model_name_or_path, subfolder="movq", torch_dtype=weight_dtype
).eval()
image_encoder = CLIPVisionModelWithProjection.from_pretrained(
args.pretrained_prior_model_name_or_path, subfolder="image_encoder", torch_dtype=weight_dtype
).eval()
unet = UNet2DConditionModel.from_pretrained(args.pretrained_decoder_model_name_or_path, subfolder="unet")
接下来,脚本包含几个图像变换和一个预处理函数,用于将变换应用于图像并返回像素值:
def preprocess_train(examples):
images = [image.convert("RGB") for image in examples[image_column]]
examples["pixel_values"] = [train_transforms(image) for image in images]
examples["clip_pixel_values"] = image_processor(images, return_tensors="pt").pixel_values
return examples
最后,训练循环处理将图像转换为潜在值、添加噪声以及预测噪声残差:
model_pred = unet(noisy_latents, timesteps, None, added_cond_kwargs=added_cond_kwargs).sample[:, :4]
5. 启动脚本
在 Naruto BLIP 标题数据集上训练以生成你自己的 Naruto 角色。如果在多个 GPU 上训练,请向 accelerate launch 命令添加 --multi_gpu 参数。
5.1 启动先验模型的训练:
export DATASET_NAME="lambdalabs/naruto-blip-captions"
accelerate launch --mixed_precision="fp16" train_text_to_image_prior.py \
--dataset_name=$DATASET_NAME \
--resolution=768 \
--train_batch_size=1 \
--gradient_accumulation_steps=4 \
--max_train_steps=15000 \
--learning_rate=1e-05 \
--max_grad_norm=1 \
--checkpoints_total_limit=3 \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--validation_prompts="A robot naruto, 4k photo" \
--report_to="wandb" \
--push_to_hub \
--output_dir="kandi2-prior-naruto-model"
5.2 启动解码器模型的训练:
export DATASET_NAME="lambdalabs/naruto-blip-captions"
accelerate launch --mixed_precision="fp16" train_text_to_image_decoder.py \
--dataset_name=$DATASET_NAME \
--resolution=768 \
--train_batch_size=1 \
--gradient_accumulation_steps=4 \
--gradient_checkpointing \
--max_train_steps=15000 \
--learning_rate=1e-05 \
--max_grad_norm=1 \
--checkpoints_total_limit=3 \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--validation_prompts="A robot naruto, 4k photo" \
--report_to="wandb" \
--push_to_hub \
--output_dir="kandi2-decoder-naruto-model"
6. 保存模型后重新加载
6.1 加载先验模型:
from diffusers import AutoPipelineForText2Image, DiffusionPipeline
import torch
prior_pipeline = DiffusionPipeline.from_pretrained(output_dir, torch_dtype=torch.float16)
prior_components = {"prior_" + k: v for k,v in prior_pipeline.components.items()}
pipeline = AutoPipelineForText2Image.from_pretrained("kandinsky-community/kandinsky-2-2-decoder", **prior_components, torch_dtype=torch.float16)
pipe.enable_model_cpu_offload()
prompt="A robot naruto, 4k photo"
image = pipeline(prompt=prompt, negative_prompt=negative_prompt).images[0]
6.2 加载编码器模型:
from diffusers import AutoPipelineForText2Image
import torch
pipeline = AutoPipelineForText2Image.from_pretrained("path/to/saved/model", torch_dtype=torch.float16)
pipeline.enable_model_cpu_offload()
prompt="A robot naruto, 4k photo"
image = pipeline(prompt=prompt).images[0]
对解码器模型,可以额外指定要加载的 UNet 模型:
from diffusers import AutoPipelineForText2Image, UNet2DConditionModel
unet = UNet2DConditionModel.from_pretrained("path/to/saved/model" + "/checkpoint-<N>/unet")
pipeline = AutoPipelineForText2Image.from_pretrained("kandinsky-community/kandinsky-2-2-decoder", unet=unet, torch_dtype=torch.float16)
pipeline.enable_model_cpu_offload()
image = pipeline(prompt="A robot naruto, 4k photo").images[0]
CogVideoX
CogVideoX 是一个文本到视频生成模型,专注于创建与提示更一致的连贯视频。它通过多种方法实现这一目标:
- 一个 3D 变分自编码器,对视频进行空间和时间压缩,提高压缩率和视频准确性。
- 一个专家 Transformer 模块,帮助对齐文本和视频,以及一个 3D 全注意力模块,用于捕捉和创建空间和时间上准确的视频。
视频指令维度的实际测试发现,CogVideoX 在主题一致性、动态信息、背景一致性、物体信息、平滑运动、颜色、场景、外观风格和时间风格方面效果良好,但在人体动作、空间关系和多物体方面无法取得良好效果。
使用 Diffusers 进行微调可以帮助弥补这些不良结果。
数据准备
训练脚本接受两种格式的数据。
第一种格式适用于小规模训练,第二种格式使用 CSV 格式,更适合大规模训练的流式数据。
1) 小格式
两个文件,其中一个文件包含行分隔的提示,另一个文件包含行分隔的视频数据路径(视频文件路径必须相对于指定 --instance_data_root 时的路径)。
假设你将 --instance_data_root 指定为 /dataset,并且这个目录包含文件:prompts.txt 和 videos.txt。
prompts.txt 文件应包含换行分隔的提示:
A black and white animated sequence featuring a rabbit, named Rabbity Ribfried, and an anthropomorphic goat in a musical, playful environment, showcasing their evolving interaction.
A black and white animated sequence on a ship's deck features a bulldog character, named Bully Bulldoger, showcasing exaggerated facial expressions and body language. The character progresses from confident to focused, then to strained and distressed, displaying a range of emotions as it navigates challenges. The ship's interior remains static in the background, with minimalistic details such as a bell and open door. The character's dynamic movements and changing expressions drive the narrative, with no camera movement to distract from its evolving reactions and physical gestures.
...
videos.txt 文件应包含换行分隔的视频文件路径。注意,路径应相对于 --instance_data_root 目录。
videos/00000.mp4
videos/00001.mp4
...
总体来说,如果你在数据集根目录上运行 tree 命令,你的数据集将如下所示:
/dataset
├── prompts.txt
├── videos.txt
├── videos
├── videos/00000.mp4
├── videos/00001.mp4
├── ...
2) 流格式
你可以使用单个 CSV 文件。在这个例子中,假设你有一个 metadata.csv 文件。期望的格式是:
<CAPTION_COLUMN>,<PATH_TO_VIDEO_COLUMN>
"""A black and white animated sequence featuring a rabbit, named Rabbity Ribfried, and an anthropomorphic goat in a musical, playful environment, showcasing their evolving interaction.""","""00000.mp4"""
"""A black and white animated sequence on a ship's deck features a bulldog character, named Bully Bulldoger, showcasing exaggerated facial expressions and body language. The character progresses from confident to focused, then to strained and distressed, displaying a range of emotions as it navigates challenges. The ship's interior remains static in the background, with minimalistic details such as a bell and open door. The character's dynamic movements and changing expressions drive the narrative, with no camera movement to distract from its evolving reactions and physical gestures.""","""00001.mp4"""
...
在这种情况下,--instance_data_root 应该是视频存储的位置,--dataset_name 应该是本地文件夹的路径或 Hub 上托管的与 load_dataset 兼容的数据集。假设你在 https://huggingface.co/datasets/my-awesome-username/minecraft-videos 有 Minecraft 游戏的视频,你需要指定 my-awesome-username/minecraft-videos。
使用这种格式时,--caption_column 必须是 <CAPTION_COLUMN>,--video_column 必须是 <PATH_TO_VIDEO_COLUMN>。
你不必严格限制在 CSV 格式。只要 load_dataset 方法支持加载基本 <PATH_TO_VIDEO_COLUMN> 和 <CAPTION_COLUMN> 的文件格式,任何格式都可以。之所以要经历这些数据集组织的复杂过程来加载视频数据,是因为 load_dataset 不完全支持所有种类的视频格式。
安装依赖
pip install diffusers transformers accelerate peft huggingface_hub # 用于所有建模和训练相关
pip install datasets decord # 用于加载视频训练数据
pip install bitsandbytes # 用于使用 8 位 Adam 或 AdamW 优化器进行内存优化训练
pip install wandb # 可选,用于监控训练日志
pip install deepspeed # 可选,用于 DeepSpeed 训练
pip install prodigyopt # 可选,如果你希望使用 Prodigy 优化器进行训练
# 安装依赖
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install -e .
# 导航到训练脚本文件夹,安装依赖
cd examples/cogvideo
pip install -r requirements.txt
初始化 Accelerate 环境
accelerate config
# 生成默认配置
accelerate config default
# 使用代码的方式,生成默认配置
from accelerate.utils import write_basic_config
write_basic_config()
登录到远程 Hub:
huggingface-cli login
# Alternatively, you could upload your model manually using:
# huggingface-cli upload my-cool-account-name/my-cool-lora-name /path/to/awesome/lora
假设您正在对 50 个相似概念的视频进行训练,我们发现 1500-2000 步效果很好。然而,官方推荐的是使用 100 个视频,总共 4000 步。假设您正在使用单个 GPU 进行训练,编号为 --train_batch_size,编号为 1:
- 在 50 个视频上进行 1500 步训练,将对应于 30 个训练周期
- 在 100 个视频上进行 4000 步训练将对应于 40 个训练周期
#!/bin/bash
GPU_IDS="0"
accelerate launch --gpu_ids $GPU_IDS examples/cogvideo/train_cogvideox_lora.py \
--pretrained_model_name_or_path THUDM/CogVideoX-2b \
--cache_dir <CACHE_DIR> \
--instance_data_root <PATH_TO_WHERE_VIDEO_FILES_ARE_STORED> \
--dataset_name my-awesome-name/my-awesome-dataset \
--caption_column <CAPTION_COLUMN> \
--video_column <PATH_TO_VIDEO_COLUMN> \
--id_token <ID_TOKEN> \
--validation_prompt "<ID_TOKEN> Spiderman swinging over buildings:::A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical atmosphere of this unique musical performance" \
--validation_prompt_separator ::: \
--num_validation_videos 1 \
--validation_epochs 10 \
--seed 42 \
--rank 64 \
--lora_alpha 64 \
--mixed_precision fp16 \
--output_dir /raid/aryan/cogvideox-lora \
--height 480 --width 720 --fps 8 --max_num_frames 49 --skip_frames_start 0 --skip_frames_end 0 \
--train_batch_size 1 \
--num_train_epochs 30 \
--checkpointing_steps 1000 \
--gradient_accumulation_steps 1 \
--learning_rate 1e-3 \
--lr_scheduler cosine_with_restarts \
--lr_warmup_steps 200 \
--lr_num_cycles 1 \
--enable_slicing \
--enable_tiling \
--optimizer Adam \
--adam_beta1 0.9 \
--adam_beta2 0.95 \
--max_grad_norm 1.0 \
--report_to wandb
设置 <ID_TOKEN> 并非必要。通过一些有限的实验,我们发现使用它效果更好(因为它类似于 Dreambooth 训练),而不使用它则效果较差。当提供 <ID_TOKEN> 时,它会被附加到每个提示的开头。因此,如果你的 <ID_TOKEN> 是 "DISNEY" 并且你的提示是 "Spiderman swinging over buildings",那么在训练中实际使用的提示将是 "DISNEY Spiderman swinging over buildings"。如果不提供,你将要么在没有额外标记的情况下进行训练,要么可以增强你的数据集,在开始训练之前将标记应用到你希望的位置。
推理
一旦您训练好了一个 lora 模型,可以通过将 lora 权重加载到 CogVideoXPipeline 来进行推理。
import torch
from diffusers import CogVideoXPipeline
from diffusers.utils import export_to_video
pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-2b", torch_dtype=torch.float16)
# pipe.load_lora_weights("/path/to/lora/weights", adapter_name="cogvideox-lora") # Or,
pipe.load_lora_weights("my-awesome-hf-username/my-awesome-lora-name", adapter_name="cogvideox-lora") # If loading from the HF Hub
pipe.to("cuda")
# Assuming lora_alpha=32 and rank=64 for training. If different, set accordingly
pipe.set_adapters(["cogvideox-lora"], [32 / 64])
prompt = "A vast, shimmering ocean flows gracefully under a twilight sky, its waves undulating in a mesmerizing dance of blues and greens. The surface glints with the last rays of the setting sun, casting golden highlights that ripple across the water. Seagulls soar above, their cries blending with the gentle roar of the waves. The horizon stretches infinitely, where the ocean meets the sky in a seamless blend of hues. Close-ups reveal the intricate patterns of the waves, capturing the fluidity and dynamic beauty of the sea in motion."
frames = pipe(prompt, guidance_scale=6, use_dynamic_cfg=True).frames[0]
export_to_video(frames, "output.mp4", fps=8)
减少内存使用
在使用 diffusers 库进行测试时,diffusers 库中包含的所有优化功能均已启用。该方案尚未在 NVIDIA A100/H100 架构以外的设备上进行实际内存使用测试。通常,该方案可适配所有 NVIDIA Ampere 架构及以上的设备。如果禁用优化,内存消耗将成倍增加,峰值内存使用量约为表格中数值的 3 倍。然而,速度将提升约 3-4 倍。您可以有选择地禁用部分优化,包括:
pipe.enable_sequential_cpu_offload()
pipe.vae.enable_slicing()
pipe.vae.enable_tiling()
说明:
- 对于多 GPU 推理,需要禁用
enable_sequential_cpu_offload()优化。 - 使用 INT8 模型会减慢推理速度,这是为了适配低内存 GPU,同时尽量减少视频质量损失,但推理速度会显著降低。
- CogVideoX-2B 模型是在 FP16 精度下训练的,所有 CogVideoX-5B 模型都是在 BF16 精度下训练的。我们建议使用模型训练时的精度进行推理。
- PytorchAO 和 Optimum-quanto 可用于量化文本编码器、Transformer 和 VAE 模块,以降低 CogVideoX 的内存需求。这使得模型能够在免费的 T4 Colabs 或内存较小的 GPU 上运行!此外,请注意 TorchAO 量化与
torch.compile完全兼容,这可以显著提高推理速度。在配备 NVIDIA H100 及更高版本的设备上,必须使用 FP8 精度,需要源安装torch、torchao、diffusers和acceleratePython 包。推荐使用 CUDA 12.4。 - 推理速度测试也使用了上述内存优化方案。没有内存优化,推理速度提高约 10%。只有 diffusers 版本的模型支持量化。
- 该模型仅支持英文输入;其他语言可以通过大型模型优化转换为英文使用。
- 模型微调的内存使用在 8 * H100 环境中进行了测试,程序会自动使用 Zero 2 优化。如果表格中标注了特定数量的 GPU,那么微调时必须使用该数量或更多的 GPU。
| 模型名称 | CogVideoX-2B | CogVideoX-5B |
|---|---|---|
| 推理精度 | FP16(推荐),BF16,FP32,FP8,INT8,不支持 INT4 | BF16(推荐),FP16,FP32,FP8*,INT8,不支持 INT4 |
| 单 GPU 推理 VRAM | FP16: 使用 diffusers 12.5GB | BF16: 使用 diffusers 20.7GB |
| INT8: 使用 torchao 的 diffusers 7.8GB | INT8: 使用 torchao 的 diffusers 11.4GB | |
| 多 GPU 推理 VRAM | FP16: 使用 diffusers 10GB* | BF16: 使用 diffusers 15GB* |
| 推理速度 | 单 A100:约 90 秒,单 H100:约 45 秒 | 单 A100:约 180 秒,单 H100:约 90 秒 |
| 微调精度 | FP16 | BF16 |
| 微调 VRAM 消耗 | 47 GB (bs=1, LORA) 61 GB (bs=2, LORA) 62GB (bs=1, SFT) | 63 GB (bs=1, LORA) 80 GB (bs=2, LORA) 75GB (bs=1, SFT) |
微调技术
Textual Inversion(文本反转)
文本反转是一种训练技术,只需几个你想要它学习的示例图像,即可个性化图像生成模型。
这种技术通过学习和更新文本嵌入(新嵌入与你在提示中必须使用的特殊单词相关联)来匹配你提供的示例图像。
如果在 vRAM 有限的 GPU 上训练,应该在训练命令中启用 gradient_checkpointing 和 mixed_precision 参数。
还可以通过使用内存高效的注意力机制(xFormers)来减少内存占用。
通过 textual_inversion.py 脚本,进行文本反转训练:
1. 安装依赖
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .
# 导航到训练脚本的文件夹,安装依赖
cd examples/textual_inversion
pip install -r requirements.txt
2. 初始化 Accelerate 环境
accelerate config
# 生成默认配置
accelerate config default
# 使用代码的方式生成默认配置
from accelerate.utils import write_basic_config
write_basic_config()
3. 脚本参数
训练脚本包含许多参数,帮助您根据需求调整训练运行。所有参数及其描述都列在 parse_args() 函数中。
示例:增加梯度累积步数(默认值为 1)
accelerate launch textual_inversion.py --gradient_accumulation_steps=4
重要参数:
| 参数 | 说明 |
|---|---|
--pretrained_model_name_or_path | Hub 上的模型名称或预训练模型的本地路径 |
--train_data_dir | 包含训练数据集(示例图像)的文件夹路径 |
--output_dir | 训练模型保存位置 |
--push_to_hub | 是否将训练模型推送到 Hub |
--checkpointing_steps | 模型训练时保存检查点的频率;如果训练因某种原因中断,可以通过在训练命令中添加 --resume_from_checkpoint 来从该检查点继续训练 |
--num_vectors | 用于学习嵌入的向量数量;增加此参数有助于模型学习更好,但会带来增加的训练成本 |
--placeholder_token | 用于将学习到的嵌入与特殊词关联的词(在推理时必须在提示中使用该词) |
--initializer_token | 粗略描述你试图训练的对象或风格的单个词 |
--learnable_property | 是否在训练模型以学习新的”风格”(例如,梵高的绘画风格)或”对象”(例如,你的狗) |
4. 训练脚本
与其他一些训练脚本不同,textual_inversion.py 拥有自定义数据集类 TextualInversionDataset 用于创建数据集。
您可以自定义图像大小、占位符标记、插值方法、是否裁剪图像等。
如果您需要更改数据集的创建方式,可以修改 TextualInversionDataset。
在 main() 函数中找到数据集预处理代码和训练循环。
加载分词器、调度器和模型:
# Load tokenizer
if args.tokenizer_name:
tokenizer = CLIPTokenizer.from_pretrained(args.tokenizer_name)
elif args.pretrained_model_name_or_path:
tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_model_name_or_path, subfolder="tokenizer")
# Load scheduler and models
noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
text_encoder = CLIPTextModel.from_pretrained(
args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
)
vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision)
unet = UNet2DConditionModel.from_pretrained(
args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision
)
特殊占位符标记被添加到分词器旁边,嵌入被重新调整以考虑新标记。
脚本从 TextualInversionDataset 创建数据集:
train_dataset = TextualInversionDataset(
data_root=args.train_data_dir,
tokenizer=tokenizer,
size=args.resolution,
placeholder_token=(" ".join(tokenizer.convert_ids_to_tokens(placeholder_token_ids))),
repeats=args.repeats,
learnable_property=args.learnable_property,
center_crop=args.center_crop,
set="train",
)
train_dataloader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.train_batch_size, shuffle=True, num_workers=args.dataloader_num_workers
)
最后,执行训练循环。
5. 启动脚本
下载一些猫玩具的图片并存储在目录中:
from huggingface_hub import snapshot_download
local_dir = "./cat"
snapshot_download(
"diffusers/cat_toy_example", local_dir=local_dir, repo_type="dataset", ignore_patterns=".gitattributes"
)
将环境变量 MODEL_NAME 设置为 Hub 上的模型 ID 或本地模型的路径,将 DATA_DIR 设置为你刚刚下载的猫图像的路径。
该脚本会创建并保存以下文件到你的仓库:
learned_embeds.bin:与你的示例图像对应的已学习嵌入向量token_identifier.txt:特殊占位符标记type_of_concept.txt:你正在训练的概念类型(“对象”或”风格”)
一个完整的训练运行在单个 V100 GPU 上需要约 1 小时。
如果对训练过程感兴趣,可以在训练过程中定期保存生成的图像。将以下参数添加到训练命令中:
--validation_prompt="A <cat-toy> train"
--num_validation_images=4
--validation_steps=100
启动命令:
export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export DATA_DIR="./cat"
accelerate launch textual_inversion.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATA_DIR \
--learnable_property="object" \
--placeholder_token="<cat-toy>" \
--initializer_token="toy" \
--resolution=512 \
--train_batch_size=1 \
--gradient_accumulation_steps=4 \
--max_train_steps=3000 \
--learning_rate=5.0e-04 \
--scale_lr \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--output_dir="textual_inversion_cat" \
--push_to_hub
6. 保存模型后重新加载
from diffusers import StableDiffusionPipeline
import torch
pipeline = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16).to("cuda")
pipeline.load_textual_inversion("sd-concepts-library/cat-toy")
image = pipeline("A <cat-toy> train", num_inference_steps=50).images[0]
image.save("cat-train.png")
DreamBooth
DreamBooth 是一种训练技术,通过仅使用主题或风格的几张图像来更新整个扩散模型。它的工作原理是在提示词中将一个特殊词汇与示例图像做关联。
如果在 vRAM 有限的 GPU 上进行训练,你应该尝试在训练命令中启用 gradient_checkpointing 和 mixed_precision 参数。
也可以通过使用内存高效的注意力机制(xFormers)来减少内存占用。
如果想要用 Flax 更快地训练,你应该配备一个内存大于 30GB 的 GPU。
通过 train_dreambooth.py 脚本,进行 DreamBooth:
1. 安装依赖
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .
# 导航到训练脚本的文件夹,安装依赖
cd examples/dreambooth
pip install -r requirements.txt
2. 初始化 Accelerate 环境
accelerate config
# 生成默认配置
accelerate config default
# 使用代码的方式生成默认配置
from accelerate.utils import write_basic_config
write_basic_config()
3. 脚本参数
DreamBooth 对训练超参数非常敏感,很容易过拟合。阅读《使用🧨 Diffusers 训练 DreamBooth 的 Stable Diffusion》博客文章,了解针对不同主题的推荐设置,以帮助你选择合适的超参数。
训练脚本提供了许多参数用于自定义您的训练过程。所有参数及其描述可以在 parse_args() 函数中找到。
示例:使用 bf16 格式进行训练:
accelerate launch train_dreambooth.py --mixed_precision="bf16"
重要参数:
| 参数 | 说明 |
|---|---|
--pretrained_model_name_or_path | Hub 上的模型名称或预训练模型的本地路径 |
--instance_data_dir | 包含训练数据集(示例图像)的文件夹路径 |
--instance_prompt | 包含示例图像特殊词的文本提示 |
--train_text_encoder | 是否也要训练文本编码器 |
--output_dir | 训练模型保存位置 |
--push_to_hub | 是否将训练模型推送到 Hub |
--checkpointing_steps | 模型训练时保存检查点的频率 |
--resume_from_checkpoint | 如果训练因某种原因中断,可以通过在训练命令中添加该参数来从该检查点继续训练 |
最小信噪比加权
最小信噪比加权策略可以通过重新平衡损失来帮助训练,实现更快的收敛。训练脚本支持预测 epsilon(噪声)或 v_prediction,但最小信噪比与这两种预测类型都兼容。
# 添加 --snr_gamma 参数并将其设置为推荐的值 5.0:
accelerate launch train_dreambooth.py --snr_gamma=5.0
先验保留损失
先验保留损失是一种方法,它使用模型自己生成的样本来帮助它学习如何生成更多样化的图像。因为这些生成的样本图像与您提供的图像属于同一类别,它们帮助模型保留它所学的关于该类别的知识,以及如何利用它已经知道的关于该类别的知识来创作新的组合。
--with_prior_preservation:是否使用先验保留损失--prior_loss_weight:控制先验保留损失对模型的影响--class_data_dir:包含生成的类别样本图像的文件夹路径--class_prompt:描述生成的样本图像类别的文本提示
accelerate launch train_dreambooth.py \
--with_prior_preservation \
--prior_loss_weight=1.0 \
--class_data_dir="path/to/class/images" \
--class_prompt="text prompt describing class"
训练文本编码器
为了提高生成输出的质量,也可以训练文本编码器,而不仅仅是 UNet。这需要额外的内存,您需要一台至少拥有 24GB 显存的 GPU。
如果您有必要的硬件,那么训练文本编码器会产生更好的结果,尤其是在生成人脸图像时。
启用此选项的方法是:
accelerate launch train_dreambooth.py --train_text_encoder
4. 训练脚本
DreamBooth 自带其数据集类:
DreamBoothDataset:对图像和类别图像进行预处理,并对提示词进行分词以用于训练PromptDataset:生成用于生成类别图像的提示嵌入
如果你启用了先验保留损失,类别图像将在此处生成:
sample_dataset = PromptDataset(args.class_prompt, num_new_images)
sample_dataloader = torch.utils.data.DataLoader(sample_dataset, batch_size=args.sample_batch_size)
sample_dataloader = accelerator.prepare(sample_dataloader)
pipeline.to(accelerator.device)
for example in tqdm(
sample_dataloader, desc="Generating class images", disable=not accelerator.is_local_main_process
):
images = pipeline(example["prompt"]).images
接下来是 main() 函数,它负责设置训练数据集以及训练循环本身。
脚本加载了分词器、调度器和模型:
# Load the tokenizer
if args.tokenizer_name:
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, revision=args.revision, use_fast=False)
elif args.pretrained_model_name_or_path:
tokenizer = AutoTokenizer.from_pretrained(
args.pretrained_model_name_or_path,
subfolder="tokenizer",
revision=args.revision,
use_fast=False,
)
# Load scheduler and models
noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
text_encoder = text_encoder_cls.from_pretrained(
args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
)
if model_has_vae(args):
vae = AutoencoderKL.from_pretrained(
args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision
)
else:
vae = None
unet = UNet2DConditionModel.from_pretrained(
args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision
)
然后,从 DreamBoothDataset 创建训练数据集和数据加载器:
train_dataset = DreamBoothDataset(
instance_data_root=args.instance_data_dir,
instance_prompt=args.instance_prompt,
class_data_root=args.class_data_dir if args.with_prior_preservation else None,
class_prompt=args.class_prompt,
class_num=args.num_class_images,
tokenizer=tokenizer,
size=args.resolution,
center_crop=args.center_crop,
encoder_hidden_states=pre_computed_encoder_hidden_states,
class_prompt_encoder_hidden_states=pre_computed_class_prompt_encoder_hidden_states,
tokenizer_max_length=args.tokenizer_max_length,
)
train_dataloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.train_batch_size,
shuffle=True,
collate_fn=lambda examples: collate_fn(examples, args.with_prior_preservation),
num_workers=args.dataloader_num_workers,
)
最后,训练循环负责处理其余步骤,例如将图像转换为潜在空间、向输入添加噪声、预测噪声残差以及计算损失。
5. 启动脚本
下载一些狗的图片,并将它们存储在一个目录中:
from huggingface_hub import snapshot_download
local_dir = "./dog"
snapshot_download(
"diffusers/dog-example",
local_dir=local_dir,
repo_type="dataset",
ignore_patterns=".gitattributes",
)
将环境变量 MODEL_NAME 设置为 Hub 上的模型 ID 或本地模型的路径,INSTANCE_DIR 设置为你刚刚下载狗图像的路径,OUTPUT_DIR 设置为你想要保存模型的路径。
使用 sks 作为特殊词汇来关联训练。
如果对跟进过程感兴趣,可以在训练过程中定期保存生成的图像。在训练命令中添加以下参数:
--validation_prompt="a photo of a sks dog"
--num_validation_images=4
--validation_steps=100
5.1 16GB
在 16GB GPU 上,你可以使用 bitsandbytes 8 位优化器和梯度检查点来帮助你训练 DreamBooth 模型。
安装 bitsandbytes:
pip install bitsandbytes
然后,将以下参数添加到您的训练命令中:
accelerate launch train_dreambooth.py \
--gradient_checkpointing \
--use_8bit_adam \
完整的启动命令如下:
export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export INSTANCE_DIR="./dog"
export OUTPUT_DIR="path_to_saved_model"
accelerate launch train_dreambooth.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--instance_data_dir=$INSTANCE_DIR \
--output_dir=$OUTPUT_DIR \
--instance_prompt="a photo of sks dog" \
--resolution=512 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--learning_rate=5e-6 \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--max_train_steps=400 \
--push_to_hub
5.2 12GB
在 12GB 的 GPU 上,你需要使用 bitsandbytes 8 位优化器、梯度检查点、xFormers,并将梯度设置为 None 而不是零以减少内存使用。
accelerate launch train_dreambooth.py \
--use_8bit_adam \
--gradient_checkpointing \
--enable_xformers_memory_efficient_attention \
--set_grads_to_none \
完整的启动命令:
export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export INSTANCE_DIR="./dog"
export OUTPUT_DIR="path_to_saved_model"
accelerate launch train_dreambooth.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--instance_data_dir=$INSTANCE_DIR \
--output_dir=$OUTPUT_DIR \
--instance_prompt="a photo of sks dog" \
--resolution=512 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--learning_rate=5e-6 \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--max_train_steps=400 \
--push_to_hub
5.3 8GB
在 8GB GPU 上,你需要使用 DeepSpeed 将一些张量从 VRAM 卸载到 CPU 或 NVME,以使用更少的 GPU 内存进行训练。
运行以下命令以配置 Accelerate 环境:
accelerate config
在配置过程中,请确认使用 DeepSpeed。现在,通过结合 DeepSpeed 阶段 2、fp16 混合精度,并将模型参数和优化器状态卸载到 CPU,应该能够在 8GB vRAM 以下进行训练。缺点是这需要更多的系统 RAM(约 25GB)。
还应该将默认的 Adam 优化器更改为 DeepSpeed 优化的 Adam deepspeed.ops.adam.DeepSpeedCPUAdam,以实现显著的加速。启用 DeepSpeedCPUAdam 需要您的系统 CUDA 工具链版本与 PyTorch 安装的版本相同。
目前,bitsandbytes 的 8 位优化器与 DeepSpeed 不兼容。
完整启动命令:
export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export INSTANCE_DIR="./dog"
export OUTPUT_DIR="path_to_saved_model"
accelerate launch train_dreambooth.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--instance_data_dir=$INSTANCE_DIR \
--output_dir=$OUTPUT_DIR \
--instance_prompt="a photo of sks dog" \
--resolution=512 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--learning_rate=5e-6 \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--max_train_steps=400 \
--push_to_hub
6. 保存模型后重新加载
from diffusers import DiffusionPipeline
import torch
pipeline = DiffusionPipeline.from_pretrained("path_to_saved_model", torch_dtype=torch.float16, use_safetensors=True).to("cuda")
image = pipeline("A photo of sks dog in a bucket", num_inference_steps=50, guidance_scale=7.5).images[0]
image.save("dog-bucket.png")
7. 与其他微调方法结合
1) LoRA
LoRA 是一种显著减少可训练参数数量的训练技术。因此,训练速度更快,并且更容易存储生成的权重,因为它们体积要小得多(约 100MB)。
使用 train_dreambooth_lora.py 脚本进行 LoRA 训练。
2) Stable Diffusion XL
Stable Diffusion XL (SDXL) 是一个强大的文生图模型,能够生成高分辨率图像,并且在其架构中增加了一个第二个文本编码器。
使用 train_dreambooth_lora_sdxl.py 脚本来训练一个带有 LoRA 的 SDXL 模型。
3) DeepFloyd IF
DeepFloyd IF 是一个具有三个阶段的级联像素扩散模型。第一阶段生成基础图像,第二和第三阶段逐步将基础图像上采样为高分辨率的 1024x1024 图像。
使用 train_dreambooth_lora.py 或 train_dreambooth.py 脚本来训练 DeepFloyd IF 模型(使用 LoRA 或完整模型)。
DeepFloyd IF 使用预测方差,但 Diffusers 训练脚本使用预测误差,因此训练好的 DeepFloyd IF 模型切换到固定方差计划。训练脚本会为你更新完全训练好的模型的调度器配置。然而,当你加载保存的 LoRA 权重时,也必须更新管道的调度器配置。
from diffusers import DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", use_safetensors=True)
pipe.load_lora_weights("<lora weights path>")
# Update scheduler config to fixed variance schedule
pipe.scheduler = pipe.scheduler.__class__.from_config(pipe.scheduler.config, variance_type="fixed_small")
阶段 2 模型需要额外的验证图像来进行放大。可以下载并使用训练图像的缩小版本来完成这一步。
from huggingface_hub import snapshot_download
local_dir = "./dog_downsized"
snapshot_download(
"diffusers/dog-example-downsized",
local_dir=local_dir,
repo_type="dataset",
ignore_patterns=".gitattributes",
)
下面的代码示例简要介绍了如何结合 DreamBooth 和 LoRA 训练 DeepFloyd IF 模型。
一些重要参数包括:
--resolution=64,由于 DeepFloyd IF 是一个像素扩散模型,且要在未压缩的像素上工作,输入图像的分辨率需要小得多--pre_compute_text_embeddings,提前计算文本嵌入以节省内存,因为 T5Model 会占用大量内存--tokenizer_max_length=77,作为文本编码器使用 T5 时,你可以使用更长的默认文本长度,但默认模型的编码过程使用较短的文本长度--text_encoder_use_attention_mask,将注意力掩码传递给文本编码器
训练配置:
a) 阶段 1 LoRA DreamBooth
DeepFloyd IF 的 LoRA 和 DreamBooth 第一阶段训练需要约 28GB 内存。
export MODEL_NAME="DeepFloyd/IF-I-XL-v1.0"
export INSTANCE_DIR="dog"
export OUTPUT_DIR="dreambooth_dog_lora"
accelerate launch train_dreambooth_lora.py \
--report_to wandb \
--pretrained_model_name_or_path=$MODEL_NAME \
--instance_data_dir=$INSTANCE_DIR \
--output_dir=$OUTPUT_DIR \
--instance_prompt="a sks dog" \
--resolution=64 \
--train_batch_size=4 \
--gradient_accumulation_steps=1 \
--learning_rate=5e-6 \
--scale_lr \
--max_train_steps=1200 \
--validation_prompt="a sks dog" \
--validation_epochs=25 \
--checkpointing_steps=100 \
--pre_compute_text_embeddings \
--tokenizer_max_length=77 \
--text_encoder_use_attention_mask
b) 阶段 2 LoRA DreamBooth
对于 DeepFloyd IF 的 LoRA 和 DreamBooth 的第二阶段,请注意这些参数:
--validation_images:验证时需要放大的图像--class_labels_conditioning=timesteps:在阶段 2 中根据需要额外条件化 UNet--learning_rate=1e-6:与阶段 1 相比,使用较低的学习率--resolution=256:上采样器的预期分辨率
export MODEL_NAME="DeepFloyd/IF-II-L-v1.0"
export INSTANCE_DIR="dog"
export OUTPUT_DIR="dreambooth_dog_upscale"
export VALIDATION_IMAGES="dog_downsized/image_1.png dog_downsized/image_2.png dog_downsized/image_3.png dog_downsized/image_4.png"
python train_dreambooth_lora.py \
--report_to wandb \
--pretrained_model_name_or_path=$MODEL_NAME \
--instance_data_dir=$INSTANCE_DIR \
--output_dir=$OUTPUT_DIR \
--instance_prompt="a sks dog" \
--resolution=256 \
--train_batch_size=4 \
--gradient_accumulation_steps=1 \
--learning_rate=1e-6 \
--max_train_steps=2000 \
--validation_prompt="a sks dog" \
--validation_epochs=100 \
--checkpointing_steps=500 \
--pre_compute_text_embeddings \
--tokenizer_max_length=77 \
--text_encoder_use_attention_mask \
--validation_images $VALIDATION_IMAGES \
--class_labels_conditioning=timesteps
c) 阶段 1 DreamBooth
对于 DeepFloyd IF 的 DreamBooth 第一阶段,请注意这些参数:
--skip_save_text_encoder:跳过保存完整 T5 文本编码器与微调模型--use_8bit_adam:使用 8 位 Adam 优化器以节省内存,因为优化器状态的大小在训练完整模型时会占用较多内存--learning_rate=1e-7:完整模型训练时应使用非常低的学习率,否则模型质量会下降(你可以使用更大的批处理大小来使用更高的学习率)
使用 8 位 Adam 优化器和 4 的批处理大小进行训练,完整模型可以在 ~48GB 的内存中训练。
export MODEL_NAME="DeepFloyd/IF-I-XL-v1.0"
export INSTANCE_DIR="dog"
export OUTPUT_DIR="dreambooth_if"
accelerate launch train_dreambooth.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--instance_data_dir=$INSTANCE_DIR \
--output_dir=$OUTPUT_DIR \
--instance_prompt="a photo of sks dog" \
--resolution=64 \
--train_batch_size=4 \
--gradient_accumulation_steps=1 \
--learning_rate=1e-7 \
--max_train_steps=150 \
--validation_prompt "a photo of sks dog" \
--validation_steps 25 \
--text_encoder_use_attention_mask \
--tokenizer_max_length 77 \
--pre_compute_text_embeddings \
--use_8bit_adam \
--set_grads_to_none \
--skip_save_text_encoder \
--push_to_hub
d) 阶段 2 DreamBooth
对于 DeepFloyd IF 的 DreamBooth 第二阶段,请注意这些参数:
--learning_rate=5e-6:使用较小的有效批处理大小和较低的学习率--resolution=256:上采样器的预期分辨率--train_batch_size=2和--gradient_accumulation_steps=6:要有效地在包含人脸的图像上进行训练需要更大的批处理大小
export MODEL_NAME="DeepFloyd/IF-II-L-v1.0"
export INSTANCE_DIR="dog"
export OUTPUT_DIR="dreambooth_dog_upscale"
export VALIDATION_IMAGES="dog_downsized/image_1.png dog_downsized/image_2.png dog_downsized/image_3.png dog_downsized/image_4.png"
accelerate launch train_dreambooth.py \
--report_to wandb \
--pretrained_model_name_or_path=$MODEL_NAME \
--instance_data_dir=$INSTANCE_DIR \
--output_dir=$OUTPUT_DIR \
--instance_prompt="a sks dog" \
--resolution=256 \
--train_batch_size=2 \
--gradient_accumulation_steps=6 \
--learning_rate=5e-6 \
--max_train_steps=2000 \
--validation_prompt="a sks dog" \
--validation_steps=150 \
--checkpointing_steps=500 \
--pre_compute_text_embeddings \
--tokenizer_max_length=77 \
--text_encoder_use_attention_mask \
--validation_images $VALIDATION_IMAGES \
--class_labels_conditioning timesteps \
--push_to_hub
8. 训练技巧
训练 DeepFloyd IF 模型可能具有挑战性,但这里有一些我们觉得有帮助的技巧:
- LoRA 足以训练阶段 1 模型,因为模型的低分辨率使得表示更精细的细节变得困难。
- 对于常见或简单的物体,你不必微调上采样器。确保传递给上采样器的提示已调整,以从实例提示中移除新标记。例如,如果你的阶段 1 提示是”a sks dog”,那么你的阶段 2 提示应该是”a dog”。
- 对于更精细的细节,如人脸,完全训练阶段 2 的上采样器比使用 LoRA 训练阶段 2 模型更好。使用较大的批量大小和较低的学习率也有帮助。
- 应使用较低的学习率来训练阶段 2 模型。
DDPMScheduler比训练脚本中使用的DPMSolver效果更好。
LoRA
LoRA(大型语言模型的低秩适配)是一种流行且轻量级的训练技术,可以显著减少可训练参数的数量。它通过向模型中插入少量新的权重,并仅训练这些权重来工作。这使得使用 LoRA 进行训练速度更快、内存效率更高,并产生更小的模型权重(几百 MB),这些权重更容易存储和分享。
LoRA 还可以与其他训练技术(如 DreamBooth)结合使用,以加快训练速度。
使用 train_text_to_image_lora.py 脚本,进行 LoRA 微调:
1. 安装依赖
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .
# 导航到训练脚本文件夹,安装依赖
cd examples/text_to_image
pip install -r requirements.txt
2. 初始化 Accelerate 环境
accelerate config
# 生成默认配置
accelerate config default
# 使用代码的方式生成默认配置
from accelerate.utils import write_basic_config
write_basic_config()
3. 脚本参数
训练脚本包含许多参数,帮助您自定义训练过程。所有参数及其描述都在 parse_args() 函数中找到。
示例:增加训练的轮数
accelerate launch train_text_to_image_lora.py --num_train_epochs=150
与 LoRA 相关的重要参数:
| 参数 | 说明 |
|---|---|
--rank | 训练低秩矩阵的内部维度;更高的秩意味着更多的可训练参数 |
--learning_rate | 默认学习率为 1e-4,但使用 LoRA 时,你可以使用更高的学习率 |
4. 训练脚本
数据集预处理代码和训练循环位于 main() 函数中。
4.1 对 UNet 微调
Diffusers 使用 PEFT 库中的 peft.LoraConfig 来设置 LoRA 适配器的参数,如秩、alpha 以及将 LoRA 权重插入哪些模块。适配器被添加到 UNet 中,并且在 lora_layers 中仅对 LoRA 层进行优化。
unet_lora_config = LoraConfig(
r=args.rank,
lora_alpha=args.rank,
init_lora_weights="gaussian",
target_modules=["to_k", "to_q", "to_v", "to_out.0"],
)
unet.add_adapter(unet_lora_config)
lora_layers = filter(lambda p: p.requires_grad, unet.parameters())
4.2 对 Text Encoder 微调
Diffusers 还支持在必要时使用 PEFT 库中的 LoRA 对文本编码器进行微调,例如微调 Stable Diffusion XL (SDXL)。~peft.LoraConfig 用于配置 LoRA 适配器的参数,这些参数随后被添加到文本编码器中,并且仅对 LoRA 层进行过滤以进行训练。
text_lora_config = LoraConfig(
r=args.rank,
lora_alpha=args.rank,
init_lora_weights="gaussian",
target_modules=["q_proj", "k_proj", "v_proj", "out_proj"],
)
text_encoder_one.add_adapter(text_lora_config)
text_encoder_two.add_adapter(text_lora_config)
text_lora_parameters_one = list(filter(lambda p: p.requires_grad, text_encoder_one.parameters()))
text_lora_parameters_two = list(filter(lambda p: p.requires_grad, text_encoder_two.parameters()))
4.3 优化器配置
优化器使用 lora_layers 进行初始化,因为这些是唯一会被优化的权重:
optimizer = optimizer_cls(
lora_layers,
lr=args.learning_rate,
betas=(args.adam_beta1, args.adam_beta2),
weight_decay=args.adam_weight_decay,
eps=args.adam_epsilon,
)
5. 启动脚本
在 Naruto BLIP 标题数据集上训练,以生成你自己的 Naruto 角色。
该脚本会创建并保存以下文件到你的仓库:
- saved model checkpoints:已保存的模型检查点
pytorch_lora_weights.safetensors:训练好的 LoRA 权重
如果在多个 GPU 上训练,请向 accelerate launch 命令添加 --multi_gpu 参数。
在配备 11GB 显存的 2080 Ti GPU 上完整训练需要约 5 小时。
export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export OUTPUT_DIR="/sddata/finetune/lora/naruto"
export HUB_MODEL_ID="naruto-lora"
export DATASET_NAME="lambdalabs/naruto-blip-captions"
accelerate launch --mixed_precision="fp16" train_text_to_image_lora.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--dataset_name=$DATASET_NAME \
--dataloader_num_workers=8 \
--resolution=512 \
--center_crop \
--random_flip \
--train_batch_size=1 \
--gradient_accumulation_steps=4 \
--max_train_steps=15000 \
--learning_rate=1e-04 \
--max_grad_norm=1 \
--lr_scheduler="cosine" \
--lr_warmup_steps=0 \
--output_dir=${OUTPUT_DIR} \
--push_to_hub \
--hub_model_id=${HUB_MODEL_ID} \
--report_to=wandb \
--checkpointing_steps=500 \
--validation_prompt="A naruto with blue eyes." \
--seed=1337
6. 保存模型后重新加载模型
from diffusers import AutoPipelineForText2Image
import torch
pipeline = AutoPipelineForText2Image.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16).to("cuda")
pipeline.load_lora_weights("path/to/lora/model", weight_name="pytorch_lora_weights.safetensors")
image = pipeline("A naruto with blue eyes").images[0]
Custom Diffusion(自定义扩散)
Custom Diffusion 是一种用于个性化图像生成模型的训练技术。与 Textual Inversion、DreamBooth 和 LoRA 类似,Custom Diffusion 仅需几个(约 4-5 个)示例图像即可使用。该技术通过仅训练交叉注意力层的权重,并使用一个特殊词汇来表示新学习到的概念。Custom Diffusion 的独特之处在于它能够同时学习多个概念。
如果在 vRAM 有限的 GPU 上进行训练,应该尝试启用 xFormers 并使用 --enable_xformers_memory_efficient_attention,以在更低的 vRAM 需求(16GB)下实现更快的训练。
为了进一步节省内存,可以在训练参数中添加 --set_grads_to_none,将梯度设置为 None 而不是零。
使用 train_custom_diffusion.py 脚本,进行自定义扩散训练:
1. 安装依赖
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .
# 导航到训练脚本文件夹,安装依赖
cd examples/custom_diffusion
pip install -r requirements.txt
pip install clip-retrieval
2. 初始化 Accelerate 环境
accelerate config
# 生成默认配置
accelerate config default
# 使用代码的方式,生成默认配置
from accelerate.utils import write_basic_config
write_basic_config()
3. 脚本参数
训练脚本包含所有帮助您自定义训练过程的参数。这些参数可以在 parse_args() 函数中找到。
示例:更改输入图像的分辨率
accelerate launch train_custom_diffusion.py --resolution=256
Custom Diffusion 重要参数:
| 参数 | 说明 |
|---|---|
--freeze_model | 冻结交叉注意力层中的键和值参数;默认值为 crossattn_kv,但您可以将其设置为 crossattn 以训练交叉注意力层中的所有参数 |
--concepts_list | 要学习多个概念,提供一个包含这些概念的 JSON 文件路径 |
--modifier_token | 用于表示已学习概念的特定词语 |
--initializer_token | 用于初始化 modifier_token 的嵌入的特定词语 |
先验保留损失
先验保留损失是一种方法,它使用模型自己生成的样本来帮助它学习如何生成更多样化的图像。因为这些生成的样本图像与您提供的图像属于同一类别,它们帮助模型保留它所学的关于该类别的知识,以及如何利用它已经知道的关于该类别的知识来创作新的组合。
许多先验保留损失的参数在 DreamBooth 训练指南中有描述。
正则化
Custom Diffusion 包括使用少量真实图像来训练目标图像,以防止过拟合。
下载 200 张带有 clip_retrieval 的真实图像。class_prompt 应与目标图像属于同一类别。这些图像存储在 class_data_dir 中。
python retrieve.py --class_prompt cat --class_data_dir real_reg/samples_cat --num_class_images 200
要启用正则化,请添加以下参数:
--with_prior_preservation:是否使用先验保留损失--prior_loss_weight:控制先验保留损失对模型的影响--real_prior:是否使用少量真实图像来防止过拟合
accelerate launch train_custom_diffusion.py \
--with_prior_preservation \
--prior_loss_weight=1.0 \
--class_data_dir="./real_reg/samples_cat" \
--class_prompt="cat" \
--real_prior=True \
4. 训练脚本
训练脚本有两个数据集类:
CustomDiffusionDataset:预处理图像、类别图像和提示词以进行训练PromptDataset:准备生成类别图像的提示词
将 modifier_token 添加到分词器中,将其转换为 token ID,并将 token 嵌入调整为适应新的 modifier_token 大小。然后使用 initializer_token 的嵌入初始化 modifier_token 嵌入。文本编码器中的所有参数都被冻结,只有 token 嵌入除外,因为模型试图学习将它们与概念关联起来。
params_to_freeze = itertools.chain(
text_encoder.text_model.encoder.parameters(),
text_encoder.text_model.final_layer_norm.parameters(),
text_encoder.text_model.embeddings.position_embedding.parameters(),
)
freeze_params(params_to_freeze)
将 Custom Diffusion 权重添加到注意力层。这是确保注意力权重形状和大小正确,以及设置每个 UNet 块中适当数量注意力处理器的关键步骤。
st = unet.state_dict()
for name, _ in unet.attn_processors.items():
cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
if name.startswith("mid_block"):
hidden_size = unet.config.block_out_channels[-1]
elif name.startswith("up_blocks"):
block_id = int(name[len("up_blocks.")])
hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
elif name.startswith("down_blocks"):
block_id = int(name[len("down_blocks.")])
hidden_size = unet.config.block_out_channels[block_id]
layer_name = name.split(".processor")[0]
weights = {
"to_k_custom_diffusion.weight": st[layer_name + ".to_k.weight"],
"to_v_custom_diffusion.weight": st[layer_name + ".to_v.weight"],
}
if train_q_out:
weights["to_q_custom_diffusion.weight"] = st[layer_name + ".to_q.weight"]
weights["to_out_custom_diffusion.0.weight"] = st[layer_name + ".to_out.0.weight"]
weights["to_out_custom_diffusion.0.bias"] = st[layer_name + ".to_out.0.bias"]
if cross_attention_dim is not None:
custom_diffusion_attn_procs[name] = attention_class(
train_kv=train_kv,
train_q_out=train_q_out,
hidden_size=hidden_size,
cross_attention_dim=cross_attention_dim,
).to(unet.device)
custom_diffusion_attn_procs[name].load_state_dict(weights)
else:
custom_diffusion_attn_procs[name] = attention_class(
train_kv=False,
train_q_out=False,
hidden_size=hidden_size,
cross_attention_dim=cross_attention_dim,
)
del st
unet.set_attn_processor(custom_diffusion_attn_procs)
custom_diffusion_layers = AttnProcsLayers(unet.attn_processors)
优化器被初始化以更新交叉注意力层参数:
optimizer = optimizer_class(
itertools.chain(text_encoder.get_input_embeddings().parameters(), custom_diffusion_layers.parameters())
if args.modifier_token is not None
else custom_diffusion_layers.parameters(),
lr=args.learning_rate,
betas=(args.adam_beta1, args.adam_beta2),
weight_decay=args.adam_weight_decay,
eps=args.adam_epsilon,
)
在训练循环中,重要的是只更新正在学习的概念的嵌入。这意味着将所有其他 token 嵌入的梯度设置为零:
if args.modifier_token is not None:
if accelerator.num_processes > 1:
grads_text_encoder = text_encoder.module.get_input_embeddings().weight.grad
else:
grads_text_encoder = text_encoder.get_input_embeddings().weight.grad
index_grads_to_zero = torch.arange(len(tokenizer)) != modifier_token_id[0]
for i in range(len(modifier_token_id[1:])):
index_grads_to_zero = index_grads_to_zero & (
torch.arange(len(tokenizer)) != modifier_token_id[i]
)
grads_text_encoder.data[index_grads_to_zero, :] = grads_text_encoder.data[
index_grads_to_zero, :
].fill_(0)
5. 启动脚本
下载并使用这些示例猫图像。
将环境变量 MODEL_NAME 设置为 Hub 上的模型 ID 或本地模型的路径,INSTANCE_DIR 设置为刚刚下载的猫图像的路径,OUTPUT_DIR 设置为你要保存模型的路径。
使用 <new1> 作为特殊词来将新学习的嵌入关联起来。
脚本会创建并保存模型检查点和 pytorch_custom_diffusion_weights.bin 文件到你的仓库。
要监控训练进度,在训练命令中添加 --report_to=wandb 参数,并使用 --validation_prompt 指定验证提示。
5.1 学习单一概念
export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export OUTPUT_DIR="path-to-save-model"
export INSTANCE_DIR="./data/cat"
accelerate launch train_custom_diffusion.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--instance_data_dir=$INSTANCE_DIR \
--output_dir=$OUTPUT_DIR \
--class_data_dir=./real_reg/samples_cat/ \
--with_prior_preservation \
--real_prior \
--prior_loss_weight=1.0 \
--class_prompt="cat" \
--num_class_images=200 \
--instance_prompt="photo of a <new1> cat" \
--resolution=512 \
--train_batch_size=2 \
--learning_rate=1e-5 \
--lr_warmup_steps=0 \
--max_train_steps=250 \
--scale_lr \
--hflip \
--modifier_token "<new1>" \
--validation_prompt="<new1> cat sitting in a bucket" \
--report_to="wandb" \
--push_to_hub
单一概念的微调模型加载:
import torch
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16,
).to("cuda")
pipeline.unet.load_attn_procs("path-to-save-model", weight_name="pytorch_custom_diffusion_weights.bin")
pipeline.load_textual_inversion("path-to-save-model", weight_name="<new1>.bin")
image = pipeline(
"<new1> cat sitting in a bucket",
num_inference_steps=100,
guidance_scale=6.0,
eta=1.0,
).images[0]
image.save("cat.png")
5.2 学习多个概念
Custom Diffusion 也可以学习多个概念,需要提供一个包含每个概念详细信息的 JSON 文件。
运行 clip-retrieval 来收集一些真实图像用于正则化:
pip install clip-retrieval
python retrieve.py --class_prompt {} --class_data_dir {} --num_class_images 200
启动脚本:
export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export OUTPUT_DIR="path-to-save-model"
accelerate launch train_custom_diffusion.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--output_dir=$OUTPUT_DIR \
--concepts_list=./concept_list.json \
--with_prior_preservation \
--real_prior \
--prior_loss_weight=1.0 \
--resolution=512 \
--train_batch_size=2 \
--learning_rate=1e-5 \
--lr_warmup_steps=0 \
--max_train_steps=500 \
--num_class_images=200 \
--scale_lr \
--hflip \
--modifier_token "<new1>+<new2>" \
--push_to_hub
多个概念的微调模型加载:
import torch
from huggingface_hub.repocard import RepoCard
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16,
).to("cuda")
model_id = "sayakpaul/custom-diffusion-cat-wooden-pot"
pipeline.unet.load_attn_procs(model_id, weight_name="pytorch_custom_diffusion_weights.bin")
pipeline.load_textual_inversion(model_id, weight_name="<new1>.bin")
pipeline.load_textual_inversion(model_id, weight_name="<new2>.bin")
image = pipeline(
"the <new1> cat sculpture in the style of a <new2> wooden pot",
num_inference_steps=100,
guidance_scale=6.0,
eta=1.0,
).images[0]
image.save("multi-subject.png")
Latent Consistency Distillation(潜在一致性蒸馏)
潜在一致性模型(LCMs)仅需几步即可生成高质量图像,这代表了一大步前进,因为许多流程至少需要 25 步以上。LCMs 是通过将潜在一致性蒸馏方法应用于任何 Stable Diffusion 模型而产生的。该方法通过在潜在空间中应用单阶段引导蒸馏,并结合跳步方法来一致地跳过时间步,以加速蒸馏过程。
如果在 vRAM 有限的 GPU 上训练,尝试启用 gradient_checkpointing、gradient_accumulation_steps 和 mixed_precision 以减少内存使用并加速训练。
通过启用 xFormers 的内存高效注意力机制和 bitsandbytes 的 8 位优化器,可以进一步减少内存使用。
通过 train_lcm_distill_sd_wds.py 脚本,进行潜在一致性蒸馏:
1. 安装依赖
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .
# 导航到训练脚本的示例文件夹,安装依赖
cd examples/consistency_distillation
pip install -r requirements.txt
2. 初始化 Accelerate 环境
accelerate config
# 生成默认配置
accelerate config default
# 使用代码的方式,生成默认配置
from accelerate.utils import write_basic_config
write_basic_config()
3. 脚本参数
训练脚本提供了许多参数来帮助你自定义训练过程。所有参数及其描述可以在 parse_args() 函数中找到。
示例:使用 fp16 格式进行混合精度训练以提高训练速度
accelerate launch train_lcm_distill_sd_wds.py --mixed_precision="fp16"
与潜在一致性蒸馏相关的参数:
| 参数 | 说明 |
|---|---|
--pretrained_teacher_model | 预训练的潜在扩散模型路径,用作教师模型 |
--pretrained_vae_model_name_or_path | 预训练的 VAE 路径;SDXL VAE 已知存在数值不稳定性,因此此参数允许您指定一个替代的 VAE |
--w_min 和 --w_max | 指导比例采样时的最小和最大指导比例值 |
--num_ddim_timesteps | DDIM 采样所需的时间步数 |
--loss_type | 用于潜在一致性蒸馏的损失类型(L2 或 Huber);通常推荐使用 Huber 损失,因为它对异常值更鲁棒 |
--huber_c | Huber 损失参数 |
4. 训练脚本
训练脚本首先创建一个数据集类 Text2ImageDataset,用于预处理图像并创建训练数据集。
def transform(example):
image = example["image"]
image = TF.resize(image, resolution, interpolation=transforms.InterpolationMode.BILINEAR)
c_top, c_left, _, _ = transforms.RandomCrop.get_params(image, output_size=(resolution, resolution))
image = TF.crop(image, c_top, c_left, resolution, resolution)
image = TF.to_tensor(image)
image = TF.normalize(image, [0.5], [0.5])
example["image"] = image
return example
为了在云端存储的大规模数据集的读写方面提升性能,该脚本使用 WebDataset 格式创建预处理管道,以应用转换并创建用于训练的数据集和数据加载器。
图像经过处理后直接输入训练循环,无需先下载完整数据集。
processing_pipeline = [
wds.decode("pil", handler=wds.ignore_and_continue),
wds.rename(image="jpg;png;jpeg;webp", text="text;txt;caption", handler=wds.warn_and_continue),
wds.map(filter_keys({"image", "text"})),
wds.map(transform),
wds.to_tuple("image", "text"),
]
在 main() 函数中,加载了所有必要的组件,如噪声调度器、分词器、文本编码器和 VAE。同时在此处加载了教师 UNet,然后你可以从教师 UNet 创建学生 UNet。
在训练过程中,学生 UNet 由优化器进行更新。
teacher_unet = UNet2DConditionModel.from_pretrained(
args.pretrained_teacher_model, subfolder="unet", revision=args.teacher_revision
)
unet = UNet2DConditionModel(**teacher_unet.config)
unet.load_state_dict(teacher_unet.state_dict(), strict=False)
unet.train()
创建优化器来更新 UNet 参数:
optimizer = optimizer_class(
unet.parameters(),
lr=args.learning_rate,
betas=(args.adam_beta1, args.adam_beta2),
weight_decay=args.adam_weight_decay,
eps=args.adam_epsilon,
)
创建数据集:
dataset = Text2ImageDataset(
train_shards_path_or_url=args.train_shards_path_or_url,
num_train_examples=args.max_train_samples,
per_gpu_batch_size=args.train_batch_size,
global_batch_size=args.train_batch_size * accelerator.num_processes,
num_workers=args.dataloader_num_workers,
resolution=args.resolution,
shuffle_buffer_size=1000,
pin_memory=True,
persistent_workers=True,
)
train_dataloader = dataset.train_dataloader
设置训练循环并实现潜在一致性蒸馏方法。
脚本的这个部分负责向潜在变量添加噪声、采样并创建指导尺度嵌入,以及从噪声中预测原始图像:
pred_x_0 = predicted_origin(
noise_pred,
start_timesteps,
noisy_model_input,
noise_scheduler.config.prediction_type,
alpha_schedule,
sigma_schedule,
)
model_pred = c_skip_start * noisy_model_input + c_out_start * pred_x_0
获取教师模型的预测结果和 LCM 的预测结果,计算损失,然后将其反向传播到 LCM。
if args.loss_type == "l2":
loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
elif args.loss_type == "huber":
loss = torch.mean(
torch.sqrt((model_pred.float() - target.float()) ** 2 + args.huber_c**2) - args.huber_c
)
5. 启动脚本
使用 --train_shards_path_or_url 来指定存储在 Hub 上的 Conceptual Captions 12M 数据集的路径。将 MODEL_DIR 环境变量设置为教师模型的名称,将 OUTPUT_DIR 设置为你要保存模型的位置。
export MODEL_DIR="stable-diffusion-v1-5/stable-diffusion-v1-5"
export OUTPUT_DIR="path/to/saved/model"
accelerate launch train_lcm_distill_sd_wds.py \
--pretrained_teacher_model=$MODEL_DIR \
--output_dir=$OUTPUT_DIR \
--mixed_precision=fp16 \
--resolution=512 \
--learning_rate=1e-6 --loss_type="huber" --ema_decay=0.95 --adam_weight_decay=0.0 \
--max_train_steps=1000 \
--max_train_samples=4000000 \
--dataloader_num_workers=8 \
--train_shards_path_or_url="pipe:curl -L -s https://huggingface.co/datasets/laion/conceptual-captions-12m-webdataset/resolve/main/data/{00000..01099}.tar?download=true" \
--validation_steps=200 \
--checkpointing_steps=200 --checkpoints_total_limit=10 \
--train_batch_size=12 \
--gradient_checkpointing --enable_xformers_memory_efficient_attention \
--gradient_accumulation_steps=1 \
--use_8bit_adam \
--resume_from_checkpoint=latest \
--report_to=wandb \
--seed=453645634 \
--push_to_hub
6. 加载 LCM 进行推理
from diffusers import UNet2DConditionModel, DiffusionPipeline, LCMScheduler
import torch
unet = UNet2DConditionModel.from_pretrained("your-username/your-model", torch_dtype=torch.float16, variant="fp16")
pipeline = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", unet=unet, torch_dtype=torch.float16, variant="fp16")
pipeline.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipeline.to("cuda")
prompt = "sushi rolls in the form of panda heads, sushi platter"
image = pipeline(prompt, num_inference_steps=4, guidance_scale=1.0).images[0]
7. 与其他微调方法结合
1) LoRA
LoRA 是一种显著减少可训练参数数量的训练技术。因此,训练速度更快,并且更容易存储生成的权重,因为它们体积更小(约 100MB)。
使用 train_lcm_distill_lora_sd_wds.py 或 train_lcm_distill_lora_sdxl.wds.py 脚本进行 LoRA 训练。
2) Stable Diffusion XL
Stable Diffusion XL (SDXL) 是一种强大的文本到图像模型,能够生成高分辨率图像,并且在其架构中添加了第二个文本编码器。
使用 train_lcm_distill_sdxl_wds.py 脚本进行 LoRA 的 SDXL 模型训练。
Reinforcement Learning Training with DDPO(使用 DDPO 进行强化学习训练)
可以通过 TRL 库和 Diffusers,在奖励函数上微调 Stable Diffusion。这是通过 Black 等人于《使用强化学习训练扩散模型》中引入的 Denoising Diffusion Policy Optimization (DDPO) 算法实现的,该算法在🤗 TRL 中以 DDPOTrainer 的形式实现。
量化方法
开始
量化专注于用更少的 bits 表示数据,同时尽量保持原始数据的精度。这通常意味着将数据类型转换为用更少的 bits 表示相同信息。例如,如果你的模型权重以 32 位浮点数存储,并且它们被量化为 16 位浮点数,这将使模型大小减半,使其更容易存储并减少内存使用。低精度也可以加速推理,因为用更少的 bits 进行计算所需时间更少。
Pipeline 级别量化
根据你对管道中每个模型的量化规格的控制程度,你可以使用 PipelineQuantizationConfig 的两种方式:
- 简单量化,只需要定义
quant_backend、quant_kwargs和components_to_quantize - 细粒度量化,提供一个
quant_mapping,它为各个模型组件提供量化规范
a. 简单量化
PipelineQuantizationConfig 参数:
quant_backend:指定要使用的量化后端。目前支持的后端包括:bitsandbytes_4bit、bitsandbytes_8bit、gguf、quanto和torchao。quant_kwargs:包含要使用的具体量化参数。components_to_quantize:指定要对管道的哪些组件进行量化。通常,你应该对计算量最大的组件进行量化,例如 transformer。如果管道有多个此类组件(例如FluxPipeline),则可以考虑对文本编码器进行量化。
示例:在保持 CLIP 模型完整的情况下,对 FluxPipeline 中的 T5 文本编码器进行了量化:
import torch
from diffusers import DiffusionPipeline
from diffusers.quantizers import PipelineQuantizationConfig
pipeline_quant_config = PipelineQuantizationConfig(
quant_backend="bitsandbytes_4bit",
quant_kwargs={"load_in_4bit": True, "bnb_4bit_quant_type": "nf4", "bnb_4bit_compute_dtype": torch.bfloat16},
components_to_quantize=["transformer", "text_encoder_2"],
)
将 pipeline_quant_config 传递给 from_pretrained() 以量化 Pipeline:
pipe = DiffusionPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
quantization_config=pipeline_quant_config,
torch_dtype=torch.bfloat16,
).to("cuda")
image = pipe("photo of a cute dog").images[0]
b. 细粒度量化
quant_mapping 参数提供了更灵活的选项,用于指定如何在管道中的每个独立组件进行量化,例如组合不同的量化后端。
初始化 PipelineQuantizationConfig 并将其传递给 quant_mapping。quant_mapping 允许你指定管道中每个组件的量化选项,例如转换器和文本编码器。
示例:分别使用 diffusers.QuantoConfig 和 transformers.BitsAndBytesConfig,对 transformer 和文本编码器进行量化:
import torch
from diffusers import DiffusionPipeline
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from diffusers.quantizers.quantization_config import QuantoConfig
from diffusers.quantizers import PipelineQuantizationConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
pipeline_quant_config = PipelineQuantizationConfig(
quant_mapping={
"transformer": QuantoConfig(weights_dtype="int8"),
"text_encoder_2": TransformersBitsAndBytesConfig(
load_in_4bit=True, compute_dtype=torch.bfloat16
),
}
)
在 Transformers 中有一个独立的 bitsandbytes 后端。你需要导入并使用 transformers.BitsAndBytesConfig 来处理来自 Transformers 的组件。
例如,FluxPipeline 中的 text_encoder_2 是一个来自 Transformers 的 T5EncoderModel,因此你需要使用 transformers.BitsAndBytesConfig 而不是 diffusers.BitsAndBytesConfig。
import torch
from diffusers import DiffusionPipeline
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from diffusers.quantizers import PipelineQuantizationConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
pipeline_quant_config = PipelineQuantizationConfig(
quant_mapping={
"transformer": DiffusersBitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16),
"text_encoder_2": TransformersBitsAndBytesConfig(
load_in_4bit=True, compute_dtype=torch.bfloat16
),
}
)
将 pipeline_quant_config 传递给 from_pretrained() 以量化管道:
pipe = DiffusionPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
quantization_config=pipeline_quant_config,
torch_dtype=torch.bfloat16,
).to("cuda")
image = pipe("photo of a cute dog").images[0]
bitsandbytes
bitsandbytes 是将模型量化为 8 位和 4 位的最佳选择。
8 位量化将 fp16 中的异常值与 int8 中的非异常值相乘,然后将非异常值转换回 fp16,最后将它们相加以返回 fp16 中的权重。这减少了异常值对模型性能的负面影响。
4 位量化进一步压缩了模型,通常与 QLoRA 一起使用,用于微调量化后的 LLM。
本指南展示了如何通过量化使 FLUX.1-dev 在少于 16GB 的 VRAM 上运行,甚至可以在免费的 Google Colab 实例上运行。
a. 安装依赖
pip install diffusers transformers accelerate bitsandbytes -U
b. 量化模型
通过将 BitsAndBytesConfig 传递给 from_pretrained() 来量化模型。
这适用于任何模态的任何模型,只要它支持使用 Accelerate 加载并包含 torch.nn.Linear 层。
1. 量化配置
8 位量化
8 位量化可以将模型内存使用量减半。
bitsandbytes 同时支持 Transformers 和 Diffusers,因此您可以量化 FluxTransformer2DModel 和 T5EncoderModel。
# CLIPTextModel 和 AutoencoderKL 没有进行量化,因为它们本身尺寸就很小,而且 AutoencoderKL 只有几层 torch.nn.Linear 层。
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
import torch
from diffusers import AutoModel
from transformers import T5EncoderModel
quant_config = TransformersBitsAndBytesConfig(load_in_8bit=True,)
text_encoder_2_8bit = T5EncoderModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="text_encoder_2",
quantization_config=quant_config,
torch_dtype=torch.float16,
)
quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True,)
transformer_8bit = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
torch_dtype=torch.float16,
)
默认情况下,其他所有模块(如 torch.nn.LayerNorm)都会被转换为 torch.float16 类型。
可以通过 torch_dtype 参数来更改这些模块的数据类型。
transformer_8bit = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
torch_dtype=torch.float32,
)
4 位量化
将模型量化为 4 位可以减少内存使用 4 倍。
bitsandbytes 在 Transformers 和 Diffusers 中都得到支持,因此你可以量化 FluxTransformer2DModel 和 T5EncoderModel。
CLIPTextModel 和 AutoencoderKL 没有进行量化,因为它们本身尺寸就很小,而且 AutoencoderKL 只有几层 torch.nn.Linear 层。
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
import torch
from diffusers import AutoModel
from transformers import T5EncoderModel
quant_config = TransformersBitsAndBytesConfig(load_in_4bit=True,)
text_encoder_2_4bit = T5EncoderModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="text_encoder_2",
quantization_config=quant_config,
torch_dtype=torch.float16,
)
quant_config = DiffusersBitsAndBytesConfig(load_in_4bit=True,)
transformer_4bit = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
torch_dtype=torch.float16,
)
默认情况下,其他所有模块(如 torch.nn.LayerNorm)都会被转换为 torch.float16 类型。你可以通过 torch_dtype 参数来更改这些模块的数据类型。
transformer_4bit = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
torch_dtype=torch.float32,
)
2. 使用量化模型
设置 device_map="auto" 会自动首先填满所有可用的 GPU,然后是 CPU,最后,如果仍然不够,才会使用硬盘。
from diffusers import FluxPipeline
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
transformer=transformer_8bit,
text_encoder_2=text_encoder_2_8bit,
torch_dtype=torch.float16,
device_map="auto",
)
pipe_kwargs = {
"prompt": "A cat holding a sign that says hello world",
"height": 1024,
"width": 1024,
"guidance_scale": 3.5,
"num_inference_steps": 50,
"max_sequence_length": 512,
}
image = pipe(**pipe_kwargs, generator=torch.manual_seed(0),).images[0]
当内存足够时,可以直接使用 .to("cuda") 将管道移动到 GPU,并通过 enable_model_cpu_offload() 来优化 GPU 内存使用。
3. 推送到 Hub
模型量化后,可以使用 push_to_hub() 方法将模型推送到 Hub。首先推送量化文件 config.json,然后推送量化后的模型权重。
可以使用 save_pretrained() 将 8 位序列化模型保存到本地。
使用 8 位和 4 位权重进行训练仅支持训练额外参数。
# 使用 get_memory_footprint 方法检查您的内存占用:
print(model.get_memory_footprint())
# 注意:这仅告诉你模型参数的内存占用,并不能估算推理所需的内存
4. 量化模型加载
量化模型可以通过 from_pretrained() 方法加载,无需指定 quantization_config 参数:
from diffusers import AutoModel, BitsAndBytesConfig
quantization_config = BitsAndBytesConfig(load_in_4bit=True)
model_4bit = AutoModel.from_pretrained(
"hf-internal-testing/flux.1-dev-nf4-pkg", subfolder="transformer"
)
8 位(LLM.int8() 算法)
异常值阈值
一个”异常值”是指大于某个阈值的隐藏状态值,这些值是在 fp16 格式下计算的。
虽然这些值通常呈正态分布([-3.5, 3.5]),但对于大型模型,这种分布可能差异很大([-60, 6] 或 [6, 60])。
8 位量化在值约为 5 时效果很好,但超过这个范围,性能损失会非常显著。一个良好的默认阈值是 6,但对于更不稳定的模型(小型模型或微调),可能需要更低的阈值。
要找到最佳阈值,建议在 BitsAndBytesConfig 中尝试 llm_int8_threshold 参数:
from diffusers import AutoModel, BitsAndBytesConfig
quantization_config = BitsAndBytesConfig(
load_in_8bit=True, llm_int8_threshold=10,
)
model_8bit = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quantization_config,
)
跳过模块转换
对于某些模型,您不需要将每个模块量化为 8 位,这实际上可能导致不稳定性。
例如,对于像 Stable Diffusion 3 这样的扩散模型,可以使用 BitsAndBytesConfig 中的 llm_int8_skip_modules 参数跳过 proj_out 模块。
from diffusers import SD3Transformer2DModel, BitsAndBytesConfig
quantization_config = BitsAndBytesConfig(
load_in_8bit=True, llm_int8_skip_modules=["proj_out"],
)
model_8bit = SD3Transformer2DModel.from_pretrained(
"stabilityai/stable-diffusion-3-medium-diffusers",
subfolder="transformer",
quantization_config=quantization_config,
)
4 位(QLoRA 算法)
计算数据类型
为了加速计算,使用 BitsAndBytesConfig 中的 bnb_4bit_compute_dtype 参数将数据类型从 float32(默认值)更改为 bf16:
import torch
from diffusers import BitsAndBytesConfig
quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)
标准浮点数 4(NF4,Normal Float 4)
NF4 是 QLoRA 论文中的一种 4 位数据类型,适用于从正态分布初始化的权重。
使用 NF4 来训练 4 位基础模型。这可以通过 BitsAndBytesConfig 中的 bnb_4bit_quant_type 参数进行配置:
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
from diffusers import AutoModel
from transformers import T5EncoderModel
quant_config = TransformersBitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
)
text_encoder_2_4bit = T5EncoderModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="text_encoder_2",
quantization_config=quant_config,
torch_dtype=torch.float16,
)
quant_config = DiffusersBitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
)
transformer_4bit = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
torch_dtype=torch.float16,
)
嵌套量化
嵌套量化是一种无需额外性能损耗即可节省额外内存的技术。
此功能对已量化的权重进行二次量化,以额外节省 0.4 比特/参数。
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
from diffusers import AutoModel
from transformers import T5EncoderModel
quant_config = TransformersBitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
)
text_encoder_2_4bit = T5EncoderModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="text_encoder_2",
quantization_config=quant_config,
torch_dtype=torch.float16,
)
quant_config = DiffusersBitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
)
transformer_4bit = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
torch_dtype=torch.float16,
)
去量化 bitsandbytes 模型
量化后,可以将模型反量化到原始精度,但这可能会导致轻微的质量损失。确保有足够的 GPU 内存来容纳反量化的模型。
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
from diffusers import AutoModel
from transformers import T5EncoderModel
quant_config = TransformersBitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
)
text_encoder_2_4bit = T5EncoderModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="text_encoder_2",
quantization_config=quant_config,
torch_dtype=torch.float16,
)
quant_config = DiffusersBitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
)
transformer_4bit = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
torch_dtype=torch.float16,
)
text_encoder_2_4bit.dequantize()
transformer_4bit.dequantize()
torch.compile
使用 torch.compile 加速推理:
8 位
torch._dynamo.config.capture_dynamic_output_shape_ops = True
quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True)
transformer_4bit = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
torch_dtype=torch.float16,
)
transformer_4bit.compile(fullgraph=True)
4 位
quant_config = DiffusersBitsAndBytesConfig(load_in_4bit=True)
transformer_4bit = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
torch_dtype=torch.float16,
)
transformer_4bit.compile(fullgraph=True)
GGUF
GGUF 文件格式通常用于存储用于 GGML 推理的模型,并支持多种块级量化选项。
Diffusers 支持通过 from_single_file 加载方式使用 Model 类加载预先量化并保存为 GGUF 格式的检查点。
目前通过 Pipelines 加载 GGUF 检查点尚不支持。
示例:使用 GGUF Q2_K 量化变体加载 FLUX.1 DEV transformer 模型
a. 安装 gguf
pip install -U gguf
b. 模型加载
由于 GGUF 是一种单文件格式,使用 FromSingleFileMixin.from_single_file 加载模型并传入 GGUFQuantizationConfig。
在使用 GGUF 检查点时,量化后的权重会保留在低内存 dtype(通常为 torch.uint8)中,并在每个模块通过模型进行前向传递时动态地解量化并转换为配置的 compute_dtype。GGUFQuantizationConfig 允许你设置 compute_dtype。
import torch
from diffusers import FluxPipeline, FluxTransformer2DModel, GGUFQuantizationConfig
ckpt_path = (
"https://huggingface.co/city96/FLUX.1-dev-gguf/blob/main/flux1-dev-Q2_K.gguf"
)
transformer = FluxTransformer2DModel.from_single_file(
ckpt_path,
quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
torch_dtype=torch.bfloat16,
)
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
transformer=transformer,
torch_dtype=torch.bfloat16,
)
pipe.enable_model_cpu_offload()
prompt = "A cat holding a sign that says hello world"
image = pipe(prompt, generator=torch.manual_seed(0)).images[0]
image.save("flux-gguf.png")
torchao
TorchAO 是一个用于 PyTorch 的架构优化库。它提供了高性能的数据类型(dtypes)、优化技术和推理与训练的内核,并具有与原生 PyTorch 功能(如 torch.compile、FullyShardedDataParallel(FSDP)等)的兼容性。
a. 安装依赖
pip install -U torch torchao
b. 量化模型加载
通过将 TorchAoConfig 传递给 from_pretrained()(也可以加载预量化模型)来量化模型。
这适用于任何模态的任何模型,只要它支持使用 Accelerate 加载并包含 torch.nn.Linear 层。
示例:将权重量化为 int8
import torch
from diffusers import FluxPipeline, AutoModel, TorchAoConfig
model_id = "black-forest-labs/FLUX.1-dev"
dtype = torch.bfloat16
quantization_config = TorchAoConfig("int8wo")
transformer = AutoModel.from_pretrained(
model_id,
subfolder="transformer",
quantization_config=quantization_config,
torch_dtype=dtype,
)
pipe = FluxPipeline.from_pretrained(
model_id,
transformer=transformer,
torch_dtype=dtype,
)
pipe.to("cuda")
# Without quantization: ~31.447 GB
# With quantization: ~20.40 GB
print(f"Pipeline memory usage: {torch.cuda.max_memory_reserved() / 1024**3:.3f} GB")
prompt = "A cat holding a sign that says hello world"
image = pipe(
prompt, num_inference_steps=50, guidance_scale=4.5, max_sequence_length=512
).images[0]
image.save("output.png")
TorchAO 完全兼容 torch.compile,这使得它区别于其他量化方法。只需一行代码即可加速推理。
transformer = torch.compile(transformer, mode="max-autotune", fullgraph=True)
torchao 还通过 autoquant 支持自动量化 API。
自动量化通过比较每种技术在选定输入类型和形状上的性能,来确定适用于模型的最佳量化策略。目前,这可以直接用于底层建模组件。
TorchAoConfig 类接受三个参数:
quant_type: 一个字符串值,指明以下量化类型之一。modules_to_not_convert: 一个模块全名或部分模块名的列表,这些模块不应执行量化。例如,要不对FluxTransformer2DModel的第一个模块执行任何量化,可以指定:modules_to_not_convert=["single_transformer_blocks.0"]。kwargs: 一个字典,包含传递给底层量化方法的键值对参数,该量化方法将根据quant_type被调用。
c. 支持的量化类型
torchao 支持对 int8、float3-float8 和 uint1-uint7 进行仅权重量化以及权重和动态激活量化。
- 仅权重量化:将模型权重存储在特定的低比特数据类型中,但使用更高精度的数据类型进行计算,例如
bfloat16。这降低了模型权重的内存需求,但保留了激活计算的内存峰值。 - 动态激活量化:将模型权重存储在低比特 dtype 中,同时实时量化激活值以节省额外内存。这降低了模型权重的内存需求,同时也降低了激活计算的内存开销。然而,有时这可能会带来质量上的折衷,因此建议对不同的模型进行彻底测试。
| 类型 | 函数名 | 别名 |
|---|---|---|
| 整数量化 | int4_weight_only | int4wo |
int8_dynamic_activation_int4_weight | int4dq | |
int8_weight_only | int8wo | |
int8_dynamic_activation_int8_weight | int8dq | |
| 浮点 8 位量化 | float8_weight_only | float8wo |
float8_dynamic_activation_float8_weight | float8wo_e5m2 | |
float8_static_activation_float8_weight | float8wo_e4m3 | |
float8dq | ||
float8dq_e4m3 | ||
float8dq_e4m3_tensor | ||
float8dq_e4m3_row | ||
| 浮点 X 位量化 | fpx_weight_only | fpX_eAwB,X 是位数(1-7),A 是指数位,B 是尾数位。约束:X == A + B + 1 |
| 无符号整数量化 | uintx_weight_only | uint1wo, uint2wo, uint3wo, uint4wo, uint5wo, uint6wo, uint7wo |
某些量化方法是别名(例如,
int8wo是int8_weight_only的常用缩写)。
d. 序列化和反序列化量化模型
示例:保存序列化量化模型
首先使用量化 dtype 加载模型,然后使用 save_pretrained() 方法保存。
import torch
from diffusers import AutoModel, TorchAoConfig
quantization_config = TorchAoConfig("int8wo")
transformer = AutoModel.from_pretrained(
"black-forest-labs/Flux.1-Dev",
subfolder="transformer",
quantization_config=quantization_config,
torch_dtype=torch.bfloat16,
)
transformer.save_pretrained("/path/to/flux_int8wo", safe_serialization=False)
示例:加载序列化的量化模型
使用 from_pretrained() 方法加载。
import torch
from diffusers import FluxPipeline, AutoModel
transformer = AutoModel.from_pretrained(
"/path/to/flux_int8wo",
torch_dtype=torch.bfloat16,
use_safetensors=False
)
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/Flux.1-Dev",
transformer=transformer,
torch_dtype=torch.bfloat16
)
pipe.to("cuda")
prompt = "A cat holding a sign that says hello world"
image = pipe(prompt, num_inference_steps=30, guidance_scale=7.0).images[0]
image.save("output.png")
示例:低版本 torch(<=2.6.0)加载问题
如果使用 torch<=2.6.0,某些量化方法(如 uint4wo)无法直接加载,并在尝试加载模型时可能导致 UnpicklingError,但在保存时工作正常。
为了解决这个问题,可以将状态字典手动加载到模型中,这需要在 torch.load 中使用 weights_only=False。
import torch
from accelerate import init_empty_weights
from diffusers import FluxPipeline, AutoModel, TorchAoConfig
# Serialize the model
transformer = AutoModel.from_pretrained(
"black-forest-labs/Flux.1-Dev",
subfolder="transformer",
quantization_config=TorchAoConfig("uint4wo"),
torch_dtype=torch.bfloat16,
)
transformer.save_pretrained(
"/path/to/flux_uint4wo",
safe_serialization=False,
max_shard_size="50GB"
)
# Load the model
state_dict = torch.load(
"/path/to/flux_uint4wo/diffusion_pytorch_model.bin",
weights_only=False,
map_location="cpu"
)
with init_empty_weights():
transformer = AutoModel.from_config("/path/to/flux_uint4wo/config.json")
transformer.load_state_dict(state_dict, strict=True, assign=True)
Quanto
Quanto 是 Optimum 的 PyTorch 量化后端:
- 所有功能都在即时模式下可用(适用于非可追踪模型)
- 支持量化感知训练
- 量化模型与
torch.compile兼容 - 量化模型设备无关(例如 CUDA、XPU、MPS、CPU)
a. 安装依赖
pip install optimum-quanto accelerate
b. 量化模型
通过将 QuantoConfig 对象传递给 from_pretrained() 方法来量化模型。
尽管 Quanto 库允许量化 nn.Conv2d 和 nn.LayerNorm 模块,但目前 Diffusers 仅支持量化模型中的 nn.Linear 层的权重。
示例:Quanto 应用 float8 量化
import torch
from diffusers import FluxTransformer2DModel, QuantoConfig
model_id = "black-forest-labs/FLUX.1-dev"
quantization_config = QuantoConfig(weights_dtype="float8")
transformer = FluxTransformer2DModel.from_pretrained(
model_id,
subfolder="transformer",
quantization_config=quantization_config,
torch_dtype=torch.bfloat16,
)
pipe = FluxPipeline.from_pretrained(
model_id, transformer=transformer, torch_dtype=torch_dtype
)
pipe.to("cuda")
prompt = "A cat holding a sign that says hello world"
image = pipe(
prompt, num_inference_steps=50, guidance_scale=4.5, max_sequence_length=512
).images[0]
image.save("output.png")
c. 跳过特定模块的量化
使用 modules_to_not_convert 参数在 QuantoConfig 中跳过对某些模块的量化。确保传递给此参数的模块与 state_dict 中模块的键匹配。
import torch
from diffusers import FluxTransformer2DModel, QuantoConfig
model_id = "black-forest-labs/FLUX.1-dev"
quantization_config = QuantoConfig(
weights_dtype="float8",
modules_to_not_convert=["proj_out"]
)
transformer = FluxTransformer2DModel.from_pretrained(
model_id,
subfolder="transformer",
quantization_config=quantization_config,
torch_dtype=torch.bfloat16,
)
d. 使用 from_single_file 加载权重
QuantoConfig 与 FromOriginalModelMixin.from_single_file 兼容。
import torch
from diffusers import FluxTransformer2DModel, QuantoConfig
ckpt_path = "https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/flux1-dev.safetensors"
quantization_config = QuantoConfig(weights_dtype="float8")
transformer = FluxTransformer2DModel.from_single_file(
ckpt_path,
quantization_config=quantization_config,
torch_dtype=torch.bfloat16
)
e. 保存量化模型
Diffusers 支持使用 ModelMixin.save_pretrained 方法序列化 Quanto 模型。
直接使用 Quanto 库量化的模型和使用 Diffusers 以 Quanto 作为后端量化的模型,其序列化和加载要求不同。目前无法使用 ModelMixin.from_pretrained 将直接使用 Quanto 量化的模型加载到 Diffusers 中。
import torch
from diffusers import FluxTransformer2DModel, QuantoConfig
model_id = "black-forest-labs/FLUX.1-dev"
quantization_config = QuantoConfig(weights_dtype="float8")
transformer = FluxTransformer2DModel.from_pretrained(
model_id,
subfolder="transformer",
quantization_config=quantization_config,
torch_dtype=torch.bfloat16,
)
# save quantized model to reuse
transformer.save_pretrained("<your quantized model save path>")
# you can reload your quantized model with
model = FluxTransformer2DModel.from_pretrained("<your quantized model save path>")
f. 使用 torch.compile
Quanto 后端支持以下量化类型中的 torch.compile。
int8 权重
import torch
from diffusers import FluxPipeline, FluxTransformer2DModel, QuantoConfig
model_id = "black-forest-labs/FLUX.1-dev"
quantization_config = QuantoConfig(weights_dtype="int8")
transformer = FluxTransformer2DModel.from_pretrained(
model_id,
subfolder="transformer",
quantization_config=quantization_config,
torch_dtype=torch.bfloat16,
)
transformer = torch.compile(transformer, mode="max-autotune", fullgraph=True)
pipe = FluxPipeline.from_pretrained(
model_id, transformer=transformer, torch_dtype=torch_dtype
)
pipe.to("cuda")
images = pipe("A cat holding a sign that says hello").images[0]
images.save("flux-quanto-compile.png")
支持的权重类型
| 权重类型 |
|---|
float8 |
int8 |
int4 |
int2 |
加速推理和减少内存开销
加速推理
扩散模型在推理时速度较慢,因为生成是一个迭代过程,噪声会在一定数量的”步骤”内逐渐被细化成图像或视频。
为了加速这一过程,你可以尝试使用不同的调度器、降低模型权重的精度以进行更快计算、使用更内存高效的注意力机制等。
模型数据类型
模型权重的精度和数据类型会影响推理速度,因为更高的精度需要更多的内存来加载,并需要更多时间进行计算。
PyTorch 默认以 float32 或全精度加载模型权重,因此更改数据类型是一种快速获得更快推理的简单方法。
bfloat16
bfloat16 与 float16 类似,但它在数值误差方面更为稳健。bfloat16 的硬件支持各不相同,但大多数现代 GPU 都能够支持 bfloat16。
import torch
from diffusers import StableDiffusionXLPipeline
pipeline = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
).to("cuda")
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]
float16
float16 与 bfloat16 相似,但可能更容易出现数值错误。
import torch
from diffusers import StableDiffusionXLPipeline
pipeline = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
).to("cuda")
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]
TF32
TensorFloat-32(TF32)模式在 NVIDIA Ampere GPU 上得到支持,它以 TF32 计算卷积和矩阵乘法操作。存储和其他操作仍保持为 float32。
当与 bfloat16 或 float16 结合使用时,这能显著提高计算速度。
PyTorch 默认仅启用卷积的 TF32 模式,而你需要显式地启用矩阵乘法的 TF32 模式。
import torch
from diffusers import StableDiffusionXLPipeline
torch.backends.cuda.matmul.allow_tf32 = True
pipeline = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
).to("cuda")
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]
缩放点积注意力
缩放点积注意力(SDPA)实现了多种注意力后端,包括 FlashAttention、xFormers 和原生 C++ 实现。它会自动选择最适合您硬件的后端。
如果你使用的是 PyTorch >= 2.0,SDPA 会默认启用,你的代码无需做任何额外修改。
不过,如果你希望自行选择注意力后端,可以尝试实验其他注意力后端。
示例:使用 torch.nn.attention.sdpa_kernel 上下文管理器来启用高效的注意力
from torch.nn.attention import SDPBackend, sdpa_kernel
import torch
from diffusers import StableDiffusionXLPipeline
pipeline = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
).to("cuda")
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
image = pipeline(prompt, num_inference_steps=30).images[0]
torch.compile
torch.compile 通过将 PyTorch 代码和操作编译成优化的内核来加速推理。Diffusers 通常会编译计算密集型模型,如 UNet、transformer 或 VAE。
为获得最大速度,启用以下编译器设置:
import torch
from diffusers import StableDiffusionXLPipeline
torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True
加载并编译 UNet 和 VAE。你可以选择几种不同的模式,但 "max-autotune" 通过编译成 CUDA 图来优化最快速度。CUDA 图通过单次 CPU 操作启动多个 GPU 操作,有效减少了开销。
将内存布局更改为 channels_last 也能优化内存和推理速度。
pipeline = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
).to("cuda")
pipeline.unet.to(memory_format=torch.channels_last)
pipeline.vae.to(memory_format=torch.channels_last)
pipeline.unet = torch.compile(
pipeline.unet, mode="max-autotune", fullgraph=True
)
pipeline.vae.decode = torch.compile(
pipeline.vae.decode,
mode="max-autotune",
fullgraph=True
)
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]
第一次编译速度较慢,但一旦编译完成,速度会显著提升。
尽量只在相同类型的推理操作上使用编译后的管道。在大小不同的图像上调用编译后的管道会重新触发编译,这既慢又低效。
区域编译(Regional Compilation)
区域编译通过仅编译模型中特定的重复区域(或块),而不是整个模型,从而减少了冷启动编译时间。编译器会重用缓存和编译好的其他块的代码。
加速器提供了 compile_regions 方法,用于自动按顺序编译 nn.Module 中的重复代码块。模型的其余部分将单独编译。
# pip install -U accelerate
import torch
from diffusers import StableDiffusionXLPipeline
from accelerate.utils import compile_regions
pipeline = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
).to("cuda")
pipeline.unet = compile_regions(pipeline.unet, mode="reduce-overhead", fullgraph=True)
图断裂(Graph Breaks)
在 torch.compile 中指定 fullgraph=True 非常重要,以确保底层模型中没有图断裂。能够在不降低性能的情况下利用 torch.compile。
对于 UNet 和 VAE,return_dict 改变了访问返回变量的方式:
# 修改前
latents = unet(
latents, timestep=timestep, encoder_hidden_states=prompt_embeds
).sample
# 修改后
latents = unet(
latents, timestep=timestep, encoder_hidden_states=prompt_embeds, return_dict=False
)[0]
GPU 同步
step() 函数在去噪器每次做出预测后都会在调度器上调用,并且索引 sigmas 变量。
当放在 GPU 上时,由于 CPU 和 GPU 之间的通信同步,会引入延迟。当去噪器已经被编译时,这种现象更加明显。
通常,sigmas 应该保持在 CPU 上,以避免通信同步和延迟。
动态量化
动态量化通过降低精度以实现更快的数学运算来提高推理速度。
这种特定的量化类型根据运行时数据来确定如何缩放激活值,而不是使用固定的缩放因子。因此,缩放因子与数据更加精确地匹配。
示例:使用 torchao 库对 UNet 和 VAE 应用动态 int8 量化
import torch
from torchao import apply_dynamic_quant
from diffusers import StableDiffusionXLPipeline
torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True
torch._inductor.config.force_fuse_int_mm_with_mul = True
torch._inductor.config.use_mixed_mm = True
使用 dynamic_quant_filter_fn 过滤掉 UNet 和 VAE 中那些从动态量化中获益不大的线性层:
def dynamic_quant_filter_fn(mod, *args):
return (
isinstance(mod, torch.nn.Linear)
and mod.in_features > 16
and (mod.in_features, mod.out_features)
not in [
(1280, 640),
(1920, 1280),
(1920, 640),
(2048, 1280),
(2048, 2560),
(2560, 1280),
(256, 128),
(2816, 1280),
(320, 640),
(512, 1536),
(512, 256),
(512, 512),
(640, 1280),
(640, 1920),
(640, 320),
(640, 5120),
(640, 640),
(960, 320),
(960, 640),
]
)
pipeline = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
).to("cuda")
apply_dynamic_quant(pipeline.unet, dynamic_quant_filter_fn)
apply_dynamic_quant(pipeline.vae, dynamic_quant_filter_fn)
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]
融合投影矩阵
在一个注意力模块中,输入被投影到由投影矩阵 Q、K 和 V 表示的三个子空间中。
这些投影通常分别计算,但你可以将它们水平组合成一个矩阵,并在一个步骤中执行投影。这增加了输入投影矩阵乘法的大小,并提高了量化的影响。
pipeline.fuse_qkv_projections()
编译和卸载量化模型
优化模型通常需要在推理速度和内存使用之间进行权衡。例如,虽然缓存可以提高推理速度,但它也会增加内存消耗,因为它需要存储中间注意力层的输出。更均衡的优化策略是结合模型量化、torch.compile 和各种卸载方法。
对于图像生成,结合量化和模型卸载通常能在质量、速度和内存之间取得最佳平衡。分组卸载对于图像生成不太有效,因为如果计算内核完成得更快,通常无法完全重叠数据传输。这会导致 CPU 和 GPU 之间的一些通信开销。
对于视频生成,结合量化和分组卸载往往效果更好,因为视频模型更受计算限制。
下表提供了 Flux 的优化策略组合及其对延迟和内存使用的影响的比较:
| 组合策略 | 延迟(s) | 内存使用(GB) |
|---|---|---|
| 量化 | 32.602 | 14.9453 |
| 量化 + torch.compile | 25.847 | 14.9448 |
| 量化 + torch.compile + 模型 CPU 卸载 | 32.312 | 12.2369 |
安装依赖
pip install -U bitsandbytes
a. 量化和 torch.compile
对模型进行量化以减少存储所需的内存,并将其编译以加速推理。
配置 Dynamo capture_dynamic_output_shape_ops = True 在编译 bitsandbytes 模型时处理动态输出。
import torch
from diffusers import DiffusionPipeline
from diffusers.quantizers import PipelineQuantizationConfig
torch._dynamo.config.capture_dynamic_output_shape_ops = True
# quantize
pipeline_quant_config = PipelineQuantizationConfig(
quant_backend="bitsandbytes_4bit",
quant_kwargs={
"load_in_4bit": True,
"bnb_4bit_quant_type": "nf4",
"bnb_4bit_compute_dtype": torch.bfloat16
},
components_to_quantize=["transformer", "text_encoder_2"],
)
pipeline = DiffusionPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
quantization_config=pipeline_quant_config,
torch_dtype=torch.bfloat16,
).to("cuda")
# compile
pipeline.transformer.to(memory_format=torch.channels_last)
pipeline.transformer.compile(mode="max-autotune", fullgraph=True)
pipeline("""
cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
""").images[0]
b. 量化、torch.compile 和卸载
除了量化和 torch.compile,如果需要进一步减少内存使用,可以尝试卸载。卸载会根据计算需求将各种层或模型组件从 CPU 移动到 GPU。
在卸载时配置 Dynamo cache_size_limit 以避免过度重新编译,并将 capture_dynamic_output_shape_ops = True 设置为在编译 bitsandbytes 模型时处理动态输出。
模型 CPU 卸载
模型 CPU 卸载将单个管道组件(如 Transformer 模型)在需要计算时移动到 GPU,否则卸载到 CPU。
import torch
from diffusers import DiffusionPipeline
from diffusers.quantizers import PipelineQuantizationConfig
torch._dynamo.config.cache_size_limit = 1000
torch._dynamo.config.capture_dynamic_output_shape_ops = True
# quantize
pipeline_quant_config = PipelineQuantizationConfig(
quant_backend="bitsandbytes_4bit",
quant_kwargs={
"load_in_4bit": True,
"bnb_4bit_quant_type": "nf4",
"bnb_4bit_compute_dtype": torch.bfloat16
},
components_to_quantize=["transformer", "text_encoder_2"],
)
pipeline = DiffusionPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
quantization_config=pipeline_quant_config,
torch_dtype=torch.bfloat16,
).to("cuda")
# model CPU offloading
pipeline.enable_model_cpu_offload()
# compile
pipeline.transformer.compile()
pipeline(
"cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California, "
"highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain"
).images[0]
分组卸载
分组卸载将单个管道组件(如 Transformer 模型)的内部层移至 GPU 进行计算,并在不需要时将其卸载。同时,它使用 CUDA 流功能预取下一层以供执行。
通过重叠计算和数据传输,它比模型 CPU 卸载更快,同时还能节省内存。
# pip install ftfy
import torch
from diffusers import AutoModel, DiffusionPipeline
from diffusers.hooks import apply_group_offloading
from diffusers.utils import export_to_video
from diffusers.quantizers import PipelineQuantizationConfig
from transformers import UMT5EncoderModel
torch._dynamo.config.cache_size_limit = 1000
torch._dynamo.config.capture_dynamic_output_shape_ops = True
# quantize
pipeline_quant_config = PipelineQuantizationConfig(
quant_backend="bitsandbytes_4bit",
quant_kwargs={
"load_in_4bit": True,
"bnb_4bit_quant_type": "nf4",
"bnb_4bit_compute_dtype": torch.bfloat16
},
components_to_quantize=["transformer", "text_encoder"],
)
text_encoder = UMT5EncoderModel.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers",
subfolder="text_encoder",
torch_dtype=torch.bfloat16
)
pipeline = DiffusionPipeline.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers",
quantization_config=pipeline_quant_config,
torch_dtype=torch.bfloat16,
).to("cuda")
# group offloading
onload_device = torch.device("cuda")
offload_device = torch.device("cpu")
pipeline.transformer.enable_group_offload(
onload_device=onload_device,
offload_device=offload_device,
offload_type="leaf_level",
use_stream=True,
non_blocking=True
)
pipeline.vae.enable_group_offload(
onload_device=onload_device,
offload_device=offload_device,
offload_type="leaf_level",
use_stream=True,
non_blocking=True
)
apply_group_offloading(
pipeline.text_encoder,
onload_device=onload_device,
offload_type="leaf_level",
use_stream=True,
non_blocking=True
)
# compile
pipeline.transformer.compile()
prompt = """
The camera rushes from far to near in a low-angle shot,
revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in
for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground.
Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts dynamic
shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
"""
negative_prompt = """
Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality,
low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured,
misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards
"""
output = pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
num_frames=81,
guidance_scale=5.0,
).frames[0]
export_to_video(output, "output.mp4", fps=16)
Pruna
Pruna 是一个模型优化框架,提供多种优化方法——量化、剪枝、缓存、编译——用于加速推理并减少内存使用。下面展示了优化方法的总体概述:
| 技术 | 描述 | 速度 | 内存 | 质量 |
|---|---|---|---|---|
| batcher | 将多个输入组合在一起同时处理,提高计算效率并减少处理时间。 | ✅ | ❌ | ➖ |
| cacher | 存储计算过程中的中间结果,以加快后续操作。 | ✅ | ➖ | ➖ |
| compiler | 针对特定硬件提供指令来优化模型。 | ✅ | ➖ | ➖ |
| distiller | 训练一个更小、更简单的模型来模仿一个更大、更复杂的模型。 | ✅ | ✅ | ❌ |
| quantizer | 降低权重和激活的精度,减少内存需求。 | ✅ | ✅ | ❌ |
| pruner | 移除不太重要或冗余的连接和神经元,形成一个更稀疏、更高效的网络。 | ✅ | ✅ | ❌ |
| recoverer | 在压缩后恢复模型的性能。 | ➖ | ➖ | ✅ |
| factorizer | 分解将多个小矩阵乘法合并为一个大的融合操作。 | ✅ | ➖ | ➖ |
| enhancer | 通过应用去噪或上采样等后处理算法来增强模型输出。 | ❌ | - | ✅ |
✅(改进),➖(大致相同),❌(恶化)
a. 安装依赖
pip install pruna
b. 优化 Diffusers 模型
以下示例使用 factorizer、compiler 和 cacher 算法组合优化 black-forest-labs/FLUX.1-dev。这种组合将推理速度提升高达 4.2 倍,并将峰值 GPU 内存使用从 34.7GB 降低到 28.0GB,同时几乎保持相同的输出质量。
首先,定义一个 SmashConfig 来使用优化算法。要优化模型,用 smash 包裹管道和 SmashConfig,然后像平常一样使用管道进行推理。
import torch
from diffusers import FluxPipeline
from pruna import PrunaModel, SmashConfig, smash
# load the model
# Try segmind/Segmind-Vega or black-forest-labs/FLUX.1-schnell with a small GPU memory
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
torch_dtype=torch.bfloat16
).to("cuda")
# define the configuration
smash_config = SmashConfig()
smash_config["factorizer"] = "qkv_diffusers"
smash_config["compiler"] = "torch_compile"
smash_config["torch_compile_target"] = "module_list"
smash_config["cacher"] = "fora"
smash_config["fora_interval"] = 2
# for the best results in terms of speed you can add these configs
# however they will increase your warmup time from 1.5 min to 10 min
# smash_config["torch_compile_mode"] = "max-autotune-no-cudagraphs"
# smash_config["quantizer"] = "torchao"
# smash_config["torchao_quant_type"] = "fp8dq"
# smash_config["torchao_excluded_modules"] = "norm+embedding"
# optimize the model
smashed_pipe = smash(pipe, smash_config)
# run the model
smashed_pipe("a knitted purple prune").images[0]
优化后,我们可以使用 Hugging Face Hub 分享和加载优化后的模型。
# save the model
smashed_pipe.save_to_hub("<username>/FLUX.1-dev-smashed")
# load the model
smashed_pipe = PrunaModel.from_hub("<username>/FLUX.1-dev-smashed")
c. 评估和基准测试 Diffusers 模型
Pruna 提供了 EvaluationAgent 来评估你优化后的模型质量。
我们可以定义关心的指标(例如总时间和吞吐量),以及要评估的数据集。我们可以定义一个模型并将其传递给 EvaluationAgent。
优化后的模型评估
我们可以通过使用 EvaluationAgent 加载并评估优化后的模型,并将其传递给 Task。
import torch
from diffusers import FluxPipeline
from pruna import PrunaModel
from pruna.data.pruna_datamodule import PrunaDataModule
from pruna.evaluation.evaluation_agent import EvaluationAgent
from pruna.evaluation.metrics import (
ThroughputMetric,
TorchMetricWrapper,
TotalTimeMetric,
)
from pruna.evaluation.task import Task
# define the device
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
# load the model
# Try PrunaAI/Segmind-Vega-smashed or PrunaAI/FLUX.1-dev-smashed with a small GPU memory
smashed_pipe = PrunaModel.from_hub("PrunaAI/FLUX.1-dev-smashed")
# Define the metrics
metrics = [
TotalTimeMetric(n_iterations=20, n_warmup_iterations=5),
ThroughputMetric(n_iterations=20, n_warmup_iterations=5),
TorchMetricWrapper("clip"),
]
# Define the datamodule
datamodule = PrunaDataModule.from_string("LAION256")
datamodule.limit_datasets(10)
# Define the task and evaluation agent
task = Task(metrics, datamodule=datamodule, device=device)
eval_agent = EvaluationAgent(task)
# Evaluate smashed model and offload it to CPU
smashed_pipe.move_to_device(device)
smashed_pipe_results = eval_agent.evaluate(smashed_pipe)
smashed_pipe.move_to_device("cpu")
评估独立模型
与其将优化后的模型与基础模型进行比较,你也可以评估独立的 diffusers 模型。如果你想在未进行优化的情况下评估模型的性能,这会很有用。我们可以通过使用 PrunaModel 包装器并在其上运行 EvaluationAgent 来实现这一点。
import torch
from diffusers import FluxPipeline
from pruna import PrunaModel
# load the model
# Try PrunaAI/Segmind-Vega-smashed or PrunaAI/FLUX.1-dev-smashed with a small GPU memory
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
torch_dtype=torch.bfloat16
).to("cpu")
wrapped_pipe = PrunaModel(model=pipe)
xFormers
安装
pip install xformers
使用
安装 xFormers 后,你可以使用 enable_xformers_memory_efficient_attention() 实现更快推理和更低内存消耗。
Token 合并(ToMe)
Token 合并(ToMe)在基于 Transformer 的网络的前向传递过程中逐步合并冗余的 token/patch,这可以加快 StableDiffusionPipeline 的推理延迟。
安装依赖
pip install tomesd
使用 tomesd 库中的 apply_patch 函数
from diffusers import StableDiffusionPipeline
import torch
import tomesd
pipeline = StableDiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True,
).to("cuda")
tomesd.apply_patch(pipeline, ratio=0.5)
image = pipeline("a photo of an astronaut riding a horse on mars").images[0]
apply_patch 函数提供了一些参数来帮助在管道推理速度和生成的 token 质量之间取得平衡。最重要的参数是 ratio,它控制在前向传递过程中合并的 token 数量。
正如论文中所报道的,ToMe 能够在大幅提升推理速度的同时,极大地保持生成图像的质量。通过增加 ratio,你可以进一步加速推理,但代价是图像质量会有所下降。
DeepCache
DeepCache 通过策略性地缓存和重用高级特征,同时利用 U-Net 架构高效更新低级特征,从而加速 StableDiffusionPipeline 和 StableDiffusionXLPipeline。
安装
pip install DeepCache
加载并启用 DeepCacheSDHelper
import torch
from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained(
'stable-diffusion-v1-5/stable-diffusion-v1-5', torch_dtype=torch.float16
).to("cuda")
from DeepCache import DeepCacheSDHelper
helper = DeepCacheSDHelper(pipe=pipe)
helper.set_params(
cache_interval=3,
cache_branch_id=0,
)
helper.enable()
image = pipe("a photo of an astronaut on a moon").images[0]
set_params 方法接受两个参数:cache_interval 和 cache_branch_id。cache_interval 表示特征缓存的频率,指定为每次缓存操作之间的步数。cache_branch_id 用于识别网络哪个分支(从最浅层到最深层的层顺序)负责执行缓存过程。选择较低的 cache_branch_id 或较大的 cache_interval 可以在牺牲图像质量的情况下提高推理速度(这两项超参数的消融实验可在论文中找到)。设置好这些参数后,使用 enable 或 disable 方法来激活或禁用 DeepCacheSDHelper。
TGATE
T-GATE 通过在收敛后跳过交叉注意力计算,加速了 Stable Diffusion、PixArt 和延迟一致性模型管道的推理。这种方法无需任何额外训练,并且可以将推理速度提高 10-50%。T-GATE 也兼容其他优化方法,如 DeepCache。
安装
pip install tgate
pip install -U torch diffusers transformers accelerate DeepCache
使用对应 pipeline 的加载器
| Pipeline | T-GATE Loader |
|---|---|
| PixArt | TgatePixArtLoader |
| Stable Diffusion XL | TgateSDXLLoader |
| Stable Diffusion XL + DeepCache | TgateSDXLDeepCacheLoader |
| Stable Diffusion | TgateSDLoader |
| Stable Diffusion + DeepCache | TgateSDDeepCacheLoader |
创建一个 TgateLoader,包含一个流水线、门控步骤(停止计算交叉注意力的时间步)以及推理步数。然后在流水线上调用 tgate 方法,传入提示、门控步骤和推理步数。
PixArt
import torch
from diffusers import PixArtAlphaPipeline
from tgate import TgatePixArtLoader
pipe = PixArtAlphaPipeline.from_pretrained(
"PixArt-alpha/PixArt-XL-2-1024-MS", torch_dtype=torch.float16
)
gate_step = 8
inference_step = 25
pipe = TgatePixArtLoader(
pipe,
gate_step=gate_step,
num_inference_steps=inference_step,
).to("cuda")
image = pipe.tgate(
"An alpaca made of colorful building blocks, cyberpunk.",
gate_step=gate_step,
num_inference_steps=inference_step,
).images[0]
Stable Diffusion XL
import torch
from diffusers import StableDiffusionXLPipeline
from diffusers import DPMSolverMultistepScheduler
from tgate import TgateSDXLLoader
pipe = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16,
variant="fp16",
use_safetensors=True,
)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
gate_step = 10
inference_step = 25
pipe = TgateSDXLLoader(
pipe,
gate_step=gate_step,
num_inference_steps=inference_step,
).to("cuda")
image = pipe.tgate(
"Astronaut in a jungle, cold color palette, muted colors, detailed, 8k.",
gate_step=gate_step,
num_inference_steps=inference_step
).images[0]
Stable Diffusion XL + DeepCache
import torch
from diffusers import StableDiffusionXLPipeline
from diffusers import DPMSolverMultistepScheduler
from tgate import TgateSDXLDeepCacheLoader
pipe = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16,
variant="fp16",
use_safetensors=True,
)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
gate_step = 10
inference_step = 25
pipe = TgateSDXLDeepCacheLoader(
pipe,
cache_interval=3,
cache_branch_id=0,
).to("cuda")
image = pipe.tgate(
"Astronaut in a jungle, cold color palette, muted colors, detailed, 8k.",
gate_step=gate_step,
num_inference_steps=inference_step
).images[0]
Latent Consistency Model
import torch
from diffusers import StableDiffusionXLPipeline
from diffusers import UNet2DConditionModel, LCMScheduler
from diffusers import DPMSolverMultistepScheduler
from tgate import TgateSDXLLoader
unet = UNet2DConditionModel.from_pretrained(
"latent-consistency/lcm-sdxl",
torch_dtype=torch.float16,
variant="fp16",
)
pipe = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
unet=unet,
torch_dtype=torch.float16,
variant="fp16",
)
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
gate_step = 1
inference_step = 4
pipe = TgateSDXLLoader(
pipe,
gate_step=gate_step,
num_inference_steps=inference_step,
lcm=True
).to("cuda")
image = pipe.tgate(
"Astronaut in a jungle, cold color palette, muted colors, detailed, 8k.",
gate_step=gate_step,
num_inference_steps=inference_step
).images[0]
xDiT
xDiT 是一款专为大规模并行部署扩散 Transformer(DiTs)而设计的推理引擎。xDiT 为扩散模型提供了一套高效的并行方法,以及 GPU 内核加速。
xDiT 支持四种并行方法,包括统一序列并行、PipeFusion、CFG 并行和数据并行。xDiT 中的四种并行方法可以混合配置,优化通信模式以最佳适应底层网络硬件。
与并行化正交的优化专注于加速单个 GPU 性能。除了利用知名的注意力优化库外,我们还利用了 torch.compile 和 onediff 等编译加速技术。
安装依赖
pip install xfuser
使用 xDiT 加速 Diffusers 模型推理
import torch
from diffusers import StableDiffusion3Pipeline
from xfuser import xFuserArgs, xDiTParallel
from xfuser.config import FlexibleArgumentParser
from xfuser.core.distributed import get_world_group
def main():
parser = FlexibleArgumentParser(description="xFuser Arguments")
args = xFuserArgs.add_cli_args(parser).parse_args()
engine_args = xFuserArgs.from_cli_args(args)
engine_config, input_config = engine_args.create_config()
local_rank = get_world_group().local_rank
pipe = StableDiffusion3Pipeline.from_pretrained(
pretrained_model_name_or_path=engine_config.model_config.model,
torch_dtype=torch.float16,
).to(f"cuda:{local_rank}")
# do anything you want with pipeline here
pipe = xDiTParallel(pipe, engine_config, input_config)
pipe(
height=input_config.height,
width=input_config.height,
prompt=input_config.prompt,
num_inference_steps=input_config.num_inference_steps,
output_type=input_config.output_type,
generator=torch.Generator(device="cuda").manual_seed(input_config.seed),
)
if input_config.output_type == "pil":
pipe.save("results", "stable_diffusion_3")
if __name__ == "__main__":
main()
使用 xDiT 的 xFuserArgs 来获取配置参数,并将这些参数连同 Diffusers 库的 pipeline 对象一起传递给 xDiTParallel,即可完成 Diffusers 中特定 pipeline 的并行化。
xDiT 运行时参数可以通过 -h 在命令行中查看,您可以参考此使用示例获取更多详细信息。
xDiT 需要使用 torchrun 来启动,以支持其多节点、多 GPU 并行功能。例如,以下命令可用于 8-GPU 并行推理:
torchrun --nproc_per_node=8 ./inference.py \
--model models/FLUX.1-dev \
--data_parallel_degree 2 \
--ulysses_degree 2 \
--ring_degree 2 \
--prompt "A snowy mountain" "A small dog" \
--num_inference_steps 50
ParaAttention
大型图像和视频生成模型,如 FLUX.1-dev 和 HunyuanVideo,由于其体积较大,对实时应用和部署可能构成推理挑战。
ParaAttention 是一个实现上下文并行和首块缓存的库,可以与其他技术(torch.compile、fp8 动态量化)结合,以加速推理。
本指南将向你展示如何在 NVIDIA L20 GPU 上对 FLUX.1-dev 和 HunyuanVideo 应用 ParaAttention。我们的基线基准测试除了 HunyuanVideo 为避免内存溢出错误外,未应用任何优化。
我们的基线基准测试显示,FLUX.1-dev 能够在 26.36 秒内通过 28 步生成 1024x1024 分辨率的图像,而 HunyuanVideo 能够在 3675.71 秒内通过 30 步生成 129 帧 720p 分辨率的视频。
a. 第一块缓存
在模型中缓存转换器块的输出并在后续推理步骤中重用,可以减少计算成本并加快推理速度。
然而,确定何时重用缓存以确保生成图像或视频的质量是困难的。ParaAttention 直接使用第一个转换器块输出的残差差异来近似模型输出之间的差异。当差异足够小时,会重用先前推理步骤的残差差异。换句话说,去噪步骤被跳过了。
在 FLUX.1-dev 和 HunyuanVideo 推理中实现了 2 倍的速度提升,质量非常好。
FLUX.1-dev
在 FLUX.1-dev 上应用第一块缓存,请按如下方式调用 apply_cache_on_pipe。0.08 是 FLUX 模型的默认残差差值。
import time
import torch
from diffusers import FluxPipeline
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
torch_dtype=torch.bfloat16,
).to("cuda")
from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
apply_cache_on_pipe(pipe, residual_diff_threshold=0.08)
# Enable memory savings
# pipe.enable_model_cpu_offload()
# pipe.enable_sequential_cpu_offload()
begin = time.time()
image = pipe(
"A cat holding a sign that says hello world",
num_inference_steps=28,
).images[0]
end = time.time()
print(f"Time: {end - begin:.2f}s")
print("Saving image to flux.png")
image.save("flux.png")
与基线相比,第一块缓存将推理速度降低至 17.01 秒,即速度提升了 1.55 倍,同时几乎保持了零质量损失。
HunyuanVideo
在 HunyuanVideo 上应用 First Block Cache,如下所示 apply_cache_on_pipe。0.06 是 HunyuanVideo 模型的默认残差差值。
import time
import torch
from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
from diffusers.utils import export_to_video
model_id = "tencent/HunyuanVideo"
transformer = HunyuanVideoTransformer3DModel.from_pretrained(
model_id,
subfolder="transformer",
torch_dtype=torch.bfloat16,
revision="refs/pr/18",
)
pipe = HunyuanVideoPipeline.from_pretrained(
model_id,
transformer=transformer,
torch_dtype=torch.float16,
revision="refs/pr/18",
).to("cuda")
from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
apply_cache_on_pipe(pipe, residual_diff_threshold=0.6)
pipe.vae.enable_tiling()
begin = time.time()
output = pipe(
prompt="A cat walks on the grass, realistic",
height=720,
width=1280,
num_frames=129,
num_inference_steps=30,
).frames[0]
end = time.time()
print(f"Time: {end - begin:.2f}s")
print("Saving video to hunyuan_video.mp4")
export_to_video(output, "hunyuan_video.mp4", fps=15)
与基线相比,第一块缓存将推理速度降低至 2271.06 秒,即速度提升了 1.62 倍,同时保持了几乎零的质量损失。
b. fp8 量化
使用动态量化的 fp8 进一步加速推理并减少内存使用。为了使用 8 位 NVIDIA Tensor Cores,激活值和权重都必须进行量化。
使用 float8_weight_only 和 float8_dynamic_activation_float8_weight 对文本编码器和 Transformer 模型进行量化。
默认的量化方法是逐张量量化,但如果你的 GPU 支持按行量化,你也可以尝试它以获得更好的精度。
1. 安装依赖
pip3 install -U torch torchao
2. 量化模型
使用 mode="max-autotune-no-cudagraphs" 或 mode="max-autotune" 的 torch.compile 可以选择最佳内核以获得最佳性能。如果模型是第一次被调用,编译可能需要很长时间,但一旦模型被编译,这是值得的。
此示例仅量化 transformer 模型,但你也可以量化文本编码器以进一步减少内存使用。
FLUX.1-dev
import time
import torch
from diffusers import FluxPipeline
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
torch_dtype=torch.bfloat16,
).to("cuda")
from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
apply_cache_on_pipe(
pipe,
residual_diff_threshold=0.12, # Use a larger value to make the cache take effect
)
from torchao.quantization import (
quantize_,
float8_dynamic_activation_float8_weight,
float8_weight_only
)
quantize_(pipe.text_encoder, float8_weight_only())
quantize_(pipe.transformer, float8_dynamic_activation_float8_weight())
pipe.transformer = torch.compile(
pipe.transformer, mode="max-autotune-no-cudagraphs",
)
# Enable memory savings
# pipe.enable_model_cpu_offload()
# pipe.enable_sequential_cpu_offload()
for i in range(2):
begin = time.time()
image = pipe(
"A cat holding a sign that says hello world",
num_inference_steps=28,
).images[0]
end = time.time()
if i == 0:
print(f"Warm up time: {end - begin:.2f}s")
else:
print(f"Time: {end - begin:.2f}s")
print("Saving image to flux.png")
image.save("flux.png")
HunyuanVideo
import time
import torch
from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
from diffusers.utils import export_to_video
model_id = "tencent/HunyuanVideo"
transformer = HunyuanVideoTransformer3DModel.from_pretrained(
model_id,
subfolder="transformer",
torch_dtype=torch.bfloat16,
revision="refs/pr/18",
)
pipe = HunyuanVideoPipeline.from_pretrained(
model_id,
transformer=transformer,
torch_dtype=torch.float16,
revision="refs/pr/18",
).to("cuda")
from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
apply_cache_on_pipe(pipe)
from torchao.quantization import (
quantize_,
float8_dynamic_activation_float8_weight,
float8_weight_only
)
quantize_(pipe.text_encoder, float8_weight_only())
quantize_(pipe.transformer, float8_dynamic_activation_float8_weight())
pipe.transformer = torch.compile(
pipe.transformer, mode="max-autotune-no-cudagraphs",
)
# Enable memory savings
pipe.vae.enable_tiling()
# pipe.enable_model_cpu_offload()
# pipe.enable_sequential_cpu_offload()
for i in range(2):
begin = time.time()
output = pipe(
prompt="A cat walks on the grass, realistic",
height=720,
width=1280,
num_frames=129,
num_inference_steps=1 if i == 0 else 30,
).frames[0]
end = time.time()
if i == 0:
print(f"Warm up time: {end - begin:.2f}s")
else:
print(f"Time: {end - begin:.2f}s")
print("Saving video to hunyuan_video.mp4")
export_to_video(output, "hunyuan_video.mp4", fps=15)
一个 NVIDIA L20 GPU 只有 48GB 内存,在编译后如果
enable_model_cpu_offload未被调用,可能会出现内存不足(OOM)错误,因为 HunyuanVideo 在以高分辨率和大帧数运行时会产生非常大的激活张量。对于内存小于 80GB 的 GPU,可以尝试降低分辨率和帧数以避免 OOM 错误。大型视频生成模型通常受注意力计算的限制,而不是全连接层。这些模型从量化和
torch.compile中获益不大。
c. 上下文并行
Context Parallelism 并行化推理并支持多 GPU 扩展。ParaAttention 的模块化设计允许你将 Context Parallelism 与 First Block Cache 和动态量化结合使用。
如果推理过程需要持久化和可服务化,建议使用 torch.multiprocessing 编写自己的推理处理器。这可以消除启动进程、加载和重新编译模型的额外开销。
FLUX.1-dev
下面的代码示例结合了 First Block Cache、fp8 动态量化、torch.compile 和上下文并行,以实现最快的推理速度。
import time
import torch
import torch.distributed as dist
from diffusers import FluxPipeline
dist.init_process_group()
torch.cuda.set_device(dist.get_rank())
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
torch_dtype=torch.bfloat16,
).to("cuda")
from para_attn.context_parallel import init_context_parallel_mesh
from para_attn.context_parallel.diffusers_adapters import parallelize_pipe
from para_attn.parallel_vae.diffusers_adapters import parallelize_vae
mesh = init_context_parallel_mesh(
pipe.device.type,
max_ring_dim_size=2,
)
parallelize_pipe(
pipe,
mesh=mesh,
)
parallelize_vae(pipe.vae, mesh=mesh._flatten())
from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
apply_cache_on_pipe(
pipe,
residual_diff_threshold=0.12, # Use a larger value to make the cache take effect
)
from torchao.quantization import (
quantize_,
float8_dynamic_activation_float8_weight,
float8_weight_only
)
quantize_(pipe.text_encoder, float8_weight_only())
quantize_(pipe.transformer, float8_dynamic_activation_float8_weight())
torch._inductor.config.reorder_for_compute_comm_overlap = True
pipe.transformer = torch.compile(
pipe.transformer, mode="max-autotune-no-cudagraphs",
)
# Enable memory savings
# pipe.enable_model_cpu_offload(gpu_id=dist.get_rank())
# pipe.enable_sequential_cpu_offload(gpu_id=dist.get_rank())
for i in range(2):
begin = time.time()
image = pipe(
"A cat holding a sign that says hello world",
num_inference_steps=28,
output_type="pil" if dist.get_rank() == 0 else "pt",
).images[0]
end = time.time()
if dist.get_rank() == 0:
if i == 0:
print(f"Warm up time: {end - begin:.2f}s")
else:
print(f"Time: {end - begin:.2f}s")
if dist.get_rank() == 0:
print("Saving image to flux.png")
image.save("flux.png")
dist.destroy_process_group()
保存到 run_flux.py,并使用 torchrun 启动:
# Use --nproc_per_node to specify the number of GPUs
torchrun --nproc_per_node=2 run_flux.py
与基线相比,使用 2 块 NVIDIA L20 GPU 时,推理速度降至 8.20 秒,或提升 3.21 倍;使用 4 块 L20 时,推理速度为 3.90 秒,或提升 6.75 倍。
HunyuanVideo
下面的代码示例结合了第一块缓存和上下文并行,以实现最快的推理速度。
import time
import torch
import torch.distributed as dist
from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
from diffusers.utils import export_to_video
dist.init_process_group()
torch.cuda.set_device(dist.get_rank())
model_id = "tencent/HunyuanVideo"
transformer = HunyuanVideoTransformer3DModel.from_pretrained(
model_id,
subfolder="transformer",
torch_dtype=torch.bfloat16,
revision="refs/pr/18",
)
pipe = HunyuanVideoPipeline.from_pretrained(
model_id,
transformer=transformer,
torch_dtype=torch.float16,
revision="refs/pr/18",
).to("cuda")
from para_attn.context_parallel import init_context_parallel_mesh
from para_attn.context_parallel.diffusers_adapters import parallelize_pipe
from para_attn.parallel_vae.diffusers_adapters import parallelize_vae
mesh = init_context_parallel_mesh(
pipe.device.type,
)
parallelize_pipe(
pipe,
mesh=mesh,
)
parallelize_vae(pipe.vae, mesh=mesh._flatten())
from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
apply_cache_on_pipe(pipe)
# from torchao.quantization import quantize_, float8_dynamic_activation_float8_weight, float8_weight_only
#
# torch._inductor.config.reorder_for_compute_comm_overlap = True
#
# quantize_(pipe.text_encoder, float8_weight_only())
# quantize_(pipe.transformer, float8_dynamic_activation_float8_weight())
# pipe.transformer = torch.compile(
# pipe.transformer, mode="max-autotune-no-cudagraphs",
# )
# Enable memory savings
pipe.vae.enable_tiling()
# pipe.enable_model_cpu_offload(gpu_id=dist.get_rank())
# pipe.enable_sequential_cpu_offload(gpu_id=dist.get_rank())
for i in range(2):
begin = time.time()
output = pipe(
prompt="A cat walks on the grass, realistic",
height=720,
width=1280,
num_frames=129,
num_inference_steps=1 if i == 0 else 30,
output_type="pil" if dist.get_rank() == 0 else "pt",
).frames[0]
end = time.time()
if dist.get_rank() == 0:
if i == 0:
print(f"Warm up time: {end - begin:.2f}s")
else:
print(f"Time: {end - begin:.2f}s")
if dist.get_rank() == 0:
print("Saving video to hunyuan_video.mp4")
export_to_video(output, "hunyuan_video.mp4", fps=15)
dist.destroy_process_group()
保存到 run_hunyuan_video.py,并使用 torchrun 启动:
# Use --nproc_per_node to specify the number of GPUs
torchrun --nproc_per_node=8 run_hunyuan_video.py
与基线相比,使用 8 块 NVIDIA L20 GPU 时,推理速度降至 649.23 秒,即速度提升了 5.66 倍。