Add support for framestacking - v2607 magpie tts

CMakeLists.txt

@@ -181,12 +181,9 @@ option(NEMO_SPEECH_CUBLAS_SHIM     "Build the in-tree drop-in cuBLAS shim (nativ
 # norm/BF16 fusions, CUDA graph and launch fixes) are transparent ggml-cuda internals.
 option(NEMO_SPEECH_GGML_PATCHED "Linked ggml has the project ASR patches applied (fused rel-pos op, F16 dw-conv)" ON)
 
-# Fused rel-pos attention CUDA op (GGML_OP_FUSED_RELPOS_ATTN). Replaces the
-# unfused 3-GEMM + rel-shift + softmax attention sequence with one kernel in the
-# non-cache-aware encoder path. Defaults ON whenever GGML_CUDA + a patched ggml
-# are present (it links a patch-only op symbol); forced OFF otherwise, where the
-# encoder runs the unfused stock-ggml path. Kept as an explicit option so it can
-# be toggled OFF for debugging / WER bisection (the kernel stores k/v/p F16).
+# Relative-position mode of the fused CUDA attention op. Replaces the unfused
+# encoder attention sequence with one kernel. Defaults ON with CUDA and the
+# patched ggml; the encoder otherwise uses stock ggml ops.
 if(GGML_CUDA AND NEMO_SPEECH_GGML_PATCHED)
     option(NEMO_SPEECH_FUSED_RELPOS_ATTN "Use the fused rel-pos attention CUDA op in the encoder" ON)
 else()

app/serve.cpp

@@ -228,6 +228,9 @@ run_server(int argc, char** argv) {
         "tts.enabled",
         [&](const std::string& value) { tts_enabled = parse_enablement("tts.enabled", value); },
         "Enable TTS: auto, true, or false");
+    parser.Register(
+        "tts.preempt", &server_config.preempt_tts,
+        "Cancel older HTTP TTS synthesis when a newer request arrives");
 #endif
     auto value = [&](int& index, const std::string& option) {
         if (++index >= argc)
@@ -503,8 +506,6 @@ run_server(int argc, char** argv) {
                 tts_config.runtime.codec_cpu = gpu < 0;
             } else {
                 tts_config.runtime.lt_backend = nemo_speech::tts::MagpieBackendPreference::Cuda;
-                tts_config.runtime.sampling_backend =
-                    nemo_speech::tts::MagpieBackendPreference::Cuda;
             }
         }
         tts_config.runtime.magpie_model = magpie_path;
@@ -518,6 +519,9 @@ run_server(int argc, char** argv) {
         config.default_voice_name = tts_config.default_voice_name;
         engines.load_tts(std::move(config));
     }
+    // The HTTP server owns request handling, so carry the existing TTS
+    // benchmark switch across from the TTS server configuration.
+    server_config.tts_benchmark = tts_config.benchmark;
 #endif
     if (!engines.ready())
         throw std::runtime_error(
@@ -642,6 +646,7 @@ print_serve_help(const char* program) {
         "  --codec-model MODEL     NanoCodec path or indexed model\n"
         "  --tokenizer-dir MODEL   TTS tokenizer directory or indexed model\n"
         "  --tn-model-dir MODEL    Optional TTS text-normalization assets\n"
+        "  --tts.preempt           Cancel older HTTP TTS synthesis for the newest request\n"
 #endif
 #if defined(NEMO_SPEECH_CLI_NMT)
         "  --nmt-model MODEL       Optional translation model\n"

app/synthesize.cpp

@@ -214,7 +214,6 @@ command_synthesize(int argc, char** argv) {
                 parsed.runtime.codec_cpu = true;
             } else if (cuda_device) {
                 parsed.runtime.lt_backend = nemo_speech::tts::MagpieBackendPreference::Cuda;
-                parsed.runtime.sampling_backend = nemo_speech::tts::MagpieBackendPreference::Cuda;
                 parsed.runtime.codec_cpu = false;
             } else {
                 parsed.runtime.lt_backend = nemo_speech::tts::MagpieBackendPreference::Cpu;

config/server.example.yaml

@@ -70,7 +70,7 @@ tts:
   steps: -1                     # -1 = model default
   top-k: -1                     # -1 = model default
 
-  chunk-frames: 3
+  chunk-frames: 4
   codec-queue-depth: 4
   codec-history-frames: -1
   codec-future-frames: 1
@@ -84,7 +84,7 @@ tts:
   codec-cpu: false
 
   lt-backend: auto              # auto | cpu | cuda
-  sampling-backend: auto        # auto | cpu | cuda
+  sampling-backend: auto        # CUDA with a CUDA Magpie/LT path; otherwise CPU
   uma-mode: auto                # auto | off | on
   longform: auto                # auto | off | on
 

config/tts.example.yaml

@@ -26,7 +26,7 @@ tts:
   # temperature: 0.6            # omitted = model default
   # cfg-scale: 2.5              # omitted = model default
 
-  chunk-frames: 3
+  chunk-frames: 4
   codec-queue-depth: 4
   codec-history-frames: -1
   codec-future-frames: 1
@@ -40,7 +40,7 @@ tts:
   codec-cpu: false
 
   lt-backend: auto              # auto | cpu | cuda
-  sampling-backend: auto        # auto | cpu | cuda
+  sampling-backend: auto        # CUDA with a CUDA Magpie/LT path; otherwise CPU
   uma-mode: auto                # auto | off | on
   longform: auto                # auto | off | on
 

conversion/tts.py

@@ -24,6 +24,7 @@
 import torch
 
 from .source import extract_archive, find_checkpoint_files, load_state_dict, read_checkpoint_config
+from .tts_tokenizer_profiles import tokenizer_profile
 
 SPECIAL_AUDIO_TOKENS = 8
 SPEAKER_NAMES = ["John", "Sofia", "Aria", "Jason", "Leo"]
@@ -79,24 +80,69 @@ def add_i32_array(writer: gguf.GGUFWriter, key: str, value: Any) -> None:
         writer.add_array(key, items)
 
 
+def _indexed_weight_indices(sd: dict[str, torch.Tensor], prefix: str) -> list[int]:
+    suffix = ".weight"
+    indices: list[int] = []
+    for name in sd:
+        if not name.startswith(prefix) or not name.endswith(suffix):
+            continue
+        raw_index = name[len(prefix) : -len(suffix)]
+        if not raw_index.isascii() or not raw_index.isdigit():
+            raise ValueError(f"{prefix} contains a non-numeric weight index: {name}")
+        indices.append(int(raw_index))
+    return indices
+
+
+def _require_contiguous_indices(label: str, indices: list[int], expected_count: int) -> None:
+    expected = list(range(expected_count))
+    actual = sorted(indices)
+    if actual != expected:
+        raise ValueError(
+            f"{label} indexes must be contiguous from 0 through {expected_count - 1}: "
+            f"found={actual}"
+        )
+
+
 def add_metadata(
     writer: gguf.GGUFWriter, cfg: dict[str, Any], sd: dict[str, torch.Tensor]
 ) -> dict[str, Any]:
+    encoder = cfg["encoder"]
+    decoder = cfg["decoder"]
+    lt_hidden = int(cfg.get("local_transformer_hidden_dim", 256))
+    frame_stacking = int(cfg.get("frame_stacking_factor", 1))
+    audio_embedding_indices = _indexed_weight_indices(sd, "audio_embeddings.")
+    n_stacked_codebooks = len(audio_embedding_indices)
+    if frame_stacking < 1 or n_stacked_codebooks == 0 or n_stacked_codebooks % frame_stacking:
+        raise ValueError(
+            "audio embedding count must be a positive multiple of frame_stacking_factor: "
+            f"embeddings={n_stacked_codebooks} frame_stacking_factor={frame_stacking}"
+        )
+    _require_contiguous_indices("audio embedding", audio_embedding_indices, n_stacked_codebooks)
+
     audio_vocab = int(sd["audio_embeddings.0.weight"].shape[0])
     codebook_size = audio_vocab - SPECIAL_AUDIO_TOKENS
     text_vocab = int(sd["text_embedding.weight"].shape[0])
     baked_t = int(sd["_baked_embedding_T"].item())
     baked_d = int(sd["_baked_embedding_D"].item())
     baked_lens = [int(x) for x in sd["baked_context_embedding_len"].tolist()]
-
-    encoder = cfg["encoder"]
-    decoder = cfg["decoder"]
-    lt_hidden = int(cfg.get("local_transformer_hidden_dim", 256))
-    frame_stacking = int(cfg.get("frame_stacking_factor", 1))
-    n_codebooks = int(
-        len([k for k in sd if k.startswith("audio_embeddings.") and k.endswith(".weight")])
+    profile = tokenizer_profile(cfg, text_vocab, frame_stacking)
+    n_codebooks = n_stacked_codebooks // frame_stacking
+    expected_logits = n_stacked_codebooks * audio_vocab
+    if int(sd["final_proj.weight"].shape[0]) != expected_logits:
+        raise ValueError(
+            "final_proj rows do not match stacked audio layout: "
+            f"rows={sd['final_proj.weight'].shape[0]} expected={expected_logits}"
+        )
+    lt_head_indices = _indexed_weight_indices(sd, "local_transformer_out_projections.")
+    n_lt_heads = len(lt_head_indices)
+    if n_lt_heads != n_stacked_codebooks:
+        raise ValueError(
+            "local transformer output heads do not match audio embeddings: "
+            f"heads={n_lt_heads} embeddings={n_stacked_codebooks}"
+        )
+    _require_contiguous_indices(
+        "local transformer output projection", lt_head_indices, n_stacked_codebooks
     )
-    n_codebooks //= frame_stacking
 
     inf = cfg.get("inference_parameters", {})
 
@@ -105,9 +151,11 @@ def add_metadata(
         "model_type": cfg.get("model_type"),
         "nemo_target": cfg.get("target"),
         "nemo_version": cfg.get("nemo_version"),
+        "tokenizer_profile": profile,
         "codec_model": cfg.get("codecmodel_path"),
         "text_vocab_size": text_vocab,
         "audio_codebooks": n_codebooks,
+        "stacked_audio_codebooks": n_stacked_codebooks,
         "audio_codebook_size": codebook_size,
         "audio_vocab_size": audio_vocab,
         "frame_stacking_factor": frame_stacking,
@@ -127,6 +175,7 @@ def add_metadata(
     )
     writer.add_string("magpietts.nemo_target", str(cfg.get("target", "")))
     writer.add_string("magpietts.nemo_version", str(cfg.get("nemo_version", "")))
+    writer.add_string("magpietts.tokenizer_profile", profile)
     writer.add_string("magpietts.model_type", str(cfg.get("model_type", "")))
     writer.add_string("magpietts.codec_model", str(cfg.get("codecmodel_path", "")))
     writer.add_string("magpietts.config_json", json.dumps(cfg, ensure_ascii=False, sort_keys=True))
@@ -135,6 +184,7 @@ def add_metadata(
 
     add_i32(writer, "magpietts.text_vocab_size", text_vocab)
     add_i32(writer, "magpietts.audio_codebooks", n_codebooks)
+    add_i32(writer, "magpietts.stacked_audio_codebooks", n_stacked_codebooks)
     add_i32(writer, "magpietts.audio_codebook_size", codebook_size)
     add_i32(writer, "magpietts.audio_vocab_size", audio_vocab)
     add_i32(writer, "magpietts.audio_bos_id", codebook_size + 0)
@@ -308,9 +358,9 @@ def convert(
 
         if cfg.get("target") != "nemo.collections.tts.models.magpietts.MagpieTTSModel":
             raise RuntimeError(f"unsupported target: {cfg.get('target')}")
-        if cfg.get("model_type") != "decoder_ce" or not cfg.get(
-            "has_baked_context_embedding", False
-        ):
+        # v2602 records this in the config while v2607 retains the baked
+        # tensor but omits the legacy config flag.
+        if cfg.get("model_type") != "decoder_ce" or "baked_context_embedding.weight" not in sd:
             raise RuntimeError(
                 "this GGML example expects MagpieTTS decoder_ce with baked context embeddings"
             )

conversion/tts_tokenizer_profiles.py

@@ -0,0 +1,170 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Known MagpieTTS tokenizer layouts shared by conversion checks."""
+
+from __future__ import annotations
+
+from typing import Any
+
+TOKENIZER_ORDERS = {
+    "v2602": (
+        "english_phoneme",
+        "spanish_phoneme",
+        "german_phoneme",
+        "mandarin_phoneme",
+        "japanese_phoneme",
+        "french_chartokenizer",
+        "hindi_chartokenizer",
+        "italian_phoneme",
+        "vietnamese_phoneme",
+        "text_ce_tokenizer",
+    ),
+    "v2607": (
+        "english_phoneme",
+        "text_ce_tokenizer",
+        "spanish_phoneme",
+        "german_phoneme",
+        "mandarin_phoneme",
+        "japanese_phoneme",
+        "portuguese_Brazilian_phoneme",
+        "hindi_phoneme",
+        "arabic_AE_chartokenizer",
+        "arabic_SA_chartokenizer",
+        "arabic_MSA_chartokenizer",
+        "french_chartokenizer",
+        "italian_chartokenizer",
+        "vietnamese_chartokenizer",
+        "korean_chartokenizer",
+    ),
+}
+
+TOKENIZER_PROFILE_DIMENSIONS = {
+    "v2602": (2362, 1),
+    "v2607": (3359, 2),
+}
+
+TOKENIZER_PROFILE_NEMO_VERSIONS = {
+    "v2602": "2.6.0rc0",
+    "v2607": "2.8.0rc0",
+}
+
+IPA_TARGET = "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.IPATokenizer"
+BYT5_TARGET = "AutoTokenizer"
+TOKENIZER_TARGETS = {
+    "v2602": {
+        "english_phoneme": IPA_TARGET,
+        "spanish_phoneme": IPA_TARGET,
+        "german_phoneme": IPA_TARGET,
+        "mandarin_phoneme": (
+            "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers."
+            "ChinesePhonemesTokenizer"
+        ),
+        "japanese_phoneme": (
+            "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers."
+            "JapanesePhonemeTokenizer"
+        ),
+        "french_chartokenizer": BYT5_TARGET,
+        "hindi_chartokenizer": (
+            "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers."
+            "HindiCharsTokenizer"
+        ),
+        "italian_phoneme": BYT5_TARGET,
+        "vietnamese_phoneme": BYT5_TARGET,
+        "text_ce_tokenizer": BYT5_TARGET,
+    },
+    "v2607": {
+        "english_phoneme": IPA_TARGET,
+        "text_ce_tokenizer": BYT5_TARGET,
+        "spanish_phoneme": IPA_TARGET,
+        "german_phoneme": IPA_TARGET,
+        "mandarin_phoneme": (
+            "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers."
+            "ChinesePhonemesTokenizer"
+        ),
+        "japanese_phoneme": (
+            "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers."
+            "JapanesePhonemeTokenizer"
+        ),
+        "portuguese_Brazilian_phoneme": IPA_TARGET,
+        "hindi_phoneme": IPA_TARGET,
+        "arabic_AE_chartokenizer": (
+            "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers."
+            "ArabicCharsTokenizer"
+        ),
+        "arabic_SA_chartokenizer": (
+            "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers."
+            "ArabicCharsTokenizer"
+        ),
+        "arabic_MSA_chartokenizer": (
+            "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers."
+            "ArabicCharsTokenizer"
+        ),
+        "french_chartokenizer": BYT5_TARGET,
+        "italian_chartokenizer": BYT5_TARGET,
+        "vietnamese_chartokenizer": BYT5_TARGET,
+        "korean_chartokenizer": BYT5_TARGET,
+    },
+}
+
+V2607_LANGUAGE_MAPPING = {
+    "en": ["english_phoneme"],
+    "de": ["german_phoneme"],
+    "es": ["spanish_phoneme"],
+    "fr": ["french_chartokenizer"],
+    "it": ["italian_chartokenizer"],
+    "vi": ["vietnamese_chartokenizer"],
+    "zh": ["mandarin_phoneme"],
+    "hi": ["hindi_phoneme"],
+    "ja": ["japanese_phoneme"],
+    "pt-BR": ["portuguese_Brazilian_phoneme"],
+    "ko": ["korean_chartokenizer"],
+    "ar-AE": ["arabic_AE_chartokenizer"],
+    "ar-SA": ["arabic_SA_chartokenizer"],
+    "ar-MSA": ["arabic_MSA_chartokenizer"],
+}
+
+
+def tokenizer_profile(cfg: dict[str, Any], text_vocab: int, frame_stacking: int) -> str:
+    tokenizers = cfg.get("text_tokenizers")
+    if not isinstance(tokenizers, dict):
+        raise ValueError("Magpie config has no text_tokenizers mapping")
+    order = tuple(tokenizers)
+    matches = [name for name, expected in TOKENIZER_ORDERS.items() if order == expected]
+    if len(matches) != 1:
+        raise ValueError(f"unsupported Magpie tokenizer layout: {list(order)}")
+    profile = matches[0]
+    if str(cfg.get("nemo_version", "")) != TOKENIZER_PROFILE_NEMO_VERSIONS[profile]:
+        raise ValueError(f"unsupported {profile} nemo_version: {cfg.get('nemo_version')!r}")
+    for name, target in TOKENIZER_TARGETS[profile].items():
+        tokenizer = tokenizers.get(name)
+        actual_target = tokenizer.get("_target_") if isinstance(tokenizer, dict) else None
+        if actual_target != target:
+            raise ValueError(
+                f"unsupported {profile} tokenizer target for {name}: " f"{actual_target!r}"
+            )
+    expected_japanese_case = "upper" if profile == "v2602" else "lower"
+    japanese_g2p = tokenizers["japanese_phoneme"].get("g2p")
+    japanese_case = (
+        japanese_g2p.get("ascii_letter_case") if isinstance(japanese_g2p, dict) else None
+    )
+    if japanese_case != expected_japanese_case:
+        raise ValueError(f"unsupported {profile} Japanese ascii_letter_case: {japanese_case!r}")
+    if profile == "v2607":
+        hindi = tokenizers["hindi_phoneme"]
+        hindi_g2p = hindi.get("g2p")
+        hindi_g2p_locale = hindi_g2p.get("locale") if isinstance(hindi_g2p, dict) else None
+        if hindi.get("locale") != "hi-IN" or hindi_g2p_locale != "hi-IN":
+            raise ValueError("unsupported v2607 Hindi IPA locale")
+        if tokenizers["portuguese_Brazilian_phoneme"].get("locale_specific_punct") is not False:
+            raise ValueError("unsupported v2607 Portuguese punctuation mode")
+        if cfg.get("language_to_tokenizer_mapping") != V2607_LANGUAGE_MAPPING:
+            raise ValueError("unsupported v2607 language_to_tokenizer_mapping")
+    expected_vocab, expected_stacking = TOKENIZER_PROFILE_DIMENSIONS[profile]
+    if (text_vocab, frame_stacking) != (expected_vocab, expected_stacking):
+        raise ValueError(
+            f"Magpie tokenizer profile {profile} requires text_vocab_size={expected_vocab} "
+            f"and frame_stacking_factor={expected_stacking}, got "
+            f"text_vocab_size={text_vocab} and frame_stacking_factor={frame_stacking}"
+        )
+    return profile