Spaces:
Running
Running
File size: 5,288 Bytes
736dcf7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | # Modified from https://github.com/1zb/functional-diffusion
import numpy as np
import torch
from torch import nn, einsum
import torch.nn.functional as F
from einops import rearrange, repeat
from timm.models.layers import DropPath
def exists(val):
return val is not None
def default(val, d):
return val if exists(val) else d
class PointEmbed(nn.Module):
def __init__(self, hidden_dim=48, dim=128):
super().__init__()
assert hidden_dim % 12 == 0
self.embedding_dim = hidden_dim
chunk_size = self.embedding_dim // 12
freq = torch.pow(2, torch.arange(chunk_size)).float() * np.pi
e = torch.zeros(6, chunk_size * 6)
for i in range(6):
start_idx = i * chunk_size
end_idx = start_idx + chunk_size
e[i, start_idx:end_idx] = freq
self.register_buffer('basis', e)
self.mlp = nn.Linear(self.embedding_dim + 6, dim)
@staticmethod
def embed(input, basis):
projections = torch.einsum('bnd,de->bne', input, basis)
embeddings = torch.cat([projections.sin(), projections.cos()], dim=2)
return embeddings
def forward(self, input):
# input: B x N x 6
x = self.embed(input, self.basis) # B,N,48
embed = self.mlp(torch.cat([x, input], dim=2)) # B x N x C
return embed
class PreNorm(nn.Module):
def __init__(self, dim, fn, context_dim = None, modulated=False):
super().__init__()
self.fn = fn
self.norm = nn.LayerNorm(dim)
self.norm_context = nn.LayerNorm(context_dim) if exists(context_dim) else None
self.modulated = modulated
if self.modulated:
self.gamma = nn.Linear(dim, dim, bias=False)
self.beta = nn.Linear(dim, dim, bias=False)
def forward(self, x, **kwargs):
x = self.norm(x)
if self.modulated:
label = kwargs.pop('label')
gamma = self.gamma(label) # b 1 c
beta = self.beta(label) # b 1 c
x = gamma * x + beta
if exists(self.norm_context):
context = kwargs['context']
normed_context = self.norm_context(context)
kwargs.update(context = normed_context)
return self.fn(x, **kwargs)
class GEGLU(nn.Module):
def forward(self, x):
x, gates = x.chunk(2, dim = -1)
return x * F.gelu(gates)
class FeedForward(nn.Module):
def __init__(self, dim, mult = 4, drop_path_rate = 0.0):
super().__init__()
self.net = nn.Sequential(
nn.Linear(dim, dim * mult * 2),
GEGLU(),
nn.Linear(dim * mult, dim)
)
self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()
def forward(self, x):
return self.drop_path(self.net(x))
class Attention(nn.Module):
def __init__(self, query_dim, context_dim = None, heads = 8, dim_head = 64, drop_path_rate = 0.0):
super().__init__()
inner_dim = dim_head * heads
context_dim = default(context_dim, query_dim)
self.scale = dim_head ** -0.5
self.heads = heads
self.to_q = nn.Linear(query_dim, inner_dim, bias = False)
self.to_kv = nn.Linear(context_dim, inner_dim * 2, bias = False)
self.to_out = nn.Linear(inner_dim, query_dim)
self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()
def forward(self, x, context = None, query_mask = None, context_mask=None, rel_pos=None):
h = self.heads
q = self.to_q(x)
context = default(context, x)
k, v = self.to_kv(context).chunk(2, dim = -1)
q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h = h), (q, k, v))
sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
if exists(rel_pos):
# rel_pos shape expected to be [b, i, j, h] or [b h, i, j]
if rel_pos.dim() == 4: # [b, i, j, h]
# Reshape to match attention heads dimension
rel_pos = rearrange(rel_pos, 'b i j h -> (b h) i j')
# Add the relative positional bias to the attention scores
sim = sim + rel_pos
if exists(query_mask): # shape (B, Nq)
query_mask = query_mask.bool()
if query_mask.dim() == 2:
query_mask = repeat(query_mask, 'b i -> (b h) i 1', h=h)
elif query_mask.dim() == 3:
query_mask = repeat(query_mask, 'b n j -> (b h) n j', h=h)
sim.masked_fill_(~query_mask, -torch.finfo(sim.dtype).max)
if exists(context_mask):
context_mask_bool = context_mask.bool()
if context_mask_bool.dim() == 2:
context_mask_bool = repeat(context_mask_bool, 'b j -> (b h) 1 j', h=h)
elif context_mask_bool.dim() == 3:
context_mask_bool = repeat(context_mask_bool, 'b n j -> (b h) n j', h=h)
sim.masked_fill_(~context_mask_bool, -torch.finfo(sim.dtype).max)
attn = sim.softmax(dim = -1)
out = einsum('b i j, b j d -> b i d', attn, v)
out = rearrange(out, '(b h) n d -> b n (h d)', h = h)
return self.drop_path(self.to_out(out))
|