[TTS][refactor] Part 2 - nemo.colletions.tts.parts

Signed-off-by: Xuesong Yang <[email protected]>

examples/tts/aligner_heteronym_disambiguation.py

@@ -22,7 +22,7 @@
 import torch
 
 from nemo.collections.tts.models import AlignerModel
-from nemo.collections.tts.torch.helpers import general_padding
+from nemo.collections.tts.parts.utils.tts_dataset_utils import general_padding
 
 
 """

nemo/collections/tts/g2p/models/heteronym_classification.py

@@ -26,8 +26,8 @@
 from nemo.collections.nlp.metrics.classification_report import ClassificationReport
 from nemo.collections.nlp.modules.common import TokenClassifier
 from nemo.collections.nlp.parts.utils_funcs import tensor2list
-from nemo.collections.tts.g2p.data.data_utils import get_heteronym_spans, get_wordid_to_phonemes, read_wordids
 from nemo.collections.tts.g2p.data.heteronym_classification_data import HeteronymClassificationDataset
+from nemo.collections.tts.parts.utils.g2p_utils import get_heteronym_spans, get_wordid_to_phonemes, read_wordids
 from nemo.core.classes.common import PretrainedModelInfo
 from nemo.utils import logging
 

nemo/collections/tts/g2p/modules.py

@@ -31,7 +31,7 @@
     english_word_tokenize,
     normalize_unicode_text,
 )
-from nemo.collections.tts.g2p.data.data_utils import GRAPHEME_CASE_MIXED, GRAPHEME_CASE_UPPER, set_grapheme_case
+from nemo.collections.tts.parts.utils.g2p_utils import GRAPHEME_CASE_MIXED, GRAPHEME_CASE_UPPER, set_grapheme_case
 from nemo.utils import logging
 from nemo.utils.decorators import experimental
 from nemo.utils.get_rank import is_global_rank_zero

nemo/collections/tts/losses/radttsloss.py

@@ -15,8 +15,8 @@
 import torch
 from torch.nn import functional as F
 
-from nemo.collections.tts.helpers.helpers import get_mask_from_lengths
 from nemo.collections.tts.losses.aligner_loss import ForwardSumLoss
+from nemo.collections.tts.parts.utils.helpers import get_mask_from_lengths
 from nemo.core.classes import Loss
 
 

nemo/collections/tts/losses/spectrogram_enhancer_losses.py

@@ -41,7 +41,7 @@
 from einops import rearrange
 from torch.autograd import grad as torch_grad
 
-from nemo.collections.tts.helpers.helpers import mask_sequence_tensor
+from nemo.collections.tts.parts.utils.helpers import mask_sequence_tensor
 
 
 class GradientPenaltyLoss(torch.nn.Module):

nemo/collections/tts/losses/tacotron2loss.py

@@ -14,7 +14,7 @@
 
 import torch
 
-from nemo.collections.tts.helpers.helpers import get_mask_from_lengths
+from nemo.collections.tts.parts.utils.helpers import get_mask_from_lengths
 from nemo.core.classes import Loss, typecheck
 from nemo.core.neural_types.elements import LengthsType, LogitsType, LossType, MelSpectrogramType
 from nemo.core.neural_types.neural_type import NeuralType

nemo/collections/tts/models/aligner.py

@@ -23,8 +23,8 @@
 from pytorch_lightning.loggers import WandbLogger
 from torch import nn
 
-from nemo.collections.tts.helpers.helpers import binarize_attention, get_mask_from_lengths, plot_alignment_to_numpy
 from nemo.collections.tts.losses.aligner_loss import BinLoss, ForwardSumLoss
+from nemo.collections.tts.parts.utils.helpers import binarize_attention, get_mask_from_lengths, plot_alignment_to_numpy
 from nemo.core.classes import ModelPT
 from nemo.core.classes.common import PretrainedModelInfo
 from nemo.utils import logging, model_utils

nemo/collections/tts/models/base.py

@@ -17,7 +17,7 @@
 
 import torch
 
-from nemo.collections.tts.helpers.helpers import OperationMode
+from nemo.collections.tts.parts.utils.helpers import OperationMode
 from nemo.core.classes import ModelPT
 from nemo.core.classes.common import PretrainedModelInfo, typecheck
 from nemo.core.neural_types.elements import AudioSignal

nemo/collections/tts/models/fastpitch.py

@@ -22,17 +22,17 @@
 from pytorch_lightning.loggers import TensorBoardLogger
 
 from nemo.collections.common.parts.preprocessing import parsers
-from nemo.collections.tts.helpers.helpers import (
+from nemo.collections.tts.losses.aligner_loss import BinLoss, ForwardSumLoss
+from nemo.collections.tts.losses.fastpitchloss import DurationLoss, EnergyLoss, MelLoss, PitchLoss
+from nemo.collections.tts.models.base import SpectrogramGenerator
+from nemo.collections.tts.modules.fastpitch import FastPitchModule
+from nemo.collections.tts.parts.utils.helpers import (
     batch_from_ragged,
     plot_alignment_to_numpy,
     plot_spectrogram_to_numpy,
     process_batch,
     sample_tts_input,
 )
-from nemo.collections.tts.losses.aligner_loss import BinLoss, ForwardSumLoss
-from nemo.collections.tts.losses.fastpitchloss import DurationLoss, EnergyLoss, MelLoss, PitchLoss
-from nemo.collections.tts.models.base import SpectrogramGenerator
-from nemo.collections.tts.modules.fastpitch import FastPitchModule
 from nemo.core.classes import Exportable
 from nemo.core.classes.common import PretrainedModelInfo, typecheck
 from nemo.core.neural_types.elements import (

nemo/collections/tts/models/fastpitch_ssl.py

@@ -20,10 +20,10 @@
 from pytorch_lightning import Trainer
 from pytorch_lightning.loggers import TensorBoardLogger
 
-from nemo.collections.tts.helpers.helpers import plot_multipitch_to_numpy, plot_spectrogram_to_numpy
 from nemo.collections.tts.losses.fastpitchloss import DurationLoss, MelLoss, PitchLoss
 from nemo.collections.tts.modules.fastpitch import FastPitchSSLModule, average_features
 from nemo.collections.tts.modules.transformer import mask_from_lens
+from nemo.collections.tts.parts.utils.helpers import plot_multipitch_to_numpy, plot_spectrogram_to_numpy
 from nemo.core.classes import ModelPT
 from nemo.core.classes.common import PretrainedModelInfo
 from nemo.utils import logging, model_utils

nemo/collections/tts/models/hifigan.py

@@ -20,10 +20,10 @@
 from omegaconf import DictConfig, OmegaConf, open_dict
 from pytorch_lightning.loggers.wandb import WandbLogger
 
-from nemo.collections.tts.helpers.helpers import get_batch_size, get_num_workers, plot_spectrogram_to_numpy
 from nemo.collections.tts.losses.hifigan_losses import DiscriminatorLoss, FeatureMatchingLoss, GeneratorLoss
 from nemo.collections.tts.models.base import Vocoder
 from nemo.collections.tts.modules.hifigan_modules import MultiPeriodDiscriminator, MultiScaleDiscriminator
+from nemo.collections.tts.parts.utils.helpers import get_batch_size, get_num_workers, plot_spectrogram_to_numpy
 from nemo.core.classes import Exportable
 from nemo.core.classes.common import PretrainedModelInfo, typecheck
 from nemo.core.neural_types.elements import AudioSignal, MelSpectrogramType

nemo/collections/tts/models/mixer_tts.py

@@ -31,15 +31,15 @@
     EnglishCharsTokenizer,
     EnglishPhonemesTokenizer,
 )
-from nemo.collections.tts.helpers.helpers import (
+from nemo.collections.tts.losses.aligner_loss import BinLoss, ForwardSumLoss
+from nemo.collections.tts.models.base import SpectrogramGenerator
+from nemo.collections.tts.modules.fastpitch import average_features, regulate_len
+from nemo.collections.tts.parts.utils.helpers import (
     binarize_attention_parallel,
     get_mask_from_lengths,
     plot_pitch_to_numpy,
     plot_spectrogram_to_numpy,
 )
-from nemo.collections.tts.losses.aligner_loss import BinLoss, ForwardSumLoss
-from nemo.collections.tts.models.base import SpectrogramGenerator
-from nemo.collections.tts.modules.fastpitch import average_features, regulate_len
 from nemo.core import Exportable
 from nemo.core.classes.common import PretrainedModelInfo, typecheck
 from nemo.core.neural_types.elements import (

nemo/collections/tts/models/radtts.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import contextlib
-import random
 
 import torch
 from hydra.utils import instantiate
@@ -21,14 +20,14 @@
 from pytorch_lightning.loggers import TensorBoardLogger
 
 from nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers import BaseTokenizer
-from nemo.collections.tts.helpers.helpers import (
+from nemo.collections.tts.losses.radttsloss import AttentionBinarizationLoss, RADTTSLoss
+from nemo.collections.tts.models.base import SpectrogramGenerator
+from nemo.collections.tts.parts.utils.helpers import (
     batch_from_ragged,
     plot_alignment_to_numpy,
     regulate_len,
     sample_tts_input,
 )
-from nemo.collections.tts.losses.radttsloss import AttentionBinarizationLoss, RADTTSLoss
-from nemo.collections.tts.models.base import SpectrogramGenerator
 from nemo.core.classes import Exportable
 from nemo.core.classes.common import typecheck
 from nemo.core.neural_types.elements import (

nemo/collections/tts/models/spectrogram_enhancer.py

@@ -49,13 +49,13 @@
 from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger
 from torch.utils.tensorboard.writer import SummaryWriter
 
-from nemo.collections.tts.helpers.helpers import mask_sequence_tensor, to_device_recursive
 from nemo.collections.tts.losses.spectrogram_enhancer_losses import (
     ConsistencyLoss,
     GeneratorLoss,
     GradientPenaltyLoss,
     HingeLoss,
 )
+from nemo.collections.tts.parts.utils.helpers import mask_sequence_tensor, to_device_recursive
 from nemo.core import Exportable, ModelPT, typecheck
 from nemo.core.neural_types import LengthsType, MelSpectrogramType, NeuralType
 from nemo.core.neural_types.elements import BoolType

nemo/collections/tts/models/tacotron2.py

@@ -24,13 +24,13 @@
 from torch import nn
 
 from nemo.collections.common.parts.preprocessing import parsers
-from nemo.collections.tts.helpers.helpers import (
+from nemo.collections.tts.losses.tacotron2loss import Tacotron2Loss
+from nemo.collections.tts.models.base import SpectrogramGenerator
+from nemo.collections.tts.parts.utils.helpers import (
     get_mask_from_lengths,
     tacotron2_log_to_tb_func,
     tacotron2_log_to_wandb_func,
 )
-from nemo.collections.tts.losses.tacotron2loss import Tacotron2Loss
-from nemo.collections.tts.models.base import SpectrogramGenerator
 from nemo.core.classes.common import PretrainedModelInfo, typecheck
 from nemo.core.neural_types.elements import (
     AudioSignal,

nemo/collections/tts/models/two_stages.py

@@ -20,8 +20,8 @@
 from hydra.utils import instantiate
 from omegaconf import MISSING, DictConfig, OmegaConf
 
-from nemo.collections.tts.helpers.helpers import OperationMode, griffin_lim
 from nemo.collections.tts.models.base import MelToSpec, Vocoder
+from nemo.collections.tts.parts.utils.helpers import OperationMode, griffin_lim
 from nemo.core.classes.common import PretrainedModelInfo
 from nemo.core.neural_types.elements import AudioSignal, MelSpectrogramType
 from nemo.core.neural_types.neural_type import NeuralType

nemo/collections/tts/models/univnet.py

@@ -21,11 +21,11 @@
 from omegaconf import DictConfig, open_dict
 from pytorch_lightning.loggers.wandb import WandbLogger
 
-from nemo.collections.tts.helpers.helpers import get_batch_size, get_num_workers, plot_spectrogram_to_numpy
 from nemo.collections.tts.losses.hifigan_losses import DiscriminatorLoss, GeneratorLoss
 from nemo.collections.tts.losses.stftlosses import MultiResolutionSTFTLoss
 from nemo.collections.tts.models.base import Vocoder
 from nemo.collections.tts.modules.univnet_modules import MultiPeriodDiscriminator, MultiResolutionDiscriminator
+from nemo.collections.tts.parts.utils.helpers import get_batch_size, get_num_workers, plot_spectrogram_to_numpy
 from nemo.core import Exportable
 from nemo.core.classes.common import PretrainedModelInfo, typecheck
 from nemo.core.neural_types.elements import AudioSignal, MelSpectrogramType

nemo/collections/tts/models/vits.py

@@ -17,18 +17,17 @@
 
 import omegaconf
 import torch
-import wandb
 from hydra.utils import instantiate
 from omegaconf import DictConfig, OmegaConf
 from pytorch_lightning import Trainer
 from pytorch_lightning.loggers import WandbLogger
 from torch.cuda.amp import autocast
 from torch.nn import functional as F
 
-from nemo.collections.tts.helpers.helpers import clip_grad_value_, plot_spectrogram_to_numpy, slice_segments
 from nemo.collections.tts.losses.vits_losses import DiscriminatorLoss, FeatureMatchingLoss, GeneratorLoss, KlLoss
 from nemo.collections.tts.models.base import TextToWaveform
 from nemo.collections.tts.modules.vits_modules import MultiPeriodDiscriminator
+from nemo.collections.tts.parts.utils.helpers import clip_grad_value_, plot_spectrogram_to_numpy, slice_segments
 from nemo.collections.tts.torch.data import DistributedBucketSampler
 from nemo.collections.tts.torch.tts_data_types import SpeakerID
 from nemo.core.classes.common import PretrainedModelInfo, typecheck

nemo/collections/tts/models/waveglow.py

@@ -18,9 +18,9 @@
 from omegaconf import DictConfig, open_dict
 from pytorch_lightning.loggers import TensorBoardLogger
 
-from nemo.collections.tts.helpers.helpers import OperationMode, waveglow_log_to_tb_func
 from nemo.collections.tts.losses.waveglowloss import WaveGlowLoss
 from nemo.collections.tts.models.base import GlowVocoder
+from nemo.collections.tts.parts.utils.helpers import OperationMode, waveglow_log_to_tb_func
 from nemo.core.classes import Exportable
 from nemo.core.classes.common import PretrainedModelInfo, typecheck
 from nemo.core.neural_types.elements import (

nemo/collections/tts/modules/aligner.py

@@ -15,8 +15,8 @@
 import torch
 from torch import nn
 
-from nemo.collections.tts.helpers.helpers import binarize_attention_parallel
 from nemo.collections.tts.modules.submodules import ConvNorm
+from nemo.collections.tts.parts.utils.helpers import binarize_attention_parallel
 
 
 class AlignmentEncoder(torch.nn.Module):

nemo/collections/tts/modules/attribute_prediction_model.py

@@ -16,10 +16,10 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from nemo.collections.tts.helpers.helpers import get_mask_from_lengths
 from nemo.collections.tts.modules.common import ConvLSTMLinear
 from nemo.collections.tts.modules.submodules import ConvNorm, MaskedInstanceNorm1d
 from nemo.collections.tts.modules.transformer import FFTransformer
+from nemo.collections.tts.parts.utils.helpers import get_mask_from_lengths
 
 
 def get_attribute_prediction_model(config):

nemo/collections/tts/modules/common.py

@@ -22,15 +22,14 @@
 from torch.cuda import amp
 from torch.cuda.amp import autocast as autocast
 from torch.nn import functional as F
-from torch.nn.utils.rnn import PackedSequence
 
-from nemo.collections.tts.helpers.helpers import get_mask_from_lengths, sort_tensor, unsort_tensor
-from nemo.collections.tts.helpers.splines import (
+from nemo.collections.tts.modules.submodules import ConvNorm, LinearNorm, MaskedInstanceNorm1d
+from nemo.collections.tts.parts.utils.helpers import get_mask_from_lengths, sort_tensor, unsort_tensor
+from nemo.collections.tts.parts.utils.splines import (
     piecewise_linear_inverse_transform,
     piecewise_linear_transform,
     unbounded_piecewise_quadratic_transform,
 )
-from nemo.collections.tts.modules.submodules import ConvNorm, LinearNorm, MaskedInstanceNorm1d
 
 
 @torch.jit.script

nemo/collections/tts/modules/fastpitch.py

@@ -44,7 +44,7 @@
 
 import torch
 
-from nemo.collections.tts.helpers.helpers import binarize_attention_parallel, regulate_len
+from nemo.collections.tts.parts.utils.helpers import binarize_attention_parallel, regulate_len
 from nemo.core.classes import NeuralModule, typecheck
 from nemo.core.neural_types.elements import (
     EncodedRepresentation,

nemo/collections/tts/modules/radtts.py

@@ -15,9 +15,6 @@
 import torch.nn.functional as F
 from torch import nn
 
-from nemo.collections.tts.helpers.helpers import get_mask_from_lengths
-from nemo.collections.tts.helpers.helpers import mas_width1 as mas
-from nemo.collections.tts.helpers.helpers import regulate_len
 from nemo.collections.tts.modules.attribute_prediction_model import get_attribute_prediction_model
 from nemo.collections.tts.modules.common import (
     AffineTransformationLayer,
@@ -29,6 +26,9 @@
     LinearNorm,
     get_radtts_encoder,
 )
+from nemo.collections.tts.parts.utils.helpers import get_mask_from_lengths
+from nemo.collections.tts.parts.utils.helpers import mas_width1 as mas
+from nemo.collections.tts.parts.utils.helpers import regulate_len
 from nemo.core.classes import Exportable, NeuralModule
 from nemo.core.neural_types.elements import Index, LengthsType, MelSpectrogramType, TokenDurationType, TokenIndex
 from nemo.core.neural_types.neural_type import NeuralType

nemo/collections/tts/modules/spectrogram_enhancer.py

@@ -46,7 +46,7 @@
 from einops import rearrange
 from kornia.filters import filter2d
 
-from nemo.collections.tts.helpers.helpers import mask_sequence_tensor
+from nemo.collections.tts.parts.utils.helpers import mask_sequence_tensor
 
 
 class Blur(torch.nn.Module):