python/tests/test_host_chain.py

#!/usr/bin/env python3
"""Host-chain unit tests for the Inflect AX TTS Python SDK.

Runs on a plain host without NPU: onnxruntime is used as a numeric stand-in
for the AXMODEL runtime (the SDK pipeline above the session layer is identical
for both backends; see README.md for what is NOT covered by this).

Runnable two ways:
    python tests/test_host_chain.py     # plain script, exit code 0 on pass
    pytest tests/                       # if pytest is installed
"""

from __future__ import annotations

import sys
from pathlib import Path

import numpy as np

SDK_ROOT = Path(__file__).resolve().parents[1]
PKG_DIR = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(SDK_ROOT))

from inflect_ax_tts import DUMMY_PHONEME_IDS, InflectTTS  # noqa: E402
from inflect_ax_tts import host_chain  # noqa: E402
from inflect_ax_tts.frontend import intersperse  # noqa: E402

CONVERT_DIR = PKG_DIR / "model_convert"
REF_NPZ = CONVERT_DIR / "export" / "reference_io.npz"
ENCODER_ONNX = CONVERT_DIR / "export" / "encoder.onnx"
DECODER_ONNX = CONVERT_DIR / "export" / "decoder.onnx"


def _cosine(a: np.ndarray, b: np.ndarray) -> float:
    a = np.asarray(a, dtype=np.float64).ravel()
    b = np.asarray(b, dtype=np.float64).ravel()
    assert a.shape == b.shape, f"shape {a.shape} vs {b.shape}"
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-30))


def _torch_reference_expand(logw, m_p, logs_p, x_len, length_scale):
    """Replicates export/export_onnx.py:host_expand with torch (ground truth)."""
    import torch

    def sequence_mask(lengths, max_length):
        return torch.arange(max_length).unsqueeze(0) < lengths.unsqueeze(1)

    x_lengths = torch.LongTensor([x_len])
    x_mask = sequence_mask(x_lengths, logw.shape[-1]).unsqueeze(1).float()
    w = torch.exp(torch.from_numpy(logw)) * x_mask * length_scale
    w_ceil = torch.ceil(w)
    y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
    y_len = int(y_lengths.item())
    y_mask = sequence_mask(y_lengths, y_len).unsqueeze(1).float()
    attn_mask = x_mask.unsqueeze(2) * y_mask.unsqueeze(-1)
    # generate_path (origin/runtime/commons.py)
    cum = torch.cumsum(w_ceil, -1).view(-1)
    path = sequence_mask(cum, y_len).float().view(1, 1, -1, y_len)
    path = path - torch.nn.functional.pad(path, (0, 0, 1, 0))[:, :, :-1]
    attn = path.transpose(2, 3) * attn_mask  # [1, 1, T', T]
    m_p_e = torch.matmul(attn.squeeze(1), torch.from_numpy(m_p).transpose(1, 2))
    logs_p_e = torch.matmul(attn.squeeze(1), torch.from_numpy(logs_p).transpose(1, 2))
    return m_p_e[0].numpy(), logs_p_e[0].numpy(), y_len


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------

def test_expand_priors_matches_torch_reference():
    """numpy host chain (durations + generate_path + matmul) vs torch."""
    ref = np.load(REF_NPZ)
    logw, m_p, logs_p = ref["logw_torch"], ref["m_p_torch"], ref["logs_p_torch"]
    x_len = int(ref["x_lengths"][0])
    for length_scale in (1.0, 0.8, 1.6):
        exp_m, exp_l, y_len = host_chain.expand_priors(
            logw, m_p, logs_p, x_len, length_scale
        )
        ref_m, ref_l, ref_y = _torch_reference_expand(
            logw, m_p, logs_p, x_len, length_scale
        )
        assert y_len == ref_y, f"T' {y_len} vs torch {ref_y}"
        assert np.allclose(exp_m, ref_m, atol=1e-5), (
            f"m_p_e max diff {np.abs(exp_m - ref_m).max():.2e} @ls={length_scale}"
        )
        assert np.allclose(exp_l, ref_l, atol=1e-5), (
            f"logs_p_e max diff {np.abs(exp_l - ref_l).max():.2e} @ls={length_scale}"
        )
    print("[PASS] expand_priors == torch host_expand (3 length scales)")


def test_encoder_padding_semantics_ort():
    """Padded static-shape encoder call == unpadded reference (first x_len frames)."""
    import onnxruntime as ort

    ref = np.load(REF_NPZ)
    tokens = np.zeros((1, 256), dtype=np.int64)
    x_len = int(ref["x_lengths"][0])
    tokens[0, :x_len] = ref["tokens"][0]
    sess = ort.InferenceSession(str(ENCODER_ONNX), providers=["CPUExecutionProvider"])
    m_p, logs_p, logw = sess.run(
        None, {"tokens": tokens, "x_lengths": np.asarray([x_len], dtype=np.int64)}
    )
    for name, got, want in (
        ("m_p", m_p, ref["m_p_torch"]),
        ("logs_p", logs_p, ref["logs_p_torch"]),
        ("logw", logw, ref["logw_torch"]),
    ):
        cos = _cosine(got[..., :x_len], want[..., :x_len])
        assert cos >= 0.999999, f"{name} cos {cos}"
    print("[PASS] encoder T=256 padding: first-x_len frames match torch ref")


def test_decoder_chunk_compose_ort():
    """SDK chunked decode (Tp=512 path) vs EXPORT reference chain wav."""
    import onnxruntime as ort

    ref = np.load(REF_NPZ)
    z_p = ref["z_p"][0]  # [192, 51]
    sess = ort.InferenceSession(str(DECODER_ONNX), providers=["CPUExecutionProvider"])

    def run_chunk(z_chunk: np.ndarray) -> np.ndarray:
        return sess.run(None, {"z_p": z_chunk[None].astype(np.float32)})[0].reshape(-1)

    wav = host_chain.decode_waveform(z_p, run_chunk)
    want = ref["wav_onnx_chain"].reshape(-1)
    assert wav.shape == want.shape, f"{wav.shape} vs {want.shape}"
    # Right zero-padding bleeds into the trailing frames (flow/dec conv
    # receptive field; SIMULATE §4 uses the same metric convention): crop
    # min(64, T'/4) frames off the right edge for the comparison.
    t_prime = z_p.shape[1]
    crop = min(64, t_prime // 4) * host_chain.HOP_LENGTH
    cos = _cosine(wav[:-crop], want[:-crop])
    assert cos >= 0.9999, f"decoder compose cos {cos}"
    head = _cosine(wav[: (t_prime // 2) * host_chain.HOP_LENGTH],
                   want[: (t_prime // 2) * host_chain.HOP_LENGTH])
    assert head >= 0.999999, f"decoder compose head cos {head}"
    print(f"[PASS] decoder chunk compose vs EXPORT chain: cos={cos:.9f} (right-edge crop), head cos={head:.9f}")


def test_decoder_multi_chunk_compose_ort():
    """T'>512 multi-chunk path: overlap crossfade vs single dynamic ORT call.

    Uses synthetic z_p (not a real prior): checks stitching correctness only,
    comparing against the dynamic-shape ONNX decoder run over the full length.
    Boundary regions are excluded from the metric (conv receptive-field bleed
    is expected; EXPORT_NOTES §5.3).
    """
    import onnxruntime as ort

    rng = np.random.default_rng(1234)
    t_prime = 512 + 300
    z_p = rng.standard_normal((192, t_prime), dtype=np.float32) * 0.5
    sess = ort.InferenceSession(str(DECODER_ONNX), providers=["CPUExecutionProvider"])
    want = sess.run(None, {"z_p": z_p[None]})[0].reshape(-1)
    got = host_chain.decode_waveform(
        z_p, lambda z: sess.run(None, {"z_p": z[None]})[0].reshape(-1)
    )
    assert got.shape == want.shape
    # Exclude 32 frames around the chunk boundary (frame 512) and the edges.
    hop = host_chain.HOP_LENGTH
    lo, hi = (512 + 32) * hop, (t_prime - 32) * hop
    interior = _cosine(got[lo:hi], want[lo:hi])
    head = _cosine(got[: 480 * hop], want[: 480 * hop])
    assert interior >= 0.999, f"multi-chunk interior cos {interior}"
    assert head >= 0.999, f"multi-chunk head cos {head}"
    print(f"[PASS] multi-chunk compose: head cos={head:.6f}, interior cos={interior:.6f}")


def test_synthesize_tokens_end_to_end_ort():
    """Full SDK pipeline on ORT stand-in: length, determinism, finiteness."""
    tts = InflectTTS(ENCODER_ONNX, DECODER_ONNX, backend="onnxruntime")
    sr, wav1 = tts.synthesize_tokens(DUMMY_PHONEME_IDS, seed=0, variation=0.667)
    _, wav2 = tts.synthesize_tokens(DUMMY_PHONEME_IDS, seed=0, variation=0.667)
    _, wav3 = tts.synthesize_tokens(DUMMY_PHONEME_IDS, seed=1, variation=0.667)
    ref = np.load(REF_NPZ)
    assert sr == 24000
    # Same Host chain as EXPORT baseline -> same T'=51 -> 13056 samples.
    assert wav1.size == int(ref["wav_onnx_chain"].size), (
        f"{wav1.size} vs baseline {ref['wav_onnx_chain'].size}"
    )
    assert np.array_equal(wav1, wav2), "same seed must be bit-reproducible"
    assert not np.array_equal(wav1, wav3), "different seed must differ"
    assert np.isfinite(wav1).all()
    assert np.abs(wav1).max() <= 1.0
    print(f"[PASS] e2e synthesize_tokens: {wav1.size} samples, deterministic, bounded")


def test_frontend_symbol_mapping():
    """Symbol table mapping: literal IPA baseline string -> EXPORT baseline ids.

    export/model_meta.json numeric_baseline was built with
    cleaned_text_to_sequence("ðə kwɪk") — a literal IPA string, NOT eSpeak
    output (eSpeak with_stress would add stress marks)."""
    from inflect_ax_tts.symbols import symbols

    symbol_to_id = {s: i for i, s in enumerate(symbols)}
    ids = [symbol_to_id[ch] for ch in "ðə kwɪk"]
    assert ids == DUMMY_PHONEME_IDS, f"{ids} vs {DUMMY_PHONEME_IDS}"
    assert intersperse(ids, 0) == [0, 81, 0, 83, 0, 16, 0, 53, 0, 65, 0, 102, 0, 53, 0]
    print("[PASS] symbol mapping: 'ðə kwɪk' -> baseline phoneme ids")


def test_frontend_smoke():
    """eSpeak frontend structural check (skipped when espeak is unavailable)."""
    try:
        from inflect_ax_tts.frontend import text_to_token_ids
        ids = text_to_token_ids("the quick")
    except Exception as exc:  # espeak missing on host -> skip, not a failure
        print(f"[SKIP] frontend (espeak unavailable: {exc})")
        return
    assert len(ids) % 2 == 1, "interspersed sequence must be odd length"
    assert ids[::2] == [0] * (len(ids) // 2 + 1), "blanks must interleave"
    from inflect_ax_tts.symbols import symbols

    assert all(0 <= i < len(symbols) for i in ids)
    print(f"[PASS] frontend smoke: 'the quick' -> {len(ids)} interspersed ids")


ALL_TESTS = [
    test_expand_priors_matches_torch_reference,
    test_encoder_padding_semantics_ort,
    test_decoder_chunk_compose_ort,
    test_decoder_multi_chunk_compose_ort,
    test_synthesize_tokens_end_to_end_ort,
    test_frontend_symbol_mapping,
    test_frontend_smoke,
]


def main() -> int:
    failures = 0
    for test in ALL_TESTS:
        try:
            test()
        except Exception as exc:  # noqa: BLE001
            failures += 1
            print(f"[FAIL] {test.__name__}: {exc}")
    print(f"{'ALL PASS' if not failures else f'{failures} FAILURES'} ({len(ALL_TESTS)} tests)")
    return 1 if failures else 0


if __name__ == "__main__":
    sys.exit(main())