mirror of https://github.com/coqui-ai/TTS.git
81 lines
2.1 KiB
Python
81 lines
2.1 KiB
Python
import torch
|
|
|
|
from TTS.vocoder.layers.wavegrad import PositionalEncoding, FiLM, UBlock, DBlock
|
|
from TTS.vocoder.models.wavegrad import Wavegrad
|
|
|
|
|
|
def test_positional_encoding():
|
|
layer = PositionalEncoding(50)
|
|
inp = torch.rand(32, 50, 100)
|
|
nl = torch.rand(32)
|
|
o = layer(inp, nl)
|
|
|
|
assert o.shape[0] == 32
|
|
assert o.shape[1] == 50
|
|
assert o.shape[2] == 100
|
|
assert isinstance(o, torch.FloatTensor)
|
|
|
|
|
|
def test_film():
|
|
layer = FiLM(50, 76)
|
|
inp = torch.rand(32, 50, 100)
|
|
nl = torch.rand(32)
|
|
shift, scale = layer(inp, nl)
|
|
|
|
assert shift.shape[0] == 32
|
|
assert shift.shape[1] == 76
|
|
assert shift.shape[2] == 100
|
|
assert isinstance(shift, torch.FloatTensor)
|
|
|
|
assert scale.shape[0] == 32
|
|
assert scale.shape[1] == 76
|
|
assert scale.shape[2] == 100
|
|
assert isinstance(scale, torch.FloatTensor)
|
|
|
|
|
|
def test_ublock():
|
|
inp1 = torch.rand(32, 50, 100)
|
|
inp2 = torch.rand(32, 50, 50)
|
|
nl = torch.rand(32)
|
|
|
|
layer_film = FiLM(50, 100)
|
|
layer = UBlock(50, 100, 2, [1, 2, 4, 8])
|
|
|
|
scale, shift = layer_film(inp1, nl)
|
|
o = layer(inp2, shift, scale)
|
|
|
|
assert o.shape[0] == 32
|
|
assert o.shape[1] == 100
|
|
assert o.shape[2] == 100
|
|
assert isinstance(o, torch.FloatTensor)
|
|
|
|
|
|
def test_dblock():
|
|
inp = torch.rand(32, 50, 130)
|
|
layer = DBlock(50, 100, 2)
|
|
o = layer(inp)
|
|
|
|
assert o.shape[0] == 32
|
|
assert o.shape[1] == 100
|
|
assert o.shape[2] == 65
|
|
assert isinstance(o, torch.FloatTensor)
|
|
|
|
|
|
def test_wavegrad_forward():
|
|
x = torch.rand(32, 1, 20 * 300)
|
|
c = torch.rand(32, 80, 20)
|
|
noise_scale = torch.rand(32)
|
|
|
|
model = Wavegrad(in_channels=80,
|
|
out_channels=1,
|
|
upsample_factors=[5, 5, 3, 2, 2],
|
|
upsample_dilations=[[1, 2, 1, 2], [1, 2, 1, 2],
|
|
[1, 2, 4, 8], [1, 2, 4, 8],
|
|
[1, 2, 4, 8]])
|
|
o = model.forward(x, c, noise_scale)
|
|
|
|
assert o.shape[0] == 32
|
|
assert o.shape[1] == 1
|
|
assert o.shape[2] == 20 * 300
|
|
assert isinstance(o, torch.FloatTensor)
|