| import torch |
|
|
| from TTS.vocoder.configs import WavegradConfig |
| from TTS.vocoder.layers.wavegrad import DBlock, FiLM, PositionalEncoding, UBlock |
| from TTS.vocoder.models.wavegrad import Wavegrad, WavegradArgs |
|
|
|
|
| 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) |
|
|
| layer.apply_weight_norm() |
| layer.remove_weight_norm() |
|
|
|
|
| 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) |
|
|
| layer.apply_weight_norm() |
| layer.remove_weight_norm() |
|
|
|
|
| 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) |
|
|
| layer.apply_weight_norm() |
| layer.remove_weight_norm() |
|
|
|
|
| def test_wavegrad_forward(): |
| x = torch.rand(32, 1, 20 * 300) |
| c = torch.rand(32, 80, 20) |
| noise_scale = torch.rand(32) |
|
|
| args = WavegradArgs( |
| 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]], |
| ) |
| config = WavegradConfig(model_params=args) |
| model = Wavegrad(config) |
| 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) |
|
|
| model.apply_weight_norm() |
| model.remove_weight_norm() |
|
|
