Bootstrap

【扩散模型（六）】Stable Diffusion 3 diffusers 源码详解1-推理代码-文本处理部分

系列文章目录

文章目录

前言

下图为《Scaling Rectified Flow Transformers for High-Resolution Image Synthesis》 (ICML 2024 )中的 SD3 架构图。
在这里插入图片描述

一、文本处理的整体流程

下面流程图只对正向提示词进行了梳理，负向提示词的流程并无差异。
在这里插入图片描述

本文分析的源代码为 diffusers 包中的 SD3 pipeline （位置在/path/to/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py），文本处理部分主要为 其中 __call__() 函数调用的 self.encode_prompt() 函数，主要涉及了 3 个 text encoder 以及对应的 3 个 tokenizer。

其输入输出如下：

 (
     prompt_embeds,
     negative_prompt_embeds,
     pooled_prompt_embeds,
     negative_pooled_prompt_embeds,
 ) = self.encode_prompt(
     prompt=prompt,
     prompt_2=prompt_2,
     prompt_3=prompt_3,
     negative_prompt=negative_prompt,
     negative_prompt_2=negative_prompt_2,
     negative_prompt_3=negative_prompt_3,
     do_classifier_free_guidance=self.do_classifier_free_guidance,
     prompt_embeds=prompt_embeds,
     negative_prompt_embeds=negative_prompt_embeds,
     pooled_prompt_embeds=pooled_prompt_embeds,
     negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
     device=device,
     clip_skip=self.clip_skip,
     num_images_per_prompt=num_images_per_prompt,
     max_sequence_length=max_sequence_length,
 )

输入：

  • 其中 prompt 和 negative_prompt 为输入的字符串
  • 其他的 prompt_2、 prompt_3、 negative_prompt_2、 negative_prompt_3、prompt_embeds、 negative_prompt_embeds、pooled_prompt_embeds、negative_pooled_prompt_embeds 均为 None
  • do_classifier_free_guidance 一般都是 True
  • max_sequence_length = 256

具体而言是在 encode_prompt 函数中，通过两次 _get_clip_prompt_embeds 和 _get_t5_prompt_embeds 来调用 3 个 Text Encoder。

prompt_embed, pooled_prompt_embed = self._get_clip_prompt_embeds(
    prompt=prompt,
    device=device,
    num_images_per_prompt=num_images_per_prompt,
    clip_skip=clip_skip,
    clip_model_index=0,
)
prompt_2_embed, pooled_prompt_2_embed = self._get_clip_prompt_embeds(
    prompt=prompt_2,
    device=device,
    num_images_per_prompt=num_images_per_prompt,
    clip_skip=clip_skip,
    clip_model_index=1,
)
clip_prompt_embeds = torch.cat([prompt_embed, prompt_2_embed], dim=-1)

t5_prompt_embed = self._get_t5_prompt_embeds(
    prompt=prompt_3,
    num_images_per_prompt=num_images_per_prompt,
    max_sequence_length=max_sequence_length,
    device=device,
)

二、Text Encoder 1、2（CLIP）

1. 模型部分

  • 根据输入的 clip_tokenizers、clip_text_encoders 序号分别选择 text_encoder (CLIP L/141) 或者 text_encoder_2 (OpenCLIP bigG/142)。
  • 从下面初始化代码可以看出，二者 text_encoder 和 text_encoder_2 采用的类一致，所以二者的区别主要是模型权重以及 config 不同。
...
def __init__(...
     text_encoder: CLIPTextModelWithProjection,
     tokenizer: CLIPTokenizer,
     text_encoder_2: CLIPTextModelWithProjection,
     tokenizer_2: CLIPTokenizer,
...

 def _get_clip_prompt_embeds(
     self,
     prompt: Union[str, List[str]],
     num_images_per_prompt: int = 1,
     device: Optional[torch.device] = None,
     clip_skip: Optional[int] = None,
     clip_model_index: int = 0,
 ):
     device = device or self._execution_device

     clip_tokenizers = [self.tokenizer, self.tokenizer_2]
     clip_text_encoders = [self.text_encoder, self.text_encoder_2]

     tokenizer = clip_tokenizers[clip_model_index]
     text_encoder = clip_text_encoders[clip_model_index]

在下载的 SD3 模型权重文件中，/path/to/stable-diffusion-3-medium-diffusers 可以找到 text_encoder 和 text_encoder_2 子目录，对比其中的 config（下图中左边为 text_encoder ，右边为 text_encoder_2 ），可以知道二者更具体的不同之处：

  1. hidden_size 不同：768 vs 1280
  2. hidden_act： quick_gelu vs gelu
  3. intermediate_size 不同：3072 vs 5120
  4. “num_attention_heads” 和 “num_hidden_layers”：12/12 vs 20/32
  5. projection_dim 不同：768 vs 1280
  • 从以上 config 中，可以明显看出 text_encoder_2 (OpenCLIP bigG/14) 确实更加 big。
  • 两个 Text Encoder 最终的输出也和上文 “一、文本处理的整体流程” 中的流程图一致，分别输出 [n, 77, 768 ] 和 [n, 77, 1280]。
    • n 为推理时的 num_images_per_prompt，每个 prompt 的出图数量。

在这里插入图片描述

2. 两个 Text Encoder 的输入和输出

  • 二者的输入是相同的 prompt，得到输出为不同的两对 prompt_embed, pooled_prompt_embed ；prompt_2_embed, pooled_prompt_2_embed。
  • 其中，
    • prompt_embed [n, 77, 768 ] 和 prompt_2_embed [n, 77, 1280]为主要的 prompt 特征，并在后续 cat 到一起，得到 clip_prompt_embeds [n, 77, 2048]。
    • pooled_prompt_embed 和 pooled_prompt_2_embed 也一样 cat，
    • 两种特质的区别：prompt_embed（prompt_2_embed）是更主要/细粒度的文本特征、而 pooled_prompt_embed（pooled_prompt_2_embed）是更粗粒度的文本特征。
    • 原文：However, as the pooled text representation retains only coarse-grained information about the text input 3, the network also requires information from the sequence representation c t x t c_{txt} ctxt​.
prompt_embed, pooled_prompt_embed = self._get_clip_prompt_embeds(
    prompt=prompt,
    device=device,
    num_images_per_prompt=num_images_per_prompt,
    clip_skip=clip_skip,
    clip_model_index=0,
)
prompt_2_embed, pooled_prompt_2_embed = self._get_clip_prompt_embeds(
    prompt=prompt_2,
    device=device,
    num_images_per_prompt=num_images_per_prompt,
    clip_skip=clip_skip,
    clip_model_index=1,
)
clip_prompt_embeds = torch.cat([prompt_embed, prompt_2_embed], dim=-1)
...
pooled_prompt_embeds = torch.cat([pooled_prompt_embed, pooled_prompt_2_embed], dim=-1)

三、Text Encoder 3（T5）

T5EncoderModel 的调用则更简洁一点，输入同样是 prompt，并且只有一个输出。

def __init__(...
	text_encoder_3: T5EncoderModel,
       tokenizer_3: T5TokenizerFast,
...

t5_prompt_embed = self._get_t5_prompt_embeds(
    prompt=prompt_3,
    num_images_per_prompt=num_images_per_prompt,
    max_sequence_length=max_sequence_length,
    device=device,
)

# 实际为 clip_prompt_embeds = torch.nn.functional.pad(
#    clip_prompt_embeds, (0, 4096-2048)
#)，即在后面 2048 个维度上 pad 全 0. 
clip_prompt_embeds = torch.nn.functional.pad(
    clip_prompt_embeds, (0, t5_prompt_embed.shape[-1] - clip_prompt_embeds.shape[-1])
)

# 在序列长度的维度（-2）上 cat 到一起，得到 77+256 = 333 的长度
prompt_embeds = torch.cat([clip_prompt_embeds, t5_prompt_embed], dim=-2)
  • 作用：增强对复杂文本的生成能力。
  • 原文：T5 对于复杂的提示词很重要，例如涉及高度细节或拼写较长的文本（第2行和第3行）。然而，对于大多数提示，作者发现在推理时删除T5仍然可以获得具有竞争力的性能。
    在这里插入图片描述

其他

强烈安利另外一位博主的文章：

  1. Stable Diffusion1.5网络结构-超详细原创
  2. Stable Diffusion XL网络结构-超详细原创

  1. Learning transferable visual models from natural language supervision, 2021. ↩︎

  2. Reproducible scaling laws for contrastive language-image learning. In 2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). IEEE, 2023. doi: 10.1109/cvpr52729.2023.00276. URL http://dx.doi.org/10.1109/CVPR52729.2 023.00276. ↩︎

  3. Sdxl: Improving latent diffusion models for high-resolution image synthesis, 2023. ↩︎

悦读

道可道，非常道；名可名，非常名。 无名，天地之始，有名，万物之母。 故常无欲，以观其妙，常有欲，以观其徼。 此两者，同出而异名，同谓之玄，玄之又玄，众妙之门。

最新收录
AR远程空间标注Vuforia+WebRTC音视频通话和空间标注功能
Java小记-Vue/ElementUI/Axios(超级无敌认真好用，万字收藏篇!!!!)
Kafka面试三道题
排序算法乱炖: 快速排序、归并排序、冒泡排序
使用HBuildX打包App
【Java 设计模式】行为型之状态模式
springboot+毕业论文管理系统 计算机毕业设计附源码42124
web前端基础（四）
c++ 11 原子操作库 （std::atomic）（三）
机器学习-常用评价指标大赏~
;