https://github.com/omerbt/TokenFlow/issues/25
https://github.com/omerbt/TokenFlow/issues/31
https://github.com/omerbt/TokenFlow/issues/32
https://github .com/eps696/SDfu
Este artigo explica principalmente como a parte Model do TokenFlow é construída. O código foi extraído de TokenFlow/tokenflow_utils.py.
A lógica de construção do modelo do Tokenflow é primeiro carregar o Stable Diffusion original e, em seguida, registrar novamente o módulo UNet que precisa ser modificado . A operação de modificação é chamada primeiro run_tokenflow.py:
self.init_method(conv_injection_t=pnp_f_t, qk_injection_t=pnp_attn_t)
Se você quiser entender o código a seguir, primeiro você precisa entender o BasicTransformerBlockcódigo-fonte SD. É melhor olhar o código-fonte PnP pnp-diffusers , porque o TokenFlow é aprimorado com base no PnP.
def init_method(self, conv_injection_t, qk_injection_t):
self.qk_injection_timesteps = self.scheduler.timesteps[:qk_injection_t] if qk_injection_t >= 0 else []
self.conv_injection_timesteps = self.scheduler.timesteps[:conv_injection_t] if conv_injection_t >= 0 else []
register_extended_attention_pnp(self, self.qk_injection_timesteps)
register_conv_injection(self, self.conv_injection_timesteps)
set_tokenflow(self.unet)
init_methodA função completou 3 coisas:
(1) register_extended_attention_pnp: substituir a autoatenção do unet (expandido para KV a partir de vários quadros, completar a injeção ao mesmo tempo.
(2) register_conv_injection: substituir o conv_injection de conv (segundo bloco resnet do UpBlock, completar a injeção.
(3)) set_tokenflow: Substitua o BasicTransformerBlock 16 do unet por TokenFlowBlock.
Entre eles, qk_injection_timestepse conv_injection_timestepsestão dois timestep list, usados para controlar a operação PnP Inject a ser executada apenas nas primeiras etapas.
Além dessas modificações no modelo UNet, a função no código-fonte também define o ID do quadro-chave batched_denoise_steppara editar o quadro-chave primeiro . register_pivotalO ID do lote é definido antes da edição de cada lote register_batch_idx. register_timeDefina a etapa t para algumas camadas da UNet antes de prever o ruído .
A seguir, analisaremos, uma por uma, em ordem, as modificações feitas pelo tokenflow no modelo de Difusão Estável original durante o processo de inferência.
registrar_extended_attention_pnp
A função da função register_extended_attention_pnp : reconstruir a função forward para estender a atenção para attn1 (16) de todas as camadas BasicTransformerBlocksa_forward da UNet , mas apenas injetar operações PnP para algumas das attn1 (8) .injection_schedule
Isso pode ser visto nos resultados de BasicTransformerBlock: Embora a implementação da classe attn1 seja CrossAttention, o contexto não é passado para KV durante a inferência e é essencialmente SelfAttention .
class BasicTransformerBlock(nn.Module):
def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True):
super().__init__()
self.attn1 = CrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout) # is a self-attention
self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
self.attn2 = CrossAttention(query_dim=dim, context_dim=context_dim,
heads=n_heads, dim_head=d_head, dropout=dropout) # cross attention
self.norm1 = nn.LayerNorm(dim)
self.norm2 = nn.LayerNorm(dim)
self.norm3 = nn.LayerNorm(dim)
def forward(self, x, context=None):
x = self.attn1(self.norm1(x)) + x
x = self.attn2(self.norm2(x), context=context) + x
x = self.ff(self.norm3(x)) + x
return x
Parâmetros de entrada : register_extended_attention_pnpOs dois parâmetros obrigatórios passados pela função são unet modele injection_schedule. Usado para controlar o intervalo de tempo de execução da injeção PnPinjection_schedule durante a inferência , porque queremos realizar operações PnP apenas nos primeiros intervalos de tempo .
Reconstruçãoextend_attention : Como o avanço do tokenflow original sa_forwardé realizar operações de matriz de atenção em todos os quadros , o que consome muitos recursos, adicionei atenção que sa_3frame_forwardcalcula apenas a atenção de 3 quadros adjacentes . A função que reconstruí register_extended_attention_pnpé mostrada abaixo: com base na função original, adicionei um is_3_frameparâmetro para escolher se deseja usá-la sa_3frame_forward.
Primeiro, vamos pular essas duas funções sa_forwarde sa_3frame_forwardver como encontrar o módulo correspondente ao UNet e modificar seu método de encaminhamento.
1. Reconstrua forward para attn1 de todas as camadas BasicTransformerBlock
Determine register_forward_funqual usar com base sa_forwarde, em seguida, percorra cada módulo de unet para determinar se o módulo modificado herda de BasicTransformerBlock. Em caso afirmativo, modifique-o forward, mas deixe-o injection_scheduleem branco (ou seja, não execute PnP).
module_names = []
register_forward_fun = sa_3frame_forward if is_3_frame else sa_forward
for module_name, module in model.unet.named_modules():
if isinstance_str(module, "BasicTransformerBlock"):
module_names.append(module_name)
# replace BasicTransformerBlock.attn1's forward with sa_forward
module.attn1.forward = register_forward_fun(module.attn1)
# set injection_schedule empty[] for BasicTransformerBlock.attn1
setattr(module.attn1, 'injection_schedule', [])
print(f"all change {
len(module_names)} layer's BasicTransformerBlock.attn1.forward() for extended_attention_pnp...")
print(module_names) # up_blocks.1.attentions.0.transformer_blocks.0
isinstance_strDetermine se a lista de tipos herdados de x contém a classe cls_name:
def isinstance_str(x: object, cls_name: str):
for _cls in x.__class__.__mro__:
if _cls.__name__ == cls_name:
return True
return False
Pela primeira vez, o avanço das seguintes 16 camadas de atenção em unet é reconstruído: 6 down_blocks, 9 up_blocks e 1 mid_block.
down_blocks.0.attentions.0.transformer_blocks.0.attn1
down_blocks.0.attentions.1.transformer_blocks.0.attn1
down_blocks.1.attentions.0.transformer_blocks.0.attn1
down_blocks.1.attentions.1.transformer_blocks.0.attn1
down_blocks.2.attentions.0.transformer_blocks.0.attn1
down_blocks.2.attentions.1.transformer_blocks.0.attn1
up_blocks.1.attentions.0.transformer_blocks.0.attn1
up_blocks.1.attentions.1.transformer_blocks.0.attn1
up_blocks.1.attentions.2.transformer_blocks.0.attn1
up_blocks.2.attentions.0.transformer_blocks.0.attn1
up_blocks.2.attentions.1.transformer_blocks.0.attn1
up_blocks.2.attentions.2.transformer_blocks.0.attn1
up_blocks.3.attentions.0.transformer_blocks.0.attn1
up_blocks.3.attentions.1.transformer_blocks.0.attn1
up_blocks.3.attentions.2.transformer_blocks.0.attn1
mid_block.attentions.0.transformer_blocks.0.attn1
2. Injetar injeção_schedule e usar a operação PnP em alguns attn1 (8)
Sob as instruções de res_dict, modifique a função de encaminhamento para o attn1 específico, registre injeção_schedule para ele e use a operação PnP.
res_dict = {
1: [1, 2], 2: [0, 1, 2], 3: [0, 1, 2]} # upblock's self-attention layers
# we are injecting attention in blocks 4 - 11 of the Unet UpBlock, so not in the first block of the lowest resolution
for res in res_dict: # res = 1
for block in res_dict[res]: # res_dict[res] = [1, 2]
module = model.unet.up_blocks[res].attentions[block].transformer_blocks[0].attn1
module.forward = sa_forward(module)
setattr(module, 'injection_schedule', injection_schedule)
A segunda reconstrução de 8 up_blocks:
model.unet.up_blocks.1.attentions.1.transformer_blocks.0.attn1
model.unet.up_blocks.1.attentions.2.transformer_blocks.0.attn1
model.unet.up_blocks.2.attentions.0.transformer_blocks.0.attn1
model.unet.up_blocks.2.attentions.1.transformer_blocks.0.attn1
model.unet.up_blocks.2.attentions.2.transformer_blocks.0.attn1
model.unet.up_blocks.3.attentions.0.transformer_blocks.0.attn1
model.unet.up_blocks.3.attentions.1.transformer_blocks.0.attn1
model.unet.up_blocks.3.attentions.2.transformer_blocks.0.attn1
3. para_encaminhar
Operação PnP : Como a entrada do TokenFlow considera ambos PnPe , a única classifer-free guidanceentrada latente original da UNet torna-se três cópias : source_latents + x + x( um x correspondeedit_prompt a um xnull_prompt e source_latents corresponde a elesource_prompt ).
latent_model_input = torch.cat([source_latents] + ([x] * 2))
Desta forma, quando Unet está inferindo, ele pode ser cortado diretamente da entrada latente x, dividido em 3 partes e source_latentsinjetado na soma (a injeção PnP é uma substituição direta ). Para autoatenção, substituímos apenas Q e K .uncond_latentscond_latents
source_latents = x[:n_frames]
uncond_latents = x[n_frames:2*n_frames]
cond_latents = x[2*n_frames:]
# source inject uncond
q[n_frames:2*n_frames] = q[:n_frames]
k[n_frames:2*n_frames] = k[:n_frames]
# source inject cond
q[2*n_frames:] = q[:n_frames]
k[2*n_frames:] = k[:n_frames]
Extend_Attention : tokenflow implementa o uso de autoatenção estendida, pois para o i-ésimo quadro, ao calcular a autoatenção, Q é o recurso do i-ésimo quadro, e KV deve vir de todos os outros quadros, portanto é necessário repeat一下K和Vfacilitar cálculos subsequentes.
T base = S oftmax ( Q i ; [ K i 1 , . . . , K ik ] d ) ⋅ [ V i 1 , . . . , V ik ] T_{base}=Softmax(\frac{Q^i; [K^{i1},...,K^{ik}]}{\sqrt{d}})\cdot[V^{i1},...,V^{ik}]Tba se=Suave máximo ( _ _ _ _dPeu ;[ Keu 1 ,... ,KEU ])⋅[ Veu 1 ,... ,VEU ]
# KV reshape and repeat for extend_attention: Softmax(Q_i_frame @ K_all_frame) @ V_all_frame
# (n_frames, seq_len, dim) -> (1, n_frames * seq_len, dim) -> (n_frames, n_frames * seq_len, dim)
k_source = k[:n_frames]
k_uncond = k[n_frames:2 * n_frames].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
k_cond = k[2 * n_frames:].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
v_source = v[:n_frames]
v_uncond = v[n_frames:2 * n_frames].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
v_cond = v[2 * n_frames:].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
Como 逐帧os cálculos são realizados e a atenção de múltiplas cabeças 逐头é calculada, um loop for duplo é construído para calcular a atenção 第 i 帧第 j 头的 attention oute, finalmente, 分别concat帧维度和头维度o resultado final da atenção é obtido:
Q @ K -> sim：
(b, 1, seq_len, dim//head) @ (b, 1, dim//head, frame*seq_len) -> (b, 1, seq_len, frame*seq_len)
sim @ V -> out：
(b, 1, seq_len, frame*seq_len) @ (b, 1, frame*seq_len, dim//head) -> (b, 1, seq_len, dim//head)
cat each head's out：
(b->n_frames, 1, seq_len, dim//head) -> (n_frames, 1, seq_len, dim//head)
cat each frame's out：
(n_frames, 1, seq_len, dim//head) -> (n_frames, heads, seq_len, dim//heads)
sa_forwardO código completo é o seguinte:
def sa_forward(self):
to_out = self.to_out # self.to_out = [linear, dropout]
if type(to_out) is torch.nn.modules.container.ModuleList:
to_out = self.to_out[0]
else:
to_out = self.to_out
def forward(x, encoder_hidden_states=None, attention_mask=None):
is_cross = encoder_hidden_states is not None # corss-attention or self-attention
h = self.heads
batch_size, sequence_length, dim = x.shape # (3*n_frames, seq_len, dim)
# batch: 前n_frames个样本为source feature, 中间n_frames个样本为uncond featur, 后n_frames个样本为cond feature
n_frames = batch_size // 3
# source_latents = x[:n_frames], uncond_latents = x[n_frames:2*n_frames], cond_latents = x[2*n_frames:]
encoder_hidden_states = encoder_hidden_states if is_cross else x
q = self.to_q(x)
k = self.to_k(encoder_hidden_states)
v = self.to_v(encoder_hidden_states)
# PnP Injection QK：只需要sample过程中的前几个timestep进行injection (判断t是否符合)，且只在UpBlock进行inject
if self.injection_schedule is not None and (self.t in self.injection_schedule or self.t == 1000):
# source inject into unconditional
q[n_frames:2 * n_frames] = q[:n_frames]
k[n_frames:2 * n_frames] = k[:n_frames]
# source inject into conditional
q[2 * n_frames:] = q[:n_frames]
k[2 * n_frames:] = k[:n_frames]
# KV reshape and repeat for extend_attention: Softmax(Q_i_frame @ K_all_frame) @ V_all_frame
# (n_frames, seq_len, dim) -> (1, n_frames * seq_len, dim) -> (n_frames, n_frames * seq_len, dim)
k_source = k[:n_frames]
k_uncond = k[n_frames:2 * n_frames].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
k_cond = k[2 * n_frames:].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
v_source = v[:n_frames]
v_uncond = v[n_frames:2 * n_frames].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
v_cond = v[2 * n_frames:].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
# project QKV's source, cond and uncond, respectively
q_source = self.reshape_heads_to_batch_dim(q[:n_frames]) # q (n_frames*heads, seq_len, dim//heads)
q_uncond = self.reshape_heads_to_batch_dim(q[n_frames:2 * n_frames])
q_cond = self.reshape_heads_to_batch_dim(q[2 * n_frames:])
k_source = self.reshape_heads_to_batch_dim(k_source) # kv (n_frames*heads, n_frames * seq_len, dim//heads)
k_uncond = self.reshape_heads_to_batch_dim(k_uncond)
k_cond = self.reshape_heads_to_batch_dim(k_cond)
v_source = self.reshape_heads_to_batch_dim(v_source)
v_uncond = self.reshape_heads_to_batch_dim(v_uncond)
v_cond = self.reshape_heads_to_batch_dim(v_cond)
# split heads
q_src = q_source.view(n_frames, h, sequence_length, dim // h)
k_src = k_source.view(n_frames, h, sequence_length, dim // h)
v_src = v_source.view(n_frames, h, sequence_length, dim // h)
q_uncond = q_uncond.view(n_frames, h, sequence_length, dim // h)
k_uncond = k_uncond.view(n_frames, h, sequence_length * n_frames, dim // h)
v_uncond = v_uncond.view(n_frames, h, sequence_length * n_frames, dim // h)
q_cond = q_cond.view(n_frames, h, sequence_length, dim // h)
k_cond = k_cond.view(n_frames, h, sequence_length * n_frames, dim // h)
v_cond = v_cond.view(n_frames, h, sequence_length * n_frames, dim // h)
out_source_all = []
out_uncond_all = []
out_cond_all = []
# each frame or single_batch frames
single_batch = n_frames<=12
b = n_frames if single_batch else 1 # b=1
# do attention for each frame respectively. frames [frame：frame=b]
for frame in range(0, n_frames, b):
out_source = []
out_uncond = []
out_cond = []
# do attention for each head respectively. head j
for j in range(h):
# do attention for source, cond and uncond respectively, (b, 1, seq_len, dim//head) @ (b, 1, dim//head, frame*seq_len) -> (b, 1, seq_len, frame*seq_len)
sim_source_b = torch.bmm(q_src[frame: frame+ b, j], k_src[frame: frame+ b, j].transpose(-1, -2)) * self.scale
sim_uncond_b = torch.bmm(q_uncond[frame: frame+ b, j], k_uncond[frame: frame+ b, j].transpose(-1, -2)) * self.scale
sim_cond = torch.bmm(q_cond[frame: frame+ b, j], k_cond[frame: frame+ b, j].transpose(-1, -2)) * self.scale
# append each head's out, (b, 1, seq_len, frame*seq_len) @ (b, 1, frame*seq_len, dim//head) -> (b, 1, seq_len, dim//head)
out_source.append(torch.bmm(sim_source_b.softmax(dim=-1), v_src[frame: frame+ b, j]))
out_uncond.append(torch.bmm(sim_uncond_b.softmax(dim=-1), v_uncond[frame: frame+ b, j]))
out_cond.append(torch.bmm(sim_cond.softmax(dim=-1), v_cond[frame: frame+ b, j]))
# cat each head's out, (b->n_frames, 1, seq_len, dim//head) -> (n_frames, 1, seq_len, dim//head)
out_source = torch.cat(out_source, dim=0)
out_uncond = torch.cat(out_uncond, dim=0)
out_cond = torch.cat(out_cond, dim=0)
if single_batch: # if use single_batch, view single_batch frame's out
out_source = out_source.view(h, n_frames,sequence_length, dim // h).permute(1, 0, 2, 3).reshape(h * n_frames, sequence_length, -1)
out_uncond = out_uncond.view(h, n_frames,sequence_length, dim // h).permute(1, 0, 2, 3).reshape(h * n_frames, sequence_length, -1)
out_cond = out_cond.view(h, n_frames,sequence_length, dim // h).permute(1, 0, 2, 3).reshape(h * n_frames, sequence_length, -1)
# append each frame's out
out_source_all.append(out_source)
out_uncond_all.append(out_uncond)
out_cond_all.append(out_cond)
# cat each frame's out, (n_frames, 1, seq_len, dim//head) -> (n_frames, heads, seq_len, dim//heads)
out_source = torch.cat(out_source_all, dim=0)
out_uncond = torch.cat(out_uncond_all, dim=0)
out_cond = torch.cat(out_cond_all, dim=0)
# cat source, cond and uncond's out, (n_frames, heads, seq_len, dim//heads) -> (3*n_frames, heads, seq_len, dim//heads)
out = torch.cat([out_source, out_uncond, out_cond], dim=0)
out = self.reshape_batch_dim_to_heads(out)
return to_out(out)
return forward
3. at_3frame_forward
Porque PnP é usado: cada vez que a autoatenção não é mais como sa_forwardrepetir n_frames na repetição de entrada, mas a autoatenção KV normal de um único quadro ésource_latent executada e a autoatenção é executada .forward_originaluncond_latentcond_latentKV 来自相邻3帧forward_extended
Cada vez que a atenção do i-ésimo quadro é calculada (window_size=3), com o i-ésimo quadro como centro, 下标=[i-1, i, i+1]3 quadros são considerados KV:
def sa_3frame_forward(self): # self attention只是扩展到连续的 3 个关键帧，而不是所有关键帧。
to_out = self.to_out
if type(to_out) is torch.nn.modules.container.ModuleList:
to_out = self.to_out[0]
else:
to_out = self.to_out
# 原始的UNet attention forward
def forward_original(q, k, v):
n_frames, seq_len, dim = q.shape
h = self.heads
head_dim = dim // h
q = self.head_to_batch_dim(q).reshape(n_frames, h, seq_len, head_dim)
k = self.head_to_batch_dim(k).reshape(n_frames, h, seq_len, head_dim)
v = self.head_to_batch_dim(v).reshape(n_frames, h, seq_len, head_dim)
out_all = []
for frame in range(n_frames):
out = []
for j in range(h):
sim = torch.matmul(q[frame, j], k[frame, j].transpose(-1, -2)) * self.scale # (seq_len, seq_len)
out.append(torch.matmul(sim.softmax(dim=-1), v[frame, j])) # h * (seq_len, head_dim)
out = torch.cat(out, dim=0).reshape(-1, seq_len, head_dim) # (h, seq_len, head_dim)
out_all.append(out) # n_frames * (h, seq_len, head_dim)
out = torch.cat(out_all, dim=0) # (n_frames * h, seq_len, head_dim)
out = self.batch_to_head_dim(out) # (n_frames, seq_len, h * head_dim)
return out
# extend UNet attention forward(all frames)
def forward_extended(q, k, v):
n_frames, seq_len, dim = q.shape
h = self.heads
head_dim = dim // h
q = self.head_to_batch_dim(q).reshape(n_frames, h, seq_len, head_dim)
k = self.head_to_batch_dim(k).reshape(n_frames, h, seq_len, head_dim)
v = self.head_to_batch_dim(v).reshape(n_frames, h, seq_len, head_dim)
out_all = []
window_size = 3
for frame in range(n_frames): # frame=32, window_size=3: window=[14, 15, 16, 17, 18]
out = []
# sliding window to improve speed. 以当前帧frame为中心，取window_size大小的帧，如frame_idx=1时, window: [0, 1, 2]
window = range(max(0, frame-window_size // 2), min(n_frames, frame+window_size//2+1))
for j in range(h):
sim_all = [] # 存当前帧frame和window内3帧的sim，len(sim_all)=3
for kframe in window: # (1, 1, seq_len, head_dim) @ (1, 1, head_dim, seq_len) -> (1, 1, seq_len, seq_len)
# 当前帧frame 依次和window内的帧kframe，计算sim存入sim_all
sim_all.append(torch.matmul(q[frame, j], k[kframe, j].transpose(-1, -2)) * self.scale) # window * (seq_len, seq_len)
sim_all = torch.cat(sim_all).reshape(len(window), seq_len, seq_len).transpose(0, 1) # (seq_len, window, seq_len)
sim_all = sim_all.reshape(seq_len, len(window) * seq_len) # (seq_len, window * seq_len)
out.append(torch.matmul(sim_all.softmax(dim=-1), v[window, j].reshape(len(window) * seq_len, head_dim))) # h * (seq_len, head_dim)
out = torch.cat(out, dim=0).reshape(-1, seq_len, head_dim) # (h, seq_len, head_dim)
out_all.append(out) # n_frames * (h, seq_len, head_dim)
out = torch.cat(out_all, dim=0) # (n_frames * h, seq_len, head_dim)
out = self.batch_to_head_dim(out) # (n_frames, seq_len, h * head_dim)
return out
def forward(x, encoder_hidden_states=None, attention_mask=None):
batch_size, sequence_length, dim = x.shape
h = self.heads
n_frames = batch_size // 3
is_cross = encoder_hidden_states is not None
encoder_hidden_states = encoder_hidden_states if is_cross else x
q = self.to_q(x)
k = self.to_k(encoder_hidden_states)
v = self.to_v(encoder_hidden_states)
if self.injection_schedule is not None and (self.t in self.injection_schedule or self.t == 1000):
# inject unconditional
q[n_frames:2 * n_frames] = q[:n_frames]
k[n_frames:2 * n_frames] = k[:n_frames]
# inject conditional
q[2 * n_frames:] = q[:n_frames]
k[2 * n_frames:] = k[:n_frames]
# source_latent 正常的self attention, uncond 和 cond进行 KV来自相邻3帧的self attention
out_source = forward_original(q[:n_frames], k[:n_frames], v[:n_frames])
out_uncond = forward_extended(q[n_frames:2 * n_frames], k[n_frames:2 * n_frames], v[n_frames:2 * n_frames])
out_cond = forward_extended(q[2 * n_frames:], k[2 * n_frames:], v[2 * n_frames:])
out = torch.cat([out_source, out_uncond, out_cond], dim=0) # (3 * n_frames, seq_len, dim)
return to_out(out)
return forward
registrar_conv_injection
Depois de registrar o encaminhamento SelfAttention para UNet, unet.up_blocks[1].resnets[1]registre um novo encaminhamento para ResnetBlock2D da UNet e registre injection_scheduleo passo de tempo de injeção de PnP de controle ao mesmo tempo.
Há conv_forwardapenas mais uma etapa da operação de injeção PnP do que o encaminhamento ResnetBlock2D comum :
if self.injection_schedule is not None and (self.t in self.injection_schedule or self.t == 1000):
source_batch_size = int(hidden_states.shape[0] // 3)
# inject unconditional
hidden_states[source_batch_size:2 * source_batch_size] = hidden_states[:source_batch_size]
# inject conditional
hidden_states[2 * source_batch_size:] = hidden_states[:source_batch_size]
set_tokenflow
__class__Significa retornar sua própria classe pai, set_tokenflowou seja, encontrar todos os módulos na UNet cuja classe pai é BasicTransformerBlock e usar BasicTransformerBlock como classe pai para eles e, em seguida, adicionar uma camada de classe TokenFlowBlock externa.
def set_tokenflow(model: torch.nn.Module):
"""
Sets the tokenflow attention blocks in a model.
"""
for _, module in model.named_modules():
if isinstance_str(module, "BasicTransformerBlock"):
# 16个 module.__class__ = <class 'diffusers.models.attention.BasicTransformerBlock'>
make_tokenflow_block_fn = make_tokenflow_attention_block
# 将BasicTransformerBlock作为父类，外面再套一层TokenFlowBlock类
module.__class__ = make_tokenflow_block_fn(module.__class__)
# Something needed for older versions of diffusers
if not hasattr(module, "use_ada_layer_norm_zero"):
module.use_ada_layer_norm = False
module.use_ada_layer_norm_zero = False
return model
make_tokenflow_attention_block
Esta função define uma classe TokenFlowBlock e retorna a classe TokenFlowBlock. TokenFlowBlock herda da classe BasicTransformerBlock e apenas reescreve forwarda função .
Primeiro pivotal_passdetermine se é um quadro-chave:
- Quadros-chave, basta salvá-los
pivot_hidden_states - Para quadros não-chave, pegue os quadros não-chave e os quadros-chave
source_latent, calcule a similaridade de cosseno cosine_sim com os quadros-chave, shape=(n_frames * seq_len, len(batch_idxs) * seq_len), encontre-o相似度最大的帧下标idxe, em seguida, empilhe 3 cópias para origem, incond e cond.- Se o lote atual não for o primeiro,
len(batch_idxs) =2salve o índice de quadros mais semelhante em idx1 e idx2, respectivamente. - Se for o primeiro lote,
len(batch_idxs) =1salve o índice de quadros mais semelhante em idx1
- Se o lote atual não for o primeiro,
batch_size, sequence_length, dim = hidden_states.shape # (batch, seq_len, dim)
n_frames = batch_size // 3 # batch = 3 * n_frames: source + uncond + cond
mid_idx = n_frames // 2
hidden_states = hidden_states.view(3, n_frames, sequence_length, dim) # (source + uncond + cond, n_frames, seq_len, dim)
norm_hidden_states = self.norm1(hidden_states).view(3, n_frames, sequence_length, dim)
if self.pivotal_pass: # is_pivotal = True # 关键帧，存下
self.pivot_hidden_states = norm_hidden_states # (3, n_frames, sequence_length, dim) ,关键帧的n_frames=5
else: # is_pivotal = False # 非关键帧，与关键帧计算source_latent的cosine_sim
idx1 = []
idx2 = []
batch_idxs = [self.batch_idx] # 每batch_size帧进行一批处理，batch_idx是第几个batch，如32帧，batch_size=8，batch_idx可以为0或1或2或3或4
if self.batch_idx > 0: # 如果不是第一个batch
batch_idxs.append(self.batch_idx - 1) # 加入前一个batch的idx，如当前batch_idx=1时，再加入0，则batch_idxs=[1,0]
# 取source_latent的非关键帧与关键帧计算cosine_sim，如果batch_idxs=[1,0]，则只拿第0个batch和第1个batch的关键帧和其norm_hidden_states计算sim
sim = batch_cosine_sim(norm_hidden_states[0].reshape(-1, dim), # (n_frames*sequence_length, dim)
self.pivot_hidden_states[0][batch_idxs].reshape(-1, dim)) # (len(batch_idxs)*sequence_length, dim)
if len(batch_idxs) == 2: # 如果不是第一个batch， 分别保存最相似的帧下标到idx1和idx2
# sim: (n_frames * seq_len, len(batch_idxs) * seq_len), len(batch_idxs)=2
sim1, sim2 = sim.chunk(2, dim=1)
idx1.append(sim1.argmax(dim=-1)) # (n_frames * seq_len) 个数,每个数在[0,76]之间
idx2.append(sim2.argmax(dim=-1)) # (n_frames * seq_len) 个数,每个数在[0,76]之间
else: # 如果是第一个batch，保存最相似的帧下标到idx1
idx1.append(sim.argmax(dim=-1))
# 为source、uncond、cond 堆叠3份
idx1 = torch.stack(idx1 * 3, dim=0) # （3, n_frames * seq_len）
idx1 = idx1.squeeze(1)
if len(batch_idxs) == 2:
idx2 = torch.stack(idx2 * 3, dim=0) # (3, n_frames * seq_len)
idx2 = idx2.squeeze(1)
Em seguida, a autoatençãoattn1 , a atenção cruzadaattn2 e o feedforwardff são realizados em sequência . Não há alteração na atenção cruzada e no feedforward. A única mudança é o processo de autoatenção :
- Para quadros-chave, calcule o resultado da autoatenção e salve-o.
- Para quadros não-chave, funda-os com os resultados de atenção dos quadros-chave. O método de fusão é a média ponderada e o peso é determinado pela distância entre o quadro e o quadro-chave. Se o quadro não-chave for um quadro do primeiro lote, o resultado de atenção do quadro-chave será usado diretamente. Se o quadro não-chave for um quadro do segundo lote, calcule os resultados de atenção com os quadros-chave do primeiro lote e os quadros-chave do segundo lote e, em seguida, execute uma média ponderada. O peso é determinado pela distância entre o quadro e dois quadros-chave. A fórmula específica é a seguinte:
peso = ∣ s − p 1 ∣ / ( ∣ s − p 1 ∣ + ∣ s − p 2 ∣ ) peso = |s - p1| / (|s - p1| + |s - p2 |)peso _ _ _ _ _=∣s _-p 1∣/ ( ∣ s-p 1∣+∣s _-p 2∣ )
onde s representa o número do quadro, p1 representa o número do primeiro quadro-chave e p2 representa o número do segundo quadro-chave.
# 1. Self-Attention
cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {
}
if self.pivotal_pass:
# norm_hidden_states.shape = 3, n_frames * seq_len, dim
self.attn_output = self.attn1(
norm_hidden_states.view(batch_size, sequence_length, dim),
encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,
**cross_attention_kwargs,
)
# 3, n_frames * seq_len, dim - > 3 * n_frames, seq_len, dim
self.kf_attn_output = self.attn_output
else:
batch_kf_size, _, _ = self.kf_attn_output.shape
self.attn_output = self.kf_attn_output.view(3, batch_kf_size // 3, sequence_length, dim)[:,
batch_idxs] # 3, n_frames, seq_len, dim --> 3, len(batch_idxs), seq_len, dim
# gather values from attn_output, using idx as indices, and get a tensor of shape 3, n_frames, seq_len, dim
if not self.pivotal_pass:
if len(batch_idxs) == 2:
attn_1, attn_2 = self.attn_output[:, 0], self.attn_output[:, 1]
attn_output1 = attn_1.gather(dim=1, index=idx1.unsqueeze(-1).repeat(1, 1, dim))
attn_output2 = attn_2.gather(dim=1, index=idx2.unsqueeze(-1).repeat(1, 1, dim))
s = torch.arange(0, n_frames).to(idx1.device) + batch_idxs[0] * n_frames
# distance from the pivot
p1 = batch_idxs[0] * n_frames + n_frames // 2
p2 = batch_idxs[1] * n_frames + n_frames // 2
d1 = torch.abs(s - p1)
d2 = torch.abs(s - p2)
# weight
w1 = d2 / (d1 + d2)
w1 = torch.sigmoid(w1)
w1 = w1.unsqueeze(0).unsqueeze(-1).unsqueeze(-1).repeat(3, 1, sequence_length, dim)
attn_output1 = attn_output1.view(3, n_frames, sequence_length, dim)
attn_output2 = attn_output2.view(3, n_frames, sequence_length, dim)
attn_output = w1 * attn_output1 + (1 - w1) * attn_output2
else:
attn_output = self.attn_output[:,0].gather(dim=1, index=idx1.unsqueeze(-1).repeat(1, 1, dim))
attn_output = attn_output.reshape(
batch_size, sequence_length, dim) # 3 * n_frames, seq_len, dim
else:
attn_output = self.attn_output
hidden_states = hidden_states.reshape(batch_size, sequence_length, dim) # 3 * n_frames, seq_len, dim
hidden_states = attn_output + hidden_states # res_connect