EMTTS Pronunciation Control

easymagpie_single_infer.py

@@ -0,0 +1,86 @@
+"""Minimal single-utterance inference script for EasyMagpieTTS."""
+import argparse
+import os
+
+os.environ['OMP_NUM_THREADS'] = '2'
+
+import soundfile as sf
+from omegaconf import open_dict
+from nemo.collections.tts.models import EasyMagpieTTSModel
+
+
+def main():
+    p = argparse.ArgumentParser(description="EasyMagpieTTS single inference")
+    p.add_argument("--transcript", required=True, help="Text to synthesize")
+    p.add_argument("--model_path", required=True, help="Path to .nemo TTS checkpoint")
+    p.add_argument("--codec_model_path", required=True, help="Path to .nemo codec model")
+    p.add_argument("--output_path", required=True, help="Output .wav file path")
+
+    ctx = p.add_mutually_exclusive_group(required=True)
+    ctx.add_argument("--context_audio_path", help="Path to context audio file")
+    ctx.add_argument("--context_text", help="Context text string")
+
+    p.add_argument("--phoneme_tokenizer_path", default=None,
+                   help="Path to phoneme tokenizer JSON (auto-detected next to this script if omitted)")
+    p.add_argument("--language", default="en", choices=["en", "zh", "es", "fr", "de", "it", "hi", "vi"])
+    p.add_argument("--temperature", type=float, default=0.7)
+    p.add_argument("--topk", type=int, default=80)
+    p.add_argument("--use_cfg", action="store_true", default=True)
+    p.add_argument("--no_cfg", dest="use_cfg", action="store_false")
+    p.add_argument("--cfg_scale", type=float, default=2.5)
+    p.add_argument("--max_steps", type=int, default=300)
+    args = p.parse_args()
+
+    phoneme_tokenizer_path = args.phoneme_tokenizer_path
+    if phoneme_tokenizer_path is None:
+        phoneme_tokenizer_path = os.path.join(
+            os.path.dirname(os.path.abspath(__file__)),
+            "bpe_ipa_tokenizer_2048_en_de_es_fr_hi_it_vi_zh.json",
+        )
+
+    # --- load model ---
+    model_cfg = EasyMagpieTTSModel.restore_from(args.model_path, return_config=True)
+    with open_dict(model_cfg):
+        model_cfg.codecmodel_path = args.codec_model_path
+        model_cfg.train_ds = None
+        model_cfg.validation_ds = None
+        if getattr(model_cfg, "phoneme_tokenizer", None) is not None:
+            model_cfg.phoneme_tokenizer.tokenizer_path = phoneme_tokenizer_path
+
+    model = EasyMagpieTTSModel.restore_from(args.model_path, override_config_path=model_cfg)
+    model.use_kv_cache_for_inference = True
+    model.eval().cuda().float()
+
+    # --- resolve context ---
+    if args.context_audio_path:
+        use_language_tag = bool(getattr(model, "add_language_to_context_text", False))
+        context_text = f"[{args.language.upper()}]" if use_language_tag else "[NO TEXT CONTEXT]"
+        context_audio_path = args.context_audio_path
+    else:
+        context_text = args.context_text
+        context_audio_path = None
+
+    transcript = args.transcript.strip()
+    if not transcript.endswith((".", "?", "!")):
+        transcript += "."
+
+    # --- infer ---
+    audio, audio_len = model.do_tts(
+        transcript=transcript,
+        context_audio_file_path=context_audio_path,
+        context_text=context_text,
+        use_cfg=args.use_cfg,
+        cfg_scale=args.cfg_scale,
+        use_local_transformer=True,
+        temperature=args.temperature,
+        topk=args.topk,
+        max_steps=args.max_steps,
+    )
+
+    audio_np = audio[0, : audio_len[0]].cpu().numpy()
+    sf.write(args.output_path, audio_np, model.output_sample_rate)
+    print(f"Saved {args.output_path}  ({audio_np.shape[0]} samples, {model.output_sample_rate} Hz)")
+
+
+if __name__ == "__main__":
+    main()

examples/tts/conf/magpietts/easy_magpietts.yaml

@@ -102,21 +102,16 @@ model:
   phoneme_corruption_timestep_ratio: 0.15
   phoneme_corruption_unk_mode_prob: 0.5
   phoneme_corruption_type: "repeat_skip_unk" # "repeat_skip_unk" or "complete_channel"
-  phoneme_as_text_prob: 0.0 # Probability of replacing training text with pronunciation-control G2P output.
+  enable_phoneme_text_input: false # Enable inline IPA spans in text as <bop>...<eop>; requires phoneme_tokenizer.
+  partial_phoneme_text_prob: 0.0 # Training-only probability of using precomputed IPA alignments with Lhotse.
+  partial_phoneme_word_prob: 0.5 # Fallback fixed per-alignment probability when min/max are not set.
+  partial_phoneme_word_prob_min: 0.25 # Sampled per-alignment probability lower bound.
+  partial_phoneme_word_prob_max: 0.75 # Sampled per-alignment probability upper bound.
+  phoneme_text_bop_marker: "<bop>"
+  phoneme_text_eop_marker: "<eop>"
   ignore_phoneme_languages: [] # Languages for which missing IPA phoneme fields are allowed during training.
   add_language_to_context_text: false # Prefix context text with language metadata for multilingual conditioning.
 
-  pronunciation_control_g2p:
-    en:
-      _target_: nemo.collections.tts.g2p.models.i18n_ipa.IpaG2p
-      phoneme_dict: "scripts/tts_dataset_files/ipa_cmudict-0.7b_nv26.07.txt"
-      heteronyms: "scripts/tts_dataset_files/heteronyms-052722"
-      phoneme_probability: 0.3
-      ignore_ambiguous_words: false
-      use_chars: true
-      use_stresses: true
-      grapheme_case: mixed
-
   phoneme_tokenizer:
     _target_: nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.IPABPETokenizer
     tokenizer_path: ???

examples/tts/conf/magpietts/easy_magpietts_lhotse.yaml

@@ -96,20 +96,15 @@ model:
   phoneme_corruption_timestep_ratio: 0.15
   phoneme_corruption_unk_mode_prob: 0.5
   phoneme_corruption_type: "repeat_skip_unk" # "repeat_skip_unk" or "complete_channel"
-  phoneme_as_text_prob: 0.0 # Probability of replacing training text with pronunciation-control G2P output.
+  enable_phoneme_text_input: false # Enable inline IPA spans in text as <bop>...<eop>; requires phoneme_tokenizer.
+  partial_phoneme_text_prob: 0.0 # Training-only probability of using precomputed IPA alignments.
+  partial_phoneme_word_prob: 0.5 # Fallback fixed per-alignment probability when min/max are not set.
+  partial_phoneme_word_prob_min: 0.25 # Sampled per-alignment probability lower bound.
+  partial_phoneme_word_prob_max: 0.75 # Sampled per-alignment probability upper bound.
+  phoneme_text_bop_marker: "<bop>"
+  phoneme_text_eop_marker: "<eop>"
   ignore_phoneme_languages: [] # Languages for which missing IPA phoneme fields are allowed during training.
   add_language_to_context_text: false # Prefix context text with language metadata for multilingual conditioning.
-
-  pronunciation_control_g2p:
-    en:
-      _target_: nemo.collections.tts.g2p.models.i18n_ipa.IpaG2p
-      phoneme_dict: "scripts/tts_dataset_files/ipa_cmudict-0.7b_nv26.07.txt"
-      heteronyms: "scripts/tts_dataset_files/heteronyms-052722"
-      phoneme_probability: 0.3
-      ignore_ambiguous_words: false
-      use_chars: true
-      use_stresses: true
-      grapheme_case: mixed
 
   phoneme_tokenizer:
     _target_: nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.IPABPETokenizer

nemo/collections/tts/data/text_to_speech_dataset.py

@@ -35,9 +35,8 @@
     get_tokenizer_for_language,
     get_weighted_sampler,
     load_audio,
-    setup_pronunciation_control_g2p,
     stack_tensors,
-    tokenize_text_with_pronunciation_control,
+    tokenize_text_with_phoneme_spans,
 )
 from nemo.core.classes import Dataset
 from nemo.utils import logging
@@ -392,8 +391,10 @@ def __init__(
         text_context_remapping: Dict[str, str] = None,
         text_context_remapping_prob: float = 0.0,
         ignore_phoneme_languages: List[str] = None,
-        phoneme_as_text_prob: float = 0.0,
-        pronunciation_control_g2p: Dict = None,
+        enable_phoneme_text_input: bool = False,
+        text_phoneme_token_offset: int = None,
+        phoneme_text_bop_marker: str = "<bop>",
+        phoneme_text_eop_marker: str = "<eop>",
         add_language_to_context_text: bool = False,
         default_tokenizer_name: str = "english_phoneme",
     ):
@@ -421,7 +422,6 @@ def __init__(
         self.tokenizer_config = tokenizer_config
         self.text_tokenizer = None  # Assigned in worker_init_fn in model file
         self.phoneme_tokenizer = None  # Assigned in worker_init_fn in model file (if any)
-        self.pronunciation_control_g2p = None
         self.load_16khz_audio = load_16khz_audio
         self.use_text_conditioning_tokenizer = use_text_conditioning_tokenizer
         self.text_conditioning_tokenizer_name = text_conditioning_tokenizer_name
@@ -431,8 +431,10 @@ def __init__(
         self.text_context_remapping = text_context_remapping
         self.text_context_remapping_prob = text_context_remapping_prob
         self.ignore_phoneme_languages = ignore_phoneme_languages or []
-        self.phoneme_as_text_prob = phoneme_as_text_prob
-        self.pronunciation_control_g2p_config = pronunciation_control_g2p
+        self.enable_phoneme_text_input = enable_phoneme_text_input
+        self.text_phoneme_token_offset = text_phoneme_token_offset
+        self.phoneme_text_bop_marker = phoneme_text_bop_marker
+        self.phoneme_text_eop_marker = phoneme_text_eop_marker
         self.add_language_to_context_text = add_language_to_context_text
         self.default_tokenizer_name = default_tokenizer_name
 
@@ -443,12 +445,6 @@ def get_num_audio_samples_to_slice(self, duration, sample_rate):
 
     def __getitem__(self, index):
         data = self.data_samples[index]
-        if (
-            self.pronunciation_control_g2p is None
-            and self.pronunciation_control_g2p_config is not None
-            and self.phoneme_as_text_prob > 0.0
-        ):
-            self.pronunciation_control_g2p = setup_pronunciation_control_g2p(self.pronunciation_control_g2p_config)
 
         def _sample_context_duration_with_available_limit(available_duration_sec: float) -> float:
             effective_duration_max = min(self.context_duration_max, available_duration_sec)
@@ -466,14 +462,15 @@ def _sample_context_duration_with_available_limit(available_duration_sec: float)
         else:
             language = 'en'
 
-        tokens = tokenize_text_with_pronunciation_control(
+        tokens = tokenize_text_with_phoneme_spans(
             text_tokenizer=self.text_tokenizer,
             text_str=data.text,
-            language=language,
             tokenizer_name=tokenizer_name,
-            dataset_type=self.dataset_type,
-            phoneme_as_text_prob=self.phoneme_as_text_prob,
-            pronunciation_control_g2p=self.pronunciation_control_g2p,
+            enable_phoneme_text_input=self.enable_phoneme_text_input,
+            phoneme_tokenizer=self.phoneme_tokenizer,
+            text_phoneme_token_offset=self.text_phoneme_token_offset,
+            bop_marker=self.phoneme_text_bop_marker,
+            eop_marker=self.phoneme_text_eop_marker,
         )
         tokens = tokens + [self.eos_id]  # Not adding BOS id
         tokens = torch.tensor(tokens, dtype=torch.int32)
@@ -651,6 +648,7 @@ def _sample_context_duration_with_available_limit(available_duration_sec: float)
             else:
                 if self.add_language_to_context_text:
                     context_text = f"[{language.upper()}]"
+                    print(f"Context text: {context_text}")
                 else:
                     context_text = "[NO TEXT CONTEXT]"
                 context_tokens = self.text_tokenizer.encode(context_text, self.text_conditioning_tokenizer_name)
@@ -1012,6 +1010,11 @@ def __getitem__(self, idx: int) -> Dict[str, Any]:
             tokenizer_name=tokenizer_name,
             text_tokenizer=self.text_tokenizer,
             eos_token_id=self.eos_id,
+            enable_phoneme_text_input=self.enable_phoneme_text_input,
+            phoneme_tokenizer=self.phoneme_tokenizer,
+            text_phoneme_token_offset=self.text_phoneme_token_offset,
+            bop_marker=self.phoneme_text_bop_marker,
+            eop_marker=self.phoneme_text_eop_marker,
         )
 
         # Handle empty text edge case

nemo/collections/tts/data/text_to_speech_dataset_lhotse.py

@@ -30,11 +30,13 @@
     IPABPETokenizer,
 )
 from nemo.collections.tts.parts.utils.tts_dataset_utils import (
+    _sample_probability_range,
     beta_binomial_prior_distribution,
+    has_phoneme_text_spans,
     normalize_volume,
-    setup_pronunciation_control_g2p,
+    partially_phonemize_text,
     stack_tensors,
-    tokenize_text_with_pronunciation_control,
+    tokenize_text_with_phoneme_spans,
 )
 from nemo.core.classes.common import safe_instantiate
 from nemo.utils import logging
@@ -183,8 +185,14 @@ def __init__(
         text_context_remapping_prob: float = 0.0,
         phoneme_tokenizer_config: DictConfig = None,
         ignore_phoneme_languages: List[str] = None,
-        phoneme_as_text_prob: float = 0.0,
-        pronunciation_control_g2p: Optional[DictConfig] = None,
+        enable_phoneme_text_input: bool = False,
+        text_phoneme_token_offset: int = None,
+        partial_phoneme_text_prob: float = 0.0,
+        partial_phoneme_word_prob: float = 0.5,
+        partial_phoneme_word_prob_min: float = None,
+        partial_phoneme_word_prob_max: float = None,
+        phoneme_text_bop_marker: str = "<bop>",
+        phoneme_text_eop_marker: str = "<eop>",
         add_language_to_context_text: bool = False,
     ):
         super().__init__()
@@ -207,13 +215,18 @@ def __init__(
         self.tokenizer_config = tokenizer_config
         self.text_tokenizer = None
         self.phoneme_tokenizer = None
-        self.pronunciation_control_g2p = None
         self.text_context_remapping = text_context_remapping
         self.text_context_remapping_prob = text_context_remapping_prob
         self.phoneme_tokenizer_config = phoneme_tokenizer_config
         self.ignore_phoneme_languages = ignore_phoneme_languages or []
-        self.phoneme_as_text_prob = phoneme_as_text_prob
-        self.pronunciation_control_g2p_config = pronunciation_control_g2p
+        self.enable_phoneme_text_input = enable_phoneme_text_input
+        self.text_phoneme_token_offset = text_phoneme_token_offset
+        self.partial_phoneme_text_prob = partial_phoneme_text_prob
+        self.partial_phoneme_word_prob = partial_phoneme_word_prob
+        self.partial_phoneme_word_prob_min = partial_phoneme_word_prob_min
+        self.partial_phoneme_word_prob_max = partial_phoneme_word_prob_max
+        self.phoneme_text_bop_marker = phoneme_text_bop_marker
+        self.phoneme_text_eop_marker = phoneme_text_eop_marker
         self.add_language_to_context_text = add_language_to_context_text
 
     def get_num_audio_samples_to_slice(self, duration, sample_rate):
@@ -237,19 +250,20 @@ def __getitem__(self, cuts: CutSet) -> Dict[str, Union[torch.Tensor, List]]:
                 all_tokenizers_config=self.tokenizer_config,
                 mode=self.dataset_type,
             )
-            self.bos_id = len(self.text_tokenizer.tokens)
+            if self.enable_phoneme_text_input and self.phoneme_tokenizer is None:
+                if self.phoneme_tokenizer_config is None:
+                    raise ValueError("`phoneme_tokenizer_config` is required when `enable_phoneme_text_input=True`.")
+                self.phoneme_tokenizer = safe_instantiate(self.phoneme_tokenizer_config)
+            base_text_vocab_size = len(self.text_tokenizer.tokens)
+            self.bos_id = base_text_vocab_size
             self.eos_id = self.bos_id + 1
+            if self.text_phoneme_token_offset is None:
+                self.text_phoneme_token_offset = self.eos_id + 2
             self.pad_id = self.text_tokenizer.pad
 
         # initialize the phoneme tokenizer once per dataset/worker when config is available.
         if self.phoneme_tokenizer is None and self.phoneme_tokenizer_config is not None:
             self.phoneme_tokenizer = safe_instantiate(self.phoneme_tokenizer_config)
-        if (
-            self.pronunciation_control_g2p is None
-            and self.pronunciation_control_g2p_config is not None
-            and self.phoneme_as_text_prob > 0.0
-        ):
-            self.pronunciation_control_g2p = setup_pronunciation_control_g2p(self.pronunciation_control_g2p_config)
 
         # define list to store batched information
         dataset_name_list = []
@@ -464,19 +478,54 @@ def _sample_context_duration_with_available_limit(available_duration_sec: float)
             else:
                 text_str = cut.supervisions[0].text
             raw_text_list.append(text_str)
+            text_for_tokens = text_str
+            if (
+                self.dataset_type == 'train'
+                and self.enable_phoneme_text_input
+                and self.partial_phoneme_text_prob > 0.0
+                and language not in self.ignore_phoneme_languages
+                and cut.supervisions[0].has_custom("ipa_alignment")
+                and not has_phoneme_text_spans(
+                    text_str,
+                    bop_marker=self.phoneme_text_bop_marker,
+                    eop_marker=self.phoneme_text_eop_marker,
+                )
+                and random.random() < self.partial_phoneme_text_prob
+            ):
+                min_word_prob = (
+                    self.partial_phoneme_word_prob
+                    if self.partial_phoneme_word_prob_min is None
+                    else self.partial_phoneme_word_prob_min
+                )
+                max_word_prob = (
+                    self.partial_phoneme_word_prob
+                    if self.partial_phoneme_word_prob_max is None
+                    else self.partial_phoneme_word_prob_max
+                )
+                sampled_word_prob = _sample_probability_range(
+                    "partial_phoneme_word_prob", min_word_prob, max_word_prob
+                )
+                text_for_tokens = partially_phonemize_text(
+                    text=text_str,
+                    ipa_alignment=cut.supervisions[0].ipa_alignment,
+                    partial_phoneme_word_prob=sampled_word_prob,
+                    bop_marker=self.phoneme_text_bop_marker,
+                    eop_marker=self.phoneme_text_eop_marker,
+                )
             if cut.has_custom("tokenizer_names"):
                 # Pick a random tokenizer from the list of tokenizers
                 tokenizer_name = random.choice(cut.tokenizer_names)
             else:
                 tokenizer_name = "english_phoneme"  # Default to english phoneme tokenizer
-            tokens = tokenize_text_with_pronunciation_control(
+            tokens = tokenize_text_with_phoneme_spans(
                 text_tokenizer=self.text_tokenizer,
-                text_str=text_str,
-                language=language,
+                text_str=text_for_tokens,
                 tokenizer_name=tokenizer_name,
-                dataset_type=self.dataset_type,
-                phoneme_as_text_prob=self.phoneme_as_text_prob,
-                pronunciation_control_g2p=self.pronunciation_control_g2p,
+                enable_phoneme_text_input=self.enable_phoneme_text_input,
+                phoneme_tokenizer=self.phoneme_tokenizer,
+                text_phoneme_token_offset=self.text_phoneme_token_offset,
+                bop_marker=self.phoneme_text_bop_marker,
+                eop_marker=self.phoneme_text_eop_marker,
             )
             tokens = tokens + [self.eos_id]  # Not adding BOS id
             tokens = torch.tensor(tokens, dtype=torch.int32)