cpp/src/host_chain.h

#pragma once

// Host-side chain between encoder.axmodel and decoder.axmodel, pure C++17
// (no Eigen, no torch). Mirrors export/export_onnx.py:host_expand, which
// replicates 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 interval logic
//   m_p'   = attn @ m_p^T ; logs_p' = attn @ logs_p^T
//
// generate_path produces a 0/1 alignment with exactly one 1 per output frame,
// so the "matmul" is implemented as an exact per-frame gather (same integer
// cumsum intervals, provably identical result).

#include <cstdint>
#include <vector>

namespace inflect {

constexpr int kSampleRate = 24000;
constexpr int kHopLength = 256;
constexpr int kHiddenChannels = 192;
constexpr int kEncoderT = 256;     // encoder static token length (0-padded)
constexpr int kDecoderTp = 512;    // decoder static frame length per chunk
constexpr int kDecoderOverlap = 64;  // frames of overlap between decoder chunks

struct ExpandedPriors {
    int t_prime = 0;           // T'
    // channel-major [kHiddenChannels * t_prime]: element (c, t) at c*t_prime+t
    std::vector<float> m_p;
    std::vector<float> logs_p;
};

// logw:   [t_total]        (encoder output logw[0,0,:])
// m_p:    [192 * t_total]  channel-major (encoder output m_p[0])
// logs_p: [192 * t_total]  channel-major
// x_len:  number of valid (unpadded) token frames
ExpandedPriors expand_priors(const float* logw, const float* m_p,
                             const float* logs_p, int t_total, int x_len,
                             float length_scale);

// z_p = m_p' + randn * exp(logs_p') * variation, channel-major [192 * T'].
// NOTE: std::mt19937_64 + normal_distribution — deterministic within this SDK
// for a given seed, NOT bit-identical to the PyTorch reference.
std::vector<float> inject_noise(const ExpandedPriors& priors, float variation,
                                uint64_t seed);

// Frame offsets of Tp-sized decoder chunks covering [0, t_prime) with
// >= kDecoderOverlap overlap (EXPORT_NOTES §5.3).
std::vector<int> decoder_chunk_starts(int t_prime);

// Crossfade-append `wav` (samples of chunk starting at `start_frame`) onto
// `out`, overlapping the previous chunk tail by ov_frames frames.
void crossfade_append(std::vector<float>& out, const std::vector<float>& wav,
                      int ov_frames);

// 5 ms raised-linear edge fade, as origin/inference.py.
void edge_fade(std::vector<float>& waveform, float milliseconds = 5.0f);

void clip_inplace(std::vector<float>& waveform, float lo = -1.0f, float hi = 1.0f);

}  // namespace inflect