VAE（变分自编码器，Variational Auto-Encoder）是基于Encoder-Decoder架构的生成模型。VAE的Encoder（编码器）结构能将输入图像转换为低维Latent特征，并作为U-Net的输入。VAE的Decoder（解码器）结构能将低维Latent特征重建还原成像素级图像。

在Stable Diffusion中，VAE模型主要起到了图像压缩和图像重建的作用

当我们输入一个尺寸为 H×W×C 的数据，VAE的Encoder模块会将其编码为一个大小为h×w×c的低维Latent特征，其中f=H/h=W/w为VAE的下采样率（Downsampling Factor）。反之，VAE的Decoder模块有一个相同的上采样率（Upsampling Factor）将低维Latent特征重建成像素级别的图像

VAE Encoder 和VAE Decoder结构图


#VAE压缩与重建代码展示
import cv2
import torch
import numpy as np
from diffusers import AutoencoderKL
 
# 加载VAE模型: VAE模型可以通过指定subfolder文件来单独加载。
# SD V1.5模型权重百度云网盘：关注Rocky的公众号WeThinkIn，后台回复：SD模型，即可获得资源链接
VAE = AutoencoderKL.from_pretrained("/本地路径/stable-diffusion-v1-5", subfolder="vae")
VAE.to("cuda", dtype=torch.float16)
 
# 用OpenCV读取和调整图像大小
raw_image = cv2.imread("catwoman.png")
raw_image = cv2.cvtColor(raw_image, cv2.COLOR_BGR2RGB)
raw_image = cv2.resize(raw_image, (1024, 1024))
 
# 将图像数据转换为浮点数并归一化
image = raw_image.astype(np.float32) / 127.5 - 1.0
 
# 调整数组维度以匹配PyTorch的格式 (N, C, H, W)
image = image.transpose(2, 0, 1)
image = image[None, :, :, :]
 
# 转换为PyTorch张量
image = torch.from_numpy(image).to("cuda", dtype=torch.float16)
 
# 压缩图像为Latent特征并重建
with torch.inference_mode():
    # 使用VAE进行压缩和重建
    latent = VAE.encode(image).latent_dist.sample()
    rec_image = VAE.decode(latent).sample
 
    # 后处理
    rec_image = (rec_image / 2 + 0.5).clamp(0, 1)
    rec_image = rec_image.cpu().permute(0, 2, 3, 1).numpy()
 
    # 反归一化
    rec_image = (rec_image * 255).round().astype("uint8")
    rec_image = rec_image[0]
 
    # 保存重建后图像
    cv2.imwrite("reconstructed_catwoman.png", cv2.cvtColor(rec_image, cv2.COLOR_RGB2BGR))

OpenAI开源的一致性解码器（consistency-decoder），支持作为Stable Diffusion 1.x和2.x的VAE模型。由于consistency-decoder模型较大（FP32：2.49G，FP16：1.2G），重建耗时会比原生的SD VAE模型大得多，并且在高分辨率（比如1024x1024）下效果并没有明显高于原生的SD VAE模型，所以最好将consistency-decoder模型作为补充储备模型之用。


import torch
from diffusers import DiffusionPipeline, ConsistencyDecoderVAE
 
# SD 1.5和consistency-decoder模型权重百度云网盘：关注Rocky的公众号WeThinkIn，后台回复：SD模型，即可获得资源链接
vae = ConsistencyDecoderVAE.from_pretrained("/本地路径/consistency-decoder", torch_dtype=pipe.torch_dtype)
pipe = StableDiffusionPipeline.from_pretrained(
    "/本地路径/stable-diffusion-v1-5", vae=vae, torch_dtype=torch.float16
).to("cuda")
 
pipe("horse", generator=torch.manual_seed(0)).images

2.Unet

U-Net模型是一个关键核心部分，能够预测噪声残差，并结合Sampling method（调度算法：DDPM、DDIM、DPM++等）对输入的特征矩阵进行重构，逐步将其从随机高斯噪声转化成图片的Latent Feature。

CrossAttention结构中，VAE模型的输出latent space用来生成Q，文字特征信息用来生成K和V

自注意力：自己特征生成q,k,v，计算注意力

交叉注意力：一个特征生成q，另一个特征生成k，v，然后计算注意力

3.CLIP Text Encoder

CLIP模型是一个基于对比学习的多模态模型，主要包含Text Encoder和Image Encoder两个模型。其中Text Encoder用来提取文本的特征，可以使用NLP中常用的text transformer模型作为Text Encoder；而Image Encoder主要用来提取图像的特征，可以使用CNN/Vision transformer模型（ResNet和ViT等）作为Image Encoder。与此同时，他直接使用4亿个图片与标签文本对数据集进行训练，来学习图片与本文内容的对应关系。

CLIP的训练：

CLIP在训练时，从训练集中随机取出一张图片和标签文本，接着CLIP模型的任务主要是通过Text Encoder和Image Encoder分别将标签文本和图片提取embedding向量，然后用余弦相似度（cosine similarity）来比较两个embedding向量的相似性，以判断随机抽取的标签文本和图片是否匹配，并进行梯度反向传播，不断进行优化训练。

上面我们讲到CLIP模型主要包含Text Encoder和Image Encoder两个部分，在Stable Diffusion中主要使用了Text Encoder部分。CLIP Text Encoder模型将输入的文本Prompt进行编码，转换成Text Embeddings（文本的语义信息），通过前面章节提到的U-Net网络的CrossAttention模块嵌入Stable Diffusion中作为Condition条件，对生成图像的内容进行一定程度上的控制与引导，目前SD模型使用的的是CLIP ViT-L/14中的Text Encoder模型。

CLIP ViT-L/14 中的Text Encoder是只包含Transformer结构的模型，一共由12个CLIPEncoderLayer模块组成，模型参数大小是123M，具体CLIP Text Encoder模型结构如下图所示。其中特征维度为768，token数量是77，所以输出的Text Embeddings的维度为77x768。


from transformers import CLIPTextModel, CLIPTokenizer
 
# 加载 CLIP Text Encoder模型和Tokenizer
# SD V1.5模型权重百度云网盘：关注Rocky的公众号WeThinkIn，后台回复：SDV1.5模型，即可获得资源链接
text_encoder = CLIPTextModel.from_pretrained("/本地路径/stable-diffusion-v1-5", subfolder="text_encoder").to("cuda")
text_tokenizer = CLIPTokenizer.from_pretrained("/本地路径/stable-diffusion-v1-5", subfolder="tokenizer")
 
# 将输入SD模型的prompt进行tokenize，得到对应的token ids特征
prompt = "1girl,beautiful"
text_token_ids = text_tokenizer(
    prompt,
    padding="max_length",
    max_length=text_tokenizer.model_max_length,
    truncation=True,
    return_tensors="pt"
).input_ids
 
print("text_token_ids' shape:",text_token_ids.shape)
print("text_token_ids:",text_token_ids)
 
# 将token ids特征输入CLIP Text Encoder模型中输出77x768的Text Embeddings特征
text_embeddings = text_encoder(text_token_ids.to("cuda"))[0] # 由于CLIP Text Encoder模型输出的是一个元组，所以需要[0]对77x768的Text Embeddings特征进行提取
print("text_embeddings' shape:",text_embeddings.shape)
print(text_embeddings)
 
---------------- 运行结果 ----------------
text_token_ids' shape: torch.Size([1, 77])
text_token_ids: tensor([[49406,   272,  1611,   267,  1215, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407]])
text_embeddings' shape: torch.Size([1, 77, 768])
tensor([[[-0.3884,  0.0229, -0.0522,  ..., -0.4899, -0.3066,  0.0675],
         [-0.8425, -1.1393,  1.2756,  ..., -0.2595,  1.6293, -0.7857],
         [ 0.1753, -0.9846,  0.1879,  ...,  0.0664, -1.4935, -1.2614],
         ...,
         [ 0.2039, -0.7296, -0.3212,  ...,  0.6748, -0.5813, -0.7323],
         [ 0.1921, -0.7344, -0.3045,  ...,  0.6803, -0.5852, -0.7230],
         [ 0.2114, -0.6436, -0.3047,  ...,  0.6624, -0.5575, -0.7586]]],
       device='cuda:0', grad_fn=<NativeLayerNormBackward0>)

一般来说，我们提取CLIP Text Encoder模型最后一层特征作为CrossAttention模块的输入，但是开源社区的不断实践为我们总结了如下经验：当我们生成二次元内容时，可以选择提取CLIP Text Encoder模型倒数第二层特征；当我们生成写实场景内容时，可以选择提取CLIP Text Encoder模型最后一层的特征。这让Rocky想起了SRGAN以及感知损失，其也是提取了VGG网络的中间层特征才达到了最好的效果，AI领域的“传承”与共性，往往在这些不经意间，让人感到人工智能的魅力与美妙。

由于CLIP训练时所采用的最大Token数是77，所以在SD模型进行前向推理时，当输入Prompt的Token数量超过77时，将通过Clip操作拉回77x768，而如果Token数不足77则会使用padding操作得到77x768。如果说全卷积网络的设计让图像输入尺寸不再受限，那么CLIP的这个设置就让输入的文本长度不再受限（可以是空文本）。无论是非常长的文本，还是空文本，最后都将得到一样维度的特征矩阵。

同时在SD模型的训练中，一般来说CLIP的整体性能是足够支撑我们的下游细分任务的，所以CLIP Text Encoder模型参数是冻结的，我们不需要对其重新训练。

4.SD训练过程：

我们进行Stable Diffusion模型训练时，VAE部分和CLIP部分都是冻结的，所以说官方在训练SD系列模型的时候，训练过程一般主要训练U-Net部分。

数据集：数据集格式为图片+对应的描述(prompt)

如上图，unet预测latent features添加的噪声noise1，计算noise1与noise的交叉熵损失，后续有N个step的添加噪声，预测噪声。


for step, batch in enumerate(train_dataloader):
    with torch.no_grad():
        # 将image转到latent空间
        latents = vae.encode(batch["image"]).latent_dist.sample()
        latents = latents * vae.config.scaling_factor # rescaling latents
        # 提取text embeddings
        text_input_ids = text_tokenizer(
            batch["text"],
            padding="max_length",
            max_length=tokenizer.model_max_length,
            truncation=True,
            return_tensors="pt"
    ).input_ids
    text_embeddings = text_encoder(text_input_ids)[0]
    
    # 随机采样噪音
    noise = torch.randn_like(latents)
    bsz = latents.shape[0]
    # 随机采样timestep
    timesteps = torch.randint(0, noise_scheduler.num_train_timesteps, (bsz,), device=latents.device)
    timesteps = timesteps.long()
 
    # 将noise添加到latent上，即扩散过程
    noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
 
    # 预测noise并计算loss
    model_pred = unet(noisy_latents, timesteps, encoder_hidden_states=text_embeddings).sample
    loss = F.mse_loss(model_pred.float(), noise.float(), reduction="mean")

5.SD推理

5.1文生图

noise充当latent feature一步步降噪，对n个step后的lantent feature使用VAE Decoder解码

5.2图生图

对latent feature一步步降噪，对n个step后的lantent feature使用VAE Decoder解码


示例1：
#读取diffuers库
from diffusers import StableDiffusionPipeline
 
#初始化SD模型，加载预训练权重
pipe = StableDiffusionPipeline.from_pretrained("/本地路径/stable-diffusion-v1-5")
 
#使用GPU加速
pipe.to("cuda")
 
#如GPU的内存少于10GB，可以加载float16精度的SD模型
pipe = StableDiffusionPipeline.from_pretrained("/本地路径/stable-diffusion-v1-5", revision="fp16", torch_dtype=torch.float16)
 
#接下来，我们就可以运行pipeline了
prompt = "a photograph of an astronaut riding a horse"
 
image = pipe(prompt).images[0]
 
 
 
 
示例2：
# 加载文生图pipeline
from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", # 其中一个版本
    torch_dtype=torch.float16).to("cuda")
# 输入text，这里text又称为prompt
prompts = [
    "a photograph of an astronaut riding a horse",
    "A cute otter in a rainbow whirlpool holding shells, watercolor",
    "An avocado armchair",
    "A white dog wearing sunglasses"]
generator = torch.Generator("cuda").manual_seed(42) # 定义随机seed，保证可重复性
# 执行推理
images = pipe(
    prompts,
    height=512,
    width=512,
    num_inference_steps=50,
    guidance_scale=7.5,
    negative_prompt=None,
    num_images_per_prompt=1,
    generator=generator).images

除了将SD模型权重整体加载，我们还可以将SD模型的不同组件权重进行单独加载：


from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler
from diffusers import LMSDiscreteScheduler
 
# 单独加载VAE模型 
vae = AutoencoderKL.from_pretrained("/本地路径/stable-diffusion-v1-5", subfolder="vae")
 
# 单独家在CLIP模型和tokenizer
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
 
# 单独加载U-Net模型
unet = UNet2DConditionModel.from_pretrained("/本地路径/stable-diffusion-v1-5", subfolder="unet")
 
# 单独加载调度算法
scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
 
 
#此部分为GPT写的，未经检验
from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import AutoencoderKL, UNet2DConditionModel, LMSDiscreteScheduler, StableDiffusionPipeline
import torch
from PIL import Image
import numpy as np
 
# 单独加载VAE模型
vae = AutoencoderKL.from_pretrained("/本地路径/stable-diffusion-v1-5", subfolder="vae")
 
# 单独加载CLIP模型和tokenizer
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
 
# 单独加载U-Net模型
unet = UNet2DConditionModel.from_pretrained("/本地路径/stable-diffusion-v1-5", subfolder="unet")
 
# 单独加载调度算法
scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
 
# 定义StableDiffusionPipeline
class CustomStableDiffusionPipeline(StableDiffusionPipeline):
    def __init__(self, vae, text_encoder, tokenizer, unet, scheduler):
        super().__init__(vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet, scheduler=scheduler)
 
# 创建自定义管道
pipe = CustomStableDiffusionPipeline(vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet, scheduler=scheduler)
 
# 设置为评估模式
pipe = pipe.to("cuda") if torch.cuda.is_available() else pipe.to("cpu")
pipe.eval()
 
# 文本输入
prompt = "A fantasy landscape with mountains and a river"
 
# 编码文本输入
inputs = tokenizer(prompt, return_tensors="pt")
input_ids = inputs["input_ids"].to(pipe.device)
 
# 使用文本编码器生成文本嵌入
text_embeddings = text_encoder(input_ids)[0]
 
# 生成初始噪声
latents = torch.randn((1, unet.in_channels, 64, 64), device=pipe.device)
 
# 调度算法初始化
scheduler.set_timesteps(50)
 
# 逐步生成图像
for t in scheduler.timesteps:
    # 预测噪声残差
    with torch.no_grad():
        noise_pred = unet(latents, t, encoder_hidden_states=text_embeddings).sample
 
    # 计算梯度下降
    latents = scheduler.step(noise_pred, t, latents).prev_sample
 
# 解码潜变量以获取图像
with torch.no_grad():
    image = vae.decode(latents).sample
 
# 将图像转换为PIL格式
image = (image / 2 + 0.5).clamp(0, 1)
image = image.cpu().permute(0, 2, 3, 1).numpy()
image = (image * 255).astype(np.uint8)
image = Image.fromarray(image[0])
 
# 显示图像
image.show()

6.SD模型的加速方法参考：

深入浅出完整解析Stable Diffusion（SD）核心基础知识 - 知乎 (zhihu.com)

第8小节

声明：本文内容由网友自发贡献，转载请注明出处：【wpsshop】