update unit test for extract tts spectrograms script

tests/test_extract_tts_spectrograms.py

@@ -5,81 +5,62 @@ import torch
 
 from tests import get_tests_input_path
 
-from TTS.tts.models.tacotron2 import Tacotron2
+from tests import get_tests_output_path, run_cli
-from TTS.tts.models.glow_tts import GlowTTS
+
+from TTS.tts.utils.generic_utils import setup_model
 
-from TTS.utils.audio import AudioProcessor
 from TTS.utils.io import load_config
-
+from TTS.tts.utils.text.symbols import phonemes, symbols
-from TTS.bin.extract_tts_spectrograms import inference
 
 torch.manual_seed(1)
-use_cuda = torch.cuda.is_available()
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-c = load_config(os.path.join(get_tests_input_path(), "test_config.json"))
-# set params from tacotron inference
-c.bidirectional_decoder = False
-c.double_decoder_consistency = False
-ap = AudioProcessor(**c.audio)
-
 
 # pylint: disable=protected-access
 class TestExtractTTSSpectrograms(unittest.TestCase):
     @staticmethod
     def test_GlowTTS():
-        input_dummy = torch.randint(0, 24, (8, 128)).long().to(device)
+        # set paths
-        input_lengths = torch.randint(100, 129, (8,)).long().to(device)
+        config_path = os.path.join(get_tests_input_path(), "test_glow_tts.json")
-        input_lengths[-1] = 128
+        checkpoint_path = os.path.join(get_tests_output_path(), 'checkpoint_test.pth.tar')
-        mel_spec = torch.rand(8, c.audio["num_mels"], 30).to(device)
+        output_path = os.path.join(get_tests_output_path(), 'output_extract_tts_spectrograms/')
-        mel_lengths = torch.randint(20, 30, (8,)).long().to(device)
+        # load config
-
+        c = load_config(config_path)
         # create model
-        model = GlowTTS(
+        num_chars = len(phonemes if c.use_phonemes else symbols)
-            num_chars=32,
+        model = setup_model(num_chars, 1, c, speaker_embedding_dim=None)
-            hidden_channels_enc=48,
+        # save model
-            hidden_channels_dec=48,
+        torch.save({"model": model.state_dict()}, checkpoint_path)
-            hidden_channels_dp=32,
+        # run test
-            out_channels=c.audio["num_mels"],
+        run_cli(f'CUDA_VISIBLE_DEVICES="" python3 TTS/bin/extract_tts_spectrograms.py --config_path "{config_path}" --checkpoint_path "{checkpoint_path}" --output_path "{output_path}"')
-            encoder_type="rel_pos_transformer",
+        run_cli(f'rm -rf "{output_path}" "{checkpoint_path}"')
-            encoder_params={
+    @staticmethod
-                "kernel_size": 3,
+    def test_Tacotron2():
-                "dropout_p": 0.1,
+        # set paths
-                "num_layers": 6,
+        config_path = os.path.join(get_tests_input_path(), "test_tacotron2_config.json")
-                "num_heads": 2,
+        checkpoint_path = os.path.join(get_tests_output_path(), 'checkpoint_test.pth.tar')
-                "hidden_channels_ffn": 16,  # 4 times the hidden_channels
+        output_path = os.path.join(get_tests_output_path(), 'output_extract_tts_spectrograms/')
-                "input_length": None,
+        # load config
-            },
+        c = load_config(config_path)
-            use_encoder_prenet=True,
+        # create model
-            num_flow_blocks_dec=12,
+        num_chars = len(phonemes if c.use_phonemes else symbols)
-            kernel_size_dec=5,
+        model = setup_model(num_chars, 1, c, speaker_embedding_dim=None)
-            dilation_rate=1,
+        # save model
-            num_block_layers=4,
+        torch.save({"model": model.state_dict()}, checkpoint_path)
-            dropout_p_dec=0.0,
+        # run test
-            num_speakers=0,
+        run_cli(f'CUDA_VISIBLE_DEVICES="" python3 TTS/bin/extract_tts_spectrograms.py --config_path "{config_path}" --checkpoint_path "{checkpoint_path}" --output_path "{output_path}"')
-            c_in_channels=0,
+        run_cli(f'rm -rf "{output_path}" "{checkpoint_path}"')
-            num_splits=4,
-            num_squeeze=1,
-            sigmoid_scale=False,
-            mean_only=False,
-        ).to(device)
-
-        model.eval()
-        _ = inference('glow_tts', model, c, ap, input_dummy, input_lengths, mel_spec.permute(0, 2, 1), mel_lengths)
-        print("GlowTTS extract tts spectrograms ok !")
-
     @staticmethod
     def test_Tacotron():
-        input_dummy = torch.randint(0, 24, (8, 128)).long().to(device)
+        # set paths
-        input_lengths = torch.randint(100, 128, (8,)).long().to(device)
+        config_path = os.path.join(get_tests_input_path(), "test_tacotron_config.json")
-        input_lengths = torch.sort(input_lengths, descending=True)[0]
+        checkpoint_path = os.path.join(get_tests_output_path(), 'checkpoint_test.pth.tar')
-        mel_spec = torch.rand(8, 30, c.audio["num_mels"]).to(device)
+        output_path = os.path.join(get_tests_output_path(), 'output_extract_tts_spectrograms/')
-        mel_lengths = torch.randint(20, 30, (8,)).long().to(device)
+        # load config
-        mel_lengths[0] = 30
+        c = load_config(config_path)
-
         # create model
-        model = Tacotron2(num_chars=24, decoder_output_dim=c.audio["num_mels"], r=c.r, num_speakers=1).to(device)
+        num_chars = len(phonemes if c.use_phonemes else symbols)
-        model.eval()
+        model = setup_model(num_chars, 1, c, speaker_embedding_dim=None)
-
+        # save model
-        _ = inference('tacotron2', model, c, ap, input_dummy, input_lengths, mel_spec, mel_lengths)
+        torch.save({"model": model.state_dict()}, checkpoint_path)
-        print("Tacotron extract tts spectrograms ok !")
+        # run test
+        run_cli(f'CUDA_VISIBLE_DEVICES="" python3 TTS/bin/extract_tts_spectrograms.py --config_path "{config_path}" --checkpoint_path "{checkpoint_path}" --output_path "{output_path}"')
+        run_cli(f'rm -rf "{output_path}" "{checkpoint_path}"')