2024不可不会的StableDiffusion之拼接各组件（五）

1. 引言

在之前的文章中，我介绍了如何安装扩散器库diffuser用以生成 AI 图像和构成stable diffusion的各个关键组件，即 CLIP 文本编码器、VAE 和 U-Net。在这篇文章中，我们将尝试把这些关键组件放在一起，并详细展示生成图像的扩散过程。

闲话少说，我们直接开始吧！

2. 概述

稳定扩散模型（Stable diffusion）采用文本输入和种子作为相关输入。然后，文本输入通过 CLIP 模型生成大小为77x768的文本嵌入（embedding），种子用于生成大小为 4x64x64 的高斯噪声，并将其作为第一个潜在图像表示。

接下来，U-Net迭代地对随机的潜在图像表示进行去噪。U-Net的输出是预测噪声，然后通过scheduler 调度器算法用于计算潜在的latents。这种去噪和基于文本的调节的过程重复 N 次（我们将使用 50 次）以便来获取更好的潜在图像表示。此过程完成后，VAE 解码器将对潜在图像表示 （4x64x64） 进行解码，用以解码生成最终的图像 （3x512x512）。

这种迭代去噪是获得良好输出图像的重要步骤。典型步骤操作数在 30-80 的范围内。然而，最近有论文声称通过使用蒸馏技术可以将其减少到4-5个操作步骤。

3. 辅助函数

接下来我们开始代码实践，本节涉及的大部分函数和所需的依赖库都已在本系列的前第三篇文章中进行了相关使用和解释。
这里我们额外添加以下两个辅助函数：

## Helper functions
def load_image(p):
    '''
    Function to load images from a defined path
    '''
    return Image.open(p).convert('RGB').resize((512,512))
    
## text encoder
def text_enc(prompts,text_encoder,maxlen=None,tokenizer=None,):
    if maxlen is None:
        maxlen = tokenizer.model_max_length
    inp = tokenizer(prompts,padding="max_length",
                    max_length=maxlen,truncation=True,
                    return_tensors="pt")
    return text_encoder(inp.input_ids.to("cuda"))[0].half()

第一个函数用来读取图像，并将其缩放至512X512的分辨率；第二个函数主要用于将输入的文本转化成文本嵌入embeding。

4. 核心函数

接下来，我们来定义函数prompt2img，用以实现 StableDiffusionPipeline.from_pretrained 函数相关功能的精简版本。

相关代码定义如下：

def prompt2img(prompts,text_encoder,tokenizer, unet, vae, scheduler,
               g=7.5, seed=100,steps=70,dim=512,save_int=False):
    bs = len(prompts)
    # convert textual prompts to embeding
    text = text_enc(prompts,text_encoder,
                    maxlen=tokenizer.model_max_length,
                    tokenizer=tokenizer)
    # adding an unconditional prompt
    uncond = text_enc( [""] * bs, text_encoder,
                      maxlen=text.shape[1] ,tokenizer=tokenizer)
    emb = torch.cat([uncond,text])
    if seed:
        torch.manual_seed(seed)

    # init random noise
    latents = torch.randn((bs,unet.in_channels,dim//8,dim//8))
    # set number of setps in scheduler
    scheduler.set_timesteps(steps)
    # add noise to the latents
    latents = latents.to("cuda").half()*scheduler.init_noise_sigma
    # iter through defined steps
    for i ,ts in enumerate(tqdm(scheduler.timesteps)):
        # we need to scale the i/p latents to match the variance
        inp = scheduler.scale_model_input(torch.cat([latents] * 2),ts)
        # predicting noise residual using UNet
        with torch.no_grad():
            u,t = unet(inp,ts,encoder_hidden_states=emb).sample.chunk(2)
        # performing Guidance
        pred = u + g* (t-u)
        # Conditioning the latents
        latents = scheduler.step(pred,ts,latents).prev_sample

        if save_int:
            os.makedirs(f"./steps",exist_ok=True)
            latent_to_pil(latents,vae)[0].save(f'steps/{i:04}.jpeg')

    return latent_to_pil(latents,vae)

5. 效果验证

我们接下来，测试上述函数的功能，即给定文本输入prompt，生成对应的图像，代码如下：

images = prompt_2_img(["A dog wearing a hat", "a photograph of an astronaut riding a horse"], save_int=False)
for img in images:
    plt.imshow(img)
    plt.show()

得到结果如下：

6. 代码讲解

可以看到，我们使用函数prompt2img实现了之前调用Stable Diffusion PipeLine类似的功能，即基于文本生成图像的功能。
接下来，我们来对函数prompt2img函数相关参数进行详细的讲解。

● prompt: 该参数是我们通过文本提示生成图像的提示词。类似于我们在第 1 部分中看到的pipe函数的prompt参数
● text_encoder: 该参数表示需要调用的文本编码器
● tokenizer：该参数表示用以将文本转为token的tokenizer
● unet： 表示基于unet的预测噪声的扩散模型
● vae：该参数表示SD中第二个组件变分自编码器
● scheduler： 该参数表示扩散过程控制每步添加噪声的权重
●g：代表guidance scale该值用于确定图像与文本提示的接近程度。这与一种称为“ Classifier free guidance”的技术有关，该技术可提高生成图像的质量。guidance scale的值越高，生成的图像就越接近文本提示
●seed ：该值决定了生成初始高斯latents的种子
● steps: 该值决定了生成最终latents所需要的去噪步数
● dim: 图像的尺寸，为简单起见，我们目前正在生成方形图像，因此只需要一个值，即可表示生成图像的长和宽
●save_int ：这是可选的bool变量，如果我们想保存中间潜在生成图像，则将其设置为True

7. 中间过程可视化

为了可视化扩散模型的中间过程，我们将参数save_int设置为True，来将中间过程进行保存。

scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, 
                                 beta_schedule="scaled_linear",
                                  num_train_timesteps=1000)
scheduler.set_timesteps(50)
prompts = ["A dog wearing a hat"]
images = prompt2img(prompts, text_encoder, tokenizer, unet, vae, scheduler, save_int=True)

并使用ffmpeg将每一个step得到的图像合并成视频进行输出：

ffmpeg -v 1 -y -f image2 -framerate 20 -i steps/%04d.jpeg -c:v libx264 -preset slow -qp 18 -pix_fmt yuv420p out.mp4

得到最终结果如下：

8. 总结

我希望这能给大家提供一个很好的概述，并将代码分解到最低限度，以便我们能够理解每个组件。现在我们已经实现了基于三个组件来实现SD文本生图的功能，并给出了详细的代码示例。

您学废了嘛？