File size: 2,974 Bytes
7934b29 |
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 97 98 99 100 |
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import tempfile
import pytest
import torch
from omegaconf import DictConfig
from nemo.collections.tts.models import WaveGlowModel
from nemo.core.classes import typecheck
mcfg = DictConfig(
{
"_target_": "nemo.collections.tts.modules.waveglow.WaveGlowModule",
"n_flows": 12,
"n_group": 8,
"n_mel_channels": 80,
"n_early_every": 4,
"n_early_size": 2,
"n_wn_channels": 512,
"n_wn_layers": 8,
"wn_kernel_size": 3,
}
)
pcfg = DictConfig(
{
"_target_": "nemo.collections.asr.parts.preprocessing.features.FilterbankFeatures",
"dither": 0.0,
"nfilt": 80,
"stft_conv": False,
}
)
wcfg = DictConfig({"waveglow": mcfg, "sigma": 1.0, "preprocessor": pcfg,})
def input_example(sz):
mel = torch.randn(1, 1, 80, sz).cuda().half()
z = torch.randn(1, 8, sz * 256 // 8, 1).cuda().half()
return (
mel,
z,
)
def taco2wg(spec, z):
spec = spec.permute(0, 3, 2, 1).contiguous()
return spec.view(spec.size(0), spec.size(1), -1), z.view(z.size(0), z.size(1), -1)
# Wrapper method to convert Jasper's Taco2 output to WG input and call inference
def forward_wrapper(self, spec, z=None):
spec, z = taco2wg(spec, z)
audio = self.waveglow.norm_dist_to_audio(spec=spec, sigma=1.0, z=z)
return audio
class TestWaveGlow:
@pytest.mark.pleasefixme
@pytest.mark.run_only_on('GPU')
@pytest.mark.unit
def test_export_to_onnx(self):
model = WaveGlowModel(wcfg)
model = model.cuda().half()
typecheck.set_typecheck_enabled(enabled=False)
with tempfile.TemporaryDirectory() as tmpdir, model.nemo_infer():
tmp_file_name = os.path.join(tmpdir, "waveglow.onnx")
n_mels = 80
# Test export.
inp = input_example(n_mels)
inp1 = taco2wg(*inp)
inp2 = inp1
res1 = model.waveglow(*inp1)
res2 = model.waveglow(*inp2)
assert torch.allclose(res1, res2, rtol=0.01, atol=0.1)
WaveGlowModel.forward_for_export = forward_wrapper
model.export(
tmp_file_name, input_example=inp, verbose=False, check_trace=False, do_constant_folding=True,
)
if __name__ == "__main__":
t = TestWaveGlow()
t.test_export_to_onnx()
|