317 lines
12 KiB
Python
317 lines
12 KiB
Python
import math
|
|
import torch
|
|
from torch import nn
|
|
from torch.nn import functional as F
|
|
|
|
from . import utils
|
|
|
|
class ResidualBlock(nn.Module):
|
|
def __init__(self, main, skip=None):
|
|
super().__init__()
|
|
self.main = nn.Sequential(*main)
|
|
self.skip = skip if skip else nn.Identity()
|
|
|
|
def forward(self, input):
|
|
return self.main(input) + self.skip(input)
|
|
|
|
|
|
# Noise level (and other) conditioning
|
|
|
|
class ConditionedModule(nn.Module):
|
|
pass
|
|
|
|
|
|
class UnconditionedModule(ConditionedModule):
|
|
def __init__(self, module):
|
|
self.module = module
|
|
|
|
def forward(self, input, cond):
|
|
return self.module(input)
|
|
|
|
|
|
class ConditionedSequential(nn.Sequential, ConditionedModule):
|
|
def forward(self, input, cond):
|
|
for module in self:
|
|
if isinstance(module, ConditionedModule):
|
|
input = module(input, cond)
|
|
else:
|
|
input = module(input)
|
|
return input
|
|
|
|
|
|
class ConditionedResidualBlock(ConditionedModule):
|
|
def __init__(self, *main, skip=None):
|
|
super().__init__()
|
|
self.main = ConditionedSequential(*main)
|
|
self.skip = skip if skip else nn.Identity()
|
|
|
|
def forward(self, input, cond):
|
|
skip = self.skip(input, cond) if isinstance(self.skip, ConditionedModule) else self.skip(input)
|
|
return self.main(input, cond) + skip
|
|
|
|
|
|
class AdaGN(ConditionedModule):
|
|
def __init__(self, feats_in, c_out, num_groups, eps=1e-5, cond_key='cond'):
|
|
super().__init__()
|
|
self.num_groups = num_groups
|
|
self.eps = eps
|
|
self.cond_key = cond_key
|
|
self.mapper = nn.Linear(feats_in, c_out * 2)
|
|
|
|
def forward(self, input, cond):
|
|
weight, bias = self.mapper(cond[self.cond_key]).chunk(2, dim=-1)
|
|
input = F.group_norm(input, self.num_groups, eps=self.eps)
|
|
return torch.addcmul(utils.append_dims(bias, input.ndim), input, utils.append_dims(weight, input.ndim) + 1)
|
|
|
|
class ResConvBlock(ResidualBlock):
|
|
def __init__(self, c_in, c_mid, c_out, is_last=False, kernel_size=5):
|
|
skip = None if c_in == c_out else nn.Conv1d(c_in, c_out, 1, bias=False)
|
|
super().__init__([
|
|
nn.Conv1d(c_in, c_mid, kernel_size, padding=kernel_size//2),
|
|
nn.GroupNorm(1, c_mid),
|
|
nn.GELU(),
|
|
nn.Conv1d(c_mid, c_out, kernel_size, padding=kernel_size//2),
|
|
nn.GroupNorm(1, c_out) if not is_last else nn.Identity(),
|
|
nn.GELU() if not is_last else nn.Identity(),
|
|
], skip)
|
|
|
|
class OutConvBlock(nn.Sequential):
|
|
def __init__(self, c_in, c_mid, c_out, is_last=False, kernel_size=5):
|
|
super().__init__(
|
|
nn.Conv1d(c_in, c_mid, kernel_size, padding=kernel_size//2),
|
|
nn.GroupNorm(1, c_mid),
|
|
nn.GELU(),
|
|
nn.Conv1d(c_mid, c_out, kernel_size, padding=kernel_size//2),
|
|
nn.GroupNorm(1, c_out) if not is_last else nn.Identity(),
|
|
nn.GELU() if not is_last else nn.Identity(),
|
|
)
|
|
|
|
class ResModConvBlock(ConditionedResidualBlock):
|
|
def __init__(self, feats_in, c_in, c_mid, c_out, group_size=32, dropout_rate=0.):
|
|
skip = None if c_in == c_out else nn.Conv1d(c_in, c_out, 1, bias=False)
|
|
super().__init__(
|
|
AdaGN(feats_in, c_in, max(1, c_in // group_size)),
|
|
nn.GELU(),
|
|
nn.Conv1d(c_in, c_mid, 3, padding=1),
|
|
nn.Dropout(dropout_rate, inplace=True),
|
|
AdaGN(feats_in, c_mid, max(1, c_mid // group_size)),
|
|
nn.GELU(),
|
|
nn.Conv1d(c_mid, c_out, 3, padding=1),
|
|
nn.Dropout(dropout_rate, inplace=True),
|
|
skip=skip)
|
|
|
|
class SelfAttention1d(nn.Module):
|
|
def __init__(self, c_in, n_head=1, dropout_rate=0.):
|
|
super().__init__()
|
|
assert c_in % n_head == 0
|
|
self.norm = nn.GroupNorm(1, c_in)
|
|
self.n_head = n_head
|
|
self.qkv_proj = nn.Conv1d(c_in, c_in * 3, 1)
|
|
self.out_proj = nn.Conv1d(c_in, c_in, 1)
|
|
self.dropout = nn.Dropout(dropout_rate, inplace=True)
|
|
|
|
def forward(self, input):
|
|
n, c, s = input.shape
|
|
qkv = self.qkv_proj(self.norm(input))
|
|
qkv = qkv.view(
|
|
[n, self.n_head * 3, c // self.n_head, s]).transpose(2, 3)
|
|
q, k, v = qkv.chunk(3, dim=1)
|
|
scale = k.shape[3]**-0.25
|
|
att = ((q * scale) @ (k.transpose(2, 3) * scale)).softmax(3)
|
|
y = (att @ v).transpose(2, 3).contiguous().view([n, c, s])
|
|
return input + self.dropout(self.out_proj(y))
|
|
|
|
class SelfAttentionMod1d(ConditionedModule):
|
|
def __init__(self, c_in, n_head, norm, dropout_rate=0.):
|
|
super().__init__()
|
|
assert c_in % n_head == 0
|
|
self.norm_in = norm(c_in)
|
|
self.n_head = n_head
|
|
self.qkv_proj = nn.Conv1d(c_in, c_in * 3, 1)
|
|
self.out_proj = nn.Conv1d(c_in, c_in, 1)
|
|
self.dropout = nn.Dropout(dropout_rate)
|
|
|
|
def forward(self, input, cond):
|
|
n, c, s = input.shape
|
|
qkv = self.qkv_proj(self.norm_in(input, cond))
|
|
qkv = qkv.view(
|
|
[n, self.n_head * 3, c // self.n_head, s]).transpose(2, 3)
|
|
q, k, v = qkv.chunk(3, dim=1)
|
|
scale = k.shape[3]**-0.25
|
|
att = ((q * scale) @ (k.transpose(2, 3) * scale)).softmax(3)
|
|
att = self.dropout(att)
|
|
y = (att @ v).transpose(2, 3).contiguous().view([n, c, s])
|
|
return input + self.out_proj(y)
|
|
|
|
class SkipBlock(nn.Module):
|
|
def __init__(self, *main):
|
|
super().__init__()
|
|
self.main = nn.Sequential(*main)
|
|
|
|
def forward(self, input):
|
|
return torch.cat([self.main(input), input], dim=1)
|
|
|
|
class FourierFeatures(nn.Module):
|
|
def __init__(self, in_features, out_features, std=1.):
|
|
super().__init__()
|
|
assert out_features % 2 == 0
|
|
self.weight = nn.Parameter(torch.randn(
|
|
[out_features // 2, in_features]) * std)
|
|
|
|
def forward(self, input):
|
|
f = 2 * math.pi * input @ self.weight.T
|
|
return torch.cat([f.cos(), f.sin()], dim=-1)
|
|
|
|
|
|
def expand_to_planes(input, shape):
|
|
return input[..., None].repeat([1, 1, shape[2]])
|
|
|
|
_kernels = {
|
|
'linear':
|
|
[1 / 8, 3 / 8, 3 / 8, 1 / 8],
|
|
'cubic':
|
|
[-0.01171875, -0.03515625, 0.11328125, 0.43359375,
|
|
0.43359375, 0.11328125, -0.03515625, -0.01171875],
|
|
'lanczos3':
|
|
[0.003689131001010537, 0.015056144446134567, -0.03399861603975296,
|
|
-0.066637322306633, 0.13550527393817902, 0.44638532400131226,
|
|
0.44638532400131226, 0.13550527393817902, -0.066637322306633,
|
|
-0.03399861603975296, 0.015056144446134567, 0.003689131001010537]
|
|
}
|
|
|
|
|
|
class Downsample1d(nn.Module):
|
|
def __init__(self, kernel='linear', pad_mode='reflect'):
|
|
super().__init__()
|
|
self.pad_mode = pad_mode
|
|
kernel_1d = torch.tensor(_kernels[kernel])
|
|
self.pad = kernel_1d.shape[0] // 2 - 1
|
|
self.register_buffer('kernel', kernel_1d)
|
|
|
|
def forward(self, x):
|
|
x = F.pad(x, (self.pad,) * 2, self.pad_mode)
|
|
weight = x.new_zeros([x.shape[1], x.shape[1], self.kernel.shape[0]])
|
|
indices = torch.arange(x.shape[1], device=x.device)
|
|
weight[indices, indices] = self.kernel.to(weight)
|
|
return F.conv1d(x, weight, stride=2)
|
|
|
|
|
|
class Upsample1d(nn.Module):
|
|
def __init__(self, kernel='linear', pad_mode='reflect'):
|
|
super().__init__()
|
|
self.pad_mode = pad_mode
|
|
kernel_1d = torch.tensor(_kernels[kernel]) * 2
|
|
self.pad = kernel_1d.shape[0] // 2 - 1
|
|
self.register_buffer('kernel', kernel_1d)
|
|
|
|
def forward(self, x):
|
|
x = F.pad(x, ((self.pad + 1) // 2,) * 2, self.pad_mode)
|
|
weight = x.new_zeros([x.shape[1], x.shape[1], self.kernel.shape[0]])
|
|
indices = torch.arange(x.shape[1], device=x.device)
|
|
weight[indices, indices] = self.kernel.to(weight)
|
|
return F.conv_transpose1d(x, weight, stride=2, padding=self.pad * 2 + 1)
|
|
|
|
def Downsample1d_2(
|
|
in_channels: int, out_channels: int, factor: int, kernel_multiplier: int = 2
|
|
) -> nn.Module:
|
|
assert kernel_multiplier % 2 == 0, "Kernel multiplier must be even"
|
|
|
|
return nn.Conv1d(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
kernel_size=factor * kernel_multiplier + 1,
|
|
stride=factor,
|
|
padding=factor * (kernel_multiplier // 2),
|
|
)
|
|
|
|
|
|
def Upsample1d_2(
|
|
in_channels: int, out_channels: int, factor: int, use_nearest: bool = False
|
|
) -> nn.Module:
|
|
|
|
if factor == 1:
|
|
return nn.Conv1d(
|
|
in_channels=in_channels, out_channels=out_channels, kernel_size=3, padding=1
|
|
)
|
|
|
|
if use_nearest:
|
|
return nn.Sequential(
|
|
nn.Upsample(scale_factor=factor, mode="nearest"),
|
|
nn.Conv1d(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
kernel_size=3,
|
|
padding=1,
|
|
),
|
|
)
|
|
else:
|
|
return nn.ConvTranspose1d(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
kernel_size=factor * 2,
|
|
stride=factor,
|
|
padding=factor // 2 + factor % 2,
|
|
output_padding=factor % 2,
|
|
)
|
|
|
|
|
|
# U-Nets
|
|
|
|
class DBlock(ConditionedSequential):
|
|
def __init__(self, n_layers, feats_in, c_in, c_mid, c_out, group_size=32, head_size=64, dropout_rate=0., downsample_ratio=1, self_attn=False):
|
|
modules = [Downsample1d_2(c_in, c_in, downsample_ratio)] if downsample_ratio > 1 else []
|
|
for i in range(n_layers):
|
|
my_c_in = c_in if i == 0 else c_mid
|
|
my_c_out = c_mid if i < n_layers - 1 else c_out
|
|
modules.append(ResModConvBlock(feats_in, my_c_in, c_mid, my_c_out, group_size, dropout_rate))
|
|
if self_attn:
|
|
norm = lambda c_in: AdaGN(feats_in, c_in, max(1, my_c_out // group_size))
|
|
modules.append(SelfAttentionMod1d(my_c_out, max(1, my_c_out // head_size), norm, dropout_rate))
|
|
super().__init__(*modules)
|
|
|
|
|
|
class UBlock(ConditionedSequential):
|
|
def __init__(self, n_layers, feats_in, c_in, c_mid, c_out, group_size=32, head_size=64, dropout_rate=0., upsample_ratio=1, self_attn=False):
|
|
modules = []
|
|
for i in range(n_layers):
|
|
my_c_in = c_in if i == 0 else c_mid
|
|
my_c_out = c_mid if i < n_layers - 1 else c_out
|
|
modules.append(ResModConvBlock(feats_in, my_c_in, c_mid, my_c_out, group_size, dropout_rate))
|
|
if self_attn:
|
|
norm = lambda c_in: AdaGN(feats_in, c_in, max(1, my_c_out // group_size))
|
|
modules.append(SelfAttentionMod1d(my_c_out, max(1, my_c_out // head_size), norm, dropout_rate))
|
|
if upsample_ratio > 1:
|
|
modules.append(Upsample1d_2(c_out, c_out, upsample_ratio))
|
|
super().__init__(*modules)
|
|
|
|
def forward(self, input, cond, skip=None):
|
|
if skip is not None:
|
|
input = torch.cat([input, skip], dim=1)
|
|
return super().forward(input, cond)
|
|
|
|
class MappingNet(nn.Sequential):
|
|
def __init__(self, feats_in, feats_out, n_layers=2):
|
|
layers = []
|
|
for i in range(n_layers):
|
|
layers.append(nn.Linear(feats_in if i == 0 else feats_out, feats_out))
|
|
layers.append(nn.GELU())
|
|
super().__init__(*layers)
|
|
for layer in self:
|
|
if isinstance(layer, nn.Linear):
|
|
nn.init.orthogonal_(layer.weight)
|
|
|
|
class UNet(ConditionedModule):
|
|
def __init__(self, d_blocks, u_blocks):
|
|
super().__init__()
|
|
self.d_blocks = nn.ModuleList(d_blocks)
|
|
self.u_blocks = nn.ModuleList(u_blocks)
|
|
|
|
def forward(self, x, cond):
|
|
skips = []
|
|
for block in self.d_blocks:
|
|
x = block(x, cond)
|
|
skips.append(x)
|
|
for i, (block, skip) in enumerate(zip(self.u_blocks, reversed(skips))):
|
|
x = block(x, cond, skip if i > 0 else None)
|
|
return x
|
|
|