File size: 2,548 Bytes
85a7d2c |
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 |
import gin
import numpy as np
import torch
import torch.nn as nn
from torch.nn.utils import weight_norm
from .pcmer import PCmer
def split_to_dict(tensor, tensor_splits):
"""Split a tensor into a dictionary of multiple tensors."""
labels = []
sizes = []
for k, v in tensor_splits.items():
labels.append(k)
sizes.append(v)
tensors = torch.split(tensor, sizes, dim=-1)
return dict(zip(labels, tensors))
class Unit2Control(nn.Module):
def __init__(
self,
input_channel,
n_spk,
output_splits):
super().__init__()
self.output_splits = output_splits
self.f0_embed = nn.Linear(1, 256)
self.phase_embed = nn.Linear(1, 256)
self.volume_embed = nn.Linear(1, 256)
self.n_spk = n_spk
if n_spk is not None and n_spk > 1:
self.spk_embed = nn.Embedding(n_spk, 256)
# conv in stack
self.stack = nn.Sequential(
nn.Conv1d(input_channel, 256, 3, 1, 1),
nn.GroupNorm(4, 256),
nn.LeakyReLU(),
nn.Conv1d(256, 256, 3, 1, 1))
# transformer
self.decoder = PCmer(
num_layers=3,
num_heads=8,
dim_model=256,
dim_keys=256,
dim_values=256,
residual_dropout=0.1,
attention_dropout=0.1)
self.norm = nn.LayerNorm(256)
# out
self.n_out = sum([v for k, v in output_splits.items()])
self.dense_out = weight_norm(
nn.Linear(256, self.n_out))
def forward(self, units, f0, phase, volume, spk_id = None, spk_mix_dict = None):
'''
input:
B x n_frames x n_unit
return:
dict of B x n_frames x feat
'''
x = self.stack(units.transpose(1,2)).transpose(1,2)
x = x + self.f0_embed((1+ f0 / 700).log()) + self.phase_embed(phase / np.pi) + self.volume_embed(volume)
if self.n_spk is not None and self.n_spk > 1:
if spk_mix_dict is not None:
for k, v in spk_mix_dict.items():
spk_id_torch = torch.LongTensor(np.array([[k]])).to(units.device)
x = x + v * self.spk_embed(spk_id_torch - 1)
else:
x = x + self.spk_embed(spk_id - 1)
x = self.decoder(x)
x = self.norm(x)
e = self.dense_out(x)
controls = split_to_dict(e, self.output_splits)
return controls
|