From 7bd503d3308b94220417755083c5b2c8b5bb70f8 Mon Sep 17 00:00:00 2001 From: Anand Joseph Date: Mon, 12 Oct 2026 13:16:19 +0000 Subject: [PATCH] Fix issue with magpie v2607 speakers being mapped to wrong id Signed-off-by: Anand Joseph --- conversion/tts.py | 68 +++++++++-- docs/tts/models.md | 27 +++++ tests/conversion/tts_speaker_names_test.py | 130 +++++++++++++++++++++ 3 files changed, 217 insertions(+), 8 deletions(-) create mode 100644 tests/conversion/tts_speaker_names_test.py diff --git a/conversion/tts.py b/conversion/tts.py index cba60d0..c9b7f23 100644 --- a/conversion/tts.py +++ b/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() diff --git a/docs/tts/models.md b/docs/tts/models.md index 3a0dc27..b40455b 100644 --- a/docs/tts/models.md +++ b/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. diff --git a/tests/conversion/tts_speaker_names_test.py b/tests/conversion/tts_speaker_names_test.py new file mode 100644 index 0000000..f9309f0 --- /dev/null +++ b/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()