90 lines
2.9 KiB
Python
90 lines
2.9 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
from torch import nn
|
|
|
|
import comfy.ops
|
|
ops = comfy.ops.disable_weight_init
|
|
from comfy.ldm.modules.attention import optimized_attention
|
|
|
|
|
|
class SSRAligner(nn.Module):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
query_dim = 768
|
|
inner_dim = 512
|
|
cross_dim = 1024
|
|
self.heads = 8
|
|
|
|
self.to_q = ops.Linear(query_dim, inner_dim, bias=False)
|
|
self.to_k = ops.Linear(cross_dim, inner_dim, bias=False)
|
|
|
|
self.to_v = nn.ModuleList([])
|
|
self.to_out = nn.ModuleList([])
|
|
for _ in range(6):
|
|
self.to_v.append(ops.Linear(cross_dim, inner_dim, bias=False))
|
|
self.to_out.append(ops.Linear(inner_dim, query_dim, bias=False))
|
|
self.norm = ops.LayerNorm(query_dim)
|
|
|
|
def forward(self, x, context_list):
|
|
q = self.to_q(x)
|
|
k = self.to_k(context_list[-1])
|
|
|
|
outs = []
|
|
for i in range(6):
|
|
v = self.to_v[i](context_list[i])
|
|
attn_out = optimized_attention(q, k, v, self.heads)
|
|
outs.append(self.to_out[i](attn_out))
|
|
|
|
return self.norm(torch.cat(outs, dim=1)) # [1, 77*6, 768]
|
|
|
|
|
|
# me being lazy
|
|
ssr_caa_dims = {
|
|
'input': [320, 320, 640, 640, 1280, 1280],
|
|
'middle': [1280],
|
|
'output': [1280, 1280, 1280, 640, 640, 640, 320, 320, 320]
|
|
}
|
|
|
|
|
|
class SSRCrossAttention(nn.Module):
|
|
def __init__(self, hidden_size) -> None:
|
|
super().__init__()
|
|
self.heads = 8
|
|
|
|
self.to_k = ops.Linear(768, hidden_size, bias=False)
|
|
self.to_v = ops.Linear(768, hidden_size, bias=False)
|
|
|
|
self.pos_contexts = [] # [{ 'scale': int, 'embed': tensor }, ...]
|
|
self.neg_contexts = []
|
|
|
|
def to(self, device):
|
|
super().to(device)
|
|
for i in range(len(self.pos_contexts)):
|
|
self.pos_contexts[i]['embed'] = self.pos_contexts[i]['embed'].to(device)
|
|
|
|
for i in range(len(self.neg_contexts)):
|
|
self.neg_contexts[i]['embed'] = self.neg_contexts[i]['embed'].to(device)
|
|
|
|
return self
|
|
|
|
def forward(self, q, k, v, extra_options):
|
|
|
|
hidden_states = optimized_attention(q, k, v, extra_options['n_heads'])
|
|
cond_or_uncond = extra_options['cond_or_uncond']
|
|
len_conds = len(cond_or_uncond)
|
|
batch_size = len(q)//len_conds
|
|
|
|
for cond_idx, cond in enumerate(cond_or_uncond):
|
|
contexts = self.pos_contexts if cond == 0 else self.neg_contexts
|
|
q_cond = q[cond_idx*batch_size:(cond_idx+1)*batch_size]
|
|
for context in contexts:
|
|
scale = context['scale']
|
|
embed = context['embed']
|
|
ssr_k = self.to_k(embed).repeat(batch_size, 1, 1)
|
|
ssr_v = self.to_v(embed).repeat(batch_size, 1, 1)
|
|
ssr_hidden_states = optimized_attention(q_cond, ssr_k, ssr_v, self.heads)
|
|
|
|
hidden_states[cond_idx*batch_size:(cond_idx+1)*batch_size] += scale * ssr_hidden_states
|
|
|
|
return hidden_states
|