tts: add word_time_offsets to python client (offline + streaming)

.gitmodules

@@ -1,4 +1,6 @@
 [submodule "common"]
 	path = common
-	url = https://github.com/nvidia-riva/common.git
-	branch = main
+	# Fork with the TTS word-timestamp proto fields (enable_word_time_offsets /
+	# WordTiming). Based on upstream main; rebase + retrack main once merged.
+	url = https://github.com/ydharavath/common.git
+	branch = feat/tts-word-timestamps

riva/client/tts.py

@@ -74,6 +74,7 @@ def synthesize(
         custom_dictionary: Optional[dict] = None,
         zero_shot_transcript: Optional[str] = None,
         custom_configuration: Optional[Dict[str, str]] = None,
+        enable_word_time_offsets: bool = False,
     ) -> Union[rtts.SynthesizeSpeechResponse, _MultiThreadedRendezvous]:
         """
         Synthesizes an entire audio for text :param:`text`.
@@ -95,6 +96,9 @@ def synthesize(
             zero_shot_transcript (:obj:`str`, `optional`): Transcript corresponding to Zero shot audio prompt.
             custom_configuration (:obj:`Dict[str, str]`, `optional`): Free-form key/value parameters forwarded
                 to the synthesizer (e.g. ``{"exaggeration_factor": "1.5"}``). Model-specific.
+            enable_word_time_offsets (:obj:`bool`, defaults to :obj:`False`): If :obj:`True`, request per-word
+                start/end timestamps, returned in ``response.meta.words`` (supported by models that produce
+                word alignment, e.g. Magpie TTS).
         Returns:
             :obj:`Union[riva.client.proto.riva_tts_pb2.SynthesizeSpeechResponse, grpc._channel._MultiThreadedRendezvous]`:
             a response with output. You may find :class:`riva.client.proto.riva_tts_pb2.SynthesizeSpeechResponse` fields
@@ -109,6 +113,7 @@ def synthesize(
         )
         if voice_name is not None:
             req.voice_name = voice_name
+        req.enable_word_time_offsets = enable_word_time_offsets
         if zero_shot_audio_prompt_file is not None:
             with zero_shot_audio_prompt_file.open('rb') as f:
                 audio_data = f.read()
@@ -139,6 +144,7 @@ def synthesize_online(
         zero_shot_quality: int = 20,
         custom_dictionary: Optional[dict] = None,
         custom_configuration: Optional[Dict[str, str]] = None,
+        enable_word_time_offsets: bool = False,
     ) -> Generator[rtts.SynthesizeSpeechResponse, None, None]:
         """
         Synthesizes and yields output audio chunks for text :param:`text` as the chunks
@@ -176,6 +182,7 @@ def synthesize_online(
         )
         if voice_name is not None:
             req.voice_name = voice_name
+        req.enable_word_time_offsets = enable_word_time_offsets
 
         if zero_shot_audio_prompt_file is not None:
             with zero_shot_audio_prompt_file.open('rb') as f:

scripts/tts/talk.py

@@ -72,6 +72,11 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument("--encoding", default="LINEAR_PCM", choices={"LINEAR_PCM", "OGGOPUS"}, help="Output audio encoding.")
     parser.add_argument("--custom-dictionary", type=str, help="A file path to a user dictionary with key-value pairs separated by double spaces.")
+    parser.add_argument(
+        "--word-time-offsets", action="store_true",
+        help="Request per-word start/end timestamps (printed from response metadata). Works for both "
+        "streaming and non-streaming synthesis.",
+    )
 
     parser.add_argument(
         "--stream",
@@ -207,6 +212,7 @@ def main() -> int:
                 zero_shot_audio_prompt_file=args.zero_shot_audio_prompt_file,
                 zero_shot_quality=(20 if args.zero_shot_quality is None else args.zero_shot_quality),
                 custom_dictionary=custom_dictionary_input,
+                enable_word_time_offsets=args.word_time_offsets,
                 **custom_configuration_kwargs,
             )
             first = True
@@ -215,6 +221,10 @@ def main() -> int:
                 if first:
                     print(f"Time to first audio: {(stop - start):.3f}s")
                     first = False
+                # Word timestamps arrive on a final metadata-only chunk (no audio).
+                if args.word_time_offsets and resp.meta.words:
+                    for word in resp.meta.words:
+                        print(f'word: "{word.word}" start: {word.start_time} ms end: {word.end_time} ms')
                 if sound_stream is not None:
                     sound_stream(resp.audio)
                 if out_f is not None:
@@ -227,10 +237,14 @@ def main() -> int:
                 zero_shot_quality=(20 if args.zero_shot_quality is None else args.zero_shot_quality),
                 custom_dictionary=custom_dictionary_input,
                 zero_shot_transcript=args.zero_shot_transcript,
+                enable_word_time_offsets=args.word_time_offsets,
                 **custom_configuration_kwargs,
             )
             stop = time.time()
             print(f"Time spent: {(stop - start):.3f}s")
+            if args.word_time_offsets:
+                for word in resp.meta.words:
+                    print(f'word: "{word.word}" start: {word.start_time} ms end: {word.end_time} ms')
             if sound_stream is not None:
                 sound_stream(resp.audio)
             if out_f is not None: