图像inpainting
最后我们要介绍的一项功能是图像inpainting，它和图生图一样也是文生图功能的一个扩展。SD的图像inpainting不是用在图像修复上，而是主要用在图像编辑 上：给定一个输入图像和想要编辑的区域mask，我们想通过文生图来编辑mask区域的内容。SD的图像inpainting原理可以参考论文Blended Latent Diffusion，其主要原理图如下所示：
它和图生图一样：首先将输入图像通过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_dtype=torch.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( prompt=prompt, image=input_image, mask_image=input_mask, num_inference_steps=50, strength=0.75, guidance_scale=7.5, num_images_per_prompt=1, generator=generator, ).images
下面是一个具体的生成效果，这里我们将输入图像的dog换成了mecha robot或者cat，从而实现了图像编辑。
要注意的是这里的参数guidance_scale也和图生图一样比较重要，要生成好的图像，需要选择合适的guidance_scale。如果guidance_scale=0.5时，生成的图像由于过于受到原图干扰而产生一些不协调，如下所示：
合适的prompt也比较重要，比如如果我们去掉prompt中的"sitting on a bench"，那么编辑的图像效果也会出现不协调：
无论是上面的图生图还是这里的图像inpainting，我们其实并没有去finetune SD模型，只是扩展了它的能力，但是这两样功能就需要精确调整参数才能得到满意的生成效果。这里，我们也给出StableDiffusionInpaintPipelineLegacy这个pipeline内部的核心代码：
import PIL import numpy as np import torch from diffusers import AutoencoderKL, UNet2DConditionModel, DDIMScheduler from transformers import CLIPTextModel, CLIPTokenizer from tqdm.auto import tqdm def preprocess_mask(mask): mask = mask.convert("L") w, h = mask.size w, h = map(lambda x: x - x % 32, (w, h)) # resize to integer multiple of 32 mask = mask.resize((w // 8, h // 8), resample=PIL.Image.NEAREST) mask = np.array(mask).astype(np.float32) / 255.0 mask = np.tile(mask, (4, 1, 1)) mask = mask[None].transpose(0, 1, 2, 3) # what does this step do? mask = 1 - mask # repaint white, keep black mask = torch.from_numpy(mask) return mask def preprocess(image): w, h = image.size w, h = map(lambda x: x - x % 32, (w, h)) # resize to integer multiple of 32 image = image.resize((w, h), resample=PIL.Image.LANCZOS) image = np.array(image).astype(np.float32) / 255.0 image = image[None].transpose(0, 3, 1, 2) image = torch.from_numpy(image) return 2.0 * image - 1.0 model_id = "runwayml/stable-diffusion-v1-5" # 1. 加载autoencoder vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae") # 2. 加载tokenizer和text encoder tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer") text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder") # 3. 加载扩散模型UNet unet = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet") # 4. 定义noise scheduler noise_scheduler = DDIMScheduler( num_train_timesteps=1000, beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, # don't clip sample, the x0 in stable diffusion not in range [-1, 1] set_alpha_to_one=False, ) # 将模型复制到GPU上 device = "cuda" vae.to(device, dtype=torch.float16) text_encoder.to(device, dtype=torch.float16) unet = unet.to(device, dtype=torch.float16) prompt = "a mecha robot sitting on a bench" strength = 0.75 guidance_scale = 7.5 batch_size = 1 num_inference_steps = 50 negative_prompt = "" generator = torch.Generator(device).manual_seed(0) with torch.no_grad(): # 获取prompt的text_embeddings text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt") text_embeddings = text_encoder(text_input.input_ids.to(device))[0] # 获取unconditional text embeddings max_length = text_input.input_ids.shape[-1] uncond_input = tokenizer( [negative_prompt] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt" ) uncond_embeddings = text_encoder(uncond_input.input_ids.to(device))[0] # 拼接batch text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) # 设置采样步数 noise_scheduler.set_timesteps(num_inference_steps, device=device) # 根据strength计算timesteps init_timestep = min(int(num_inference_steps * strength), num_inference_steps) t_start = max(num_inference_steps - init_timestep, 0) timesteps = noise_scheduler.timesteps[t_start:] # 预处理init_image init_input = preprocess(input_image) init_latents = vae.encode(init_input.to(device, dtype=torch.float16)).latent_dist.sample(generator) init_latents = 0.18215 * init_latents init_latents = torch.cat([init_latents] * batch_size, dim=0) init_latents_orig = init_latents # 处理mask mask_image = preprocess_mask(input_mask) mask_image = mask_image.to(device=device, dtype=init_latents.dtype) mask = torch.cat([mask_image] * batch_size) # 给init_latents加噪音 noise = torch.randn(init_latents.shape, generator=generator, device=device, dtype=init_latents.dtype) init_latents = noise_scheduler.add_noise(init_latents, noise, timesteps[:1]) latents = init_latents # 作为初始latents # Do denoise steps for t in tqdm(timesteps): # 这里latens扩展2份，是为了同时计算unconditional prediction latent_model_input = torch.cat([latents] * 2) latent_model_input = noise_scheduler.scale_model_input(latent_model_input, t) # for DDIM, do nothing # 预测噪音 noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample # CFG noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # 计算上一步的noisy latents：x_t -> x_t-1 latents = noise_scheduler.step(noise_pred, t, latents).prev_sample # 将unmask区域替换原始图像的nosiy latents init_latents_proper = noise_scheduler.add_noise(init_latents_orig, noise, torch.tensor([t])) latents = (init_latents_proper * mask) + (latents * (1 - mask)) # 注意要对latents进行scale latents = 1 / 0.18215 * latents image = vae.decode(latents).sample
另外，runwayml在发布SD 1.5版本的同时还发布了一个inpainting模型：runwayml/stable-diffusion-inpainting，与前面所讲不同的是，这是一个在SD 1.2上finetune的模型 。原来SD的UNet的输入是64x64x4，为了实现inpainting，现在给UNet的第一个卷机层增加5个channels，分别为masked图像的latents（经过autoencoder编码，64x64x4）和mask图像（直接下采样8x，64x64x1），增加的权重填零初始化。在diffusers中，可以使用StableDiffusionInpaintPipeline来调用这个模型，具体代码如下：
import torch from diffusers import StableDiffusionInpaintPipeline from PIL import Image from tqdm.auto import tqdm import PIL # Load pipeline model_id = "runwayml/stable-diffusion-inpainting/" pipe = StableDiffusionInpaintPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda") prompt = ["a mecha robot sitting on a bench", "a dog sitting on a bench", "a bench"] generator = torch.Generator("cuda").manual_seed(2023) input_image = Image.open("overture-creations-5sI6fQgYIuo.png").resize((512, 512)) input_mask = Image.open("overture-creations-5sI6fQgYIuo_mask.png").resize((512, 512)) images = pipe( prompt=prompt, image=input_image, mask_image=input_mask, num_inference_steps=50, generator=generator, ).images
其生成的效果图如下所示：
经过finetune的inpainting在生成细节上可能会更好，但是有可能会丧失部分文生图的能力，而且也比较难迁移其它finetune的SD模型。
SD 2.0
SD 2.0
Stability AI公司在2022年11月（stable-diffusion-v2-release）放出了SD 2.0版本，这里我们也简单介绍一下相比SD 1.x版本SD 2.0的具体改进点。SD 2.0相比SD 1.x版本的主要变动在于模型结构 和训练数据 两个部分。
首先是模型结构方面，SD 1.x版本的text encoder采用的是OpenAI的CLIP ViT-L/14模型，其模型参数量为123.65M；而SD 2.0采用了更大的text encoder：基于OpenCLIP在laion-2b数据集上训练的CLIP ViT-H/14模型，其参数量为354.03M，相比原来的text encoder模型大了约3倍。两个CLIP模型的对比如下所示：
可以看到CLIP ViT-H/14模型相比原来的OpenAI的L/14模型，在imagenet1K上分类准确率和mscoco多模态检索任务上均有明显的提升，这也意味着对应的text encoder更强，能够抓住更准确的文本语义信息。另外是一个小细节是SD 2.0提取的是text encoder倒数第二层的特征，而SD 1.x提取的是倒数第一层的特征。由于倒数第一层的特征之后就是CLIP的对比学习任务，所以倒数第一层的特征可能部分丢失细粒度语义信息，Imagen论文（见论文D.1部分）和novelai（见novelai blog）均采用了倒数第二层特征。对于UNet模型，SD 2.0相比SD 1.x几乎没有改变，唯一的一个小的变动是：SD 2.0不同stage的attention模块是固定attention head dim为64，而SD 1.0则是不同stage的attention模块采用固定attention head数量，明显SD 2.0的这种设定更常用，但是这个变动不会影响模型参数。然后是训练数据，前面说过SD 1.x版本其实最后主要采用laion-2B中美学评分为5以上的子集来训练，而SD 2.0版本采用评分在4.5以上的子集，相当于扩大了训练数据集，具体的训练细节见model card。另外SD 2.0除了512x512版本的模型，还包括768x768版本的模型（https://huggingface.co/stabilityai/stable-diffusion-2），所谓的768x768模型是在512x512模型基础上用图像分辨率大于768x768的子集继续训练的，不过优化目标不再是noise_prediction，而是采用Progressive Distillation for Fast Sampling of Diffusion Models论文中所提出的 v-objective。下图为SD 2.0和SD 1.x版本在COCO2017验证集上评测的对比，可以看到2.0相比1.5，CLIP score有一个明显的提升，同时FID也有一定的提升。但是正如前面所讨论的，FID和CLIP score这两个指标均有一定的局限性，所以具体效果还是上手使用来对比。
Stability AI在发布SD 2.0的同时，还发布了另外3个模型：stable-diffusion-x4-upscaler ，stable-diffusion-2-inpainting 和stable-diffusion-2-depth 。stable-diffusion-x4-upscaler是一个基于扩散模型的4x超分模型，它也是基于latent diffusion，不过这里采用的autoencoder是基于VQ-reg的，下采样率为。在实现上，它是将低分辨率图像直接和noisy latent拼接在一起送入UNet，因为autoencoder将高分辨率图像压缩为原来的1/4，而低分辨率图像也为高分辨率图像的1/4，所以低分辨率图像的空间维度和latent是一致的。另外，这个超分模型也采用了Cascaded Diffusion Models for High Fidelity Image Generation所提出的noise conditioning augmentation，简单来说就是在训练过程中给低分辨率图像加上高斯噪音，可以通过扩散过程来实现，注意这里的扩散过程的scheduler与主扩散模型的scheduler可以不一样，同时也将对应的noise_level（对应扩散模型的time step）通过class labels的方式送入UNet，让UNet知道加入噪音的程度。stable-diffusion-x4-upscaler是使用LAION中>2048x2048大小的子集（10M）训练的，训练过程中采用512x512的crops来训练（降低显存消耗）。SD模型可以用来生成512x512图像，加上这个超分模型，就可以得到2048x2048大小的图像。
在diffusers库中，可以如下使用这个超分模型（这里的noise level是指推理时对低分辨率图像加入噪音的程度）：
import requests from PIL import Image from io import BytesIO from diffusers import StableDiffusionUpscalePipeline import torch # load model and scheduler model_id = "stabilityai/stable-diffusion-x4-upscaler" pipeline = StableDiffusionUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16) pipeline = pipeline.to("cuda") # let's download an image url = "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale/low_res_cat.png" response = requests.get(url) low_res_img = Image.open(BytesIO(response.content)).convert("RGB") low_res_img = low_res_img.resize((128, 128)) prompt = "a white cat" upscaled_image = pipeline(prompt=prompt, image=low_res_img, noise_level=20).images[0] upscaled_image.save("upsampled_cat.png")
stable-diffusion-2-inpainting是图像inpainting模型，和前面所说的runwayml/stable-diffusion-inpainting基本一样，不过它是在SD 2.0的512x512版本上finetune的。
stable-diffusion-2-depth 是也是在SD 2.0的512x512版本上finetune的模型，它是额外增加了图像的深度图作为condition，这里是直接将深度图下采样8x，然后和nosiy latent拼接在一起送入UNet模型中。深度图可以作为一种结构控制，下图展示了加入深度图后生成的图像效果：
你可以调用diffusers库中的StableDiffusionDepth2ImgPipeline来实现基于深度图控制的文生图：
import torch import requests from PIL import Image from diffusers import StableDiffusionDepth2ImgPipeline pipe = StableDiffusionDepth2ImgPipeline.from_pretrained( "stabilityai/stable-diffusion-2-depth", torch_dtype=torch.float16, ).to("cuda") url = "http://images.cocodataset.org/val2017/000000039769.jpg" init_image = Image.open(requests.get(url, stream=True).raw) prompt = "two tigers" n_propmt = "bad, deformed, ugly, bad anotomy" image = pipe(prompt=prompt, image=init_image, negative_prompt=n_propmt, strength=0.7).images[0]
除此之外，Stability AI公司还开源了两个加强版的autoencoder：ft-EMA和ft-MSE（前者使用L1 loss后者使用MSE loss），前面已经说过，它们是在LAION数据集继续finetune decoder来增强重建效果。
SD 2.1
在SD 2.0版本发布几周后，Stability AI又发布了SD 2.1。SD 2.0在训练过程中采用NSFW检测器过滤掉了可能包含色情的图像（punsafe=0.1），但是也同时过滤了很多人像图片，这导致SD 2.0在人像生成上效果可能较差，所以SD 2.1是在SD 2.0的基础上放开了限制（punsafe=0.98）继续finetune，所以增强了人像的生成效果。
和SD 2.0一样，SD 2.1也包含两个版本：512x512版本和768x768版本。
