网站式小程序,哪个网站可以做代码题目,百度站长反馈中心,做盗版电影网站赚钱AIGC系列博文#xff1a; 【AIGC系列】1#xff1a;自编码器#xff08;AutoEncoder, AE#xff09; 【AIGC系列】2#xff1a;DALLE 2模型介绍#xff08;内含扩散模型介绍#xff09; 【AIGC系列】3#xff1a;Stable Diffusion模型原理介绍 【AIGC系列】4#xff1… AIGC系列博文 【AIGC系列】1自编码器AutoEncoder, AE 【AIGC系列】2DALL·E 2模型介绍内含扩散模型介绍 【AIGC系列】3Stable Diffusion模型原理介绍 【AIGC系列】4Stable Diffusion应用实践和代码分析 【AIGC系列】5视频生成模型数据处理和预训练流程介绍Sora、MovieGen、HunyuanVideo 目录 1 AutoEncoder2 CLIP text encoder3 UNet4 应用4.1 文生图4.2 图生图4.3 图像inpainting 5 其他 上一篇博文我们学习了Stable Diffusion的原理这一篇我们继续深入了解Stable Diffusion的应用实践和代码分析。
1 AutoEncoder
SD采用基于KL-reg的autoencoder当输入图像为512x512时将得到64x64x4大小的latent。autoencoder模型是在OpenImages数据集上基于256x256大小训练的但是由于模型是全卷积结构的基于ResnetBlock所以可以扩展应用在尺寸256的图像上。
下面我们使用diffusers库来加载autoencoder模型实现图像的压缩和重建代码如下
import torch
from diffusers import AutoencoderKL
import numpy as np
from PIL import Imageprint(torch.cuda.is_available())#加载模型: autoencoder可以通过SD权重指定subfolder来单独加载
print(Start...)
autoencoder AutoencoderKL.from_pretrained(runwayml/stable-diffusion-v1-5, subfoldervae)
autoencoder.to(cuda, dtypetorch.float16)
print(Get weight successfully)# 读取图像并预处理
# raw_image Image.open(liuyifei.jpg).convert(RGB).resize((256, 256))
raw_image Image.open(liuyifei.jpg).convert(RGB)
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, dtypetorch.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.save(liuyifei_re.jpg)重建效果如下所示对比手表上的文字可以看出autoencoder将图片压缩到latent后再重建其实是有损的。 为了改善这种畸变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/) 2 CLIP text encoder
SD采用CLIP text encoder来对输入的文本生成text embeddings采用的CLIP模型是clip-vit-large-patch14该模型的text encoder层数为12特征维度为768模型参数大小是123M。文本输入text encoder后得到最后的hidden states特征维度大小为77x76877是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, subfoldertext_encoder).to(cuda)
# text_encoder CLIPTextModel.from_pretrained(openai/clip-vit-large-patch14).to(cuda)
tokenizer CLIPTokenizer.from_pretrained(runwayml/stable-diffusion-v1-5, subfoldertokenizer)
# 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 tokenizer( prompt, paddingmax_length, max_lengthtokenizer.model_max_length, truncationTrue, return_tensorspt
).input_idsprint(f \n\n text_input_ids: {text_input_ids} \n\n)# 将token ids送入text model得到77x768的特征
text_embeddings text_encoder(text_input_ids.to(cuda))[0]
print(f \n\n text_embeddings: {text_embeddings} \n\n)
输出如下 text_input_ids: tensor([[49406, 320, 8853, 539, 550, 18376, 6765, 320, 4558, 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: tensor([[[-0.3884, 0.0229, -0.0522, ..., -0.4899, -0.3066, 0.0675],[ 0.0290, -1.3258, 0.3085, ..., -0.5257, 0.9768, 0.6652],[ 0.4595, 0.5617, 1.6663, ..., -1.9515, -1.2307, 0.0104],...,[-3.0421, -0.0656, -0.1793, ..., 0.3943, -0.0190, 0.7664],[-3.0551, -0.1036, -0.1936, ..., 0.4236, -0.0189, 0.7575],[-2.9854, -0.0832, -0.1715, ..., 0.4355, 0.0095, 0.7485]]],devicecuda:0, grad_fnNativeLayerNormBackward0)
值得注意的是这里的tokenizer最大长度为77CLIP训练时所采用的设置当输入text的tokens数量超过77后将进行截断如果不足则进行paddings这样将保证无论输入任何长度的文本甚至是空文本都得到77x768大小的特征。在上面的例子里输入的tokens数量少于77所以后面都padding了id为49407的token。
在训练SD的过程中CLIP text encoder模型是冻结的。在早期的工作中比如OpenAI的GLIDE和latent diffusion中的LDM均采用一个随机初始化的tranformer模型来提取text的特征但是最新的工作都是采用预训练好的text model。比如谷歌的Imagen采用纯文本模型T5 encoder来提出文本特征而SD则采用CLIP text encoder预训练好的模型往往已经在大规模数据集上进行了训练它们要比直接采用一个从零训练好的模型要好。
3 UNet
SD的扩散模型是一个860M的UNet其主要结构如下图所示其中encoder部分包括3个CrossAttnDownBlock2D模块和1个DownBlock2D模块而decoder部分包括1个UpBlock2D模块和3个CrossAttnUpBlock2D模块中间还有一个UNetMidBlock2DCrossAttn模块。
encoder和decoder两个部分是完全对应的中间有skip connection。3个CrossAttnDownBlock2D模块最后均有一个2x的downsample操作而DownBlock2D模块是不包含下采样的。 其中CrossAttnDownBlock2D模块的主要结构如下图所示text condition将通过CrossAttention模块嵌入进来此时Attention的query是UNet的中间特征而key和value则是text embeddings。 SD和DDPM一样采用预测noise的方法来训练UNet其训练损失也和DDPM一样。基于diffusers库我们可以实现SD的训练其核心代码如下
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
from diffusers import AutoencoderKL, UNet2DConditionModel, DDPMScheduler
from transformers import CLIPTextModel, CLIPTokenizer
import torch.nn.functional as F# 自定义Dataset类
class CustomImageTextDataset(Dataset):def __init__(self, image_paths, text_descriptions, transformNone):self.image_paths image_pathsself.text_descriptions text_descriptionsself.transform transformdef __len__(self):return len(self.image_paths)def __getitem__(self, idx):image_path self.image_paths[idx]text_description self.text_descriptions[idx]# 加载图像image Image.open(image_path).convert(RGB)if self.transform:image self.transform(image)return {image: image,text: text_description}# 数据准备
image_paths [path/to/image1.jpg, path/to/image2.jpg] # 替换为实际的图像路径
text_descriptions [description for image1, description for image2] # 替换为实际的文本描述# 图像转换预处理
transform transforms.Compose([transforms.Resize((256, 256)), # 调整大小transforms.ToTensor(), # 转换为张量
])# 创建数据集实例
dataset CustomImageTextDataset(image_pathsimage_paths, text_descriptionstext_descriptions, transformtransform)# 创建DataLoader
train_dataloader DataLoader(dataset, batch_size4, shuffleTrue, num_workers0)# 加载autoencoder
vae AutoencoderKL.from_pretrained(runwayml/stable-diffusion-v1-5, subfoldervae)
# 加载text encoder
text_encoder CLIPTextModel.from_pretrained(runwayml/stable-diffusion-v1-5, subfoldertext_encoder)
tokenizer CLIPTokenizer.from_pretrained(runwayml/stable-diffusion-v1-5, subfoldertokenizer)model_config {sample_size: 32,in_channels: 4,out_channels: 4,down_block_types: (DownBlock2D, CrossAttnDownBlock2D, CrossAttnDownBlock2D, CrossAttnDownBlock2D),up_block_types: (UpBlock2D, CrossAttnUpBlock2D, CrossAttnUpBlock2D, CrossAttnUpBlock2D),block_out_channels: (320, 640, 1280, 1280),layers_per_block: 2,cross_attention_dim: 768,attention_head_dim: 8,
}
# 初始化UNet
unet UNet2DConditionModel(**model_config)# 定义scheduler
noise_scheduler DDPMScheduler(beta_start0.00085,beta_end0.012,beta_schedulescaled_linear,num_train_timesteps1000
)# 冻结vae和text_encoder
vae.requires_grad_(False)
text_encoder.requires_grad_(False)opt torch.optim.AdamW(unet.parameters(), lr1e-4)# 训练循环
device torch.device(cuda if torch.cuda.is_available() else cpu)
unet.to(device)
vae.to(device)
text_encoder.to(device)for epoch in range(10): # 假设训练10个epochunet.train()for step, batch in enumerate(train_dataloader):with torch.no_grad():# 将image转到latent空间latents vae.encode(batch[image].to(device)).latent_dist.sample()# rescaling latentslatents latents * vae.config.scaling_factor# 提取text embeddingstext_input_ids tokenizer(batch[text],paddingmax_length,max_lengthtokenizer.model_max_length,truncationTrue,return_tensorspt).input_ids.to(device)text_embeddings text_encoder(text_input_ids)[0]# 随机采样噪音noise torch.randn_like(latents)bsz latents.shape[0]# 随机采样timesteptimesteps torch.randint(0, noise_scheduler.num_train_timesteps, (bsz,), devicedevice).long()# 将noise添加到latent上即扩散过程noisy_latents noise_scheduler.add_noise(latents, noise, timesteps)# 预测noise并计算lossmodel_pred unet(noisy_latents, timesteps, encoder_hidden_statestext_embeddings).sampleloss F.mse_loss(model_pred.float(), noise.float(), reductionmean)opt.zero_grad()loss.backward()opt.step()if step % 10 0:print(fEpoch {epoch}, Step {step}, Loss: {loss.item()})# 在训练完成后保存模型
model_save_path path/to/your/unet_model.pth
torch.save(unet.state_dict(), model_save_path)
print(fModel has been saved to {model_save_path})optimizer_save_path path/to/your/optimizer.pth
torch.save(opt.state_dict(), optimizer_save_path)
print(fOptimizer state has been saved to {optimizer_save_path})# 加载模型进行推理或继续训练
unet_load_path path/to/your/unet_model.pth
unet_loaded UNet2DConditionModel(**model_config) # 创建一个与原模型结构相同的实例
unet_loaded.load_state_dict(torch.load(unet_load_path))
unet_loaded.to(device)
unet_loaded.eval() # 设置为评估模式# 恢复优化器的状态以继续训练
opt_load_path path/to/your/optimizer.pth
opt_loaded torch.optim.AdamW(unet_loaded.parameters(), lr1e-4) # 创建一个新的优化器实例
opt_loaded.load_state_dict(torch.load(opt_load_path))# 使用unet_loaded进行推理或者用opt_loaded继续训练。注意的是SD的noise scheduler虽然也是采用一个1000步长的scheduler但是不是linear的而是scaled linear具体的计算如下所示
betas torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtypetorch.float32) ** 2 在训练条件扩散模型时往往会采用Classifier-Free Guidance CFG即在训练条件扩散模型的同时也训练一个无条件的扩散模型同时在采样阶段将条件控制下预测的噪音和无条件下的预测噪音组合在一起来确定最终的噪音CFG对于提升条件扩散模型的图像生成效果是至关重要的。
4 应用
4.1 文生图
根据文本生成图像是文生图的最核心的功能SD的文生图的推理流程图首先根据输入text用text encoder提取text embeddings同时初始化一个随机噪音noiselatent上的512x512图像对应的noise维度为64x64x4然后将text embeddings和noise送入扩散模型UNet中生成去噪后的latent最后送入autoencoder的decoder模块得到生成的图像。
使用diffusers库我们可以直接调用StableDiffusionPipeline来实现文生图具体代码如下所示
import torch
from diffusers import StableDiffusionPipeline
from PIL import Image # 组合图像生成grid
def image_grid(imgs, rows, cols): assert len(imgs) rows*cols w, h imgs[0].size grid Image.new(RGB, size(cols*w, rows*h)) grid_w, grid_h grid.size for i, img in enumerate(imgs): grid.paste(img, box(i%cols*w, i//cols*h)) return grid # 加载文生图pipeline
pipe StableDiffusionPipeline.from_pretrained( runwayml/stable-diffusion-v1-5, # 或者使用 SD v1.4: CompVis/stable-diffusion-v1-4 torch_dtypetorch.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, height512, width512, num_inference_steps50, guidance_scale7.5, negative_promptNone, num_images_per_prompt1, generatorgenerator
).images # 保存每个单独的图片
for idx, img in enumerate(images):img.save(fimage_{idx}.png)# 创建并保存组合后的网格图
grid image_grid(images, rows1, colslen(prompts))
grid.save(combined_images.png)
print(所有图片已保存到本地。)生成的结果如下 重要参数说明 指定width和height来决定生成图像的大小前面说过SD最后是在512x512尺度上训练的所以生成512x512尺寸效果是最好的但是实际上SD可以生成任意尺寸的图片一方面autoencoder支持任意尺寸的图片的编码和解码另外一方面扩散模型UNet也是支持任意尺寸的latents生成的UNet是卷积attention的混合结构。但是生成512x512以外的图片会存在一些问题比如生成低分辨率图像时图像的质量大幅度下降等等。 num_inference_steps指推理过程中的去噪步数或者采样步数。SD在训练过程采用的是步数为1000的noise scheduler但是在推理时往往采用速度更快的scheduler只需要少量的采样步数就能生成不错的图像比如SD默认采用PNDM scheduler它只需要采样50步就可以出图。当然我们也可以换用其它类型的scheduler比如DDIM scheduler和DPM-Solver scheduler。我们可以在diffusers中直接替换scheduler比如我们想使用DDIM
from diffusers import DDIMScheduler # 注意这里的clip_sample要关闭否则生成图像存在问题因为不能对latent进行clip
pipe.scheduler DDIMScheduler.from_config(pipe.scheduler.config, clip_sampleFalse) guidance_scale当CFG的guidance_scale越大时生成的图像应该会和输入文本更一致。SD默认采用的guidance_scale为7.5。但是过大的guidance_scale也会出现问题主要是由于训练和测试的不一致过大的guidance_scale会导致生成的样本超出范围。 negative_prompt这个参数和CFG有关去噪过程的噪音预测不仅仅依赖条件扩散模型也依赖无条件扩散模型这里的negative_prompt便是无条件扩散模型的text输入前面说过训练过程中我们将text置为空字符串来实现无条件扩散模型所以这里negative_prompt None 。但是有时候我们可以使用不为空的negative_prompt来避免模型生成的图像包含不想要的东西因为从上述公式可以看到这里的无条件扩散模型是我们想远离的部分。
4.2 图生图
图生图image2image是对文生图功能的一个扩展这个功能来源于SDEdit这个工作其核心思路也非常简单给定一个笔画的色块图像可以先给它加一定的高斯噪音执行扩散过程得到噪音图像然后基于扩散模型对这个噪音图像进行去噪就可以生成新的图像但是这个图像在结构和布局和输入图像基本一致。
相比文生图流程来说这里的初始latent不再是一个随机噪音而是由初始图像经过autoencoder编码之后的latent加高斯噪音得到这里的加噪过程就是扩散过程。要注意的是去噪过程的步数要和加噪过程的步数一致就是说你加了多少噪音就应该去掉多少噪音这样才能生成想要的无噪音图像。
在diffusers中我们可以使用StableDiffusionImg2ImgPipeline来实现文生图具体代码如下所示
import torch
from diffusers import StableDiffusionImg2ImgPipeline
from PIL import Image# 加载图生图pipeline
model_id runwayml/stable-diffusion-v1-5
pipe StableDiffusionImg2ImgPipeline.from_pretrained(model_id, torch_dtypetorch.float16).to(cuda)# 读取初始图片
init_image Image.open(liuyifei.jpg).convert(RGB).resize((512, 512))
print(init_image.size)
init_image.save(liuyifei_512.jpg)# 推理
prompt A girl wearing a hat on her head.
generator torch.Generator(devicecuda).manual_seed(2023)image pipe(promptprompt,imageinit_image,strength0.8,guidance_scale7.5,generatorgenerator
).images[0]# 保存生成的图像
output_path generated_liuyifei.jpg
image.save(output_path)
print(fGenerated image saved to {output_path})原始图片 效果如下 相比文生图的pipeline图生图的pipeline还多了一个参数strength这个参数介于0-1之间表示对输入图片加噪音的程度这个值越大加的噪音越多对原始图片的破坏也就越大当strength1时其实就变成了一个随机噪音此时就相当于纯粹的文生图pipeline了。
4.3 图像inpainting
图像inpainting和图生图一样也是文生图功能的一个扩展。SD的图像inpainting不是用在图像修复上而是主要用在图像编辑上给定一个输入图像和想要编辑的区域mask我们想通过文生图来编辑mask区域的内容。
它和图生图一样首先将输入图像通过autoencoder编码为latent然后加入一定的高斯噪音生成noisy latent再进行去噪生成图像但是这里为了保证mask以外的区域不发生变化在去噪过程的每一步都将扩散模型预测的noisy latent用真实图像同level的nosiy latent替换。
在diffusers中使用StableDiffusionInpaintPipelineLegacy可以实现文本引导下的图像inpainting具体代码如下所示
import torch
from diffusers import StableDiffusionInpaintPipelineLegacy
from PIL import Image # 加载inpainting pipeline
model_id runwayml/stable-diffusion-v1-5
pipe StableDiffusionInpaintPipelineLegacy.from_pretrained(model_id, torch_dtypetorch.float16).to(cuda) # 读取输入图像和输入mask
input_image Image.open(overture-creations-5sI6fQgYIuo.png).resize((512, 512))
input_mask Image.open(overture-creations-5sI6fQgYIuo_mask.png).resize((512, 512)) # 执行推理
prompt [a mecha robot sitting on a bench, a cat sitting on a bench]
generator torch.Generator(cuda).manual_seed(0) with torch.autocast(cuda): images pipe( promptprompt, imageinput_image, mask_imageinput_mask, num_inference_steps50, strength0.75, guidance_scale7.5, num_images_per_prompt1, generatorgenerator, ).images # 保存每个单独的图片
for idx, img in enumerate(images):img.save(fimage_{idx}.png)print(所有图片已保存到本地。)5 其他
Colab上开源的Stable Diffusion 2.1 GUIstable_diffusion_2_0.ipynb。
最强大且模块化的具有图形/节点界面的稳定扩散GUIComfyUI。
Huggingface模型库https://huggingface.co/stabilityai。
Huggingface的Diffuser库https://github.com/huggingface/diffusers。