Spaces:
Running
Running
| # 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) | |
| 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)) | |