47 lines
1.8 KiB
Python
47 lines
1.8 KiB
Python
from einops import rearrange
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
def flow_attention(img_attn, num_heads, dim_head, transformer_options):
|
|
flow = transformer_options.get('FLOW', None)
|
|
if flow is None:
|
|
return img_attn
|
|
|
|
hidden_states = img_attn
|
|
cond_list = transformer_options['cond_or_uncond']
|
|
batch_size = len(hidden_states) // len(cond_list)
|
|
flow = None
|
|
flows = []
|
|
flows = transformer_options['FLOW']['forward_flows']
|
|
|
|
for possible_flow in flows:
|
|
if possible_flow['forward_trajectory'].shape[2] == hidden_states.shape[1]:
|
|
flow = possible_flow
|
|
break
|
|
|
|
if not flow:
|
|
return hidden_states
|
|
|
|
backward_trajectory = flow['backward_trajectory'].to(hidden_states.device)
|
|
forward_trajectory = flow['forward_trajectory'].to(hidden_states.device)
|
|
attn_mask = flow['attn_masks'].to(hidden_states.device)
|
|
|
|
hidden_states = rearrange(hidden_states, "(b f) d c -> f (b c) d", f=batch_size)
|
|
|
|
hidden_states = torch.gather(hidden_states, 2, forward_trajectory.expand(-1,hidden_states.shape[1],-1))
|
|
|
|
hidden_states = rearrange(hidden_states, "f (b c) d -> (b d) f c", b=len(cond_list))
|
|
hidden_states = hidden_states.view(-1, batch_size, num_heads, dim_head).transpose(1, 2).detach()
|
|
|
|
hidden_states = F.scaled_dot_product_attention(
|
|
hidden_states, hidden_states, hidden_states,
|
|
attn_mask = (attn_mask.repeat(len(cond_list),1,1,1))
|
|
)
|
|
|
|
hidden_states = rearrange(hidden_states, "(b d) h f c -> f (b h c) d", b=len(cond_list))
|
|
hidden_states = torch.gather(hidden_states, 2, backward_trajectory.expand(-1,hidden_states.shape[1],-1)).detach()
|
|
hidden_states = rearrange(hidden_states, "f (b h c) d -> (b f) h d c", b=len(cond_list), h=num_heads)
|
|
hidden_states = rearrange(hidden_states, 'b d q f -> b q (d f)')
|
|
|
|
return hidden_states |