236 lines
7.6 KiB
Python
236 lines
7.6 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
from models.diffusion.nn import normalization, conv_nd, zero_module
|
|
from models.diffusion.unet_diffusion import Downsample, AttentionBlock, QKVAttention, QKVAttentionLegacy, Upsample
|
|
|
|
# Combined resnet & full-attention encoder for converting an audio clip into an embedding.
|
|
from trainer.networks import register_model
|
|
from utils.util import checkpoint, opt_get
|
|
|
|
|
|
class ResBlock(nn.Module):
|
|
def __init__(
|
|
self,
|
|
channels,
|
|
dropout,
|
|
out_channels=None,
|
|
use_conv=False,
|
|
use_scale_shift_norm=False,
|
|
dims=2,
|
|
up=False,
|
|
down=False,
|
|
kernel_size=3,
|
|
do_checkpoint=True,
|
|
):
|
|
super().__init__()
|
|
self.channels = channels
|
|
self.dropout = dropout
|
|
self.out_channels = out_channels or channels
|
|
self.use_conv = use_conv
|
|
self.use_scale_shift_norm = use_scale_shift_norm
|
|
self.do_checkpoint = do_checkpoint
|
|
padding = 1 if kernel_size == 3 else 2
|
|
|
|
self.in_layers = nn.Sequential(
|
|
normalization(channels),
|
|
nn.SiLU(),
|
|
conv_nd(dims, channels, self.out_channels, kernel_size, padding=padding),
|
|
)
|
|
|
|
self.updown = up or down
|
|
|
|
if up:
|
|
self.h_upd = Upsample(channels, False, dims)
|
|
self.x_upd = Upsample(channels, False, dims)
|
|
elif down:
|
|
self.h_upd = Downsample(channels, False, dims)
|
|
self.x_upd = Downsample(channels, False, dims)
|
|
else:
|
|
self.h_upd = self.x_upd = nn.Identity()
|
|
|
|
self.out_layers = nn.Sequential(
|
|
normalization(self.out_channels),
|
|
nn.SiLU(),
|
|
nn.Dropout(p=dropout),
|
|
zero_module(
|
|
conv_nd(dims, self.out_channels, self.out_channels, kernel_size, padding=padding)
|
|
),
|
|
)
|
|
|
|
if self.out_channels == channels:
|
|
self.skip_connection = nn.Identity()
|
|
elif use_conv:
|
|
self.skip_connection = conv_nd(
|
|
dims, channels, self.out_channels, kernel_size, padding=padding
|
|
)
|
|
else:
|
|
self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
|
|
|
|
def forward(self, x):
|
|
if self.do_checkpoint:
|
|
return checkpoint(
|
|
self._forward, x
|
|
)
|
|
else:
|
|
return self._forward(x)
|
|
|
|
def _forward(self, x):
|
|
if self.updown:
|
|
in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
|
|
h = in_rest(x)
|
|
h = self.h_upd(h)
|
|
x = self.x_upd(x)
|
|
h = in_conv(h)
|
|
else:
|
|
h = self.in_layers(x)
|
|
h = self.out_layers(h)
|
|
return self.skip_connection(x) + h
|
|
|
|
|
|
class AudioMiniEncoder(nn.Module):
|
|
def __init__(self,
|
|
spec_dim,
|
|
embedding_dim,
|
|
base_channels=128,
|
|
depth=2,
|
|
resnet_blocks=2,
|
|
attn_blocks=4,
|
|
num_attn_heads=4,
|
|
dropout=0,
|
|
downsample_factor=2,
|
|
kernel_size=3,
|
|
do_checkpointing=False):
|
|
super().__init__()
|
|
self.init = nn.Sequential(
|
|
conv_nd(1, spec_dim, base_channels, 3, padding=1)
|
|
)
|
|
ch = base_channels
|
|
res = []
|
|
for l in range(depth):
|
|
for r in range(resnet_blocks):
|
|
res.append(ResBlock(ch, dropout, dims=1, do_checkpoint=False, kernel_size=kernel_size))
|
|
res.append(Downsample(ch, use_conv=True, dims=1, out_channels=ch*2, factor=downsample_factor))
|
|
ch *= 2
|
|
self.res = nn.Sequential(*res)
|
|
self.final = nn.Sequential(
|
|
normalization(ch),
|
|
nn.SiLU(),
|
|
conv_nd(1, ch, embedding_dim, 1)
|
|
)
|
|
attn = []
|
|
for a in range(attn_blocks):
|
|
attn.append(AttentionBlock(embedding_dim, num_attn_heads, do_checkpoint=False))
|
|
self.attn = nn.Sequential(*attn)
|
|
self.dim = embedding_dim
|
|
self.do_checkpointing = do_checkpointing
|
|
|
|
def forward(self, x):
|
|
h = self.init(x)
|
|
h = self.res(h)
|
|
h = self.final(h)
|
|
if self.do_checkpointing:
|
|
h = checkpoint(self.attn, h)
|
|
else:
|
|
h = self.attn(h)
|
|
return h[:, :, 0]
|
|
|
|
|
|
class AudioMiniEncoderWithClassifierHead(nn.Module):
|
|
def __init__(self, classes, **kwargs):
|
|
super().__init__()
|
|
self.enc = AudioMiniEncoder(**kwargs)
|
|
self.head = nn.Linear(self.enc.dim, classes)
|
|
|
|
def forward(self, x):
|
|
h = self.enc(x)
|
|
return self.head(h)
|
|
|
|
|
|
class QueryProvidedAttentionBlock(nn.Module):
|
|
"""
|
|
An attention block that provides a separate signal for the query vs the keys/parameters.
|
|
"""
|
|
def __init__(
|
|
self,
|
|
channels,
|
|
num_heads=1,
|
|
num_head_channels=-1,
|
|
use_new_attention_order=False,
|
|
):
|
|
super().__init__()
|
|
self.channels = channels
|
|
if num_head_channels == -1:
|
|
self.num_heads = num_heads
|
|
else:
|
|
assert (
|
|
channels % num_head_channels == 0
|
|
), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
|
|
self.num_heads = channels // num_head_channels
|
|
self.norm = normalization(channels)
|
|
self.q = nn.Linear(channels, channels)
|
|
self.qnorm = nn.LayerNorm(channels)
|
|
self.kv = conv_nd(1, channels, channels*2, 1)
|
|
if use_new_attention_order:
|
|
# split qkv before split heads
|
|
self.attention = QKVAttention(self.num_heads)
|
|
else:
|
|
# split heads before split qkv
|
|
self.attention = QKVAttentionLegacy(self.num_heads)
|
|
|
|
self.proj_out = zero_module(conv_nd(1, channels, channels, 1))
|
|
|
|
def forward(self, qx, kvx, mask=None):
|
|
return checkpoint(self._forward, qx, kvx, mask)
|
|
|
|
def _forward(self, qx, kvx, mask=None):
|
|
q = self.q(self.qnorm(qx)).unsqueeze(1).repeat(1, kvx.shape[1], 1).permute(0,2,1)
|
|
kv = self.kv(self.norm(kvx.permute(0,2,1)))
|
|
qkv = torch.cat([q, kv], dim=1)
|
|
h = self.attention(qkv, mask)
|
|
h = self.proj_out(h)
|
|
return kvx + h.permute(0,2,1)
|
|
|
|
|
|
# Next up: combine multiple embeddings given a conditioning signal into a single embedding.
|
|
class EmbeddingCombiner(nn.Module):
|
|
def __init__(self, embedding_dim, attn_blocks=3, num_attn_heads=2, cond_provided=True):
|
|
super().__init__()
|
|
block = QueryProvidedAttentionBlock if cond_provided else AttentionBlock
|
|
self.attn = nn.ModuleList([block(embedding_dim, num_attn_heads) for _ in range(attn_blocks)])
|
|
self.cond_provided = cond_provided
|
|
|
|
# x_s: (b,n,d); b=batch_sz, n=number of embeddings, d=embedding_dim
|
|
# cond: (b,d) or None
|
|
def forward(self, x_s, attn_mask=None, cond=None):
|
|
assert cond is not None and self.cond_provided or cond is None and not self.cond_provided
|
|
y = x_s
|
|
for blk in self.attn:
|
|
if self.cond_provided:
|
|
y = blk(cond, y, mask=attn_mask)
|
|
else:
|
|
y = blk(y, mask=attn_mask)
|
|
return y[:, 0]
|
|
|
|
|
|
@register_model
|
|
def register_mini_audio_encoder_classifier(opt_net, opt):
|
|
return AudioMiniEncoderWithClassifierHead(**opt_get(opt_net, ['kwargs'], {}))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
'''
|
|
x = torch.randn(2, 80, 223)
|
|
cond = torch.randn(2, 512)
|
|
encs = [AudioMiniEncoder(80, 512) for _ in range(5)]
|
|
combiner = EmbeddingCombiner(512)
|
|
|
|
e = torch.stack([e(x) for e in encs], dim=2)
|
|
|
|
print(combiner(e, cond).shape)
|
|
'''
|
|
x = torch.randn(2, 80, 223)
|
|
m = AudioMiniEncoderWithClassifierHead(4, 80, 512)
|
|
print(m(x).shape)
|