Add Nemotron 3 Diarization Core ML (fp16 storage, fp32 compute)

8976c3c verified about 1 month ago

7.54 kB

#!/usr/bin/env python3
"""Export Nemotron-3-Diarization (Sortformer) to CoreML via forward_for_export.
The export boundary is NVIDIA's own streaming graph:
    (chunk, chunk_lengths, spkcache, spkcache_lengths, fifo, fifo_lengths)
        -> (spkcache_fifo_chunk_preds, chunk_pre_encode_embs, chunk_pre_encode_lengths)
Mel features are deliberately NOT part of the graph: the host computes them, so
front-end drift stays isolated from conversion error (gate 1 vs gates 2-3).
The speaker cache / FIFO are host-owned state passed in per chunk, which keeps
the CoreML model stateless and matches how FluidAudio drives Sortformer today.
"""
from __future__ import annotations

import argparse
import warnings
from pathlib import Path

import numpy as np
import torch

warnings.filterwarnings("ignore")

import coremltools as ct  # noqa: E402
from nemo.collections.asr.models import SortformerEncLabelModel  # noqa: E402

import coreml_patches  # noqa: E402

coreml_patches.apply()

HERE = Path(__file__).resolve().parent
NEMO_PATH = HERE.parent / "Nemotron-3-Diarization-preview.nemo"
# The final release ships bare weights; its graph is identical to the preview's.
FINAL_WEIGHTS = Path.home() / "Nemotron-3-Diarization-safetensors" / "model.safetensors"
# Auxiliary 3-class activity head: feeds only the training loss upstream, never the speaker preds.
TRAINING_ONLY_PREFIX = "sortformer_modules.activity_head."

# Latency profiles from the model card. Frames are 80 ms; the input-buffer
# latency is (chunk_len + chunk_right_context) * 80 ms. spkcache_len is 264 in
# every profile, so only the chunk/FIFO geometry changes.
PROFILES = {
    # 30.4 s buffer -- batch use, one call covers half a minute of audio.
    "offline": dict(chunk_len=340, chunk_right_context=40, fifo_len=40,
                    spkcache_update_period=300),
    # 1.04 s buffer -- the usable streaming profile.
    "low": dict(chunk_len=9, chunk_right_context=4, fifo_len=264,
                spkcache_update_period=222),
    # 0.64 s buffer.
    "verylow": dict(chunk_len=6, chunk_right_context=2, fifo_len=264,
                    spkcache_update_period=222),
    # 0.32 s buffer -- lowest latency the card documents.
    "ultralow": dict(chunk_len=3, chunk_right_context=1, fifo_len=264,
                     spkcache_update_period=222),
}
PROFILE_SECONDS = {
    name: (cfg["chunk_len"] + cfg["chunk_right_context"]) * 0.08
    for name, cfg in PROFILES.items()
}

OFFLINE_PROFILE = PROFILES["offline"]

OUTPUT_NAMES = ["spkcache_fifo_chunk_preds", "chunk_pre_encode_embs",
                "chunk_pre_encode_lengths"]


class StreamingWrapper(torch.nn.Module):
    """Trace-friendly shim.
    Lengths cross the CoreML boundary as int32 (CoreML has no int64 inputs) and
    are widened back to int64 here, where the NeMo graph expects them.
    """

    def __init__(self, model: SortformerEncLabelModel):
        super().__init__()
        self.model = model

    def forward(self, chunk, chunk_lengths, spkcache, spkcache_lengths, fifo, fifo_lengths):
        return self.model.forward_for_export(
            chunk,
            chunk_lengths.to(torch.int64),
            spkcache,
            spkcache_lengths.to(torch.int64),
            fifo,
            fifo_lengths.to(torch.int64),
        )


def load_model(profile: dict, weights: str = "preview") -> SortformerEncLabelModel:
    model = SortformerEncLabelModel.restore_from(
        restore_path=str(NEMO_PATH), map_location="cpu", strict=False
    )
    if weights == "final":
        from safetensors.torch import load_file

        state = {
            key: value.float()
            for key, value in load_file(str(FINAL_WEIGHTS)).items()
            if not key.startswith(TRAINING_ONLY_PREFIX)
        }
        missing, unexpected = model.load_state_dict(state, strict=False)
        if missing or unexpected:
            raise SystemExit(f"weight mismatch: missing={missing[:5]} unexpected={unexpected[:5]}")
    model.eval()
    for key, value in profile.items():
        setattr(model.sortformer_modules, key, value)
    return model


def build_examples(model: SortformerEncLabelModel):
    """Model-sized examples, but with lengths as int32 to match the CoreML iface."""
    chunk, chunk_len, spkcache, spkcache_len, fifo, fifo_len = model.streaming_input_examples(1)
    return (
        chunk,
        chunk_len.to(torch.int32),
        spkcache,
        spkcache_len.to(torch.int32),
        fifo,
        fifo_len.to(torch.int32),
    )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--precision", choices=["fp16", "fp32"], default="fp16")
    parser.add_argument("--profile", choices=sorted(PROFILES), default="offline")
    parser.add_argument("--weights", choices=["preview", "final"], default="preview")
    parser.add_argument("--out", type=Path, default=None)
    args = parser.parse_args()

    suffix = "" if args.profile == "offline" else f"_{args.profile}"
    prefix = "nemotron_diar" if args.weights == "preview" else "nemotron_diar_final"
    out_path = args.out or (HERE / "artifacts" /
                            f"{prefix}_{args.precision}{suffix}.mlpackage")
    out_path.parent.mkdir(parents=True, exist_ok=True)

    profile = PROFILES[args.profile]
    print(f"profile: {args.profile} ({PROFILE_SECONDS[args.profile]:.2f}s buffer) {profile}")
    model = load_model(profile, args.weights)
    wrapper = StreamingWrapper(model).eval()
    examples = build_examples(model)

    for name, tensor in zip(
        ["chunk", "chunk_lengths", "spkcache", "spkcache_lengths", "fifo", "fifo_lengths"],
        examples,
    ):
        print(f"  input {name:<18} {tuple(tensor.shape)} {tensor.dtype}")

    with torch.no_grad():
        reference = wrapper(*examples)
    for name, tensor in zip(OUTPUT_NAMES, reference):
        print(f"  output {name:<28} {tuple(tensor.shape)} {tensor.dtype}")

    print("\ntracing...")
    with torch.no_grad():
        traced = torch.jit.trace(wrapper, examples, strict=False)
    # Every shape in this graph is static (batch 1, fixed chunk/cache/FIFO), but the
    # trace still carries ~200 aten::Int shape computations. Freezing constant-folds
    # them; coremltools chokes on the ones that carry (1,)-shaped tensors.
    traced = torch.jit.freeze(traced)
    traced = torch.jit.run_frozen_optimizations(traced) or traced

    print(f"converting to CoreML ({args.precision})...")
    precision = ct.precision.FLOAT16 if args.precision == "fp16" else ct.precision.FLOAT32
    mlmodel = ct.convert(
        traced,
        convert_to="mlprogram",
        inputs=[
            ct.TensorType(name="chunk", shape=examples[0].shape, dtype=np.float32),
            ct.TensorType(name="chunk_lengths", shape=examples[1].shape, dtype=np.int32),
            ct.TensorType(name="spkcache", shape=examples[2].shape, dtype=np.float32),
            ct.TensorType(name="spkcache_lengths", shape=examples[3].shape, dtype=np.int32),
            ct.TensorType(name="fifo", shape=examples[4].shape, dtype=np.float32),
            ct.TensorType(name="fifo_lengths", shape=examples[5].shape, dtype=np.int32),
        ],
        outputs=[ct.TensorType(name=name) for name in OUTPUT_NAMES],
        compute_precision=precision,
        minimum_deployment_target=ct.target.macOS14,
    )

    mlmodel.save(str(out_path))
    size_mb = sum(f.stat().st_size for f in out_path.rglob("*") if f.is_file()) / 1e6
    print(f"\nsaved {out_path}  ({size_mb:.1f} MB)")


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""Export Nemotron-3-Diarization (Sortformer) to CoreML via forward_for_export.

	The export boundary is NVIDIA's own streaming graph:

	(chunk, chunk_lengths, spkcache, spkcache_lengths, fifo, fifo_lengths)
	-> (spkcache_fifo_chunk_preds, chunk_pre_encode_embs, chunk_pre_encode_lengths)

	Mel features are deliberately NOT part of the graph: the host computes them, so
	front-end drift stays isolated from conversion error (gate 1 vs gates 2-3).
	The speaker cache / FIFO are host-owned state passed in per chunk, which keeps
	the CoreML model stateless and matches how FluidAudio drives Sortformer today.
	"""
	from __future__ import annotations

	import argparse
	import warnings
	from pathlib import Path

	import numpy as np
	import torch

	warnings.filterwarnings("ignore")

	import coremltools as ct # noqa: E402
	from nemo.collections.asr.models import SortformerEncLabelModel # noqa: E402

	import coreml_patches # noqa: E402

	coreml_patches.apply()

	HERE = Path(__file__).resolve().parent
	NEMO_PATH = HERE.parent / "Nemotron-3-Diarization-preview.nemo"
	# The final release ships bare weights; its graph is identical to the preview's.
	FINAL_WEIGHTS = Path.home() / "Nemotron-3-Diarization-safetensors" / "model.safetensors"
	# Auxiliary 3-class activity head: feeds only the training loss upstream, never the speaker preds.
	TRAINING_ONLY_PREFIX = "sortformer_modules.activity_head."

	# Latency profiles from the model card. Frames are 80 ms; the input-buffer
	# latency is (chunk_len + chunk_right_context) * 80 ms. spkcache_len is 264 in
	# every profile, so only the chunk/FIFO geometry changes.
	PROFILES = {
	# 30.4 s buffer -- batch use, one call covers half a minute of audio.
	"offline": dict(chunk_len=340, chunk_right_context=40, fifo_len=40,
	spkcache_update_period=300),
	# 1.04 s buffer -- the usable streaming profile.
	"low": dict(chunk_len=9, chunk_right_context=4, fifo_len=264,
	spkcache_update_period=222),
	# 0.64 s buffer.
	"verylow": dict(chunk_len=6, chunk_right_context=2, fifo_len=264,
	spkcache_update_period=222),
	# 0.32 s buffer -- lowest latency the card documents.
	"ultralow": dict(chunk_len=3, chunk_right_context=1, fifo_len=264,
	spkcache_update_period=222),
	}
	PROFILE_SECONDS = {
	name: (cfg["chunk_len"] + cfg["chunk_right_context"]) * 0.08
	for name, cfg in PROFILES.items()
	}

	OFFLINE_PROFILE = PROFILES["offline"]

	OUTPUT_NAMES = ["spkcache_fifo_chunk_preds", "chunk_pre_encode_embs",
	"chunk_pre_encode_lengths"]


	class StreamingWrapper(torch.nn.Module):
	"""Trace-friendly shim.

	Lengths cross the CoreML boundary as int32 (CoreML has no int64 inputs) and
	are widened back to int64 here, where the NeMo graph expects them.
	"""

	def __init__(self, model: SortformerEncLabelModel):
	super().__init__()
	self.model = model

	def forward(self, chunk, chunk_lengths, spkcache, spkcache_lengths, fifo, fifo_lengths):
	return self.model.forward_for_export(
	chunk,
	chunk_lengths.to(torch.int64),
	spkcache,
	spkcache_lengths.to(torch.int64),
	fifo,
	fifo_lengths.to(torch.int64),
	)


	def load_model(profile: dict, weights: str = "preview") -> SortformerEncLabelModel:
	model = SortformerEncLabelModel.restore_from(
	restore_path=str(NEMO_PATH), map_location="cpu", strict=False
	)
	if weights == "final":
	from safetensors.torch import load_file

	state = {
	key: value.float()
	for key, value in load_file(str(FINAL_WEIGHTS)).items()
	if not key.startswith(TRAINING_ONLY_PREFIX)
	}
	missing, unexpected = model.load_state_dict(state, strict=False)
	if missing or unexpected:
	raise SystemExit(f"weight mismatch: missing={missing[:5]} unexpected={unexpected[:5]}")
	model.eval()
	for key, value in profile.items():
	setattr(model.sortformer_modules, key, value)
	return model


	def build_examples(model: SortformerEncLabelModel):
	"""Model-sized examples, but with lengths as int32 to match the CoreML iface."""
	chunk, chunk_len, spkcache, spkcache_len, fifo, fifo_len = model.streaming_input_examples(1)
	return (
	chunk,
	chunk_len.to(torch.int32),
	spkcache,
	spkcache_len.to(torch.int32),
	fifo,
	fifo_len.to(torch.int32),
	)


	def main() -> None:
	parser = argparse.ArgumentParser(description=__doc__)
	parser.add_argument("--precision", choices=["fp16", "fp32"], default="fp16")
	parser.add_argument("--profile", choices=sorted(PROFILES), default="offline")
	parser.add_argument("--weights", choices=["preview", "final"], default="preview")
	parser.add_argument("--out", type=Path, default=None)
	args = parser.parse_args()

	suffix = "" if args.profile == "offline" else f"_{args.profile}"
	prefix = "nemotron_diar" if args.weights == "preview" else "nemotron_diar_final"
	out_path = args.out or (HERE / "artifacts" /
	f"{prefix}_{args.precision}{suffix}.mlpackage")
	out_path.parent.mkdir(parents=True, exist_ok=True)

	profile = PROFILES[args.profile]
	print(f"profile: {args.profile} ({PROFILE_SECONDS[args.profile]:.2f}s buffer) {profile}")
	model = load_model(profile, args.weights)
	wrapper = StreamingWrapper(model).eval()
	examples = build_examples(model)

	for name, tensor in zip(
	["chunk", "chunk_lengths", "spkcache", "spkcache_lengths", "fifo", "fifo_lengths"],
	examples,
	):
	print(f" input {name:<18} {tuple(tensor.shape)} {tensor.dtype}")

	with torch.no_grad():
	reference = wrapper(*examples)
	for name, tensor in zip(OUTPUT_NAMES, reference):
	print(f" output {name:<28} {tuple(tensor.shape)} {tensor.dtype}")

	print("\ntracing...")
	with torch.no_grad():
	traced = torch.jit.trace(wrapper, examples, strict=False)
	# Every shape in this graph is static (batch 1, fixed chunk/cache/FIFO), but the
	# trace still carries ~200 aten::Int shape computations. Freezing constant-folds
	# them; coremltools chokes on the ones that carry (1,)-shaped tensors.
	traced = torch.jit.freeze(traced)
	traced = torch.jit.run_frozen_optimizations(traced) or traced

	print(f"converting to CoreML ({args.precision})...")
	precision = ct.precision.FLOAT16 if args.precision == "fp16" else ct.precision.FLOAT32
	mlmodel = ct.convert(
	traced,
	convert_to="mlprogram",
	inputs=[
	ct.TensorType(name="chunk", shape=examples[0].shape, dtype=np.float32),
	ct.TensorType(name="chunk_lengths", shape=examples[1].shape, dtype=np.int32),
	ct.TensorType(name="spkcache", shape=examples[2].shape, dtype=np.float32),
	ct.TensorType(name="spkcache_lengths", shape=examples[3].shape, dtype=np.int32),
	ct.TensorType(name="fifo", shape=examples[4].shape, dtype=np.float32),
	ct.TensorType(name="fifo_lengths", shape=examples[5].shape, dtype=np.int32),
	],
	outputs=[ct.TensorType(name=name) for name in OUTPUT_NAMES],
	compute_precision=precision,
	minimum_deployment_target=ct.target.macOS14,
	)

	mlmodel.save(str(out_path))
	size_mb = sum(f.stat().st_size for f in out_path.rglob("*") if f.is_file()) / 1e6
	print(f"\nsaved {out_path} ({size_mb:.1f} MB)")


	if __name__ == "__main__":
	main()