强大到离谱!硬核解读Stable Diffusion(完整版)

这篇具有很好参考价值的文章主要介绍了强大到离谱!硬核解读Stable Diffusion(完整版)。希望对大家有所帮助。如果存在错误或未考虑完全的地方,请大家不吝赐教,您也可以点击"举报违法"按钮提交疑问。

原文链接:

硬核解读Stable Diffusion(完整版)

2022年可谓是AIGC(AI Generated Content)元年,上半年有文生图大模型DALL-E2Stable Diffusion,下半年有OpenAI的文本对话大模型ChatGPT问世,这让冷却的AI又沸腾起来了,因为AIGC能让更多的人真真切切感受到AI的力量。这篇文章将介绍比较火的文生图模型Stable Diffusion(简称SD),Stable Diffusion不仅是一个完全开源的模型(代码,数据,模型全部开源),而且是它的参数量只有1B左右,大部分人可以在普通的显卡上进行推理甚至精调模型。毫不夸张的说,Stable Diffusion的出现和开源对AIGC的火热和发展是有巨大推动作用的,因为它让更多的人能快地上手AI作画。这里将基于Hugging Face的diffusers库深入讲解SD的技术原理以及部分的实现细节,然后也会介绍SD的常用功能,注意本文主要以SD V1.5版本为例,在最后也会简单介绍 SD 2.0版本以及基于SD的扩展应用。

SD模型原理

SD是CompVis、Stability AI和LAION等公司研发的一个文生图模型,它的模型和代码是开源的,而且训练数据LAION-5B也是开源的。SD在开源90天github仓库就收获了33K的stars,可见这个模型是多受欢迎。

强大到离谱!硬核解读Stable Diffusion(完整版)

SD是一个基于latent的扩散模型,它在UNet中引入text condition来实现基于文本生成图像。SD的核心来源于Latent Diffusion这个工作,常规的扩散模型是基于pixel的生成模型,而Latent Diffusion是基于latent的生成模型,它先采用一个autoencoder将图像压缩到latent空间,然后用扩散模型来生成图像的latents,最后送入autoencoder的decoder模块就可以得到生成的图像。基于latent的扩散模型的优势在于计算效率更高效,因为图像的latent空间要比图像pixel空间要小,这也是SD的核心优势。文生图模型往往参数量比较大,基于pixel的方法往往限于算力只生成64x64大小的图像,比如OpenAI的DALL-E2和谷歌的Imagen,然后再通过超分辨模型将图像分辨率提升至256x256和1024x1024;而基于latent的SD是在latent空间操作的,它可以直接生成256x256和512x512甚至更高分辨率的图像。

SD模型的主体结构如下图所示,主要包括三个模型:

  • autoencoder:encoder将图像压缩到latent空间,而decoder将latent解码为图像;

  • CLIP text encoder:提取输入text的text embeddings,通过cross attention方式送入扩散模型的UNet中作为condition;

  • UNet:扩散模型的主体,用来实现文本引导下的latent生成。

强大到离谱!硬核解读Stable Diffusion(完整版)

对于SD模型,其autoencoder模型参数大小为84M,CLIP text encoder模型大小为123M,而UNet参数大小为860M,所以SD模型的总参数量约为1B

autoencoder

autoencoder是一个基于encoder-decoder架构的图像压缩模型,对于一个大小为的输入图像,encoder模块将其编码为一个大小为的latent,其中为下采样率(downsampling factor)。在训练autoencoder过程中,除了采用L1重建损失外,还增加了感知损失(perceptual loss,即LPIPS,具体见论文The Unreasonable Effectiveness of Deep Features as a Perceptual Metric)以及基于patch的对抗训练。辅助loss主要是为了确保重建的图像局部真实性以及避免模糊,具体损失函数见latent diffusion的loss部分。同时为了防止得到的latent的标准差过大,采用了两种正则化方法:第一种是KL-reg,类似VAE增加一个latent和标准正态分布的KL loss,不过这里为了保证重建效果,采用比较小的权重(~10e-6);第二种是VQ-reg,引入一个VQ (vector quantization)layer,此时的模型可以看成是一个VQ-GAN,不过VQ层是在decoder模块中,这里VQ的codebook采样较高的维度(8192)来降低正则化对重建效果的影响。latent diffusion论文中实验了不同参数下的autoencoder模型,如下表所示,可以看到当较小和较大时,重建效果越好(PSNR越大),这也比较符合预期,毕竟此时压缩率小。

强大到离谱!硬核解读Stable Diffusion(完整版)

论文进一步将不同的autoencoder在扩散模型上进行实验,在ImageNet数据集上训练同样的步数(2M steps),其训练过程的生成质量如下所示,可以看到过小的(比如1和2)下收敛速度慢,此时图像的感知压缩率较小,扩散模型需要较长的学习;而过大的其生成质量较差,此时压缩损失过大。

强大到离谱!硬核解读Stable Diffusion(完整版)

当在4~16时,可以取得相对好的效果。SD采用基于KL-reg的autoencoder,其中下采样率,特征维度为,当输入图像为512x512大小时将得到64x64x4大小的latent。autoencoder模型时在OpenImages数据集上基于256x256大小训练的,但是由于autoencoder的模型是全卷积结构的(基于ResnetBlock),所以它可以扩展应用在尺寸>256的图像上。下面我们给出使用diffusers库来加载autoencoder模型,并使用autoencoder来实现图像的压缩和重建,代码如下所示:

import torch
from diffusers import AutoencoderKL
import numpy as np
from PIL import Image

#加载模型: autoencoder可以通过SD权重指定subfolder来单独加载
autoencoder = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")
autoencoder.to("cuda", dtype=torch.float16)

# 读取图像并预处理
raw_image = Image.open("boy.png").convert("RGB").resize((256, 256))
image = np.array(raw_image).astype(np.float32) / 127.5 - 1.0
image = image[None].transpose(0, 3, 1, 2)
image = torch.from_numpy(image)

# 压缩图像为latent并重建
with torch.inference_mode():
    latent = autoencoder.encode(image.to("cuda", dtype=torch.float16)).latent_dist.sample()
    rec_image = autoencoder.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 = Image.fromarray(rec_image[0])
rec_image

这里我们给出了两张图片在256x256和512x512下的重建效果对比,如下所示,第一列为原始图片,第二列为512x512尺寸下的重建图,第三列为256x256尺寸下的重建图。对比可以看出,autoencoder将图片压缩到latent后再重建其实是有损的,比如会出现文字和人脸的畸变,在256x256分辨率下是比较明显的,512x512下效果会好很多。

强大到离谱!硬核解读Stable Diffusion(完整版)

这种有损压缩肯定是对SD的生成图像质量是有一定影响的,不过好在SD模型基本上是在512x512以上分辨率下使用的。为了改善这种畸变,stabilityai在发布SD 2.0时同时发布了两个在LAION子数据集上精调的autoencoder,注意这里只精调autoencoder的decoder部分,SD的UNet在训练过程只需要encoder部分,所以这样精调后的autoencoder可以直接用在先前训练好的UNet上(这种技巧还是比较通用的,比如谷歌的Parti也是在训练好后自回归生成模型后,扩大并精调ViT-VQGAN的decoder模块来提升生成质量)。我们也可以直接在diffusers中使用这些autoencoder,比如mse版本(采用mse损失来finetune的模型):

autoencoder = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse/")

对于同样的两张图,这个mse版本的重建效果如下所示,可以看到相比原始版本的autoencoder,畸变是有一定改善的。

强大到离谱!硬核解读Stable Diffusion(完整版)

由于SD采用的autoencoder是基于KL-reg的,所以这个autoencoder在编码图像时其实得到的是一个高斯分布DiagonalGaussianDistribution(分布的均值和标准差),然后通过调用sample方法来采样一个具体的latent(调用mode方法可以得到均值)。由于KL-reg的权重系数非常小,实际得到latent的标准差还是比较大的,latent diffusion论文中提出了一种rescaling方法:首先计算出第一个batch数据中的latent的标准差,然后采用的系数来rescale latent,这样就尽量保证latent的标准差接近1(防止扩散过程的SNR较高,影响生成效果,具体见latent diffusion论文的D1部分讨论),然后扩散模型也是应用在rescaling的latent上,在解码时只需要将生成的latent除以,然后再送入autoencoder的decoder即可。对于SD所使用的autoencoder,这个rescaling系数为0.18215。

CLIP text encoder

SD采用CLIP text encoder来对输入text提取text embeddings,具体的是采用目前OpenAI所开源的最大CLIP模型:clip-vit-large-patch14,这个CLIP的text encoder是一个transformer模型(只有encoder模块):层数为12,特征维度为768,模型参数大小是123M。对于输入text,送入CLIP text encoder后得到最后的hidden states(即最后一个transformer block得到的特征),其特征维度大小为77x768(77是token的数量),这个细粒度的text embeddings将以cross attention的方式送入UNet中。在transofmers库中,可以如下使用CLIP text encoder:

from transformers import CLIPTextModel, CLIPTokenizer

text_encoder = CLIPTextModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="text_encoder").to("cuda")
# text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to("cuda")
tokenizer = CLIPTokenizer.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="tokenizer")
# tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")

# 对输入的text进行tokenize,得到对应的token ids
prompt = "a photograph of an astronaut riding a horse"
text_input_ids = text_tokenizer(
    prompt,
    padding="max_length",
    max_length=tokenizer.model_max_length,
    truncation=True,
    return_tensors="pt"
).input_ids

# 将token ids送入text model得到77x768的特征
text_embeddings = text_encoder(text_input_ids.to("cuda"))[0]

值得注意的是,这里的tokenizer最大长度为77(CLIP训练时所采用的设置),当输入text的tokens数量超过77后,将进行截断,如果不足则进行paddings,这样将保证无论输入任何长度的文本(甚至是空文本)都得到77x768大小的特征。在训练SD的过程中,CLIP text encoder模型是冻结的。在早期的工作中,比如OpenAI的GLIDE和latent diffusion中的LDM均采用一个随机初始化的tranformer模型来提取text的特征,但是最新的工作都是采用预训练好的text model。比如谷歌的Imagen采用纯文本模型T5 encoder来提出文本特征,而SD则采用CLIP text encoder,预训练好的模型往往已经在大规模数据集上进行了训练,它们要比直接采用一个从零训练好的模型要好。

UNet

SD的扩散模型是一个860M的UNet,其主要结构如下图所示(这里以输入的latent为64x64x4维度为例),其中encoder部分包括3个CrossAttnDownBlock2D模块和1个DownBlock2D模块,而decoder部分包括1个UpBlock2D模块和3个CrossAttnUpBlock2D模块,中间还有一个UNetMidBlock2DCrossAttn模块。encoder和decoder两个部分是完全对应的,中间存在skip connection。注意3个CrossAttnDownBlock2D模块最后均有一个2x的downsample操作,而DownBlock2D模块是不包含下采样的。

强大到离谱!硬核解读Stable Diffusion(完整版)

其中CrossAttnDownBlock2D模块的主要结构如下图所示,text condition将通过CrossAttention模块嵌入进来,此时Attention的query是UNet的中间特征,而key和value则是text embeddings。SD和DDPM一样采用预测noise的方法来训练UNet,其训练损失也和DDPM一样:这里的为text embeddings,此时的模型是一个条件扩散模型。基于diffusers库,我们可以很快实现SD的训练,其核心代码如下所示(这里参考diffusers库下examples中的finetune代码):文章来源地址https://www.toymoban.com/news/detail-446631.html

import torch
from diffusers import AutoencoderKL, UNet2DConditionModel, DDPMScheduler
from transformers import CLIPTextModel, CLIPTokenizer
import torch.nn.functional as F

# 加载autoencoder
vae = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")
# 加载text encoder
text_encoder = CLIPTextModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="text_encoder")
tokenizer = CLIPTokenizer.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="tokenizer")
# 初始化UNet
unet = UNet2DConditionModel(**model_config) # model_config为模型参数配置
# 定义scheduler
noise_scheduler = DDPMScheduler(
    beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000
)

# 冻结vae和text_encoder
vae.requires_grad_(False)
text_encoder.requires_grad_(False)

opt = torch.optim.AdamW(unet.parameters(), lr=1e-4)

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.ra

到了这里,关于强大到离谱!硬核解读Stable Diffusion(完整版)的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处: 如若内容造成侵权/违法违规/事实不符,请点击违法举报进行投诉反馈,一经查实,立即删除!

领支付宝红包 赞助服务器费用

相关文章

  • Stable Diffusion云服务器部署完整版教程

    2023年07月04日 22:30 3607浏览 · 18喜欢 · 22评论 薯片_AI 粉丝: 1513 文章: 1 设置分组 取消关注 已关注         文本旨在将stable diffusion部署在云服务器上,利用云服务器的优势让我们更好的体验AI绘图。 本文的教程是作者一步步实践所总结出来的,完整的按照作者的步骤执

    2024年02月06日
    浏览(41)
  • Fooocus:一个简单且功能强大的Stable Diffusion webUI

    Stable Diffusion是一个强大的图像生成AI模型,但它通常需要大量调整和提示工程。Fooocus的目标是改变这种状况。 Fooocus的创始人Lvmin Zhang(也是 ControlNet论文的作者)将这个项目描述为对“Stable Diffusion”和“ Midjourney”设计的重新设计。Fooocus就像是Midjourney的免费离线版本,但是

    2024年02月11日
    浏览(43)
  • 详细解读ControlNet Stable Diffusion代码

    Diffusion models代码解读:入门与实战 前言 :ControlNet的原理部分可以看我之前写的博客《Image-to-Image任务的终结者:详解ControlNet原理与代码》 ,这篇博客主要讲解一下diffusers版本的ControlNet代码,从代码中洞悉原理,加深理解。 目录 原理详解

    2024年02月12日
    浏览(49)
  • 超有趣的stable diffusion干货解读

    Diffusion模型尝试使用一种更加简单的方法生成图像。大家是否记得老式电视机信号不好时屏幕上闪烁的雪花? 这些雪花是随机、无序、混乱的,因而被称为噪声。当电视机信号不好的时候,屏幕上就会出现这些噪声点,信号越差就会出现越多的噪声点,直到最后屏幕完全被随

    2024年02月11日
    浏览(40)
  • 一文解读:Stable Diffusion 3究竟厉害在哪里?

    知乎原文:叫我Alonzo就好了 最近一段时间,正当所有人都在为OpenAI发布Sora狂欢时,Stability AI更是推出了Stable Diffusion 3的技术报告。**这两项技术不约而同都采用了Diffusion Transformer的架构设计,之前我也在我的文章中进行了解读:Diffusion Transformer究竟好在哪里?感兴趣的朋友可

    2024年04月16日
    浏览(44)
  • 代码解读:使用Stable Diffusion完成相似图像生成任务

    Diffusion models代码解读:入门与实战 前言 :作为内容生产重要的一部分,生成相似图像是一项有意义的工作,例如很多内容创作分享平台单纯依赖用户贡献的图片已经不够了,最省力的方法就是利用已有的图片生成相似的图片作为补充。这篇博客详细解读基于Stable Diffusion生成

    2024年04月25日
    浏览(55)
  • 解读Stable Video Diffusion:详细解读视频生成任务中的数据清理技术

    Diffusion Models视频生成-博客汇总 前言 :Stable Video Diffusion已经开源一周多了,技术报告《Stable Video Diffusion: Scaling Latent Video Diffusion Models to Large Datasets》对数据清洗的部分描述非常详细,虽然没有开源源代码,但是博主正在尝试复现其中的操作。这篇博客先梳理一下Stable Video

    2024年02月04日
    浏览(77)
  • 万字长文解读Stable Diffusion的核心插件—ControlNet

    目录 一、介绍 二、使用方法 三、ControlNet结构 1.整体结构 2.ControlLDM 3.Timestep Embedding 4.HintBlock 5.ResBlock 6.SpatialTransformer 7.SD Encoder Block 8.SD Decoder Block 9.ControlNet Encoder Block 10.Stable Diffusion 四、训练 1.准备数据集 2.生成ControlNet模型 3.执行训练 五、其它 1.损失函数 2.随机替换提示

    2023年04月14日
    浏览(45)
  • 一文读懂Stable Diffusion 论文原理+代码超详细解读

    Stable diffusion是一个基于Latent Diffusion Models(LDMs)实现的的文图生成(text-to-image)模型。 2022年8月,游戏设计师Jason Allen凭借AI绘画作品《太空歌剧院(Théâtre D’opéra Spatial)》获得美国科罗拉多州博览会“数字艺术/数码摄影“竞赛单元一等奖,“AI绘画”引发全球热议。得力

    2024年01月19日
    浏览(57)
  • 【Stable Diffusion XL】huggingface diffusers 官方教程解读

    相关链接: GitHub: https://github.com/huggingface/diffusers 官方教程:https://huggingface.co/docs/diffusers/tutorials/tutorial_overview StableDiffuson: https://huggingface.co/blog/stable_diffusion#how-does-stable-diffusion-work Diffusers被设计成一个用户友好和灵活的工具箱,用于构建适合您用例的扩散系统。工具箱的核

    2024年02月06日
    浏览(54)

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

博客赞助

微信扫一扫打赏

请作者喝杯咖啡吧~博客赞助

支付宝扫一扫领取红包,优惠每天领

二维码1

领取红包

二维码2

领红包