Fix issue with magpie v2607 speakers being mapped to wrong id

conversion/tts.py

@@ -25,11 +25,16 @@
 import torch
 from gguf.quants import quantize
 
-from .source import extract_archive, find_checkpoint_files, load_state_dict, read_checkpoint_config
+from .source import (
+    extract_archive,
+    find_checkpoint_files,
+    load_state_dict,
+    read_checkpoint_config,
+    read_nemo_config,
+)
 from .tts_tokenizer_profiles import tokenizer_profile
 
 SPECIAL_AUDIO_TOKENS = 8
-SPEAKER_NAMES = ["John", "Sofia", "Aria", "Jason", "Leo"]
 
 # 2-D projection weights stored as Q8_0 by ``--outtype q8_0``: the attention and feed-forward
 # projections (per-tap Conv1d slices) of the text encoder, the decoder and the local
@@ -53,7 +58,15 @@ def extract_nemo(path: Path) -> tuple[Path, tempfile.TemporaryDirectory[str] | N
 
     tmp = tempfile.TemporaryDirectory(prefix="magpietts-nemo-")
     out = Path(tmp.name)
-    extract_archive(path, out, basenames={"model_config.yaml", "model_weights.ckpt"})
+    try:
+        basenames = {"model_config.yaml", "config.yaml", "model_weights.ckpt"}
+        speaker_map = read_nemo_config(path).get("speaker_map")
+        if isinstance(speaker_map, str) and speaker_map:
+            basenames.add(Path(speaker_map.removeprefix("nemo:")).name)
+        extract_archive(path, out, basenames=basenames)
+    except Exception:
+        tmp.cleanup()
+        raise
     return out, tmp
 
 
@@ -68,6 +81,38 @@ def read_state_dict(root: Path) -> dict[str, torch.Tensor]:
     return load_state_dict(checkpoint, weights_only=True)
 
 
+def read_speaker_names(
+    cfg: dict[str, Any], speaker_count: int, root: Path | None = None
+) -> list[str]:
+    """Return checkpoint names in baked speaker ID order, never JSON key order."""
+    speaker_map = cfg.get("speaker_map")
+    if isinstance(speaker_map, str) and speaker_map:
+        if root is None:
+            raise ValueError("speaker_map requires the extracted checkpoint directory")
+        # NeMo packages artifacts under a generated basename and records a
+        # nemo: URI; original absolute paths may also survive in older configs.
+        basename = Path(speaker_map.removeprefix("nemo:")).name
+        candidates = sorted(root.rglob(basename))
+        if len(candidates) != 1:
+            raise ValueError(
+                f"speaker_map artifact {basename!r}: expected one file in {root}, "
+                f"found {len(candidates)}"
+            )
+        with candidates[0].open(encoding="utf-8") as stream:
+            speaker_map = json.load(stream)
+    if not isinstance(speaker_map, dict) or not speaker_map:
+        raise ValueError("checkpoint must provide a non-empty speaker_map; cannot infer voice IDs")
+    names: set[str] = set()
+    for name, index in speaker_map.items():
+        if not isinstance(name, str) or not name.strip() or name.casefold() in names:
+            raise ValueError("speaker_map names must be non-empty and unique ignoring case")
+        names.add(name.casefold())
+        if type(index) is not int:
+            raise ValueError(f"speaker_map ID for {name!r} must be an integer")
+    _require_contiguous_indices("speaker_map", list(speaker_map.values()), speaker_count)
+    return sorted(speaker_map, key=speaker_map.__getitem__)
+
+
 def cfg_get(cfg: dict[str, Any], path: str, default: Any = None) -> Any:
     cur: Any = cfg
     for part in path.split("."):
@@ -121,7 +166,10 @@ def _require_contiguous_indices(label: str, indices: list[int], expected_count:
 
 
 def add_metadata(
-    writer: gguf.GGUFWriter, cfg: dict[str, Any], sd: dict[str, torch.Tensor]
+    writer: gguf.GGUFWriter,
+    cfg: dict[str, Any],
+    sd: dict[str, torch.Tensor],
+    checkpoint_root: Path | None = None,
 ) -> dict[str, Any]:
     encoder = cfg["encoder"]
     decoder = cfg["decoder"]
@@ -160,6 +208,10 @@ def add_metadata(
     _require_contiguous_indices(
         "local transformer output projection", lt_head_indices, n_stacked_codebooks
     )
+    speaker_count = int(sd["baked_context_embedding.weight"].shape[0])
+    speaker_names = read_speaker_names(cfg, speaker_count, checkpoint_root)
+    if len(baked_lens) != speaker_count:
+        raise ValueError("baked context length count does not match baked speaker count")
 
     inf = cfg.get("inference_parameters", {})
 
@@ -180,7 +232,7 @@ def add_metadata(
         "encoder_layers": int(encoder["n_layers"]),
         "decoder_layers": int(decoder["n_layers"]),
         "context_length": int(decoder["max_length_causal_mask"]),
-        "speaker_names": SPEAKER_NAMES,
+        "speaker_names": speaker_names,
         "baked_context_length": baked_t,
         "baked_context_dim": baked_d,
         "baked_context_lens": baked_lens,
@@ -196,7 +248,7 @@ def add_metadata(
     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))
-    writer.add_array("magpietts.speaker_names", SPEAKER_NAMES)
+    writer.add_array("magpietts.speaker_names", speaker_names)
     writer.add_array("magpietts.baked_context_lens", baked_lens)
 
     add_i32(writer, "magpietts.text_vocab_size", text_vocab)
@@ -235,7 +287,7 @@ def add_metadata(
     )
     add_i32(writer, "magpietts.baked_context_length", baked_t)
     add_i32(writer, "magpietts.baked_context_dim", baked_d)
-    add_i32(writer, "magpietts.baked_speakers", int(sd["baked_context_embedding.weight"].shape[0]))
+    add_i32(writer, "magpietts.baked_speakers", speaker_count)
     add_i32(writer, "magpietts.inference.max_decoder_steps", inf.get("max_decoder_steps", 500))
     add_i32(writer, "magpietts.inference.topk", inf.get("topk", 80))
     add_i32(writer, "magpietts.inference.min_generated_frames", inf.get("min_generated_frames", 4))
@@ -398,7 +450,7 @@ def convert(
 
         output.parent.mkdir(parents=True, exist_ok=True)
         writer = gguf.GGUFWriter(output, "magpietts")
-        summary = add_metadata(writer, cfg, sd)
+        summary = add_metadata(writer, cfg, sd, root)
         n_written, skipped = add_tensors(writer, sd, outtype, local_transformer_outtype)
 
         writer.write_header_to_file()

docs/tts/models.md

@@ -112,6 +112,33 @@ environment setup.
 python3 convert_model.py custom-magpie.nemo --outfile custom-magpie.q8_0.gguf
 ```
 
+Voice names are read from the checkpoint's `speaker_map` (including its bundled
+`speakers.json`) and stored in `magpietts.speaker_names` in numeric speaker ID
+order. Conversion rejects missing maps or IDs that do not match the baked
+speaker count. Extracted checkpoints must retain the referenced speaker file.
+The mappings differ between releases:
+
+| Speaker ID | v2602 | v2607 |
+|---|---|---|
+| 0 | John | Aria |
+| 1 | Sofia | Jason |
+| 2 | Aria | John |
+| 3 | Jason | Leo |
+| 4 | Leo | Sofia |
+
+Older converters wrote the v2602 names into every GGUF, causing v2607 named
+voices to select the wrong speaker. Reconvert the v2607 `.nemo` with the updated
+converter and select the new GGUF, then restart the runtime:
+
+```bash
+python3 convert_model.py magpie_tts_multilingual_357m.nemo \
+    --outtype f16 --outfile magpie_tts_multilingual_357m.v2607.fixed.f16.gguf
+```
+
+Updating the converter does not change existing GGUFs or published downloads.
+Until the affected GGUF is replaced, requesting the numeric voice `"2"` selects
+John on v2607. v2602's John remains at ID 0.
+
 Conversion does not require `nemo_toolkit`. The optional
 `scripts/tts/tokenize-magpietts.py` debugging helper does.
 

tests/conversion/tts_speaker_names_test.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import io
+import json
+import sys
+import tarfile
+import tempfile
+import unittest
+from pathlib import Path
+from unittest import mock
+
+import gguf
+import torch
+import yaml
+
+ROOT = Path(__file__).resolve().parents[2]
+sys.path.insert(0, str(ROOT))
+
+from conversion.tts import convert, read_speaker_names  # noqa: E402
+
+
+class TtsSpeakerNamesTest(unittest.TestCase):
+    def test_checkpoint_ids_determine_order(self) -> None:
+        # Deliberately use key orders different from speaker ID order.
+        for mapping, expected in (
+            (
+                {"Aria": 2, "Jason": 3, "John": 0, "Leo": 4, "Sofia": 1},
+                ["John", "Sofia", "Aria", "Jason", "Leo"],
+            ),
+            (
+                {"Sofia": 4, "John": 2, "Leo": 3, "Aria": 0, "Jason": 1},
+                ["Aria", "Jason", "John", "Leo", "Sofia"],
+            ),
+        ):
+            with self.subTest(mapping=mapping):
+                self.assertEqual(read_speaker_names({"speaker_map": mapping}, 5), expected)
+
+    def test_invalid_maps_are_rejected(self) -> None:
+        for mapping in (
+            None,
+            {},
+            {"John": 0, "Aria": 0},
+            {"John": 0, "Aria": 2},
+            {"John": -1, "Aria": 0},
+            {"John": True, "Aria": 0},
+            {"John": "0", "Aria": 1},
+            {"": 0, "Aria": 1},
+            {"John": 0, "john": 1},
+            {"John": 0},
+        ):
+            with self.subTest(mapping=mapping), self.assertRaises(ValueError):
+                read_speaker_names({"speaker_map": mapping}, 2)
+
+    def test_missing_or_ambiguous_artifact_is_rejected(self) -> None:
+        with tempfile.TemporaryDirectory() as temporary:
+            root = Path(temporary)
+            config = {"speaker_map": "nemo:speakers.json"}
+            with self.assertRaisesRegex(ValueError, "found 0"):
+                read_speaker_names(config, 1, root)
+            for directory in (root / "a", root / "b"):
+                directory.mkdir()
+                (directory / "speakers.json").write_text('{"John": 0}', encoding="utf-8")
+            with self.assertRaisesRegex(ValueError, "found 2"):
+                read_speaker_names(config, 1, root)
+
+    def test_archive_conversion_preserves_speaker_ids_in_gguf_and_summary(self) -> None:
+        layouts = {
+            "v2602": {"Aria": 2, "Jason": 3, "John": 0, "Leo": 4, "Sofia": 1},
+            "v2607": {"Aria": 0, "Jason": 1, "John": 2, "Leo": 3, "Sofia": 4},
+        }
+        for profile, mapping in layouts.items():
+            with self.subTest(profile=profile), tempfile.TemporaryDirectory() as temporary:
+                root = Path(temporary)
+                artifact = "nested/hash_magpie.nemo.speakers.json"
+                config = {
+                    "target": "nemo.collections.tts.models.magpietts.MagpieTTSModel",
+                    "model_type": "decoder_ce",
+                    "speaker_map": "nemo:" + Path(artifact).name,
+                    "encoder": {"n_layers": 1, "sa_n_heads": 1, "kernel_size": 1},
+                    "decoder": {
+                        "d_model": 1,
+                        "d_ffn": 1,
+                        "n_layers": 1,
+                        "max_length_causal_mask": 8,
+                        "sa_n_heads": 1,
+                        "xa_n_heads": 1,
+                        "xa_d_head": 1,
+                        "kernel_size": 1,
+                    },
+                }
+                state = {
+                    "audio_embeddings.0.weight": torch.zeros((16, 1)),
+                    "text_embedding.weight": torch.zeros((1, 1)),
+                    "_baked_embedding_T": torch.tensor(1),
+                    "_baked_embedding_D": torch.tensor(1),
+                    "baked_context_embedding_len": torch.ones(5, dtype=torch.int32),
+                    "baked_context_embedding.weight": torch.zeros((5, 1)),
+                    "final_proj.weight": torch.zeros((16, 1)),
+                    "local_transformer_out_projections.0.weight": torch.zeros((1, 1)),
+                    "local_transformer.position_embeddings.weight": torch.zeros((1, 1)),
+                }
+                archive = root / "model.nemo"
+                with tarfile.open(archive, "w") as tar:
+                    for name, payload in (
+                        ("model_config.yaml", yaml.safe_dump(config).encode()),
+                        ("model_weights.ckpt", b"mocked weights"),
+                        (artifact, json.dumps(mapping).encode()),
+                    ):
+                        member = tarfile.TarInfo(name)
+                        member.size = len(payload)
+                        tar.addfile(member, io.BytesIO(payload))
+                output, summary_path = root / "model.gguf", root / "summary.json"
+                with (
+                    mock.patch("conversion.tts.read_state_dict", return_value=state),
+                    mock.patch("conversion.tts.tokenizer_profile", return_value=profile),
+                    mock.patch("conversion.tts.add_tensors", return_value=(0, [])),
+                ):
+                    convert(archive, output, metadata_json=summary_path)
+                expected = sorted(mapping, key=mapping.__getitem__)
+                reader = gguf.GGUFReader(output)
+                self.assertEqual(reader.get_field("magpietts.speaker_names").contents(), expected)
+                self.assertEqual(json.loads(summary_path.read_text())["speaker_names"], expected)
+
+
+if __name__ == "__main__":
+    unittest.main()