diffusers 原始碼解析（十一）
# 版權所有 2024 HunyuanDiT 作者，Qixun Wang 和 HuggingFace 團隊。保留所有權利。
#
# 根據 Apache 許可證第 2.0 版（"許可證"）進行許可；
# 除非符合許可證，否則您不得使用此檔案。
# 您可以在以下網址獲取許可證副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非適用法律或書面同意，否則根據許可證分發的軟體均按 "原樣" 基礎提供，
# 不提供任何種類的保證或條件，無論是明示或暗示的。
# 有關許可證的具體條款和條件，請參閱許可證。
from typing import Dict, Optional, Union # 匯入字典、可選和聯合型別定義
import torch # 匯入 PyTorch 庫
from torch import nn # 從 PyTorch 匯入神經網路模組
from ...configuration_utils import ConfigMixin, register_to_config # 從配置工具匯入混合類和註冊功能
from ...utils import logging # 從工具包匯入日誌記錄功能
from ...utils.torch_utils import maybe_allow_in_graph # 匯入可能允許圖形內操作的功能
from ..attention import FeedForward # 從注意力模組匯入前饋網路
from ..attention_processor import Attention, AttentionProcessor, FusedHunyuanAttnProcessor2_0, HunyuanAttnProcessor2_0 # 匯入注意力處理器
from ..embeddings import ( # 匯入嵌入模組
HunyuanCombinedTimestepTextSizeStyleEmbedding, # 組合時間步、文字、大小和樣式的嵌入
PatchEmbed, # 影像補丁嵌入
PixArtAlphaTextProjection, # 畫素藝術文字投影
)
from ..modeling_outputs import Transformer2DModelOutput # 匯入 2D 變換器模型輸出型別
from ..modeling_utils import ModelMixin # 匯入模型混合類
from ..normalization import AdaLayerNormContinuous, FP32LayerNorm # 匯入自適應層歸一化和 FP32 層歸一化
logger = logging.get_logger(__name__) # 建立當前模組的日誌記錄器，禁用 pylint 警告
class AdaLayerNormShift(nn.Module): # 定義自適應層歸一化偏移類，繼承自 nn.Module
r""" # 類文件字串，描述類的功能
Norm layer modified to incorporate timestep embeddings. # 歸一化層，修改以包含時間步嵌入
Parameters: # 引數說明
embedding_dim (`int`): The size of each embedding vector. # 嵌入向量的大小
num_embeddings (`int`): The size of the embeddings dictionary. # 嵌入字典的大小
"""
def __init__(self, embedding_dim: int, elementwise_affine=True, eps=1e-6): # 初始化方法
super().__init__() # 呼叫父類初始化方法
self.silu = nn.SiLU() # 定義 SiLU 啟用函式
self.linear = nn.Linear(embedding_dim, embedding_dim) # 定義線性層，輸入輸出維度均為嵌入維度
self.norm = FP32LayerNorm(embedding_dim, elementwise_affine=elementwise_affine, eps=eps) # 定義層歸一化
def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: # 定義前向傳播方法
shift = self.linear(self.silu(emb.to(torch.float32)).to(emb.dtype)) # 計算偏移量
x = self.norm(x) + shift.unsqueeze(dim=1) # 對輸入進行歸一化並加上偏移
return x # 返回處理後的張量
@maybe_allow_in_graph # 裝飾器，可能允許在計算圖中使用
class HunyuanDiTBlock(nn.Module): # 定義 Hunyuan-DiT 模型中的變換器塊類
r""" # 類文件字串，描述類的功能
Transformer block used in Hunyuan-DiT model (https://github.com/Tencent/HunyuanDiT). Allow skip connection and # Hunyuan-DiT 模型中的變換器塊，允許跳過連線和
QKNorm # QKNorm 功能
# 引數說明部分，定義各引數的型別和作用
Parameters:
dim (`int`): # 輸入和輸出的通道數
The number of channels in the input and output.
num_attention_heads (`int`): # 多頭注意力機制中使用的頭數
The number of heads to use for multi-head attention.
cross_attention_dim (`int`, *optional*): # 跨注意力的編碼器隱藏狀態向量的大小
The size of the encoder_hidden_states vector for cross attention.
dropout (`float`, *optional*, defaults to 0.0): # 用於正則化的丟棄機率
The dropout probability to use.
activation_fn (`str`, *optional*, defaults to `"geglu"`): # 前饋網路中使用的啟用函式
Activation function to be used in feed-forward.
norm_elementwise_affine (`bool`, *optional*, defaults to `True`): # 是否使用可學習的元素逐個仿射引數進行歸一化
Whether to use learnable elementwise affine parameters for normalization.
norm_eps (`float`, *optional*, defaults to 1e-6): # 加到歸一化層分母的小常數，以防止除以零
A small constant added to the denominator in normalization layers to prevent division by zero.
final_dropout (`bool`, *optional*, defaults to False): # 在最後的前饋層後是否應用最終丟棄
Whether to apply a final dropout after the last feed-forward layer.
ff_inner_dim (`int`, *optional*): # 前饋塊中隱藏層的大小，預設為 None
The size of the hidden layer in the feed-forward block. Defaults to `None`.
ff_bias (`bool`, *optional*, defaults to `True`): # 前饋塊中是否使用偏置
Whether to use bias in the feed-forward block.
skip (`bool`, *optional*, defaults to `False`): # 是否使用跳過連線，預設為下塊和中塊的 False
Whether to use skip connection. Defaults to `False` for down-blocks and mid-blocks.
qk_norm (`bool`, *optional*, defaults to `True`): # 在 QK 計算中是否使用歸一化，預設為 True
Whether to use normalization in QK calculation. Defaults to `True`.
"""
# 建構函式的定義，初始化各引數
def __init__(
self,
dim: int, # 輸入和輸出的通道數
num_attention_heads: int, # 多頭注意力機制中使用的頭數
cross_attention_dim: int = 1024, # 預設的跨注意力維度
dropout=0.0, # 預設的丟棄機率
activation_fn: str = "geglu", # 預設的啟用函式
norm_elementwise_affine: bool = True, # 預設使用可學習的仿射引數
norm_eps: float = 1e-6, # 預設的歸一化小常數
final_dropout: bool = False, # 預設不應用最終丟棄
ff_inner_dim: Optional[int] = None, # 預設的前饋塊隱藏層大小
ff_bias: bool = True, # 預設使用偏置
skip: bool = False, # 預設不使用跳過連線
qk_norm: bool = True, # 預設在 QK 計算中使用歸一化
):
# 呼叫父類建構函式
super().__init__()
# 定義三個塊，每個塊都有自己的歸一化層。
# 注意：新版本釋出時，檢查 norm2 和 norm3
# 1. 自注意力機制
self.norm1 = AdaLayerNormShift(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps)
# 建立自注意力機制的例項
self.attn1 = Attention(
query_dim=dim, # 查詢向量的維度
cross_attention_dim=None, # 交叉注意力的維度，未使用
dim_head=dim // num_attention_heads, # 每個頭的維度
heads=num_attention_heads, # 注意力頭的數量
qk_norm="layer_norm" if qk_norm else None, # 查詢和鍵的歸一化方法
eps=1e-6, # 數值穩定性常數
bias=True, # 是否使用偏置
processor=HunyuanAttnProcessor2_0(), # 注意力處理器的例項
)
# 2. 交叉注意力機制
self.norm2 = FP32LayerNorm(dim, norm_eps, norm_elementwise_affine)
# 建立交叉注意力機制的例項
self.attn2 = Attention(
query_dim=dim, # 查詢向量的維度
cross_attention_dim=cross_attention_dim, # 交叉注意力的維度
dim_head=dim // num_attention_heads, # 每個頭的維度
heads=num_attention_heads, # 注意力頭的數量
qk_norm="layer_norm" if qk_norm else None, # 查詢和鍵的歸一化方法
eps=1e-6, # 數值穩定性常數
bias=True, # 是否使用偏置
processor=HunyuanAttnProcessor2_0(), # 注意力處理器的例項
)
# 3. 前饋網路
self.norm3 = FP32LayerNorm(dim, norm_eps, norm_elementwise_affine)
# 建立前饋網路的例項
self.ff = FeedForward(
dim, # 輸入維度
dropout=dropout, # dropout 比例
activation_fn=activation_fn, # 啟用函式
final_dropout=final_dropout, # 最終 dropout 比例
inner_dim=ff_inner_dim, # 內部維度，通常是 dim 的倍數
bias=ff_bias, # 是否使用偏置
)
# 4. 跳躍連線
if skip: # 如果啟用跳躍連線
self.skip_norm = FP32LayerNorm(2 * dim, norm_eps, elementwise_affine=True) # 建立歸一化層
self.skip_linear = nn.Linear(2 * dim, dim) # 建立線性層
else: # 如果不啟用跳躍連線
self.skip_linear = None # 設定為 None
# 將塊大小預設為 None
self._chunk_size = None # 初始化塊大小
self._chunk_dim = 0 # 初始化塊維度
# 從 diffusers.models.attention.BasicTransformerBlock 複製的設定塊前饋方法
def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0):
# 設定塊前饋
self._chunk_size = chunk_size # 設定塊大小
self._chunk_dim = dim # 設定塊維度
def forward(
self,
hidden_states: torch.Tensor, # 輸入的隱藏狀態
encoder_hidden_states: Optional[torch.Tensor] = None, # 編碼器的隱藏狀態
temb: Optional[torch.Tensor] = None, # 額外的嵌入
image_rotary_emb=None, # 影像旋轉嵌入
skip=None, # 跳躍連線標誌
) -> torch.Tensor:
# 注意：以下程式碼塊中的計算總是在歸一化之後進行。
# 0. 長跳躍連線
# 如果 skip_linear 不為 None，執行跳躍連線
if self.skip_linear is not None:
# 將當前的隱藏狀態與跳躍連線的輸出在最後一維上拼接
cat = torch.cat([hidden_states, skip], dim=-1)
# 對拼接後的結果進行歸一化處理
cat = self.skip_norm(cat)
# 透過線性層處理歸一化後的結果，更新隱藏狀態
hidden_states = self.skip_linear(cat)
# 1. 自注意力
# 對當前隱藏狀態進行歸一化，準備進行自注意力計算
norm_hidden_states = self.norm1(hidden_states, temb) ### checked: self.norm1 is correct
# 計算自注意力的輸出
attn_output = self.attn1(
norm_hidden_states,
image_rotary_emb=image_rotary_emb,
)
# 將自注意力的輸出加到隱藏狀態上，形成新的隱藏狀態
hidden_states = hidden_states + attn_output
# 2. 交叉注意力
# 將交叉注意力的輸出加到當前的隱藏狀態上
hidden_states = hidden_states + self.attn2(
self.norm2(hidden_states), # 先進行歸一化
encoder_hidden_states=encoder_hidden_states, # 使用編碼器的隱藏狀態
image_rotary_emb=image_rotary_emb, # 傳遞旋轉嵌入
)
# 前饋網路層 ### TODO: 在狀態字典中切換 norm2 和 norm3
# 對當前的隱藏狀態進行歸一化處理，準備進入前饋網路
mlp_inputs = self.norm3(hidden_states)
# 透過前饋網路處理歸一化後的輸入，更新隱藏狀態
hidden_states = hidden_states + self.ff(mlp_inputs)
# 返回最終的隱藏狀態
return hidden_states
# 定義 HunyuanDiT2DModel 類，繼承自 ModelMixin 和 ConfigMixin
class HunyuanDiT2DModel(ModelMixin, ConfigMixin):
"""
HunYuanDiT: 基於 Transformer 的擴散模型。
繼承 ModelMixin 和 ConfigMixin 以與 diffusers 的取樣器 StableDiffusionPipeline 相容。
引數:
num_attention_heads (`int`, *可選*, 預設為 16):
多頭注意力的頭數。
attention_head_dim (`int`, *可選*, 預設為 88):
每個頭的通道數。
in_channels (`int`, *可選*):
輸入和輸出的通道數（如果輸入為 **連續**，需指定）。
patch_size (`int`, *可選*):
輸入的補丁大小。
activation_fn (`str`, *可選*, 預設為 `"geglu"`):
前饋網路中使用的啟用函式。
sample_size (`int`, *可選*):
潛在影像的寬度。訓練期間固定使用，以學習位置嵌入的數量。
dropout (`float`, *可選*, 預設為 0.0):
使用的 dropout 機率。
cross_attention_dim (`int`, *可選*):
clip 文字嵌入中的維度數量。
hidden_size (`int`, *可選*):
條件嵌入層中隱藏層的大小。
num_layers (`int`, *可選*, 預設為 1):
使用的 Transformer 塊的層數。
mlp_ratio (`float`, *可選*, 預設為 4.0):
隱藏層大小與輸入大小的比率。
learn_sigma (`bool`, *可選*, 預設為 `True`):
是否預測方差。
cross_attention_dim_t5 (`int`, *可選*):
t5 文字嵌入中的維度數量。
pooled_projection_dim (`int`, *可選*):
池化投影的大小。
text_len (`int`, *可選*):
clip 文字嵌入的長度。
text_len_t5 (`int`, *可選*):
T5 文字嵌入的長度。
use_style_cond_and_image_meta_size (`bool`, *可選*):
是否使用風格條件和影像後設資料大小。版本 <=1.1 為 True，版本 >= 1.2 為 False
"""
# 註冊到配置中
@register_to_config
def __init__(
# 多頭注意力的頭數，預設為 16
self,
num_attention_heads: int = 16,
# 每個頭的通道數，預設為 88
attention_head_dim: int = 88,
# 輸入和輸出的通道數，預設為 None
in_channels: Optional[int] = None,
# 輸入的補丁大小，預設為 None
patch_size: Optional[int] = None,
# 啟用函式，預設為 "gelu-approximate"
activation_fn: str = "gelu-approximate",
# 潛在影像的寬度，預設為 32
sample_size=32,
# 條件嵌入層中隱藏層的大小，預設為 1152
hidden_size=1152,
# 使用的 Transformer 塊的層數，預設為 28
num_layers: int = 28,
# 隱藏層大小與輸入大小的比率，預設為 4.0
mlp_ratio: float = 4.0,
# 是否預測方差，預設為 True
learn_sigma: bool = True,
# clip 文字嵌入中的維度數量，預設為 1024
cross_attention_dim: int = 1024,
# 正則化型別，預設為 "layer_norm"
norm_type: str = "layer_norm",
# t5 文字嵌入中的維度數量，預設為 2048
cross_attention_dim_t5: int = 2048,
# 池化投影的大小，預設為 1024
pooled_projection_dim: int = 1024,
# clip 文字嵌入的長度，預設為 77
text_len: int = 77,
# T5 文字嵌入的長度，預設為 256
text_len_t5: int = 256,
# 是否使用風格條件和影像後設資料大小，預設為 True
use_style_cond_and_image_meta_size: bool = True,
):
# 呼叫父類的初始化方法
super().__init__()
# 根據是否學習 sigma 決定輸出通道數
self.out_channels = in_channels * 2 if learn_sigma else in_channels
# 設定注意力頭的數量
self.num_heads = num_attention_heads
# 計算內部維度，等於注意力頭數量乘以每個頭的維度
self.inner_dim = num_attention_heads * attention_head_dim
# 初始化文字嵌入器，用於將輸入特徵投影到更高維空間
self.text_embedder = PixArtAlphaTextProjection(
# 輸入特徵維度
in_features=cross_attention_dim_t5,
# 隱藏層大小為輸入特徵的四倍
hidden_size=cross_attention_dim_t5 * 4,
# 輸出特徵維度
out_features=cross_attention_dim,
# 啟用函式設定為"siluf_fp32"
act_fn="silu_fp32",
)
# 初始化文字嵌入的填充引數，使用隨機正態分佈初始化
self.text_embedding_padding = nn.Parameter(
torch.randn(text_len + text_len_t5, cross_attention_dim, dtype=torch.float32)
)
# 初始化位置嵌入，構建影像的補丁嵌入
self.pos_embed = PatchEmbed(
# 補丁的高度
height=sample_size,
# 補丁的寬度
width=sample_size,
# 輸入通道數
in_channels=in_channels,
# 嵌入維度
embed_dim=hidden_size,
# 補丁大小
patch_size=patch_size,
# 位置嵌入型別設定為 None
pos_embed_type=None,
)
# 初始化時間和風格嵌入，結合時間步和文字大小
self.time_extra_emb = HunyuanCombinedTimestepTextSizeStyleEmbedding(
# 隱藏層大小
hidden_size,
# 池化投影維度
pooled_projection_dim=pooled_projection_dim,
# 輸入序列長度
seq_len=text_len_t5,
# 交叉注意力維度
cross_attention_dim=cross_attention_dim_t5,
# 是否使用風格條件和影像後設資料大小
use_style_cond_and_image_meta_size=use_style_cond_and_image_meta_size,
)
# 初始化 HunyuanDiT 塊列表
self.blocks = nn.ModuleList(
[
# 為每一層建立 HunyuanDiTBlock
HunyuanDiTBlock(
# 內部維度
dim=self.inner_dim,
# 注意力頭數量
num_attention_heads=self.config.num_attention_heads,
# 啟用函式
activation_fn=activation_fn,
# 前饋網路內部維度
ff_inner_dim=int(self.inner_dim * mlp_ratio),
# 交叉注意力維度
cross_attention_dim=cross_attention_dim,
# 查詢-鍵歸一化開啟
qk_norm=True, # 詳情見 http://arxiv.org/abs/2302.05442
# 如果當前層數大於層數的一半，則跳過
skip=layer > num_layers // 2,
)
# 遍歷層數
for layer in range(num_layers)
]
)
# 初始化輸出的自適應層歸一化
self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
# 初始化輸出的線性層，將內部維度對映到輸出通道數
self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
# 從 diffusers.models.unets.unet_2d_condition.UNet2DConditionModel 中複製的程式碼，用於融合 QKV 投影，更新為 FusedHunyuanAttnProcessor2_0
def fuse_qkv_projections(self):
"""
啟用融合的 QKV 投影。對於自注意力模組，所有投影矩陣（即查詢、鍵、值）都被融合。
對於交叉注意力模組，鍵和值投影矩陣被融合。
<Tip warning={true}>
該 API 是 🧪 實驗性的。
</Tip>
"""
# 初始化原始注意力處理器為 None
self.original_attn_processors = None
# 遍歷所有注意力處理器
for _, attn_processor in self.attn_processors.items():
# 檢查注意力處理器類名中是否包含 "Added"
if "Added" in str(attn_processor.__class__.__name__):
# 如果包含，則丟擲錯誤，表示不支援融合 QKV 投影
raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
# 儲存當前的注意力處理器以備後用
self.original_attn_processors = self.attn_processors
# 遍歷當前模組中的所有子模組
for module in self.modules():
# 檢查模組是否為 Attention 型別
if isinstance(module, Attention):
# 對於 Attention 模組，啟用投影融合
module.fuse_projections(fuse=True)
# 設定融合的注意力處理器
self.set_attn_processor(FusedHunyuanAttnProcessor2_0())
# 從 diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections 複製
def unfuse_qkv_projections(self):
"""
如果已啟用，則禁用融合的 QKV 投影。
<Tip warning={true}>
該 API 是 🧪 實驗性的。
</Tip>
"""
# 檢查是否有原始注意力處理器
if self.original_attn_processors is not None:
# 恢復為原始注意力處理器
self.set_attn_processor(self.original_attn_processors)
@property
# 從 diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors 複製
def attn_processors(self) -> Dict[str, AttentionProcessor]:
r"""
返回:
`dict` 型別的注意力處理器：一個字典，包含模型中使用的所有注意力處理器，按權重名稱索引。
"""
# 初始化處理器字典
processors = {}
# 定義遞迴新增處理器的函式
def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
# 檢查模組是否具有 get_processor 方法
if hasattr(module, "get_processor"):
# 將處理器新增到字典中
processors[f"{name}.processor"] = module.get_processor()
# 遍歷子模組
for sub_name, child in module.named_children():
# 遞迴呼叫以新增子模組的處理器
fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
return processors
# 遍歷當前模組的所有子模組
for name, module in self.named_children():
# 呼叫遞迴函式新增處理器
fn_recursive_add_processors(name, module, processors)
# 返回處理器字典
return processors
# 從 diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor 複製
# 定義設定注意力處理器的方法，接收一個注意力處理器或處理器字典
def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
r"""
設定用於計算注意力的注意力處理器。
引數：
processor（`dict` of `AttentionProcessor` 或 `AttentionProcessor`）:
已例項化的處理器類或將作為處理器設定的處理器類字典
用於**所有** `Attention` 層。
如果 `processor` 是字典，則鍵需要定義相應的交叉注意力處理器路徑。
當設定可訓練的注意力處理器時，這強烈建議。
"""
# 計算當前注意力處理器的數量
count = len(self.attn_processors.keys())
# 檢查傳入的處理器是否為字典，並驗證其長度與注意力層數量是否一致
if isinstance(processor, dict) and len(processor) != count:
raise ValueError(
# 丟擲錯誤，提示處理器數量與注意力層數量不匹配
f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
)
# 定義遞迴設定注意力處理器的內部函式
def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
# 檢查模組是否具有設定處理器的方法
if hasattr(module, "set_processor"):
# 如果處理器不是字典，直接設定處理器
if not isinstance(processor, dict):
module.set_processor(processor)
else:
# 從字典中取出處理器並設定
module.set_processor(processor.pop(f"{name}.processor"))
# 遍歷模組的子模組，遞迴呼叫設定處理器的方法
for sub_name, child in module.named_children():
fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
# 遍歷當前物件的所有子模組，呼叫遞迴設定處理器的方法
for name, module in self.named_children():
fn_recursive_attn_processor(name, module, processor)
# 定義設定預設注意力處理器的方法
def set_default_attn_processor(self):
"""
禁用自定義注意力處理器，並設定預設的注意力實現。
"""
# 呼叫設定注意力處理器的方法，使用預設的 HunyuanAttnProcessor2_0
self.set_attn_processor(HunyuanAttnProcessor2_0())
# 定義前向傳播的方法，接收多個輸入引數
def forward(
self,
hidden_states,
timestep,
encoder_hidden_states=None,
text_embedding_mask=None,
encoder_hidden_states_t5=None,
text_embedding_mask_t5=None,
image_meta_size=None,
style=None,
image_rotary_emb=None,
controlnet_block_samples=None,
return_dict=True,
# 從 diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking 複製的程式碼
# 定義一個方法以啟用前饋層的分塊處理，引數為分塊大小和維度
def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
"""
設定注意力處理器使用 [前饋分塊處理](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers)。
引數:
chunk_size (`int`, *可選*):
前饋層的分塊大小。如果未指定，將對每個維度為 `dim` 的張量單獨執行前饋層。
dim (`int`, *可選*, 預設為 `0`):
前饋計算應分塊的維度。選擇 dim=0（批處理）或 dim=1（序列長度）。
"""
# 檢查維度是否為 0 或 1
if dim not in [0, 1]:
# 丟擲值錯誤，提示維度設定不當
raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}")
# 預設分塊大小為 1
chunk_size = chunk_size or 1
# 定義遞迴函式以設定前饋層的分塊處理
def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int):
# 如果模組具有設定分塊前饋的屬性，則進行設定
if hasattr(module, "set_chunk_feed_forward"):
module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
# 遍歷模組的所有子模組並遞迴呼叫
for child in module.children():
fn_recursive_feed_forward(child, chunk_size, dim)
# 遍歷當前物件的所有子模組並應用遞迴函式
for module in self.children():
fn_recursive_feed_forward(module, chunk_size, dim)
# 從 diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.disable_forward_chunking 複製
def disable_forward_chunking(self):
# 定義遞迴函式以禁用前饋層的分塊處理
def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int):
# 如果模組具有設定分塊前饋的屬性，則進行設定
if hasattr(module, "set_chunk_feed_forward"):
module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
# 遍歷模組的所有子模組並遞迴呼叫
for child in module.children():
fn_recursive_feed_forward(child, chunk_size, dim)
# 遍歷當前物件的所有子模組並應用遞迴函式，禁用分塊處理
for module in self.children():
fn_recursive_feed_forward(module, None, 0)
# 版權宣告，表明該程式碼的版權屬於 Latte 團隊和 HuggingFace 團隊
# Copyright 2024 the Latte Team and The HuggingFace Team. All rights reserved.
#
# 根據 Apache 許可證 2.0 版（"許可證"）進行授權；
# 除非遵循該許可證，否則您不得使用此檔案。
# 您可以在以下網址獲取許可證的副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非法律要求或書面同意，否則根據許可證分發的軟體是按 "原樣" 基礎進行的，
# 不提供任何形式的擔保或條件，無論是明示的還是暗示的。
# 請參見許可證以瞭解有關特定語言的許可權和限制。
from typing import Optional # 從 typing 模組匯入 Optional 型別，用於指示可選引數
import torch # 匯入 PyTorch 庫
from torch import nn # 從 PyTorch 匯入神經網路模組
# 從配置工具匯入 ConfigMixin 和註冊配置的功能
from ...configuration_utils import ConfigMixin, register_to_config
# 匯入與影像嵌入相關的類和函式
from ...models.embeddings import PixArtAlphaTextProjection, get_1d_sincos_pos_embed_from_grid
# 匯入基礎變換器塊的定義
from ..attention import BasicTransformerBlock
# 匯入影像塊嵌入的定義
from ..embeddings import PatchEmbed
# 匯入 Transformer 2D 模型輸出的定義
from ..modeling_outputs import Transformer2DModelOutput
# 匯入模型混合功能的定義
from ..modeling_utils import ModelMixin
# 匯入自適應層歸一化的定義
from ..normalization import AdaLayerNormSingle
# 定義一個 3D Transformer 模型類，繼承自 ModelMixin 和 ConfigMixin
class LatteTransformer3DModel(ModelMixin, ConfigMixin):
# 設定支援梯度檢查點的標誌為 True
_supports_gradient_checkpointing = True
"""
一個用於影片類資料的 3D Transformer 模型，相關論文連結：
https://arxiv.org/abs/2401.03048，官方程式碼地址：
https://github.com/Vchitect/Latte
"""
# 引數說明
Parameters:
# 使用的多頭注意力頭數，預設為16
num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention.
# 每個頭的通道數，預設為88
attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head.
# 輸入通道數
in_channels (`int`, *optional*):
The number of channels in the input.
# 輸出通道數
out_channels (`int`, *optional*):
The number of channels in the output.
# Transformer塊的層數，預設為1
num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use.
# dropout機率，預設為0.0
dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
# 用於cross attention的編碼器隱藏狀態維度
cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
# 配置TransformerBlocks的注意力是否包含偏置引數
attention_bias (`bool`, *optional*):
Configure if the `TransformerBlocks` attention should contain a bias parameter.
# 潛在影像的寬度（如果輸入為離散型別，需指定）
sample_size (`int`, *optional*):
The width of the latent images (specify if the input is **discrete**).
# 在訓練期間固定，用於學習位置嵌入數量。
patch_size (`int`, *optional*):
# 在補丁嵌入層中使用的補丁大小。
The size of the patches to use in the patch embedding layer.
# 前饋中的啟用函式，預設為"geglu"
activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to use in feed-forward.
# 訓練期間使用的擴散步驟數
num_embeds_ada_norm ( `int`, *optional*):
The number of diffusion steps used during training. Pass if at least one of the norm_layers is
`AdaLayerNorm`. This is fixed during training since it is used to learn a number of embeddings that are
added to the hidden states. During inference, you can denoise for up to but not more steps than
`num_embeds_ada_norm`.
# 使用的歸一化型別，選項為"layer_norm"或"ada_layer_norm"
norm_type (`str`, *optional*, defaults to `"layer_norm"`):
The type of normalization to use. Options are `"layer_norm"` or `"ada_layer_norm"`.
# 是否在歸一化層中使用逐元素仿射，預設為True
norm_elementwise_affine (`bool`, *optional*, defaults to `True`):
Whether or not to use elementwise affine in normalization layers.
# 歸一化層中使用的epsilon值，預設為1e-5
norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon value to use in normalization layers.
# 標註嵌入的通道數
caption_channels (`int`, *optional*):
The number of channels in the caption embeddings.
# 影片類資料中的幀數
video_length (`int`, *optional*):
The number of frames in the video-like data.
"""
# 註冊配置的裝飾器
@register_to_config
# 初始化方法，用於設定模型的引數
def __init__(
# 注意力頭的數量，預設為16
num_attention_heads: int = 16,
# 每個注意力頭的維度，預設為88
attention_head_dim: int = 88,
# 輸入通道數，預設為None
in_channels: Optional[int] = None,
# 輸出通道數，預設為None
out_channels: Optional[int] = None,
# 網路層數，預設為1
num_layers: int = 1,
# dropout比率，預設為0.0
dropout: float = 0.0,
# 跨注意力的維度，預設為None
cross_attention_dim: Optional[int] = None,
# 是否使用注意力偏置，預設為False
attention_bias: bool = False,
# 樣本大小，預設為64
sample_size: int = 64,
# 每個patch的大小，預設為None
patch_size: Optional[int] = None,
# 啟用函式，預設為"geglu"
activation_fn: str = "geglu",
# 自適應歸一化的嵌入數量，預設為None
num_embeds_ada_norm: Optional[int] = None,
# 歸一化型別，預設為"layer_norm"
norm_type: str = "layer_norm",
# 歸一化是否進行逐元素仿射，預設為True
norm_elementwise_affine: bool = True,
# 歸一化的epsilon值，預設為1e-5
norm_eps: float = 1e-5,
# caption的通道數，預設為None
caption_channels: int = None,
# 影片長度，預設為16
video_length: int = 16,
# 設定梯度檢查點的函式，接收一個模組和一個布林值
def _set_gradient_checkpointing(self, module, value=False):
# 將梯度檢查點設定為給定的布林值
self.gradient_checkpointing = value
# 前向傳播方法，定義模型的前向計算
def forward(
# 輸入的隱藏狀態，型別為torch.Tensor
hidden_states: torch.Tensor,
# 可選的時間步長，型別為torch.LongTensor
timestep: Optional[torch.LongTensor] = None,
# 可選的編碼器隱藏狀態，型別為torch.Tensor
encoder_hidden_states: Optional[torch.Tensor] = None,
# 可選的編碼器注意力掩碼，型別為torch.Tensor
encoder_attention_mask: Optional[torch.Tensor] = None,
# 是否啟用時間注意力，預設為True
enable_temporal_attentions: bool = True,
# 是否返回字典形式的輸出，預設為True
return_dict: bool = True,
# 版權宣告，表明此程式碼的版權歸 2024 Alpha-VLLM 作者及 HuggingFace 團隊所有
#
# 根據 Apache 2.0 許可證（"許可證"）進行許可；
# 除非遵循許可證，否則您不得使用此檔案。
# 您可以在以下網址獲取許可證副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非適用法律或書面協議另有規定，軟體在許可證下分發是按"原樣"基礎，
# 不提供任何形式的明示或暗示的擔保或條件。
# 請參閱許可證以獲取有關許可權和
# 限制的具體資訊。
from typing import Any, Dict, Optional # 匯入型別提示相關的模組
import torch # 匯入 PyTorch 庫
import torch.nn as nn # 匯入 PyTorch 的神經網路模組
# 從配置和工具模組匯入必要的類和函式
from ...configuration_utils import ConfigMixin, register_to_config
from ...utils import logging
from ..attention import LuminaFeedForward # 匯入自定義的前饋網路
from ..attention_processor import Attention, LuminaAttnProcessor2_0 # 匯入注意力處理器
from ..embeddings import (
LuminaCombinedTimestepCaptionEmbedding, # 匯入組合時間步長的嵌入
LuminaPatchEmbed, # 匯入補丁嵌入
)
from ..modeling_outputs import Transformer2DModelOutput # 匯入模型輸出類
from ..modeling_utils import ModelMixin # 匯入模型混合類
from ..normalization import LuminaLayerNormContinuous, LuminaRMSNormZero, RMSNorm # 匯入不同的歸一化方法
logger = logging.get_logger(__name__) # 獲取當前模組的日誌記錄器，禁用 pylint 的名稱警告
class LuminaNextDiTBlock(nn.Module): # 定義一個名為 LuminaNextDiTBlock 的類，繼承自 nn.Module
"""
LuminaNextDiTBlock 用於 LuminaNextDiT2DModel。
引數：
dim (`int`): 輸入特徵的嵌入維度。
num_attention_heads (`int`): 注意力頭的數量。
num_kv_heads (`int`):
鍵和值特徵中的注意力頭數量（如果使用 GQA），
或設定為 None 以與查詢相同。
multiple_of (`int`): 前饋網路層的倍數。
ffn_dim_multiplier (`float`): 前饋網路層維度的乘數因子。
norm_eps (`float`): 歸一化層的 epsilon 值。
qk_norm (`bool`): 查詢和鍵的歸一化。
cross_attention_dim (`int`): 輸入文字提示的跨注意力嵌入維度。
norm_elementwise_affine (`bool`, *可選*, 預設為 True)，
"""
def __init__( # 初始化方法
self,
dim: int, # 輸入特徵的維度
num_attention_heads: int, # 注意力頭的數量
num_kv_heads: int, # 鍵和值特徵的頭數量
multiple_of: int, # 前饋網路層的倍數
ffn_dim_multiplier: float, # 前饋網路維度的乘數
norm_eps: float, # 歸一化的 epsilon 值
qk_norm: bool, # 是否對查詢和鍵進行歸一化
cross_attention_dim: int, # 跨注意力嵌入的維度
norm_elementwise_affine: bool = True, # 是否使用逐元素仿射歸一化，預設值為 True
) -> None: # 定義方法的返回型別為 None，表示不返回任何值
super().__init__() # 呼叫父類的建構函式，初始化父類的屬性
self.head_dim = dim // num_attention_heads # 計算每個注意力頭的維度
self.gate = nn.Parameter(torch.zeros([num_attention_heads])) # 建立一個可學習的引數，初始化為零，大小為注意力頭的數量
# Self-attention # 定義自注意力機制
self.attn1 = Attention( # 建立第一個注意力層
query_dim=dim, # 查詢的維度
cross_attention_dim=None, # 交叉注意力的維度，此處為 None 表示不使用
dim_head=dim // num_attention_heads, # 每個頭的維度
qk_norm="layer_norm_across_heads" if qk_norm else None, # 如果 qk_norm 為真，使用跨頭層歸一化
heads=num_attention_heads, # 注意力頭的數量
kv_heads=num_kv_heads, # 鍵值對的頭數量
eps=1e-5, # 數值穩定性引數
bias=False, # 不使用偏置項
out_bias=False, # 輸出層不使用偏置項
processor=LuminaAttnProcessor2_0(), # 使用指定的注意力處理器
)
self.attn1.to_out = nn.Identity() # 輸出層使用恆等對映
# Cross-attention # 定義交叉注意力機制
self.attn2 = Attention( # 建立第二個注意力層
query_dim=dim, # 查詢的維度
cross_attention_dim=cross_attention_dim, # 交叉注意力的維度
dim_head=dim // num_attention_heads, # 每個頭的維度
qk_norm="layer_norm_across_heads" if qk_norm else None, # 如果 qk_norm 為真，使用跨頭層歸一化
heads=num_attention_heads, # 注意力頭的數量
kv_heads=num_kv_heads, # 鍵值對的頭數量
eps=1e-5, # 數值穩定性引數
bias=False, # 不使用偏置項
out_bias=False, # 輸出層不使用偏置項
processor=LuminaAttnProcessor2_0(), # 使用指定的注意力處理器
)
self.feed_forward = LuminaFeedForward( # 建立前饋神經網路層
dim=dim, # 輸入維度
inner_dim=4 * dim, # 內部維度，通常為輸入維度的四倍
multiple_of=multiple_of, # 確保內部維度是某個數字的倍數
ffn_dim_multiplier=ffn_dim_multiplier, # 前饋網路維度的乘數
)
self.norm1 = LuminaRMSNormZero( # 建立第一個 RMS 歸一化層
embedding_dim=dim, # 歸一化的嵌入維度
norm_eps=norm_eps, # 歸一化的 epsilon 引數
norm_elementwise_affine=norm_elementwise_affine, # 是否使用元素級仿射變換
)
self.ffn_norm1 = RMSNorm(dim, eps=norm_eps, elementwise_affine=norm_elementwise_affine) # 建立前饋網路的 RMS 歸一化層
self.norm2 = RMSNorm(dim, eps=norm_eps, elementwise_affine=norm_elementwise_affine) # 建立第二個 RMS 歸一化層
self.ffn_norm2 = RMSNorm(dim, eps=norm_eps, elementwise_affine=norm_elementwise_affine) # 建立前饋網路的第二個 RMS 歸一化層
self.norm1_context = RMSNorm(cross_attention_dim, eps=norm_eps, elementwise_affine=norm_elementwise_affine) # 建立上下文的 RMS 歸一化層
def forward( # 定義前向傳播方法
self,
hidden_states: torch.Tensor, # 輸入的隱藏狀態張量
attention_mask: torch.Tensor, # 注意力掩碼張量
image_rotary_emb: torch.Tensor, # 影像旋轉嵌入張量
encoder_hidden_states: torch.Tensor, # 編碼器的隱藏狀態張量
encoder_mask: torch.Tensor, # 編碼器的掩碼張量
temb: torch.Tensor, # 位置編碼或時間編碼張量
cross_attention_kwargs: Optional[Dict[str, Any]] = None, # 可選的交叉注意力引數字典
):
"""
執行 LuminaNextDiTBlock 的前向傳遞。
引數:
hidden_states (`torch.Tensor`): LuminaNextDiTBlock 的輸入隱藏狀態。
attention_mask (`torch.Tensor): 對應隱藏狀態的注意力掩碼。
image_rotary_emb (`torch.Tensor`): 預計算的餘弦和正弦頻率。
encoder_hidden_states: (`torch.Tensor`): 透過 Gemma 編碼器處理的文字提示的隱藏狀態。
encoder_mask (`torch.Tensor`): 文字提示的隱藏狀態注意力掩碼。
temb (`torch.Tensor`): 帶有文字提示嵌入的時間步嵌入。
cross_attention_kwargs (`Dict[str, Any]`): 交叉注意力的引數。
"""
# 儲存輸入的隱藏狀態，以便後續使用
residual = hidden_states
# 自注意力
# 對隱藏狀態進行歸一化，並計算門控機制的輸出
norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb)
# 計算自注意力的輸出
self_attn_output = self.attn1(
hidden_states=norm_hidden_states,
encoder_hidden_states=norm_hidden_states,
attention_mask=attention_mask,
query_rotary_emb=image_rotary_emb,
key_rotary_emb=image_rotary_emb,
**cross_attention_kwargs,
)
# 交叉注意力
# 對編碼器的隱藏狀態進行歸一化
norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states)
# 計算交叉注意力的輸出
cross_attn_output = self.attn2(
hidden_states=norm_hidden_states,
encoder_hidden_states=norm_encoder_hidden_states,
attention_mask=encoder_mask,
query_rotary_emb=image_rotary_emb,
key_rotary_emb=None,
**cross_attention_kwargs,
)
# 將交叉注意力的輸出進行縮放
cross_attn_output = cross_attn_output * self.gate.tanh().view(1, 1, -1, 1)
# 將自注意力和交叉注意力的輸出混合
mixed_attn_output = self_attn_output + cross_attn_output
# 將混合輸出展平，以便後續處理
mixed_attn_output = mixed_attn_output.flatten(-2)
# 線性投影
# 透過線性層處理混合輸出，得到新的隱藏狀態
hidden_states = self.attn2.to_out[0](mixed_attn_output)
# 更新隱藏狀態，加入殘差連線和門控機制
hidden_states = residual + gate_msa.unsqueeze(1).tanh() * self.norm2(hidden_states)
# 透過前饋網路計算輸出
mlp_output = self.feed_forward(self.ffn_norm1(hidden_states) * (1 + scale_mlp.unsqueeze(1)))
# 更新隱藏狀態，加入前饋網路輸出和門控機制
hidden_states = hidden_states + gate_mlp.unsqueeze(1).tanh() * self.ffn_norm2(mlp_output)
# 返回最終的隱藏狀態
return hidden_states
# 定義一個名為 LuminaNextDiT2DModel 的類，繼承自 ModelMixin 和 ConfigMixin
class LuminaNextDiT2DModel(ModelMixin, ConfigMixin):
"""
LuminaNextDiT: 使用 Transformer 主幹的擴散模型。
繼承 ModelMixin 和 ConfigMixin 以相容 diffusers 的 StableDiffusionPipeline 取樣器。
引數：
sample_size (`int`): 潛在影像的寬度。此值在訓練期間固定，因為
它用於學習位置嵌入的數量。
patch_size (`int`, *optional*, (`int`, *optional*, defaults to 2):
影像中每個補丁的大小。此引數定義輸入到模型中的補丁的解析度。
in_channels (`int`, *optional*, defaults to 4):
模型的輸入通道數量。通常，這與輸入影像的通道數量匹配。
hidden_size (`int`, *optional*, defaults to 4096):
模型隱藏層的維度。此引數決定了模型隱藏表示的寬度。
num_layers (`int`, *optional*, default to 32):
模型中的層數。此值定義了神經網路的深度。
num_attention_heads (`int`, *optional*, defaults to 32):
每個注意力層中的注意力頭數量。此引數指定使用多少個獨立的注意力機制。
num_kv_heads (`int`, *optional*, defaults to 8):
注意力機制中的鍵值頭數量，如果與注意力頭數量不同。如果為 None，則預設值為 num_attention_heads。
multiple_of (`int`, *optional*, defaults to 256):
隱藏大小應該是一個倍數的因子。這可以幫助最佳化某些硬體
配置。
ffn_dim_multiplier (`float`, *optional*):
前饋網路維度的乘數。如果為 None，則使用基於
模型配置的預設值。
norm_eps (`float`, *optional*, defaults to 1e-5):
新增到歸一化層的分母中的一個小值，用於數值穩定性。
learn_sigma (`bool`, *optional*, defaults to True):
模型是否應該學習 sigma 引數，該引數可能與預測中的不確定性或方差相關。
qk_norm (`bool`, *optional*, defaults to True):
指示注意力機制中的查詢和鍵是否應該被歸一化。
cross_attention_dim (`int`, *optional*, defaults to 2048):
文字嵌入的維度。此引數定義了用於模型的文字表示的大小。
scaling_factor (`float`, *optional*, defaults to 1.0):
應用於模型某些引數或層的縮放因子。此引數可用於調整模型操作的整體規模。
"""
# 註冊到配置
@register_to_config
def __init__(
# 樣本大小，預設值為128
self,
sample_size: int = 128,
# 補丁大小，預設為2，表示影像切割塊的大小
patch_size: Optional[int] = 2,
# 輸入通道數，預設為4，表示輸入資料的特徵通道
in_channels: Optional[int] = 4,
# 隱藏層大小，預設為2304
hidden_size: Optional[int] = 2304,
# 網路層數，預設為32
num_layers: Optional[int] = 32,
# 注意力頭數量，預設為32
num_attention_heads: Optional[int] = 32,
# KV頭的數量，預設為None
num_kv_heads: Optional[int] = None,
# 數量的倍數，預設為256
multiple_of: Optional[int] = 256,
# FFN維度乘數，預設為None
ffn_dim_multiplier: Optional[float] = None,
# 歸一化的epsilon值，預設為1e-5
norm_eps: Optional[float] = 1e-5,
# 是否學習方差，預設為True
learn_sigma: Optional[bool] = True,
# 是否進行QK歸一化，預設為True
qk_norm: Optional[bool] = True,
# 交叉注意力維度，預設為2048
cross_attention_dim: Optional[int] = 2048,
# 縮放因子，預設為1.0
scaling_factor: Optional[float] = 1.0,
) -> None:
# 呼叫父類初始化方法
super().__init__()
# 設定樣本大小屬性
self.sample_size = sample_size
# 設定補丁大小屬性
self.patch_size = patch_size
# 設定輸入通道數屬性
self.in_channels = in_channels
# 根據是否學習方差設定輸出通道數
self.out_channels = in_channels * 2 if learn_sigma else in_channels
# 設定隱藏層大小屬性
self.hidden_size = hidden_size
# 設定注意力頭數量屬性
self.num_attention_heads = num_attention_heads
# 計算並設定每個注意力頭的維度
self.head_dim = hidden_size // num_attention_heads
# 設定縮放因子屬性
self.scaling_factor = scaling_factor
# 建立補丁嵌入層，並初始化其引數
self.patch_embedder = LuminaPatchEmbed(
patch_size=patch_size, in_channels=in_channels, embed_dim=hidden_size, bias=True
)
# 建立一個可學習的填充標記，初始化為空張量
self.pad_token = nn.Parameter(torch.empty(hidden_size))
# 建立時間和標題的組合嵌入層
self.time_caption_embed = LuminaCombinedTimestepCaptionEmbedding(
hidden_size=min(hidden_size, 1024), cross_attention_dim=cross_attention_dim
)
# 建立包含多個層的模組列表
self.layers = nn.ModuleList(
[
# 在模組列表中新增多個下一代塊
LuminaNextDiTBlock(
hidden_size,
num_attention_heads,
num_kv_heads,
multiple_of,
ffn_dim_multiplier,
norm_eps,
qk_norm,
cross_attention_dim,
)
for _ in range(num_layers) # 根據層數迴圈
]
)
# 建立層歸一化輸出層
self.norm_out = LuminaLayerNormContinuous(
embedding_dim=hidden_size,
conditioning_embedding_dim=min(hidden_size, 1024),
elementwise_affine=False,
eps=1e-6,
bias=True,
out_dim=patch_size * patch_size * self.out_channels,
)
# 註釋掉的最終層的初始化（若需要可取消註釋）
# self.final_layer = LuminaFinalLayer(hidden_size, patch_size, self.out_channels)
# 確保隱藏層大小與注意力頭數量的關係，保證為4的倍數
assert (hidden_size // num_attention_heads) % 4 == 0, "2d rope needs head dim to be divisible by 4"
# 前向傳播函式定義
def forward(
# 隱藏狀態的輸入張量
self,
hidden_states: torch.Tensor,
# 時間步的輸入張量
timestep: torch.Tensor,
# 編碼器的隱藏狀態張量
encoder_hidden_states: torch.Tensor,
# 編碼器的掩碼張量
encoder_mask: torch.Tensor,
# 影像的旋轉嵌入張量
image_rotary_emb: torch.Tensor,
# 交叉注意力的其他引數，預設為None
cross_attention_kwargs: Dict[str, Any] = None,
# 是否返回字典形式的結果，預設為True
return_dict=True,
# LuminaNextDiT 的前向傳播函式
) -> torch.Tensor:
"""
前向傳播的 LuminaNextDiT 模型。
引數:
hidden_states (torch.Tensor): 輸入張量，形狀為 (N, C, H, W)。
timestep (torch.Tensor): 擴散時間步的張量，形狀為 (N,).
encoder_hidden_states (torch.Tensor): 描述特徵的張量，形狀為 (N, D)。
encoder_mask (torch.Tensor): 描述特徵掩碼的張量，形狀為 (N, L)。
"""
# 透過補丁嵌入器處理隱藏狀態，獲取掩碼、影像大小和影像旋轉嵌入
hidden_states, mask, img_size, image_rotary_emb = self.patch_embedder(hidden_states, image_rotary_emb)
# 將影像旋轉嵌入轉移到與隱藏狀態相同的裝置上
image_rotary_emb = image_rotary_emb.to(hidden_states.device)
# 生成時間嵌入，結合時間步和編碼器隱藏狀態
temb = self.time_caption_embed(timestep, encoder_hidden_states, encoder_mask)
# 將編碼器掩碼轉換為布林值
encoder_mask = encoder_mask.bool()
# 對每一層進行遍歷，更新隱藏狀態
for layer in self.layers:
hidden_states = layer(
hidden_states,
mask,
image_rotary_emb,
encoder_hidden_states,
encoder_mask,
temb=temb,
cross_attention_kwargs=cross_attention_kwargs,
)
# 對隱藏狀態進行歸一化處理
hidden_states = self.norm_out(hidden_states, temb)
# 反補丁操作
height_tokens = width_tokens = self.patch_size # 獲取補丁大小
height, width = img_size[0] # 從影像大小中提取高度和寬度
batch_size = hidden_states.size(0) # 獲取批次大小
sequence_length = (height // height_tokens) * (width // width_tokens) # 計算序列長度
# 調整隱藏狀態的形狀，以適應輸出要求
hidden_states = hidden_states[:, :sequence_length].view(
batch_size, height // height_tokens, width // width_tokens, height_tokens, width_tokens, self.out_channels
)
# 調整維度以獲得最終輸出
output = hidden_states.permute(0, 5, 1, 3, 2, 4).flatten(4, 5).flatten(2, 3)
# 如果不需要返回字典，則返回輸出元組
if not return_dict:
return (output,)
# 返回 Transformer2DModelOutput 的結果
return Transformer2DModelOutput(sample=output)
# 版權所有 2024 HuggingFace 團隊。保留所有權利。
#
# 根據 Apache 許可證第 2.0 版（“許可證”）進行許可；
# 除非遵守許可證，否則您不得使用此檔案。
# 您可以在以下網址獲取許可證副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非適用法律或書面協議另有規定，軟體
# 按“原樣”提供，沒有任何形式的明示或暗示的擔保或條件。
# 請參閱許可證以瞭解有關許可權和
# 限制的具體條款。
from typing import Any, Dict, Optional, Union # 匯入型別提示相關的模組
import torch # 匯入 PyTorch 庫
from torch import nn # 從 PyTorch 匯入神經網路模組
from ...configuration_utils import ConfigMixin, register_to_config # 匯入配置相關的混合類和註冊函式
from ...utils import is_torch_version, logging # 匯入工具函式：檢查 PyTorch 版本和日誌記錄
from ..attention import BasicTransformerBlock # 匯入基礎 Transformer 塊
from ..attention_processor import Attention, AttentionProcessor, FusedAttnProcessor2_0 # 匯入注意力相關的處理器
from ..embeddings import PatchEmbed, PixArtAlphaTextProjection # 匯入嵌入相關的模組
from ..modeling_outputs import Transformer2DModelOutput # 匯入模型輸出相關的類
from ..modeling_utils import ModelMixin # 匯入模型混合類
from ..normalization import AdaLayerNormSingle # 匯入自適應層歸一化類
logger = logging.get_logger(__name__) # 獲取當前模組的日誌記錄器；pylint 禁用命名檢查
class PixArtTransformer2DModel(ModelMixin, ConfigMixin): # 定義 PixArt 2D Transformer 模型類，繼承自 ModelMixin 和 ConfigMixin
r""" # 文件字串：描述模型及其來源
A 2D Transformer model as introduced in PixArt family of models (https://arxiv.org/abs/2310.00426,
https://arxiv.org/abs/2403.04692).
"""
_supports_gradient_checkpointing = True # 設定支援梯度檢查點
_no_split_modules = ["BasicTransformerBlock", "PatchEmbed"] # 指定不進行分割的模組
@register_to_config # 使用裝飾器將初始化函式註冊到配置中
def __init__( # 定義初始化函式
self,
num_attention_heads: int = 16, # 注意力頭的數量，預設為 16
attention_head_dim: int = 72, # 每個注意力頭的維度，預設為 72
in_channels: int = 4, # 輸入通道數，預設為 4
out_channels: Optional[int] = 8, # 輸出通道數，預設為 8，可選
num_layers: int = 28, # 層數，預設為 28
dropout: float = 0.0, # dropout 比例，預設為 0.0
norm_num_groups: int = 32, # 歸一化的組數，預設為 32
cross_attention_dim: Optional[int] = 1152, # 交叉注意力的維度，預設為 1152，可選
attention_bias: bool = True, # 是否使用注意力偏置，預設為 True
sample_size: int = 128, # 樣本尺寸，預設為 128
patch_size: int = 2, # 每個補丁的尺寸，預設為 2
activation_fn: str = "gelu-approximate", # 啟用函式型別，預設為近似 GELU
num_embeds_ada_norm: Optional[int] = 1000, # 自適應歸一化的嵌入數量，預設為 1000，可選
upcast_attention: bool = False, # 是否提高注意力精度，預設為 False
norm_type: str = "ada_norm_single", # 歸一化型別，預設為單一自適應歸一化
norm_elementwise_affine: bool = False, # 是否使用逐元素仿射變換，預設為 False
norm_eps: float = 1e-6, # 歸一化的 epsilon 值，預設為 1e-6
interpolation_scale: Optional[int] = None, # 插值尺度，可選
use_additional_conditions: Optional[bool] = None, # 是否使用額外條件，可選
caption_channels: Optional[int] = None, # 說明通道數，可選
attention_type: Optional[str] = "default", # 注意力型別，預設為預設型別
):
# 初始化函式引數設定
...
def _set_gradient_checkpointing(self, module, value=False): # 定義設定梯度檢查點的方法
if hasattr(module, "gradient_checkpointing"): # 檢查模組是否具有梯度檢查點屬性
module.gradient_checkpointing = value # 設定梯度檢查點的值
@property # 定義一個屬性
# 從 diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors 複製的屬性
# 定義一個方法，返回模型中所有注意力處理器的字典，鍵為權重名稱
def attn_processors(self) -> Dict[str, AttentionProcessor]:
r"""
Returns:
`dict` of attention processors: A dictionary containing all attention processors used in the model with
indexed by its weight name.
"""
# 建立一個空字典，用於儲存注意力處理器
processors = {}
# 定義一個遞迴函式，用於新增處理器到字典中
def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
# 如果模組具有獲取處理器的方法，則將其新增到字典中
if hasattr(module, "get_processor"):
processors[f"{name}.processor"] = module.get_processor()
# 遍歷子模組，遞迴呼叫該函式以新增處理器
for sub_name, child in module.named_children():
fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
# 返回更新後的處理器字典
return processors
# 遍歷當前模組的所有子模組，呼叫遞迴函式以填充處理器字典
for name, module in self.named_children():
fn_recursive_add_processors(name, module, processors)
# 返回所有注意力處理器的字典
return processors
# 從 diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor 複製
# 定義一個方法，用於設定計算注意力的處理器
def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
r"""
Sets the attention processor to use to compute attention.
Parameters:
processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
The instantiated processor class or a dictionary of processor classes that will be set as the processor
for **all** `Attention` layers.
If `processor` is a dict, the key needs to define the path to the corresponding cross attention
processor. This is strongly recommended when setting trainable attention processors.
"""
# 獲取當前注意力處理器的數量
count = len(self.attn_processors.keys())
# 檢查傳入的處理器字典的長度是否與注意力層的數量匹配
if isinstance(processor, dict) and len(processor) != count:
raise ValueError(
f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
)
# 定義一個遞迴函式，用於設定處理器
def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
# 如果模組具有設定處理器的方法，則進行設定
if hasattr(module, "set_processor"):
if not isinstance(processor, dict):
module.set_processor(processor) # 設定單一處理器
else:
module.set_processor(processor.pop(f"{name}.processor")) # 從字典中移除並設定處理器
# 遍歷子模組，遞迴呼叫以設定處理器
for sub_name, child in module.named_children():
fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
# 遍歷當前模組的所有子模組，呼叫遞迴函式以設定處理器
for name, module in self.named_children():
fn_recursive_attn_processor(name, module, processor)
# 從 diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections 複製
# 定義融合 QKV 投影的函式
def fuse_qkv_projections(self):
# 啟用融合的 QKV 投影，對自注意力模組進行融合查詢、鍵、值矩陣
# 對交叉注意力模組則僅融合鍵和值投影矩陣
"""
Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
are fused. For cross-attention modules, key and value projection matrices are fused.
<Tip warning={true}>
This API is 🧪 experimental.
</Tip>
"""
# 初始化原始注意力處理器為 None
self.original_attn_processors = None
# 遍歷所有注意力處理器
for _, attn_processor in self.attn_processors.items():
# 檢查處理器類名中是否包含 "Added"
if "Added" in str(attn_processor.__class__.__name__):
# 如果存在，則丟擲錯誤，說明不支援此融合操作
raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
# 儲存當前的注意力處理器
self.original_attn_processors = self.attn_processors
# 遍歷模型中的所有模組
for module in self.modules():
# 如果模組是 Attention 型別
if isinstance(module, Attention):
# 執行投影融合
module.fuse_projections(fuse=True)
# 設定新的融合注意力處理器
self.set_attn_processor(FusedAttnProcessor2_0())
# 從 UNet2DConditionModel 中複製的函式，用於取消融合 QKV 投影
def unfuse_qkv_projections(self):
# 禁用已啟用的融合 QKV 投影
"""Disables the fused QKV projection if enabled.
<Tip warning={true}>
This API is 🧪 experimental.
</Tip>
"""
# 檢查原始注意力處理器是否存在
if self.original_attn_processors is not None:
# 恢復到原始注意力處理器
self.set_attn_processor(self.original_attn_processors)
# 定義前向傳播函式
def forward(
# 輸入隱藏狀態的張量
hidden_states: torch.Tensor,
# 編碼器隱藏狀態（可選）
encoder_hidden_states: Optional[torch.Tensor] = None,
# 時間步長（可選）
timestep: Optional[torch.LongTensor] = None,
# 新增的條件關鍵字引數（字典型別，可選）
added_cond_kwargs: Dict[str, torch.Tensor] = None,
# 交叉注意力關鍵字引數（字典型別，可選）
cross_attention_kwargs: Dict[str, Any] = None,
# 注意力掩碼（可選）
attention_mask: Optional[torch.Tensor] = None,
# 編碼器注意力掩碼（可選）
encoder_attention_mask: Optional[torch.Tensor] = None,
# 是否返回字典（預設值為 True）
return_dict: bool = True,
# 從 dataclasses 模組匯入 dataclass 裝飾器
from dataclasses import dataclass
# 匯入字典、可選值和聯合型別的定義
from typing import Dict, Optional, Union
# 匯入 PyTorch 及其功能模組
import torch
import torch.nn.functional as F
# 從 PyTorch 匯入神經網路模組
from torch import nn
# 匯入配置和註冊功能的相關類
from ...configuration_utils import ConfigMixin, register_to_config
# 匯入 PeftAdapter 和 UNet2DConditionLoader 的相關類
from ...loaders import PeftAdapterMixin, UNet2DConditionLoadersMixin
# 匯入基礎輸出工具類
from ...utils import BaseOutput
# 匯入基本變換器塊
from ..attention import BasicTransformerBlock
# 匯入注意力處理器的相關元件
from ..attention_processor import (
ADDED_KV_ATTENTION_PROCESSORS,
CROSS_ATTENTION_PROCESSORS,
AttentionProcessor,
AttnAddedKVProcessor,
AttnProcessor,
)
# 匯入時間步嵌入和時間步類
from ..embeddings import TimestepEmbedding, Timesteps
# 匯入模型混合工具類
from ..modeling_utils import ModelMixin
# 定義 PriorTransformerOutput 資料類，繼承自 BaseOutput
@dataclass
class PriorTransformerOutput(BaseOutput):
"""
[`PriorTransformer`] 的輸出。
Args:
predicted_image_embedding (`torch.Tensor` 的形狀為 `(batch_size, embedding_dim)`):
基於 CLIP 文字嵌入輸入的預測 CLIP 影像嵌入。
"""
# 定義預測的影像嵌入屬性
predicted_image_embedding: torch.Tensor
# 定義 PriorTransformer 類，繼承多個混合類
class PriorTransformer(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin, PeftAdapterMixin):
"""
一種 Prior Transformer 模型。
# 引數說明部分
Parameters:
# 用於多頭注意力的頭數量，預設為 32
num_attention_heads (`int`, *optional*, defaults to 32): The number of heads to use for multi-head attention.
# 每個頭的通道數量，預設為 64
attention_head_dim (`int`, *optional*, defaults to 64): The number of channels in each head.
# Transformer 塊的層數，預設為 20
num_layers (`int`, *optional*, defaults to 20): The number of layers of Transformer blocks to use.
# 模型輸入 `hidden_states` 的維度，預設為 768
embedding_dim (`int`, *optional*, defaults to 768): The dimension of the model input `hidden_states`
# 模型輸入 `hidden_states` 的嵌入數量，預設為 77
num_embeddings (`int`, *optional*, defaults to 77):
The number of embeddings of the model input `hidden_states`
# 附加令牌的數量，預設為 4，追加到投影的 `hidden_states`
additional_embeddings (`int`, *optional*, defaults to 4): The number of additional tokens appended to the
projected `hidden_states`. The actual length of the used `hidden_states` is `num_embeddings +
additional_embeddings`.
# 用於 dropout 的機率，預設為 0.0
dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
# 建立時間步嵌入時使用的啟用函式，預設為 'silu'
time_embed_act_fn (`str`, *optional*, defaults to 'silu'):
The activation function to use to create timestep embeddings.
# 在傳遞給 Transformer 塊之前應用的歸一化層，預設為 None
norm_in_type (`str`, *optional*, defaults to None): The normalization layer to apply on hidden states before
passing to Transformer blocks. Set it to `None` if normalization is not needed.
# 輸入 `proj_embedding` 上應用的歸一化層，預設為 None
embedding_proj_norm_type (`str`, *optional*, defaults to None):
The normalization layer to apply on the input `proj_embedding`. Set it to `None` if normalization is not
needed.
# 輸入 `encoder_hidden_states` 上應用的投影層，預設為 `linear`
encoder_hid_proj_type (`str`, *optional*, defaults to `linear`):
The projection layer to apply on the input `encoder_hidden_states`. Set it to `None` if
`encoder_hidden_states` is `None`.
# 條件模型的附加嵌入型別，預設為 `prd`
added_emb_type (`str`, *optional*, defaults to `prd`): Additional embeddings to condition the model.
Choose from `prd` or `None`. if choose `prd`, it will prepend a token indicating the (quantized) dot
product between the text embedding and image embedding as proposed in the unclip paper
https://arxiv.org/abs/2204.06125 If it is `None`, no additional embeddings will be prepended.
# 時間步嵌入的維度，預設為 None，如果為 None，則設定為 `num_attention_heads * attention_head_dim`
time_embed_dim (`int, *optional*, defaults to None): The dimension of timestep embeddings.
If None, will be set to `num_attention_heads * attention_head_dim`
# `proj_embedding` 的維度，預設為 None，如果為 None，則設定為 `embedding_dim`
embedding_proj_dim (`int`, *optional*, default to None):
The dimension of `proj_embedding`. If None, will be set to `embedding_dim`.
# 輸出的維度，預設為 None，如果為 None，則設定為 `embedding_dim`
clip_embed_dim (`int`, *optional*, default to None):
The dimension of the output. If None, will be set to `embedding_dim`.
"""
# 註冊到配置中
@register_to_config
# 初始化類的建構函式，設定預設引數
def __init__(
# 注意力頭的數量，預設值為32
self,
num_attention_heads: int = 32,
# 每個注意力頭的維度，預設值為64
attention_head_dim: int = 64,
# 層的數量，預設值為20
num_layers: int = 20,
# 嵌入的維度，預設值為768
embedding_dim: int = 768,
# 嵌入的數量，預設值為77
num_embeddings=77,
# 額外嵌入的數量，預設值為4
additional_embeddings=4,
# dropout的比率，預設值為0.0
dropout: float = 0.0,
# 時間嵌入啟用函式的型別，預設值為"silu"
time_embed_act_fn: str = "silu",
# 輸入歸一化型別，預設為None
norm_in_type: Optional[str] = None, # layer
# 嵌入投影歸一化型別，預設為None
embedding_proj_norm_type: Optional[str] = None, # layer
# 編碼器隱藏投影型別，預設值為"linear"
encoder_hid_proj_type: Optional[str] = "linear", # linear
# 新增的嵌入型別，預設值為"prd"
added_emb_type: Optional[str] = "prd", # prd
# 時間嵌入維度，預設為None
time_embed_dim: Optional[int] = None,
# 嵌入投影維度，預設為None
embedding_proj_dim: Optional[int] = None,
# 裁剪嵌入維度，預設為None
clip_embed_dim: Optional[int] = None,
# 定義一個屬性，獲取注意力處理器
@property
# 從diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors複製而來
def attn_processors(self) -> Dict[str, AttentionProcessor]:
r"""
返回值：
`dict`型別的注意力處理器：一個字典，包含模型中使用的所有注意力處理器，並按其權重名稱索引。
"""
# 建立一個空字典用於儲存處理器
processors = {}
# 定義遞迴函式，新增處理器到字典
def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
# 如果模組有獲取處理器的方法，新增到字典中
if hasattr(module, "get_processor"):
processors[f"{name}.processor"] = module.get_processor()
# 遍歷模組的所有子模組，遞迴呼叫自身
for sub_name, child in module.named_children():
fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
# 返回更新後的處理器字典
return processors
# 遍歷當前類的所有子模組，呼叫遞迴函式新增處理器
for name, module in self.named_children():
fn_recursive_add_processors(name, module, processors)
# 返回最終的處理器字典
return processors
# 從diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor複製而來
# 設定用於計算注意力的處理器
def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
r"""
設定用於計算注意力的處理器。
引數：
processor（`dict` 或 `AttentionProcessor`）：
例項化的處理器類或處理器類的字典，將作為 **所有** `Attention` 層的處理器。
如果 `processor` 是一個字典，鍵需要定義相應的交叉注意力處理器的路徑。建議在設定可訓練注意力處理器時使用此方法。
"""
# 計算當前注意力處理器的數量
count = len(self.attn_processors.keys())
# 如果傳入的是字典且其長度與注意力層數量不匹配，丟擲錯誤
if isinstance(processor, dict) and len(processor) != count:
raise ValueError(
f"傳入了處理器的字典，但處理器的數量 {len(processor)} 與注意力層的數量 {count} 不匹配。請確保傳入 {count} 個處理器類。"
)
# 定義遞迴設定注意力處理器的函式
def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
# 如果模組具有 set_processor 方法，則設定處理器
if hasattr(module, "set_processor"):
if not isinstance(processor, dict):
module.set_processor(processor)
else:
module.set_processor(processor.pop(f"{name}.processor"))
# 遍歷子模組並遞迴呼叫
for sub_name, child in module.named_children():
fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
# 遍歷所有子模組，呼叫遞迴函式設定處理器
for name, module in self.named_children():
fn_recursive_attn_processor(name, module, processor)
# 從 diffusers.models.unets.unet_2d_condition.UNet2DConditionModel 複製的設定預設注意力處理器的方法
def set_default_attn_processor(self):
"""
禁用自定義注意力處理器並設定預設的注意力實現。
"""
# 如果所有處理器都是新增的 KV 注意力處理器，則設定為新增的 KV 處理器
if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
processor = AttnAddedKVProcessor()
# 如果所有處理器都是交叉注意力處理器，則設定為普通的注意力處理器
elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
processor = AttnProcessor()
else:
raise ValueError(
f"當注意力處理器的型別為 {next(iter(self.attn_processors.values()))} 時，無法呼叫 `set_default_attn_processor`"
)
# 呼叫設定處理器的方法
self.set_attn_processor(processor)
# 前向傳播方法定義
def forward(
self,
hidden_states,
timestep: Union[torch.Tensor, float, int],
proj_embedding: torch.Tensor,
encoder_hidden_states: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.BoolTensor] = None,
return_dict: bool = True,
):
# 處理傳入的潛在變數
def post_process_latents(self, prior_latents):
# 將潛在變數進行標準化處理
prior_latents = (prior_latents * self.clip_std) + self.clip_mean
return prior_latents
# 版權宣告，註明版權歸屬
# Copyright 2024 Stability AI and The HuggingFace Team. All rights reserved.
#
# 根據 Apache License 2.0 授權協議進行許可
# Licensed under the Apache License, Version 2.0 (the "License");
# 只有在遵守許可證的情況下，您才能使用此檔案
# you may not use this file except in compliance with the License.
# 您可以在以下網址獲取許可證副本
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非法律要求或書面同意，軟體按“原樣”分發
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 不提供任何形式的保證或條件，明示或暗示
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# 請參見許可證以獲取有關許可權和限制的具體資訊
# See the License for the specific language governing permissions and
# limitations under the License.
# 匯入型別註解
from typing import Any, Dict, Optional, Union
# 匯入 NumPy 庫
import numpy as np
# 匯入 PyTorch 庫
import torch
# 匯入 PyTorch 的神經網路模組
import torch.nn as nn
# 匯入 PyTorch 的檢查點工具
import torch.utils.checkpoint
# 從配置工具匯入相關類
from ...configuration_utils import ConfigMixin, register_to_config
# 從注意力模組匯入前饋網路
from ...models.attention import FeedForward
# 從注意力處理器模組匯入多個類
from ...models.attention_processor import (
Attention,
AttentionProcessor,
StableAudioAttnProcessor2_0,
)
# 從建模工具匯入模型混合類
from ...models.modeling_utils import ModelMixin
# 從變換器模型匯入輸出類
from ...models.transformers.transformer_2d import Transformer2DModelOutput
# 匯入實用工具
from ...utils import is_torch_version, logging
# 匯入可能允許圖形中的工具函式
from ...utils.torch_utils import maybe_allow_in_graph
# 獲取當前模組的日誌記錄器
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
class StableAudioGaussianFourierProjection(nn.Module):
"""用於噪聲級別的高斯傅立葉嵌入。"""
# 從 diffusers.models.embeddings.GaussianFourierProjection.__init__ 複製的內容
def __init__(
self, embedding_size: int = 256, scale: float = 1.0, set_W_to_weight=True, log=True, flip_sin_to_cos=False
):
super().__init__() # 呼叫父類的建構函式
# 初始化權重為隨機值，且不需要計算梯度
self.weight = nn.Parameter(torch.randn(embedding_size) * scale, requires_grad=False)
self.log = log # 是否對輸入取對數的標誌
self.flip_sin_to_cos = flip_sin_to_cos # 是否翻轉正弦和餘弦的順序
if set_W_to_weight: # 如果設定將 W 賦值給權重
# 之後將刪除此行
del self.weight # 刪除原有權重
# 初始化 W 為隨機值，並不計算梯度
self.W = nn.Parameter(torch.randn(embedding_size) * scale, requires_grad=False)
self.weight = self.W # 將 W 賦值給權重
del self.W # 刪除 W
def forward(self, x):
# 如果 log 為 True，則對輸入進行對數變換
if self.log:
x = torch.log(x)
# 計算投影，使用 2π 乘以輸入和權重的外積
x_proj = 2 * np.pi * x[:, None] @ self.weight[None, :]
if self.flip_sin_to_cos: # 如果翻轉正弦和餘弦
# 連線餘弦和正弦，形成輸出
out = torch.cat([torch.cos(x_proj), torch.sin(x_proj)], dim=-1)
else:
# 連線正弦和餘弦，形成輸出
out = torch.cat([torch.sin(x_proj), torch.cos(x_proj)], dim=-1)
return out # 返回輸出
@maybe_allow_in_graph # 可能允許在計算圖中使用
class StableAudioDiTBlock(nn.Module):
r"""
用於穩定音訊模型的變換器塊 (https://github.com/Stability-AI/stable-audio-tools)。允許跳躍連線和 QKNorm
# 引數說明
Parameters:
dim (`int`): 輸入和輸出的通道數。
num_attention_heads (`int`): 查詢狀態所使用的頭數。
num_key_value_attention_heads (`int`): 鍵和值狀態所使用的頭數。
attention_head_dim (`int`): 每個頭中的通道數。
dropout (`float`, *optional*, defaults to 0.0): 使用的丟棄機率。
cross_attention_dim (`int`, *optional*): 跨注意力的 encoder_hidden_states 向量的大小。
upcast_attention (`bool`, *optional*):
是否將注意力計算上升到 float32。這對混合精度訓練很有用。
"""
# 初始化函式
def __init__(
self,
dim: int, # 輸入和輸出的通道數
num_attention_heads: int, # 查詢狀態的頭數
num_key_value_attention_heads: int, # 鍵和值狀態的頭數
attention_head_dim: int, # 每個頭的通道數
dropout=0.0, # 丟棄機率，預設為0
cross_attention_dim: Optional[int] = None, # 跨注意力的維度，可選
upcast_attention: bool = False, # 是否上升到 float32，預設為 False
norm_eps: float = 1e-5, # 歸一化層的小常數
ff_inner_dim: Optional[int] = None, # 前饋層內部維度，可選
):
super().__init__() # 呼叫父類建構函式
# 定義三個模組。每個模組都有自己的歸一化層。
# 1. 自注意力層
self.norm1 = nn.LayerNorm(dim, elementwise_affine=True, eps=norm_eps) # 自注意力的歸一化層
self.attn1 = Attention( # 自注意力模組
query_dim=dim, # 查詢維度
heads=num_attention_heads, # 頭數
dim_head=attention_head_dim, # 每個頭的維度
dropout=dropout, # 丟棄機率
bias=False, # 不使用偏置
upcast_attention=upcast_attention, # 是否上升到 float32
out_bias=False, # 不使用輸出偏置
processor=StableAudioAttnProcessor2_0(), # 使用的處理器
)
# 2. 跨注意力層
self.norm2 = nn.LayerNorm(dim, norm_eps, True) # 跨注意力的歸一化層
self.attn2 = Attention( # 跨注意力模組
query_dim=dim, # 查詢維度
cross_attention_dim=cross_attention_dim, # 跨注意力維度
heads=num_attention_heads, # 頭數
dim_head=attention_head_dim, # 每個頭的維度
kv_heads=num_key_value_attention_heads, # 鍵和值的頭數
dropout=dropout, # 丟棄機率
bias=False, # 不使用偏置
upcast_attention=upcast_attention, # 是否上升到 float32
out_bias=False, # 不使用輸出偏置
processor=StableAudioAttnProcessor2_0(), # 使用的處理器
) # 如果 encoder_hidden_states 為 None，則為自注意力
# 3. 前饋層
self.norm3 = nn.LayerNorm(dim, norm_eps, True) # 前饋層的歸一化層
self.ff = FeedForward( # 前饋神經網路模組
dim, # 輸入維度
dropout=dropout, # 丟棄機率
activation_fn="swiglu", # 啟用函式
final_dropout=False, # 最後是否丟棄
inner_dim=ff_inner_dim, # 內部維度
bias=True, # 使用偏置
)
# 將塊大小預設設定為 None
self._chunk_size = None # 塊大小
self._chunk_dim = 0 # 塊維度
def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0):
# 設定塊前饋
self._chunk_size = chunk_size # 設定塊大小
self._chunk_dim = dim # 設定塊維度
# 定義前向傳播方法，接收隱藏狀態和可選的注意力掩碼等引數
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
encoder_hidden_states: Optional[torch.Tensor] = None,
encoder_attention_mask: Optional[torch.Tensor] = None,
rotary_embedding: Optional[torch.FloatTensor] = None,
) -> torch.Tensor:
# 注意：在後續計算中，歸一化總是應用於實際計算之前。
# 0. 自注意力
# 對輸入的隱藏狀態進行歸一化處理
norm_hidden_states = self.norm1(hidden_states)
# 計算自注意力輸出
attn_output = self.attn1(
norm_hidden_states,
attention_mask=attention_mask,
rotary_emb=rotary_embedding,
)
# 將自注意力輸出與原始隱藏狀態相加
hidden_states = attn_output + hidden_states
# 2. 跨注意力
# 對更新後的隱藏狀態進行歸一化處理
norm_hidden_states = self.norm2(hidden_states)
# 計算跨注意力輸出
attn_output = self.attn2(
norm_hidden_states,
encoder_hidden_states=encoder_hidden_states,
attention_mask=encoder_attention_mask,
)
# 將跨注意力輸出與更新後的隱藏狀態相加
hidden_states = attn_output + hidden_states
# 3. 前饋網路
# 對隱藏狀態進行歸一化處理
norm_hidden_states = self.norm3(hidden_states)
# 計算前饋網路輸出
ff_output = self.ff(norm_hidden_states)
# 將前饋網路輸出與當前隱藏狀態相加
hidden_states = ff_output + hidden_states
# 返回最終的隱藏狀態
return hidden_states
# 定義一個名為 StableAudioDiTModel 的類，繼承自 ModelMixin 和 ConfigMixin
class StableAudioDiTModel(ModelMixin, ConfigMixin):
"""
Stable Audio 中引入的擴散變換器模型。
參考文獻：https://github.com/Stability-AI/stable-audio-tools
引數：
sample_size ( `int`, *可選*, 預設值為 1024)：輸入樣本的大小。
in_channels (`int`, *可選*, 預設值為 64)：輸入中的通道數。
num_layers (`int`, *可選*, 預設值為 24)：使用的變換器塊的層數。
attention_head_dim (`int`, *可選*, 預設值為 64)：每個頭的通道數。
num_attention_heads (`int`, *可選*, 預設值為 24)：用於查詢狀態的頭數。
num_key_value_attention_heads (`int`, *可選*, 預設值為 12)：
用於鍵和值狀態的頭數。
out_channels (`int`, 預設值為 64)：輸出通道的數量。
cross_attention_dim ( `int`, *可選*, 預設值為 768)：交叉注意力投影的維度。
time_proj_dim ( `int`, *可選*, 預設值為 256)：時間步內投影的維度。
global_states_input_dim ( `int`, *可選*, 預設值為 1536)：
全域性隱藏狀態投影的輸入維度。
cross_attention_input_dim ( `int`, *可選*, 預設值為 768)：
交叉注意力投影的輸入維度。
"""
# 支援梯度檢查點
_supports_gradient_checkpointing = True
# 註冊到配置的建構函式
@register_to_config
def __init__(
# 輸入樣本的大小，預設為1024
self,
sample_size: int = 1024,
# 輸入的通道數，預設為64
in_channels: int = 64,
# 變換器塊的層數，預設為24
num_layers: int = 24,
# 每個頭的通道數，預設為64
attention_head_dim: int = 64,
# 查詢狀態的頭數，預設為24
num_attention_heads: int = 24,
# 鍵和值狀態的頭數，預設為12
num_key_value_attention_heads: int = 12,
# 輸出通道的數量，預設為64
out_channels: int = 64,
# 交叉注意力投影的維度，預設為768
cross_attention_dim: int = 768,
# 時間步內投影的維度，預設為256
time_proj_dim: int = 256,
# 全域性隱藏狀態投影的輸入維度，預設為1536
global_states_input_dim: int = 1536,
# 交叉注意力投影的輸入維度，預設為768
cross_attention_input_dim: int = 768,
):
# 呼叫父類的初始化方法
super().__init__()
# 設定樣本大小
self.sample_size = sample_size
# 設定輸出通道數
self.out_channels = out_channels
# 計算內部維度，等於注意力頭數量乘以每個頭的維度
self.inner_dim = num_attention_heads * attention_head_dim
# 建立穩定音訊高斯傅立葉投影物件，embedding_size 為時間投影維度的一半
self.time_proj = StableAudioGaussianFourierProjection(
embedding_size=time_proj_dim // 2,
flip_sin_to_cos=True, # 是否翻轉正弦和餘弦
log=False, # 是否使用對數
set_W_to_weight=False, # 是否將 W 設定為權重
)
# 時間步投影的神經網路序列，包含兩個線性層和一個啟用函式
self.timestep_proj = nn.Sequential(
nn.Linear(time_proj_dim, self.inner_dim, bias=True), # 輸入為 time_proj_dim，輸出為 inner_dim
nn.SiLU(), # 使用 SiLU 啟用函式
nn.Linear(self.inner_dim, self.inner_dim, bias=True), # 再次投影到 inner_dim
)
# 全域性狀態投影的神經網路序列，包含兩個線性層和一個啟用函式
self.global_proj = nn.Sequential(
nn.Linear(global_states_input_dim, self.inner_dim, bias=False), # 輸入為 global_states_input_dim，輸出為 inner_dim
nn.SiLU(), # 使用 SiLU 啟用函式
nn.Linear(self.inner_dim, self.inner_dim, bias=False), # 再次投影到 inner_dim
)
# 交叉注意力投影的神經網路序列，包含兩個線性層和一個啟用函式
self.cross_attention_proj = nn.Sequential(
nn.Linear(cross_attention_input_dim, cross_attention_dim, bias=False), # 輸入為 cross_attention_input_dim，輸出為 cross_attention_dim
nn.SiLU(), # 使用 SiLU 啟用函式
nn.Linear(cross_attention_dim, cross_attention_dim, bias=False), # 再次投影到 cross_attention_dim
)
# 一維卷積層，用於預處理，卷積核大小為 1，不使用偏置
self.preprocess_conv = nn.Conv1d(in_channels, in_channels, 1, bias=False)
# 輸入線性層，將輸入通道數投影到 inner_dim，不使用偏置
self.proj_in = nn.Linear(in_channels, self.inner_dim, bias=False)
# 建立一個模組列表，包含多個 StableAudioDiTBlock
self.transformer_blocks = nn.ModuleList(
[
StableAudioDiTBlock(
dim=self.inner_dim, # 輸入維度為 inner_dim
num_attention_heads=num_attention_heads, # 注意力頭數量
num_key_value_attention_heads=num_key_value_attention_heads, # 鍵值注意力頭數量
attention_head_dim=attention_head_dim, # 每個注意力頭的維度
cross_attention_dim=cross_attention_dim, # 交叉注意力維度
)
for i in range(num_layers) # 根據層數建立相應數量的塊
]
)
# 輸出線性層，將 inner_dim 投影到輸出通道數，不使用偏置
self.proj_out = nn.Linear(self.inner_dim, self.out_channels, bias=False)
# 一維卷積層，用於後處理，卷積核大小為 1，不使用偏置
self.postprocess_conv = nn.Conv1d(self.out_channels, self.out_channels, 1, bias=False)
# 初始化梯度檢查點標誌，預設為 False
self.gradient_checkpointing = False
@property
# 從 UNet2DConditionModel 複製的屬性
def attn_processors(self) -> Dict[str, AttentionProcessor]:
r"""
Returns:
`dict` of attention processors: A dictionary containing all attention processors used in the model with
indexed by its weight name.
"""
# 建立一個空字典，用於儲存注意力處理器
processors = {}
# 定義遞迴函式，用於新增註意力處理器
def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
# 如果模組有 get_processor 方法，獲取處理器並新增到字典
if hasattr(module, "get_processor"):
processors[f"{name}.processor"] = module.get_processor()
# 遍歷子模組，遞迴呼叫新增處理器函式
for sub_name, child in module.named_children():
fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
return processors
# 遍歷當前物件的所有子模組，呼叫遞迴函式
for name, module in self.named_children():
fn_recursive_add_processors(name, module, processors)
# 返回包含所有處理器的字典
return processors
# 從 diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor 複製而來
def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
r"""
設定用於計算注意力的注意力處理器。
引數：
processor (`dict` of `AttentionProcessor` 或 `AttentionProcessor`):
例項化的處理器類或將作為處理器設定為 **所有** `Attention` 層的處理器類字典。
如果 `processor` 是一個字典，鍵需要定義對應的交叉注意力處理器的路徑。
當設定可訓練的注意力處理器時，這一點強烈推薦。
"""
# 獲取當前注意力處理器的數量
count = len(self.attn_processors.keys())
# 如果傳入的是字典且字典長度與注意力層數量不匹配，則丟擲異常
if isinstance(processor, dict) and len(processor) != count:
raise ValueError(
f"傳入了處理器字典，但處理器數量 {len(processor)} 與"
f" 注意力層數量 {count} 不匹配。請確保傳入 {count} 個處理器類。"
)
# 定義遞迴設定注意力處理器的函式
def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
# 如果模組有 set_processor 方法，則設定處理器
if hasattr(module, "set_processor"):
if not isinstance(processor, dict):
# 如果處理器不是字典，則直接設定
module.set_processor(processor)
else:
# 從字典中彈出對應的處理器並設定
module.set_processor(processor.pop(f"{name}.processor"))
# 遍歷模組的子模組，遞迴呼叫設定處理器的函式
for sub_name, child in module.named_children():
fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
# 遍歷當前物件的子模組，呼叫遞迴設定處理器的函式
for name, module in self.named_children():
fn_recursive_attn_processor(name, module, processor)
# 從 diffusers.models.transformers.hunyuan_transformer_2d.HunyuanDiT2DModel.set_default_attn_processor 複製而來，將 Hunyuan 替換為 StableAudio
def set_default_attn_processor(self):
"""
禁用自定義注意力處理器，並設定預設的注意力實現。
"""
# 呼叫設定注意力處理器的方法，使用 StableAudioAttnProcessor2_0 例項
self.set_attn_processor(StableAudioAttnProcessor2_0())
# 設定梯度檢查點的私有方法
def _set_gradient_checkpointing(self, module, value=False):
# 如果模組有 gradient_checkpointing 屬性，則設定其值
if hasattr(module, "gradient_checkpointing"):
module.gradient_checkpointing = value
# 前向傳播方法定義
def forward(
# 輸入的隱藏狀態張量
hidden_states: torch.FloatTensor,
# 時間步張量，預設為 None
timestep: torch.LongTensor = None,
# 編碼器的隱藏狀態張量，預設為 None
encoder_hidden_states: torch.FloatTensor = None,
# 全域性隱藏狀態張量，預設為 None
global_hidden_states: torch.FloatTensor = None,
# 旋轉嵌入張量，預設為 None
rotary_embedding: torch.FloatTensor = None,
# 是否返回字典格式，預設為 True
return_dict: bool = True,
# 注意力掩碼，預設為 None
attention_mask: Optional[torch.LongTensor] = None,
# 編碼器的注意力掩碼，預設為 None
encoder_attention_mask: Optional[torch.LongTensor] = None,
# 版權所有 2024 The HuggingFace Team. 保留所有權利。
#
# 根據 Apache 許可證第 2.0 版（“許可證”）進行許可；
# 除非遵守許可證，否則您不得使用此檔案。
# 您可以在以下網址獲取許可證副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非適用法律或書面協議另有規定，依據許可證分發的軟體
# 是按“原樣”基礎分發的，沒有任何形式的保證或條件，
# 無論是明示還是暗示。有關許可證下特定語言的許可權和
# 限制，請參閱許可證。
# 匯入數學模組以執行數學運算
import math
# 從 typing 匯入可選型別和元組
from typing import Optional, Tuple
# 匯入 PyTorch 庫
import torch
# 從 torch 匯入神經網路模組
from torch import nn
# 匯入配置工具和註冊功能
from ...configuration_utils import ConfigMixin, register_to_config
# 匯入注意力處理器
from ..attention_processor import Attention
# 匯入獲取時間步嵌入的功能
from ..embeddings import get_timestep_embedding
# 匯入模型工具的基類
from ..modeling_utils import ModelMixin
class T5FilmDecoder(ModelMixin, ConfigMixin):
r"""
T5 風格的解碼器，具有 FiLM 條件。
引數：
input_dims (`int`, *可選*, 預設為 `128`):
輸入維度的數量。
targets_length (`int`, *可選*, 預設為 `256`):
目標的長度。
d_model (`int`, *可選*, 預設為 `768`):
輸入隱藏狀態的大小。
num_layers (`int`, *可選*, 預設為 `12`):
使用的 `DecoderLayer` 數量。
num_heads (`int`, *可選*, 預設為 `12`):
使用的注意力頭的數量。
d_kv (`int`, *可選*, 預設為 `64`):
鍵值投影向量的大小。
d_ff (`int`, *可選*, 預設為 `2048`):
`DecoderLayer` 中間前饋層的維度數量。
dropout_rate (`float`, *可選*, 預設為 `0.1`):
丟棄機率。
"""
# 使用裝飾器註冊初始化函式到配置
@register_to_config
def __init__(
# 輸入維度，預設為128
self,
input_dims: int = 128,
# 目標長度，預設為256
targets_length: int = 256,
# 最大解碼噪聲時間，預設為2000.0
max_decoder_noise_time: float = 2000.0,
# 隱藏狀態的維度，預設為768
d_model: int = 768,
# 解碼層的數量，預設為12
num_layers: int = 12,
# 注意力頭的數量，預設為12
num_heads: int = 12,
# 鍵值維度大小，預設為64
d_kv: int = 64,
# 中間前饋層的維度，預設為2048
d_ff: int = 2048,
# 丟棄率，預設為0.1
dropout_rate: float = 0.1,
# 初始化父類
):
super().__init__()
# 建立條件嵌入層，包含兩層線性變換和啟用函式
self.conditioning_emb = nn.Sequential(
# 第一個線性層，輸入維度為 d_model，輸出維度為 d_model * 4，不使用偏置
nn.Linear(d_model, d_model * 4, bias=False),
# 使用 SiLU 啟用函式
nn.SiLU(),
# 第二個線性層，輸入維度為 d_model * 4，輸出維度同樣為 d_model * 4，不使用偏置
nn.Linear(d_model * 4, d_model * 4, bias=False),
# 使用 SiLU 啟用函式
nn.SiLU(),
)
# 建立位置編碼嵌入，大小為 (targets_length, d_model)
self.position_encoding = nn.Embedding(targets_length, d_model)
# 禁止位置編碼的權重更新
self.position_encoding.weight.requires_grad = False
# 建立連續輸入的線性投影層，輸入維度為 input_dims，輸出維度為 d_model，不使用偏置
self.continuous_inputs_projection = nn.Linear(input_dims, d_model, bias=False)
# 建立 dropout 層，丟棄率為 dropout_rate
self.dropout = nn.Dropout(p=dropout_rate)
# 建立解碼器層的模組列表
self.decoders = nn.ModuleList()
# 迴圈建立 num_layers 個解碼器層
for lyr_num in range(num_layers):
# 初始化 FiLM 條件 T5 解碼器層
lyr = DecoderLayer(d_model=d_model, d_kv=d_kv, num_heads=num_heads, d_ff=d_ff, dropout_rate=dropout_rate)
# 將解碼器層新增到列表中
self.decoders.append(lyr)
# 建立解碼器層的歸一化層
self.decoder_norm = T5LayerNorm(d_model)
# 建立後續 dropout 層，丟棄率為 dropout_rate
self.post_dropout = nn.Dropout(p=dropout_rate)
# 建立輸出層，將 d_model 的輸出對映回 input_dims，不使用偏置
self.spec_out = nn.Linear(d_model, input_dims, bias=False)
# 定義編碼器-解碼器掩碼函式
def encoder_decoder_mask(self, query_input: torch.Tensor, key_input: torch.Tensor) -> torch.Tensor:
# 計算查詢和鍵輸入的掩碼，進行逐元素相乘
mask = torch.mul(query_input.unsqueeze(-1), key_input.unsqueeze(-2))
# 返回掩碼並擴充套件維度
return mask.unsqueeze(-3)
# 前向傳播方法，接受編碼及掩碼、解碼器輸入和噪聲時間
def forward(self, encodings_and_masks, decoder_input_tokens, decoder_noise_time):
# 獲取批次大小和解碼器輸入的形狀
batch, _, _ = decoder_input_tokens.shape
# 確保噪聲時間的形狀與批次一致
assert decoder_noise_time.shape == (batch,)
# 將 decoder_noise_time 重新縮放到期望的時間範圍
time_steps = get_timestep_embedding(
decoder_noise_time * self.config.max_decoder_noise_time,
embedding_dim=self.config.d_model,
max_period=self.config.max_decoder_noise_time,
).to(dtype=self.dtype)
# 使用時間步長生成條件嵌入，並擴充套件維度
conditioning_emb = self.conditioning_emb(time_steps).unsqueeze(1)
# 確保條件嵌入的形狀正確
assert conditioning_emb.shape == (batch, 1, self.config.d_model * 4)
# 獲取解碼器輸入的序列長度
seq_length = decoder_input_tokens.shape[1]
# 如果使用相對位置，基於編碼和掩碼的長度偏移序列
decoder_positions = torch.broadcast_to(
torch.arange(seq_length, device=decoder_input_tokens.device),
(batch, seq_length),
)
# 計算位置編碼
position_encodings = self.position_encoding(decoder_positions)
# 對解碼器輸入進行連續輸入投影
inputs = self.continuous_inputs_projection(decoder_input_tokens)
# 將位置編碼新增到輸入中
inputs += position_encodings
# 應用 dropout 操作
y = self.dropout(inputs)
# 建立解碼器掩碼，沒有填充
decoder_mask = torch.ones(
decoder_input_tokens.shape[:2], device=decoder_input_tokens.device, dtype=inputs.dtype
)
# 將編碼掩碼轉換為編碼器-解碼器掩碼
encodings_and_encdec_masks = [(x, self.encoder_decoder_mask(decoder_mask, y)) for x, y in encodings_and_masks]
# 交叉注意力風格：拼接編碼
encoded = torch.cat([x[0] for x in encodings_and_encdec_masks], dim=1)
encoder_decoder_mask = torch.cat([x[1] for x in encodings_and_encdec_masks], dim=-1)
# 對每一層解碼器進行迴圈處理
for lyr in self.decoders:
y = lyr(
y,
conditioning_emb=conditioning_emb,
encoder_hidden_states=encoded,
encoder_attention_mask=encoder_decoder_mask,
)[0]
# 對輸出進行歸一化
y = self.decoder_norm(y)
# 應用 dropout 後處理
y = self.post_dropout(y)
# 生成最終的頻譜輸出
spec_out = self.spec_out(y)
# 返回頻譜輸出
return spec_out
# T5 解碼器層的定義
class DecoderLayer(nn.Module):
r"""
T5 decoder layer. # T5解碼器層的文件說明
Args: # 引數說明
d_model (`int`): # 輸入隱藏狀態的大小
Size of the input hidden states. # 輸入隱藏狀態的大小
d_kv (`int`): # 鍵值投影向量的大小
Size of the key-value projection vectors. # 鍵值投影向量的大小
num_heads (`int`): # 注意力頭的數量
Number of attention heads. # 注意力頭的數量
d_ff (`int`): # 中間前饋層的大小
Size of the intermediate feed-forward layer. # 中間前饋層的大小
dropout_rate (`float`): # 丟棄機率
Dropout probability. # 丟棄機率
layer_norm_epsilon (`float`, *optional*, defaults to `1e-6`): # 數值穩定性的小值
A small value used for numerical stability to avoid dividing by zero. # 數值穩定性的小值
"""
# 初始化方法，定義各個引數
def __init__(
self, d_model: int, d_kv: int, num_heads: int, d_ff: int, dropout_rate: float, layer_norm_epsilon: float = 1e-6
):
super().__init__() # 呼叫父類建構函式
self.layer = nn.ModuleList() # 初始化模組列表以儲存層
# 條件自注意力：第 0 層
self.layer.append(
T5LayerSelfAttentionCond(d_model=d_model, d_kv=d_kv, num_heads=num_heads, dropout_rate=dropout_rate) # 新增條件自注意力層
)
# 交叉注意力：第 1 層
self.layer.append(
T5LayerCrossAttention(
d_model=d_model, # 輸入隱藏狀態的大小
d_kv=d_kv, # 鍵值投影向量的大小
num_heads=num_heads, # 注意力頭的數量
dropout_rate=dropout_rate, # 丟棄機率
layer_norm_epsilon=layer_norm_epsilon, # 數值穩定性的小值
)
)
# Film Cond MLP + 丟棄：最後一層
self.layer.append(
T5LayerFFCond(d_model=d_model, d_ff=d_ff, dropout_rate=dropout_rate, layer_norm_epsilon=layer_norm_epsilon) # 新增條件前饋層
)
# 前向傳播方法
def forward(
self,
hidden_states: torch.Tensor, # 輸入隱藏狀態張量
conditioning_emb: Optional[torch.Tensor] = None, # 條件嵌入（可選）
attention_mask: Optional[torch.Tensor] = None, # 注意力掩碼（可選）
encoder_hidden_states: Optional[torch.Tensor] = None, # 編碼器隱藏狀態（可選）
encoder_attention_mask: Optional[torch.Tensor] = None, # 編碼器注意力掩碼（可選）
encoder_decoder_position_bias=None, # 編碼器-解碼器位置偏置
) -> Tuple[torch.Tensor]: # 返回張量的元組
hidden_states = self.layer[0]( # 透過第一層處理輸入隱藏狀態
hidden_states,
conditioning_emb=conditioning_emb, # 使用條件嵌入
attention_mask=attention_mask, # 使用注意力掩碼
)
# 如果存在編碼器隱藏狀態
if encoder_hidden_states is not None:
# 擴充套件編碼器注意力掩碼
encoder_extended_attention_mask = torch.where(encoder_attention_mask > 0, 0, -1e10).to(
encoder_hidden_states.dtype # 轉換為編碼器隱藏狀態的資料型別
)
hidden_states = self.layer[1]( # 透過第二層處理隱藏狀態
hidden_states,
key_value_states=encoder_hidden_states, # 使用編碼器隱藏狀態作為鍵值
attention_mask=encoder_extended_attention_mask, # 使用擴充套件的注意力掩碼
)
# 應用 Film 條件前饋層
hidden_states = self.layer[-1](hidden_states, conditioning_emb) # 透過最後一層處理隱藏狀態，使用條件嵌入
return (hidden_states,) # 返回處理後的隱藏狀態元組
# T5樣式的自注意力層，帶條件
class T5LayerSelfAttentionCond(nn.Module):
r"""
T5 style self-attention layer with conditioning. # T5樣式的自注意力層，帶條件說明
# 函式引數說明
Args:
d_model (`int`): # 輸入隱藏狀態的大小
Size of the input hidden states.
d_kv (`int`): # 鍵值投影向量的大小
Size of the key-value projection vectors.
num_heads (`int`): # 注意力頭的數量
Number of attention heads.
dropout_rate (`float`): # 丟棄機率
Dropout probability.
"""
# 初始化方法，設定類的基本引數
def __init__(self, d_model: int, d_kv: int, num_heads: int, dropout_rate: float):
super().__init__() # 呼叫父類建構函式
# 建立層歸一化層，輸入大小為 d_model
self.layer_norm = T5LayerNorm(d_model)
# 建立 FiLM 層，輸入特徵為 d_model * 4，輸出特徵為 d_model
self.FiLMLayer = T5FiLMLayer(in_features=d_model * 4, out_features=d_model)
# 建立注意力層，設定查詢維度、頭數、鍵值維度等引數
self.attention = Attention(query_dim=d_model, heads=num_heads, dim_head=d_kv, out_bias=False, scale_qk=False)
# 建立丟棄層，設定丟棄機率
self.dropout = nn.Dropout(dropout_rate)
# 前向傳播方法
def forward(
self,
hidden_states: torch.Tensor, # 輸入的隱藏狀態張量
conditioning_emb: Optional[torch.Tensor] = None, # 可選的條件嵌入
attention_mask: Optional[torch.Tensor] = None, # 可選的注意力掩碼
) -> torch.Tensor:
# 對輸入的隱藏狀態進行層歸一化
normed_hidden_states = self.layer_norm(hidden_states)
# 如果有條件嵌入，應用 FiLM 層
if conditioning_emb is not None:
normed_hidden_states = self.FiLMLayer(normed_hidden_states, conditioning_emb)
# 自注意力模組，獲取注意力輸出
attention_output = self.attention(normed_hidden_states)
# 將注意力輸出與原隱藏狀態相加，並應用丟棄層
hidden_states = hidden_states + self.dropout(attention_output)
# 返回更新後的隱藏狀態
return hidden_states
# T5風格的交叉注意力層
class T5LayerCrossAttention(nn.Module):
r"""
T5風格的交叉注意力層。
引數：
d_model (`int`):
輸入隱藏狀態的大小。
d_kv (`int`):
鍵值投影向量的大小。
num_heads (`int`):
注意力頭的數量。
dropout_rate (`float`):
丟棄機率。
layer_norm_epsilon (`float`):
用於數值穩定性的小值，避免除以零。
"""
# 初始化方法，設定模型引數
def __init__(self, d_model: int, d_kv: int, num_heads: int, dropout_rate: float, layer_norm_epsilon: float):
# 呼叫父類初始化方法
super().__init__()
# 建立注意力層
self.attention = Attention(query_dim=d_model, heads=num_heads, dim_head=d_kv, out_bias=False, scale_qk=False)
# 建立層歸一化層
self.layer_norm = T5LayerNorm(d_model, eps=layer_norm_epsilon)
# 建立丟棄層
self.dropout = nn.Dropout(dropout_rate)
# 前向傳播方法
def forward(
self,
hidden_states: torch.Tensor,
key_value_states: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
# 對隱藏狀態進行層歸一化
normed_hidden_states = self.layer_norm(hidden_states)
# 計算注意力輸出
attention_output = self.attention(
normed_hidden_states,
encoder_hidden_states=key_value_states,
attention_mask=attention_mask.squeeze(1),
)
# 計算層輸出，新增丟棄
layer_output = hidden_states + self.dropout(attention_output)
# 返回層輸出
return layer_output
# T5風格的前饋條件層
class T5LayerFFCond(nn.Module):
r"""
T5風格的前饋條件層。
引數：
d_model (`int`):
輸入隱藏狀態的大小。
d_ff (`int`):
中間前饋層的大小。
dropout_rate (`float`):
丟棄機率。
layer_norm_epsilon (`float`):
用於數值穩定性的小值，避免除以零。
"""
# 初始化方法，設定模型引數
def __init__(self, d_model: int, d_ff: int, dropout_rate: float, layer_norm_epsilon: float):
# 呼叫父類初始化方法
super().__init__()
# 建立帶門啟用的前饋層
self.DenseReluDense = T5DenseGatedActDense(d_model=d_model, d_ff=d_ff, dropout_rate=dropout_rate)
# 建立條件層
self.film = T5FiLMLayer(in_features=d_model * 4, out_features=d_model)
# 建立層歸一化層
self.layer_norm = T5LayerNorm(d_model, eps=layer_norm_epsilon)
# 建立丟棄層
self.dropout = nn.Dropout(dropout_rate)
# 前向傳播方法
def forward(self, hidden_states: torch.Tensor, conditioning_emb: Optional[torch.Tensor] = None) -> torch.Tensor:
# 對隱藏狀態進行層歸一化
forwarded_states = self.layer_norm(hidden_states)
# 如果存在條件嵌入，則應用條件層
if conditioning_emb is not None:
forwarded_states = self.film(forwarded_states, conditioning_emb)
# 應用前饋層
forwarded_states = self.DenseReluDense(forwarded_states)
# 更新隱藏狀態，新增丟棄
hidden_states = hidden_states + self.dropout(forwarded_states)
# 返回更新後的隱藏狀態
return hidden_states
# T5風格的前饋層，具有門控啟用和丟棄
class T5DenseGatedActDense(nn.Module):
r"""
T5風格的前饋層，具有門控啟用和丟棄。
# 引數說明部分
Args:
d_model (`int`): # 輸入隱藏狀態的尺寸
Size of the input hidden states.
d_ff (`int`): # 中間前饋層的尺寸
Size of the intermediate feed-forward layer.
dropout_rate (`float`): # 丟棄機率
Dropout probability.
"""
# 初始化方法，接受模型引數
def __init__(self, d_model: int, d_ff: int, dropout_rate: float):
super().__init__() # 呼叫父類的初始化方法
# 定義第一線性變換層，不使用偏置，輸入維度為d_model，輸出維度為d_ff
self.wi_0 = nn.Linear(d_model, d_ff, bias=False)
# 定義第二線性變換層，不使用偏置，輸入維度為d_model，輸出維度為d_ff
self.wi_1 = nn.Linear(d_model, d_ff, bias=False)
# 定義輸出線性變換層，不使用偏置，輸入維度為d_ff，輸出維度為d_model
self.wo = nn.Linear(d_ff, d_model, bias=False)
# 定義丟棄層，使用指定的丟棄機率
self.dropout = nn.Dropout(dropout_rate)
# 初始化自定義啟用函式
self.act = NewGELUActivation()
# 前向傳播方法，接受輸入的隱藏狀態
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
# 透過第一線性層和啟用函式得到隱層狀態
hidden_gelu = self.act(self.wi_0(hidden_states))
# 透過第二線性層得到隱層狀態
hidden_linear = self.wi_1(hidden_states)
# 將兩個隱層狀態進行逐元素相乘
hidden_states = hidden_gelu * hidden_linear
# 應用丟棄層
hidden_states = self.dropout(hidden_states)
# 透過輸出線性層得到最終的隱層狀態
hidden_states = self.wo(hidden_states)
# 返回最終的隱層狀態
return hidden_states
# T5風格的層歸一化模組
class T5LayerNorm(nn.Module):
r"""
T5風格的層歸一化模組。
Args:
hidden_size (`int`):
輸入隱藏狀態的大小。
eps (`float`, `optional`, defaults to `1e-6`):
用於數值穩定性的小值，以避免除以零。
"""
# 初始化函式，接受隱藏狀態大小和epsilon
def __init__(self, hidden_size: int, eps: float = 1e-6):
"""
構造一個T5風格的層歸一化模組。沒有偏置，也不減去均值。
"""
# 呼叫父類建構函式
super().__init__()
# 初始化權重為全1的可學習引數
self.weight = nn.Parameter(torch.ones(hidden_size))
# 儲存epsilon值
self.variance_epsilon = eps
# 前向傳播函式，接受隱藏狀態
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
# 計算隱藏狀態的方差，使用平方和的均值，保持維度
variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
# 按照方差進行歸一化，同時考慮到epsilon
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
# 如果權重為半精度，轉換隱藏狀態為相應型別
if self.weight.dtype in [torch.float16, torch.bfloat16]:
hidden_states = hidden_states.to(self.weight.dtype)
# 返回歸一化後的結果乘以權重
return self.weight * hidden_states
# 實現GELU啟用函式的模組
class NewGELUActivation(nn.Module):
"""
實現與Google BERT庫中相同的GELU啟用函式（與OpenAI GPT相同）。也可以參考
Gaussian Error Linear Units論文：https://arxiv.org/abs/1606.08415
"""
# 前向傳播函式，接受輸入張量
def forward(self, input: torch.Tensor) -> torch.Tensor:
# 計算GELU啟用值
return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (input + 0.044715 * torch.pow(input, 3.0))))
# T5風格的FiLM層
class T5FiLMLayer(nn.Module):
"""
T5風格的FiLM層。
Args:
in_features (`int`):
輸入特徵的數量。
out_features (`int`):
輸出特徵的數量。
"""
# 初始化函式，接受輸入和輸出特徵數量
def __init__(self, in_features: int, out_features: int):
# 呼叫父類建構函式
super().__init__()
# 定義線性層，用於生成縮放和偏移引數
self.scale_bias = nn.Linear(in_features, out_features * 2, bias=False)
# 前向傳播函式，接受輸入張量和條件嵌入
def forward(self, x: torch.Tensor, conditioning_emb: torch.Tensor) -> torch.Tensor:
# 透過線性層計算縮放和偏移
emb = self.scale_bias(conditioning_emb)
# 將結果分成縮放和偏移兩個部分
scale, shift = torch.chunk(emb, 2, -1)
# 進行縮放和偏移操作
x = x * (1 + scale) + shift
# 返回處理後的結果
return x