Spaces:
Sleeping
Sleeping
| from inspect import isfunction | |
| import math | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import nn, einsum | |
| from einops import rearrange, repeat | |
| # import configigure | |
| # from ldm.modules.diffusionmodules.util import checkpoint, FourierEmbedder | |
| from torch.utils import checkpoint | |
| import os | |
| from torchvision.utils import save_image | |
| iter_att = 0 | |
| def exists(val): | |
| return val is not None | |
| def uniq(arr): | |
| return{el: True for el in arr}.keys() | |
| def default(val, d): | |
| if exists(val): | |
| return val | |
| return d() if isfunction(d) else d | |
| def max_neg_value(t): | |
| return -torch.finfo(t.dtype).max | |
| def init_(tensor): | |
| dim = tensor.shape[-1] | |
| std = 1 / math.sqrt(dim) | |
| tensor.uniform_(-std, std) | |
| return tensor | |
| # feedforward | |
| class GEGLU(nn.Module): | |
| def __init__(self, dim_in, dim_out): | |
| super().__init__() | |
| self.proj = nn.Linear(dim_in, dim_out * 2) | |
| def forward(self, x): | |
| x, gate = self.proj(x).chunk(2, dim=-1) | |
| return x * F.gelu(gate) | |
| class FeedForward(nn.Module): | |
| def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.): | |
| super().__init__() | |
| inner_dim = int(dim * mult) | |
| dim_out = default(dim_out, dim) | |
| project_in = nn.Sequential( | |
| nn.Linear(dim, inner_dim), | |
| nn.GELU() | |
| ) if not glu else GEGLU(dim, inner_dim) | |
| self.net = nn.Sequential( | |
| project_in, | |
| nn.Dropout(dropout), | |
| nn.Linear(inner_dim, dim_out) | |
| ) | |
| def forward(self, x): | |
| return self.net(x) | |
| def zero_module(module): | |
| """ | |
| Zero out the parameters of a module and return it. | |
| """ | |
| for p in module.parameters(): | |
| p.detach().zero_() | |
| return module | |
| def Normalize(in_channels): | |
| return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) | |
| class LinearAttention(nn.Module): | |
| def __init__(self, dim, heads=4, dim_head=32): | |
| super().__init__() | |
| self.heads = heads | |
| hidden_dim = dim_head * heads | |
| self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False) | |
| self.to_out = nn.Conv2d(hidden_dim, dim, 1) | |
| def forward(self, x): | |
| b, c, h, w = x.shape | |
| qkv = self.to_qkv(x) | |
| q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads, qkv=3) | |
| k = k.softmax(dim=-1) | |
| context = torch.einsum('bhdn,bhen->bhde', k, v) | |
| out = torch.einsum('bhde,bhdn->bhen', context, q) | |
| out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w) | |
| return self.to_out(out) | |
| class CrossAttention(nn.Module): | |
| def __init__(self, query_dim, key_dim, value_dim, heads=8, dim_head=64, dropout=0): | |
| super().__init__() | |
| inner_dim = dim_head * heads | |
| self.scale = dim_head ** -0.5 | |
| self.heads = heads | |
| self.to_q = nn.Linear(query_dim, inner_dim, bias=False) | |
| self.to_k = nn.Linear(key_dim, inner_dim, bias=False) | |
| self.to_v = nn.Linear(value_dim, inner_dim, bias=False) | |
| self.to_out = nn.Sequential( nn.Linear(inner_dim, query_dim), nn.Dropout(dropout) ) | |
| def fill_inf_from_mask(self, sim, mask): | |
| if mask is not None: | |
| B,M = mask.shape | |
| mask = mask.unsqueeze(1).repeat(1,self.heads,1).reshape(B*self.heads,1,-1) | |
| max_neg_value = -torch.finfo(sim.dtype).max | |
| sim.masked_fill_(~mask, max_neg_value) | |
| return sim | |
| # def scaled_dot_product(q, k, v, mask=None): | |
| # d_k = q.size()[-1] | |
| # attn_logits = torch.matmul(q, k.transpose(-2, -1)) | |
| # attn_logits = attn_logits / math.sqrt(d_k) | |
| # if mask is not None: | |
| # attn_logits = attn_logits.masked_fill(mask == 0, -9e15) | |
| # attention = F.softmax(attn_logits, dim=-1) | |
| # values = torch.matmul(attention, v) | |
| # return values, attention | |
| def forward(self, x, key, value, mask=None): | |
| # import pdb; pdb.set_trace() | |
| q = self.to_q(x) # B*N*(H*C) | |
| k = self.to_k(key) # B*M*(H*C) | |
| v = self.to_v(value) # B*M*(H*C) | |
| B, N, HC = q.shape | |
| _, M, _ = key.shape | |
| H = self.heads | |
| C = HC // H | |
| q = q.view(B,N,H,C).permute(0,2,1,3).reshape(B*H,N,C) # (B*H)*N*C | |
| k = k.view(B,M,H,C).permute(0,2,1,3).reshape(B*H,M,C) # (B*H)*M*C | |
| v = v.view(B,M,H,C).permute(0,2,1,3).reshape(B*H,M,C) # (B*H)*M*C | |
| sim = torch.einsum('b i d, b j d -> b i j', q, k) * self.scale # (B*H)*N*M | |
| self.fill_inf_from_mask(sim, mask) | |
| attn = sim.softmax(dim=-1) # (B*H)*N*M | |
| # import pdb; pdb.set_trace() | |
| # if attn.shape[1] == 4096: | |
| # self.visual_att(attn) | |
| out = torch.einsum('b i j, b j d -> b i d', attn, v) # (B*H)*N*C | |
| out = out.view(B,H,N,C).permute(0,2,1,3).reshape(B,N,(H*C)) # B*N*(H*C) | |
| return self.to_out(out), attn | |
| def visual_att(self, att): | |
| global iter_att | |
| ll = [0,2,7] | |
| for i in range(12): | |
| kk = torch.sum(att[:,:,i], axis=0) | |
| kk = kk.reshape(64,64) | |
| save_image( (kk-kk.min()) / (kk.max() - kk.min()) , os.path.join('att', str(iter_att) + '_' +str(i) + '.png')) | |
| iter_att += 1 | |
| class SelfAttention(nn.Module): | |
| def __init__(self, query_dim, heads=8, dim_head=64, dropout=0.): | |
| super().__init__() | |
| inner_dim = dim_head * heads | |
| self.scale = dim_head ** -0.5 | |
| self.heads = heads | |
| self.to_q = nn.Linear(query_dim, inner_dim, bias=False) | |
| self.to_k = nn.Linear(query_dim, inner_dim, bias=False) | |
| self.to_v = nn.Linear(query_dim, inner_dim, bias=False) | |
| self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout) ) | |
| def forward(self, x, gated=False): | |
| q = self.to_q(x) # B*N*(H*C) | |
| k = self.to_k(x) # B*N*(H*C) | |
| v = self.to_v(x) # B*N*(H*C) | |
| B, N, HC = q.shape | |
| H = self.heads | |
| C = HC // H | |
| # if gated: import pdb; pdb.set_trace() | |
| # import pdb; pdb.set_trace() | |
| q = q.view(B,N,H,C).permute(0,2,1,3).reshape(B*H,N,C) # (B*H)*N*C | |
| k = k.view(B,N,H,C).permute(0,2,1,3).reshape(B*H,N,C) # (B*H)*N*C | |
| v = v.view(B,N,H,C).permute(0,2,1,3).reshape(B*H,N,C) # (B*H)*N*C | |
| sim = torch.einsum('b i c, b j c -> b i j', q, k) * self.scale # (B*H)*N*N | |
| attn = sim.softmax(dim=-1) # (B*H)*N*N | |
| # if gated and attn.shape[1] == 4126: | |
| # self.visual_att(attn) | |
| out = torch.einsum('b i j, b j c -> b i c', attn, v) # (B*H)*N*C | |
| out = out.view(B,H,N,C).permute(0,2,1,3).reshape(B,N,(H*C)) # B*N*(H*C) | |
| return self.to_out(out), attn | |
| def visual_att(self, att): | |
| global iter_att | |
| ll = [0,2,7] | |
| for i in range(): | |
| kk = torch.sum(att[i],axis=0) | |
| kk = kk[:4096].reshape(64,64) | |
| save_image( (kk-kk.min()) / (kk.max() - kk.min()) , os.path.join('att', str(iter_att) + '_' +str(i) + '.png')) | |
| iter_att += 1 | |
| class GatedCrossAttentionDense(nn.Module): | |
| def __init__(self, query_dim, key_dim, value_dim, n_heads, d_head): | |
| super().__init__() | |
| self.attn = CrossAttention(query_dim=query_dim, key_dim=key_dim, value_dim=value_dim, heads=n_heads, dim_head=d_head) | |
| self.ff = FeedForward(query_dim, glu=True) | |
| self.norm1 = nn.LayerNorm(query_dim) | |
| self.norm2 = nn.LayerNorm(query_dim) | |
| self.register_parameter('alpha_attn', nn.Parameter(torch.tensor(0.)) ) | |
| self.register_parameter('alpha_dense', nn.Parameter(torch.tensor(0.)) ) | |
| # this can be useful: we can externally change magnitude of tanh(alpha) | |
| # for example, when it is set to 0, then the entire model is same as original one | |
| self.scale = 1 | |
| def forward(self, x, objs): | |
| x = x + self.scale*torch.tanh(self.alpha_attn) * self.attn( self.norm1(x), objs, objs) | |
| x = x + self.scale*torch.tanh(self.alpha_dense) * self.ff( self.norm2(x) ) | |
| return x | |
| class GatedSelfAttentionDense(nn.Module): | |
| def __init__(self, query_dim, context_dim, n_heads, d_head): | |
| super().__init__() | |
| # we need a linear projection since we need cat visual feature and obj feature | |
| self.linear = nn.Linear(context_dim, query_dim) | |
| self.attn = SelfAttention(query_dim=query_dim, heads=n_heads, dim_head=d_head) | |
| self.ff = FeedForward(query_dim, glu=True) | |
| self.norm1 = nn.LayerNorm(query_dim) | |
| self.norm2 = nn.LayerNorm(query_dim) | |
| self.register_parameter('alpha_attn', nn.Parameter(torch.tensor(0.)) ) | |
| self.register_parameter('alpha_dense', nn.Parameter(torch.tensor(0.)) ) | |
| # this can be useful: we can externally change magnitude of tanh(alpha) | |
| # for example, when it is set to 0, then the entire model is same as original one | |
| self.scale = 1 | |
| def forward(self, x, objs,t): | |
| # if t >300: | |
| # self.scale = 1 | |
| # elif t > 200: | |
| # self.scale = 0.9 | |
| # else: | |
| # self.scale = 0.6 | |
| # if t >700: | |
| # self.scale = 1 | |
| # elif t > 300: | |
| # self.scale = 0.7 | |
| # else: | |
| # self.scale = 0.4 | |
| # self.scale = 0 | |
| N_visual = x.shape[1] | |
| objs = self.linear(objs) | |
| out, grounding_att = self.attn( self.norm1(torch.cat([x,objs],dim=1)), True ) | |
| out = out[:,0:N_visual,:] | |
| x = x + self.scale*torch.tanh(self.alpha_attn) * out | |
| x = x + self.scale*torch.tanh(self.alpha_dense) * self.ff( self.norm2(x) ) | |
| return x , grounding_att | |
| class GatedSelfAttentionDense2(nn.Module): | |
| def __init__(self, query_dim, context_dim, n_heads, d_head): | |
| super().__init__() | |
| # we need a linear projection since we need cat visual feature and obj feature | |
| self.linear = nn.Linear(context_dim, query_dim) | |
| self.attn = SelfAttention(query_dim=query_dim, heads=n_heads, dim_head=d_head) | |
| self.ff = FeedForward(query_dim, glu=True) | |
| self.norm1 = nn.LayerNorm(query_dim) | |
| self.norm2 = nn.LayerNorm(query_dim) | |
| self.register_parameter('alpha_attn', nn.Parameter(torch.tensor(0.)) ) | |
| self.register_parameter('alpha_dense', nn.Parameter(torch.tensor(0.)) ) | |
| # this can be useful: we can externally change magnitude of tanh(alpha) | |
| # for example, when it is set to 0, then the entire model is same as original one | |
| self.scale = 1 | |
| def forward(self, x, objs): | |
| B, N_visual, _ = x.shape | |
| B, N_ground, _ = objs.shape | |
| objs = self.linear(objs) | |
| # sanity check | |
| size_v = math.sqrt(N_visual) | |
| size_g = math.sqrt(N_ground) | |
| assert int(size_v) == size_v, "Visual tokens must be square rootable" | |
| assert int(size_g) == size_g, "Grounding tokens must be square rootable" | |
| size_v = int(size_v) | |
| size_g = int(size_g) | |
| # select grounding token and resize it to visual token size as residual | |
| out = self.attn( self.norm1(torch.cat([x,objs],dim=1)) )[:,N_visual:,:] | |
| out = out.permute(0,2,1).reshape( B,-1,size_g,size_g ) | |
| out = torch.nn.functional.interpolate(out, (size_v,size_v), mode='bicubic') | |
| residual = out.reshape(B,-1,N_visual).permute(0,2,1) | |
| # add residual to visual feature | |
| x = x + self.scale*torch.tanh(self.alpha_attn) * residual | |
| x = x + self.scale*torch.tanh(self.alpha_dense) * self.ff( self.norm2(x) ) | |
| return x | |
| class BasicTransformerBlock(nn.Module): | |
| def __init__(self, query_dim, key_dim, value_dim, n_heads, d_head, fuser_type, use_checkpoint=True): | |
| super().__init__() | |
| self.attn1 = SelfAttention(query_dim=query_dim, heads=n_heads, dim_head=d_head) | |
| self.ff = FeedForward(query_dim, glu=True) | |
| self.attn2 = CrossAttention(query_dim=query_dim, key_dim=key_dim, value_dim=value_dim, heads=n_heads, dim_head=d_head) | |
| self.norm1 = nn.LayerNorm(query_dim) | |
| self.norm2 = nn.LayerNorm(query_dim) | |
| self.norm3 = nn.LayerNorm(query_dim) | |
| self.use_checkpoint = use_checkpoint | |
| if fuser_type == "gatedSA": | |
| # note key_dim here actually is context_dim | |
| self.fuser = GatedSelfAttentionDense(query_dim, key_dim, n_heads, d_head) | |
| elif fuser_type == "gatedSA2": | |
| # note key_dim here actually is context_dim | |
| self.fuser = GatedSelfAttentionDense2(query_dim, key_dim, n_heads, d_head) | |
| elif fuser_type == "gatedCA": | |
| self.fuser = GatedCrossAttentionDense(query_dim, key_dim, value_dim, n_heads, d_head) | |
| else: | |
| assert False | |
| def forward(self, x, context, objs,t): | |
| # return checkpoint(self._forward, (x, context, objs), self.parameters(), self.use_checkpoint) | |
| # import pdb; pdb.set_trace() | |
| # if self.use_checkpoint and x.requires_grad: | |
| # return checkpoint.checkpoint(self._forward, x, context, objs,t) | |
| # else: | |
| return self._forward(x, context, objs,t) | |
| def _forward(self, x, context, objs,t): | |
| # self_att_grounding = [] | |
| out, self_prob = self.attn1( self.norm1(x) ) | |
| x = x + out | |
| x, self_prob_grounding = self.fuser(x, objs,t) # identity mapping in the beginning | |
| x_1, prob = self.attn2(self.norm2(x), context, context) | |
| x = x + x_1 | |
| x = self.ff(self.norm3(x)) + x | |
| # self_att_grounding.append(self_prob) | |
| # self_att_grounding.append(self_prob_grounding) | |
| return x, prob, self_prob | |
| class SpatialTransformer(nn.Module): | |
| def __init__(self, in_channels, key_dim, value_dim, n_heads, d_head, depth=1, fuser_type=None, use_checkpoint=True): | |
| super().__init__() | |
| self.in_channels = in_channels | |
| query_dim = n_heads * d_head | |
| self.norm = Normalize(in_channels) | |
| self.proj_in = nn.Conv2d(in_channels, | |
| query_dim, | |
| kernel_size=1, | |
| stride=1, | |
| padding=0) | |
| self.transformer_blocks = nn.ModuleList( | |
| [BasicTransformerBlock(query_dim, key_dim, value_dim, n_heads, d_head, fuser_type, use_checkpoint=use_checkpoint) | |
| for d in range(depth)] | |
| ) | |
| self.proj_out = zero_module(nn.Conv2d(query_dim, | |
| in_channels, | |
| kernel_size=1, | |
| stride=1, | |
| padding=0)) | |
| def forward(self, x, context, objs,t): | |
| b, c, h, w = x.shape | |
| x_in = x | |
| x = self.norm(x) | |
| x = self.proj_in(x) | |
| x = rearrange(x, 'b c h w -> b (h w) c') | |
| probs = [] | |
| self_prob_list = [] | |
| for block in self.transformer_blocks: | |
| x, prob, self_prob = block(x, context, objs,t) | |
| probs.append(prob) | |
| self_prob_list.append(self_prob) | |
| x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w) | |
| x = self.proj_out(x) | |
| return x + x_in, probs, self_prob_list |