SDXL base部分的权重:https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/tree/main
diffusers库中的SDXL代码pipelines:
https://github.com/huggingface/diffusers/tree/main/src/diffusers/pipelines/stable_diffusion_xl
参考:深入浅出完整解析Stable Diffusion XL(SDXL)核心基础知识 - 知乎 (zhihu.com)
Stable Diffusion XL是一个二阶段的级联扩散模型(Latent Diffusion Model),包括Base模型和Refiner模型。其中Base模型的主要工作和Stable Diffusion 1.x-2.x一致,具备文生图(txt2img)、图生图(img2img)、图像inpainting等能力。在Base模型之后,级联了Refiner模型,对Base模型生成的图像Latent特征进行精细化提升,其本质上是在做图生图的工作。
SDXL Base模型由U-Net、VAE以及CLIP Text Encoder(两个)三个模块组成
SDXL Refiner模型同样由U-Net、VAE和CLIP Text Encoder(一个)三个模块
1.VAE
VAE Encoder与VAE Decoder结构图
VAE官方开源权重:https://huggingface.co/stabilityai/sdxl-vae
Stable Diffusion XL VAE模型与之前的Stable Diffusion系列并不兼容。如果在SDXL上使用之前系列的VAE,会生成充满噪声的图片。
Stable Diffusion XL VAE采用FP16精度时会出现数值溢出成NaNs的情况,导致重建的图像是一个黑图,所以必须使用FP32精度进行推理重建。
import cv2
import torch
import numpy as np
from diffusers import AutoencoderKL# 加载SDXL VAE模型: SDXL VAE模型可以通过指定subfolder文件来单独加载。
# SDXL VAE模型权重百度云网盘:关注Rocky的公众号WeThinkIn,后台回复:SDXL模型,即可获得资源链接
VAE = AutoencoderKL.from_pretrained("/本地路径/sdxl-vae")
VAE.to("cuda") # 用OpenCV读取和调整图像大小
raw_image = cv2.imread("test_vae.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")# 压缩图像为Latent特征并重建
with torch.inference_mode():# 使用SDXL 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_sdxl.png", cv2.cvtColor(rec_image, cv2.COLOR_RGB2BGR))
2.Unet
SDXL Base部分的 U-Net的完整结构图
Stable Diffusion XL中的Text Condition信息由两个Text Encoder提供(OpenCLIP ViT-bigG和OpenAI CLIP ViT-L),将两个Text Encoder提取的Token Embedding进行Contact,通过Cross Attention组件嵌入,作为K Matrix和V Matrix。与此同时,图片的Latent Feature作为Q Matrix
3.Text Encoder模型
Stable Diffusion XL分别提取两个Text Encoder的倒数第二层特征,并进行concat操作作为文本条件(Text Conditioning)。其中OpenCLIP ViT-bigG的特征维度为77x1280,而OpenAI CLIP ViT-L/14的特征维度是77x768,所以输入总的特征维度是77x2048(77是最大的token数,2048是SDXL的context dim),再通过Cross Attention模块将文本信息传入Stable Diffusion XL的训练过程与推理过程中。
Stable Diffusion XL与之前的系列相比使用了两个CLIP Text Encoder,分别是OpenCLIP ViT-bigG(694M)和OpenAI CLIP ViT-L/14(123.65M),从而大大增强了Stable Diffusion XL对文本的提取和理解能力,同时提高了输入文本和生成图片的一致性。
SDXL OpenCLIP ViT-bigG的完整结构图
SDXL OpenCLIP ViT-bigG的文本编码过程:
from transformers import CLIPTextModel, CLIPTokenizer# 加载 OpenCLIP ViT-bigG Text Encoder模型和Tokenizer
# SDXL模型权重百度云网盘:关注Rocky的公众号WeThinkIn,后台回复:SDXL模型,即可获得资源链接
text_encoder = CLIPTextModel.from_pretrained("/本地路径/stable-diffusion-xl-base-1.0", subfolder="text_encoder_2").to("cuda")
text_tokenizer = CLIPTokenizer.from_pretrained("/本地路径/stable-diffusion-xl-base-1.0", subfolder="tokenizer_2")# 将输入SDXL模型的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_idsprint("text_token_ids' shape:",text_token_ids.shape)
print("text_token_ids:",text_token_ids)# 将token ids特征输入OpenCLIP ViT-bigG Text Encoder模型中输出77x1280的Text Embeddings特征
text_embeddings = text_encoder(text_token_ids.to("cuda"))[0] # 由于Text Encoder模型输出的是一个元组,所以需要[0]对77x1280的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, 0, 0, 0, 0,0, 0, 0, 0, 0, 0, 0, 0, 0, 0,0, 0, 0, 0, 0, 0, 0, 0, 0, 0,0, 0, 0, 0, 0, 0, 0, 0, 0, 0,0, 0, 0, 0, 0, 0, 0, 0, 0, 0,0, 0, 0, 0, 0, 0, 0, 0, 0, 0,0, 0, 0, 0, 0, 0, 0, 0, 0, 0,0, 0, 0, 0, 0, 0, 0]])
text_embeddings' shape: torch.Size([1, 77, 1280])
tensor([[[-0.1025, -0.3104, 0.1660, ..., -0.1596, -0.0680, -0.0180],[ 0.7724, 0.3004, 0.5225, ..., 0.4482, 0.8743, -1.0429],[-0.3963, 0.0041, -0.3626, ..., 0.1841, 0.2224, -1.9317],...,[-0.8887, -0.2579, 1.3508, ..., -0.4421, 0.2193, 1.2736],[-0.9659, -0.0447, 1.4424, ..., -0.4350, -0.1186, 1.2042],[-0.5213, -0.0255, 1.8161, ..., -0.7231, -0.3752, 1.0876]]],device='cuda:0', grad_fn=<NativeLayerNormBackward0>)
SDXL OpenAI CLIP ViT-L/14的完整结构图
SDXL OpenAI CLIP ViT-L/14的文本编码过程:
from transformers import CLIPTextModel, CLIPTokenizer# 加载 OpenAI CLIP ViT-L/14 Text Encoder模型和Tokenizer
# SDXL模型权重百度云网盘:关注Rocky的公众号WeThinkIn,后台回复:SDXL模型,即可获得资源链接
text_encoder = CLIPTextModel.from_pretrained("/本地路径/stable-diffusion-xl-base-1.0", subfolder="text_encoder").to("cuda")
text_tokenizer = CLIPTokenizer.from_pretrained("/本地路径/stable-diffusion-xl-base-1.0", subfolder="tokenizer")# 将输入SDXL模型的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_idsprint("text_token_ids' shape:",text_token_ids.shape)
print("text_token_ids:",text_token_ids)# 将token ids特征输入OpenAI CLIP ViT-L/14 Text Encoder模型中输出77x768的Text Embeddings特征
text_embeddings = text_encoder(text_token_ids.to("cuda"))[0] # 由于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.3885, 0.0230, -0.0521, ..., -0.4901, -0.3065, 0.0674],[-0.8424, -1.1387, 1.2767, ..., -0.2598, 1.6289, -0.7855],[ 0.1751, -0.9847, 0.1881, ..., 0.0657, -1.4940, -1.2612],...,[ 0.2039, -0.7298, -0.3206, ..., 0.6751, -0.5814, -0.7320],[ 0.1921, -0.7345, -0.3039, ..., 0.6806, -0.5852, -0.7228],[ 0.2112, -0.6438, -0.3042, ..., 0.6628, -0.5576, -0.7583]]],device='cuda:0', grad_fn=<NativeLayerNormBackward0>)
以上都为SDXL的base模型
4.Refiner模型
由于已经有U-Net(Base)模型生成了图像的Latent特征,所以Refiner模型的主要工作是在Latent特征进行小噪声去除和细节质量提升。
Refiner模型和Base模型一样是基于Latent的扩散模型,也采用了Encoder-Decoder结构,和U-Net兼容同一个VAE模型。不过在Text Encoder部分,Refiner模型只使用了OpenCLIP ViT-bigG的Text Encoder,同样提取了倒数第二层特征以及进行了pooled text embedding的嵌入。
refine模型中的Unet结构:
单独使用Stable Diffusion XL中的Base模型来生成图像:
# 加载diffusers和torch依赖库
from diffusers import DiffusionPipeline
import torch# 加载Stable Diffusion XL Base模型(stable-diffusion-xl-base-1.0或stable-diffusion-xl-base-0.9)
pipe = DiffusionPipeline.from_pretrained("/本地路径/stable-diffusion-xl-base-1.0",torch_dtype=torch.float16, variant="fp16")
# "/本地路径/stable-diffusion-xl-base-1.0"表示我们需要加载的Stable Diffusion XL Base模型路径
# 大家可以关注Rocky的公众号WeThinkIn,后台回复:SDXL模型,即可获得SDXL模型权重资源链接
# "fp16"代表启动fp16精度。比起fp32,fp16可以使模型显存占用减半# 使用GPU进行Pipeline的推理
pipe.to("cuda")# 输入提示词
prompt = "Watercolor painting of a desert landscape, with sand dunes, mountains, and a blazing sun, soft and delicate brushstrokes, warm and vibrant colors"# 输入负向提示词,表示我们不想要生成的特征
negative_prompt = "(EasyNegative),(watermark), (signature), (sketch by bad-artist), (signature), (worst quality), (low quality), (bad anatomy), NSFW, nude, (normal quality)"# 设置seed,可以固定生成图像中的构图
seed = torch.Generator("cuda").manual_seed(42)# SDXL Base Pipeline进行推理
image = pipe(prompt, negative_prompt=negative_prompt,generator=seed).images[0]
# Pipeline生成的images包含在一个list中:[<PIL.Image.Image image mode=RGB size=1024x1024>]
#所以需要使用images[0]来获取list中的PIL图像# 保存生成图像
image.save("SDXL-Base.png")
将SDXL Base模型和SDXL Refiner模型级联来生成图像:
from diffusers import DiffusionPipeline
import torch# 下面的五行代码不变
pipe = DiffusionPipeline.from_pretrained("/本地路径/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16")pipe.to("cuda")prompt = "Watercolor painting of a desert landscape, with sand dunes, mountains, and a blazing sun, soft and delicate brushstrokes, warm and vibrant colors"negative_prompt = "(EasyNegative),(watermark), (signature), (sketch by bad-artist), (signature), (worst quality), (low quality), (bad anatomy), NSFW, nude, (normal quality)"seed = torch.Generator("cuda").manual_seed(42)# 运行SDXL Base模型的Pipeline,设置输出格式为output_type="latent"
image = pipe(prompt=prompt, negative_prompt=negative_prompt, generator=seed, output_type="latent").images# 加载Stable Diffusion XL Refiner模型(stable-diffusion-xl-refiner-1.0或stable-diffusion-xl-refiner-0.9)
pipe = DiffusionPipeline.from_pretrained("/本地路径/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, variant="fp16")
# "本地路径/stable-diffusion-xl-refiner-1.0"表示我们需要加载的Stable Diffusion XL Refiner模型,
# 大家可以关注Rocky的公众号WeThinkIn,后台回复:SDXL模型,即可获得SDXL模型权重资源链接pipe.to("cuda")# SDXL Refiner Pipeline进行推理
images = pipe(prompt=prompt, negative_prompt=negative_prompt, generator=seed, image=image).images# 保存生成图像
images[0].save("SDXL-Base-Refiner.png")