File size: 4,416 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 101 102 103 104 105 106 107 108 109 110 |
# 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 OmegaConf
from nemo.collections.tts.models import FastPitchModel, HifiGanModel, RadTTSModel
from nemo.utils.app_state import AppState
@pytest.fixture()
def fastpitch_model():
model = FastPitchModel.from_pretrained(model_name="tts_en_fastpitch")
model.export_config['enable_volume'] = True
model.export_config['enable_ragged_batches'] = True
return model
@pytest.fixture()
def hifigan_model():
model = HifiGanModel.from_pretrained(model_name="tts_en_hifigan")
return model
@pytest.fixture()
def radtts_model():
this_test_dir = os.path.dirname(os.path.abspath(__file__))
cfg = OmegaConf.load(os.path.join(this_test_dir, '../../../examples/tts/conf/rad-tts_feature_pred.yaml'))
cfg.model.init_from_ptl_ckpt = None
cfg.model.train_ds.dataset.manifest_filepath = "dummy.json"
cfg.model.train_ds.dataset.sup_data_path = "dummy.json"
cfg.model.validation_ds.dataset.manifest_filepath = "dummy.json"
cfg.model.validation_ds.dataset.sup_data_path = "dummy.json"
cfg.pitch_mean = 212.35
cfg.pitch_std = 68.52
app_state = AppState()
app_state.is_model_being_restored = True
model = RadTTSModel(cfg=cfg.model)
app_state.is_model_being_restored = False
model.eval()
model.export_config['enable_ragged_batches'] = True
model.export_config['enable_volume'] = True
return model
class TestExportable:
@pytest.mark.run_only_on('GPU')
@pytest.mark.unit
def test_FastPitchModel_export_to_onnx(self, fastpitch_model):
model = fastpitch_model.cuda()
with tempfile.TemporaryDirectory() as tmpdir:
filename = os.path.join(tmpdir, 'fp.onnx')
model.export(output=filename, verbose=True, onnx_opset_version=14, check_trace=True)
@pytest.mark.with_downloads()
@pytest.mark.run_only_on('GPU')
@pytest.mark.unit
def test_HifiGanModel_export_to_onnx(self, hifigan_model):
model = hifigan_model.cuda()
assert hifigan_model.generator is not None
with tempfile.TemporaryDirectory() as tmpdir:
filename = os.path.join(tmpdir, 'hfg.onnx')
model.export(output=filename, verbose=True, check_trace=True)
@pytest.mark.pleasefixme
@pytest.mark.run_only_on('GPU')
@pytest.mark.unit
def test_RadTTSModel_export_to_torchscript(self, radtts_model):
model = radtts_model.cuda()
with tempfile.TemporaryDirectory() as tmpdir:
filename = os.path.join(tmpdir, 'rad.ts')
with torch.cuda.amp.autocast(enabled=True, cache_enabled=False, dtype=torch.float16):
input_example1 = model.input_module.input_example(max_batch=13, max_dim=777)
input_example2 = model.input_module.input_example(max_batch=19, max_dim=999)
model.export(output=filename, verbose=True, input_example=input_example1, check_trace=[input_example2])
@pytest.mark.pleasefixme
@pytest.mark.run_only_on('GPU')
@pytest.mark.unit
def test_RadTTSModel_export_to_onnx(self, radtts_model):
model = radtts_model.cuda()
with tempfile.TemporaryDirectory() as tmpdir:
filename = os.path.join(tmpdir, 'rad.onnx')
with torch.cuda.amp.autocast(enabled=True, cache_enabled=False, dtype=torch.float16):
input_example1 = model.input_module.input_example(max_batch=13, max_dim=777)
input_example2 = model.input_module.input_example(max_batch=19, max_dim=999)
model.export(
output=filename,
input_example=input_example1,
verbose=True,
onnx_opset_version=14,
check_trace=[input_example2],
)
|