File size: 2,396 Bytes
45ee559
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
88
89
90
91
92
93
94
95
96
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()