python/inflect_ax_tts/host_chain.py

"""Host-side chain between encoder.axmodel and decoder.axmodel (pure numpy).

Replicates, in numpy, the torch host chain of export/export_onnx.py:host_expand
(which itself mirrors origin/runtime/models.py:546-558):

    w      = exp(logw) * x_mask * length_scale
    w_ceil = ceil(w);  T' = max(sum(w_ceil), 1)
    attn   = generate_path(w_ceil, mask)          # pure integer logic
    m_p'   = attn @ m_p^T ;  logs_p' = attn @ logs_p^T
    z_p    = m_p' + randn * exp(logs_p') * variation   (seed-controlled, host)

Also contains the fixed-shape plumbing: encoder T=256 zero padding and decoder
Tp=512 chunking with overlap crossfade (EXPORT_NOTES §5.3 recommends
>=64-frame overlap between chunks; single-chunk T'<=512 is the path validated
by SIMULATE).
"""

from __future__ import annotations

import numpy as np

SAMPLE_RATE = 24000
HOP_LENGTH = 256
HIDDEN_CHANNELS = 192
ENCODER_T = 256
DECODER_TP = 512
DECODER_OVERLAP = 64  # frames of overlap between decoder chunks (EXPORT_NOTES §5.3)


# ---------------------------------------------------------------------------
# Duration rounding + alignment (generate_path)
# ---------------------------------------------------------------------------

def expand_priors(
    logw: np.ndarray,
    m_p: np.ndarray,
    logs_p: np.ndarray,
    x_len: int,
    length_scale: float,
) -> tuple[np.ndarray, np.ndarray, int]:
    """Duration rounding + generate_path + prior expansion.

    Args:
        logw:   [1, 1, T]  log-durations (encoder output, static T)
        m_p:    [1, C, T]  prior means
        logs_p: [1, C, T]  prior log-stds
        x_len:  number of valid (unpadded) token frames
        length_scale: 1.0 / speed

    Returns:
        (m_p_e [T', C], logs_p_e [T', C], T')
    """
    t_total = logw.shape[-1]
    if not 1 <= x_len <= t_total:
        raise ValueError(f"x_len {x_len} out of range [1, {t_total}]")
    x_mask = (np.arange(t_total) < x_len).astype(np.float32)
    w = np.exp(logw[0, 0].astype(np.float32)) * x_mask * np.float32(length_scale)
    w_ceil = np.ceil(w).astype(np.int64)
    y_len = max(int(w_ceil.sum()), 1)

    # generate_path (origin/runtime/commons.py): with the masked durations the
    # alignment reduces to pure integer interval logic — frame t of the output
    # copies input frame i for cum[i-1] <= t < cum[i].  O(T) loop, no matmul
    # of the [T', T] path needed; we still build attn explicitly for the prior
    # expansion to stay a literal port (T=256, T'<=~2000: cheap).
    cum = np.cumsum(w_ceil)
    attn = np.zeros((y_len, t_total), dtype=np.float32)
    lo = 0
    for i in range(t_total):
        hi = min(int(cum[i]), y_len)
        if hi > lo:
            attn[lo:hi, i] = 1.0
        lo = hi
    m_p_e = attn @ m_p[0].T.astype(np.float32)
    logs_p_e = attn @ logs_p[0].T.astype(np.float32)
    return m_p_e, logs_p_e, y_len


def inject_noise(
    m_p_e: np.ndarray,
    logs_p_e: np.ndarray,
    variation: float,
    seed: int,
) -> np.ndarray:
    """z_p = m_p' + randn * exp(logs_p') * variation, seed-controlled.

    NOTE: numpy PCG64 is used, not torch's MT19937 — a given `seed` is
    deterministic and reproducible within this SDK but is NOT bit-identical to
    the PyTorch reference (origin/inference.py).
    """
    rng = np.random.default_rng(seed)
    noise = rng.standard_normal(m_p_e.shape, dtype=np.float32)
    z_p = m_p_e + noise * np.exp(logs_p_e.astype(np.float32)) * np.float32(variation)
    return z_p.T.copy()  # [C, T']


# ---------------------------------------------------------------------------
# Decoder chunk composition (Tp=512 static shape, overlap crossfade)
# ---------------------------------------------------------------------------

def decoder_chunk_starts(t_prime: int, chunk: int = DECODER_TP,
                         overlap: int = DECODER_OVERLAP) -> list[int]:
    """Frame offsets of decoder chunks covering [0, t_prime) with >=overlap."""
    if t_prime <= chunk:
        return [0]
    stride = chunk - overlap
    starts = list(range(0, t_prime - chunk + 1, stride))
    last = t_prime - chunk
    if starts[-1] < last:
        starts.append(last)
    return starts


def decode_waveform(
    z_p: np.ndarray,
    run_decoder,
    chunk: int = DECODER_TP,
    overlap: int = DECODER_OVERLAP,
) -> np.ndarray:
    """Run the fixed-shape decoder over z_p and stitch chunks.

    Args:
        z_p: [C, T'] float32 prior with noise injected.
        run_decoder: callable taking z_p_chunk [C, chunk] float32 (already
            right-padded with zeros) and returning wav [chunk*HOP] float32.
        chunk/overlap: see DECODER_TP / DECODER_OVERLAP.

    Returns:
        wav [T'*HOP] float32 (tail trimmed; chunk overlaps crossfaded).
    """
    t_prime = z_p.shape[1]
    hop = HOP_LENGTH
    starts = decoder_chunk_starts(t_prime, chunk, overlap)

    out: np.ndarray | None = None
    prev_end_frame = 0
    for start in starts:
        take = min(chunk, t_prime - start)
        z_chunk = np.zeros((z_p.shape[0], chunk), dtype=np.float32)
        z_chunk[:, :take] = z_p[:, start:start + take]
        wav = np.asarray(run_decoder(z_chunk), dtype=np.float32)
        if out is None:
            out = wav
        else:
            ov_frames = prev_end_frame - start
            ov = ov_frames * hop
            fade_in = np.linspace(0.0, 1.0, ov, endpoint=True, dtype=np.float32)
            out = np.concatenate(
                [out[:-ov], out[-ov:] * (1.0 - fade_in) + wav[:ov] * fade_in, wav[ov:]]
            )
        prev_end_frame = start + chunk
    assert out is not None
    return out[: t_prime * hop]


# ---------------------------------------------------------------------------
# Waveform post-processing (origin/inference.py)
# ---------------------------------------------------------------------------

def edge_fade(waveform: np.ndarray, sample_rate: int = SAMPLE_RATE,
              milliseconds: float = 5.0) -> np.ndarray:
    frames = min(round(sample_rate * milliseconds / 1000.0), waveform.size // 2)
    if frames <= 0:
        return waveform
    output = waveform.copy()
    ramp = np.linspace(0.0, 1.0, frames, endpoint=True, dtype=np.float32)
    output[:frames] *= ramp
    output[-frames:] *= ramp[::-1]
    return output


# ---------------------------------------------------------------------------
# Text chunking (origin/inference.py: split_text / boundary_pause_seconds)
# ---------------------------------------------------------------------------

def split_text(text: str, limit: int = 280) -> list[str]:
    import re

    normalized = " ".join(text.split())
    sentences = [
        part.strip()
        for part in re.split(r"(?<=[.!?;:])\s+", normalized)
        if part.strip()
    ]
    chunks: list[str] = []
    for sentence in sentences or [normalized]:
        while len(sentence) > limit:
            search = sentence[: limit + 1]
            punctuation = max(search.rfind(mark) for mark in (",", ";", ":"))
            split_at = (
                punctuation + 1
                if punctuation >= limit // 2
                else sentence.rfind(" ", 0, limit + 1)
            )
            if split_at < limit // 2:
                split_at = limit
            chunks.append(sentence[:split_at].strip())
            sentence = sentence[split_at:].strip()
        if sentence:
            chunks.append(sentence)
    return chunks


def boundary_pause_seconds(chunk: str) -> float:
    ending = chunk.rstrip()[-1:] if chunk.strip() else ""
    return {
        "?": 0.28,
        "!": 0.24,
        ".": 0.22,
        ";": 0.16,
        ":": 0.13,
        ",": 0.09,
    }.get(ending, 0.08)


def seconds_to_samples(seconds: float) -> int:
    return round(seconds * SAMPLE_RATE)