深入理解Diffusers: 从基础到Stable Diffusion

Diffusers是一个功能强大且灵活的扩散模型工具箱,它让构建和部署扩散系统变得简单直观。本文将带你深入了解Diffusers的核心组件,并通过实例展示如何从零构建扩散系统。

核心组件

Diffusers工具箱的核心是模型(Model)和调度器(Scheduler):

模型: 负责预测噪声残差
调度器: 负责管理去噪过程
Pipeline: 将模型和调度器组合在一起,提供便捷的推理接口

从基础Pipeline开始

让我们先看一个简单的例子 - DDPM Pipeline:

python 复制代码

from diffusers import DDPMPipeline

ddpm = DDPMPipeline.from_pretrained("google/ddpm-cat-256").to("cuda")
image = ddpm(num_inference_steps=25).images[0]
image.save("ddpm_generated_cat.png")

这个Pipeline内部包含了:

UNet2DModel: 用于预测噪声残差
DDPMScheduler: 用于管理去噪过程

拆解Pipeline

让我们看看如何手动构建这个系统:

加载模型和调度器:

python 复制代码

from diffusers import DDPMScheduler, UNet2DModel

scheduler = DDPMScheduler.from_pretrained("google/ddpm-cat-256")
model = UNet2DModel.from_pretrained("google/ddpm-cat-256").to("cuda")

设置时间步:

python 复制代码

scheduler.set_timesteps(50)

创建随机噪声:

python 复制代码

import torch
sample_size = model.config.sample_size
noise = torch.randn((1, 3, sample_size, sample_size), device="cuda")

实现去噪循环:

python 复制代码

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

# 将生成的图像保存
image = (input / 2 + 0.5).clamp(0, 1)
image = image.detach().cpu().permute(0, 2, 3, 1).numpy()[0]
image = (image * 255).round().astype("uint8")
image = Image.fromarray(image)
image.save("manual_ddpm_generated_cat.png")

生成的图片：

Stable Diffusion Pipeline

Stable Diffusion是一个更复杂的文本到图像潜空间扩散模型。它包含以下组件:

VAE (变分自编码器)
CLIP文本编码器
UNet条件模型
调度器

让我们看看如何构建Stable Diffusion系统:

加载所有组件:

python 复制代码

from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import AutoencoderKL, UNet2DConditionModel, UniPCMultistepScheduler

torch_device = "cuda" if torch.cuda.is_available() else "cpu"
vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae").to(torch_device)
tokenizer = CLIPTokenizer.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="tokenizer")
text_encoder = CLIPTextModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="text_encoder").to(torch_device)
unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet").to(torch_device)
scheduler = UniPCMultistepScheduler.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="scheduler")

设置参数和生成文本嵌入:

python 复制代码

prompt = ["a photograph of an astronaut riding a horse"]
batch_size = 1
height = 512
width = 512
guidance_scale = 7.5
generator = torch.manual_seed(0)

text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, return_tensors="pt")
text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]

创建潜空间噪声:

python 复制代码

latents = torch.randn(
    (batch_size, unet.config.in_channels, height // 8, width // 8),
    generator=generator,
    device=torch_device,
)

实现去噪循环:

python 复制代码

scheduler.set_timesteps(50)

for t in scheduler.timesteps:
    latent_model_input = torch.cat([latents] * 2)
    latent_model_input = scheduler.scale_model_input(latent_model_input, timestep=t)
    
    with torch.no_grad():
        noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
    
    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
    
    latents = scheduler.step(noise_pred, t, latents).prev_sample

解码和保存图像:

python 复制代码

latents = 1 / 0.18215 * latents
with torch.no_grad():
    image = vae.decode(latents).sample

image = (image / 2 + 0.5).clamp(0, 1)
image = image.detach().cpu().permute(0, 2, 3, 1).numpy()[0]
image = (image * 255).round().astype("uint8")
image = Image.fromarray(image)
image.save("stable_diffusion_generated.png")

总结

Diffusers的设计理念是让构建扩散系统变得简单直观。通过理解模型和调度器的工作原理,你可以:

构建自己的扩散系统
自定义和优化现有Pipeline
尝试不同的模型和调度器组合

无论你是想构建简单的图像生成系统,还是复杂的文本到图像模型,Diffusers都能为你提供所需的工具和灵活性。

下一步

探索Diffusers库中的其他Pipeline
尝试构建自己的Pipeline
实验不同的模型和调度器组合
优化推理性能

Diffusers正在不断发展,欢迎加入社区,一起推动扩散模型的发展!