Blame for scripts/word_stream_piper.cpp

#include "word_stream_piper.hpp"
#include <algorithm>
#include <chrono>
#include <cstdio>
#include <cstring>
#include <random>
#include <stdexcept>
using namespace MNN;
using namespace MNN::Express;

static const int MASK_SCALES[4] = {1, 8, 64, 256};
static const int NDP = 4096, NZ = 4096;  // noise table lengths (positions wrap)

static double msSince(std::chrono::steady_clock::time_point t) {
  return std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - t).count();
}

WordStreamPiper::WordStreamPiper(const std::string& encoderPath, const std::string& decoderPath, int backend, int threads) {
  gpu_ = backend == 1;
  auto type = gpu_ ? MNN_FORWARD_OPENCL : MNN_FORWARD_CPU;
  // OpenCL reads numThread as mode flags, not a thread count.
  int mode = gpu_ ? (MNN_GPU_MEMORY_BUFFER | MNN_GPU_TUNING_NONE) : threads;
  BackendConfig bc;
  exec_ = Executor::newExecutor(type, bc, mode);
  if (!exec_) throw std::runtime_error(gpu_ ? "MNN OpenCL executor unavailable" : "MNN CPU executor unavailable");
  ExecutorScope scope(exec_);
  ScheduleConfig sc; sc.type = type; sc.backupType = MNN_FORWARD_CPU; sc.numThread = mode;
  BackendConfig rbc; sc.backendConfig = &rbc;
  rt_.reset(Executor::RuntimeManager::createRuntimeManager(sc), Executor::RuntimeManager::destroy);
  if (!rt_) throw std::runtime_error("MNN runtime creation failed");
  if (gpu_) rt_->setCache(decoderPath + ".opencl-cache");
  Module::Config cfg; cfg.shapeMutable = false;  // both graphs have fixed shapes
  enc_.reset(Module::load({"input", "input_lengths", "scales", "noise_dp", "noise_z"}, {"z", "durations", "frames"},
                          encoderPath.c_str(), rt_, &cfg));
  if (!enc_) throw std::runtime_error("Encoder load failed");
  dec_.reset(Module::load({"z", "valid_mask_1", "valid_mask_8", "valid_mask_64", "valid_mask_256"}, {"output"},
                          decoderPath.c_str(), rt_, &cfg));
  if (!dec_) throw std::runtime_error("Decoder load failed");
  inIds_ = _Input({1, N_IDS}, NCHW, halide_type_of<int>());
  inLen_ = _Input({1}, NCHW, halide_type_of<int>());
  inScales_ = _Input({3}, NCHW, halide_type_of<float>());
  inNoiseDp_ = _Input({1, 2, N_IDS}, NCHW, halide_type_of<float>());
  inNoiseZ_ = _Input({1, C, N_FRAMES}, NCHW, halide_type_of<float>());
  decZ_ = _Input({1, C, DEC_FRAMES}, NCHW, halide_type_of<float>());
  for (int i = 0; i < 4; ++i) decMask_[i] = _Input({1, 1, DEC_FRAMES * MASK_SCALES[i]}, NCHW, halide_type_of<float>());
  reset(0.667f, 1.f, 0.8f, 2, 1, 0, 1234);
}

void WordStreamPiper::reset(float noise, float length, float noisew, int ctxWords, int lookahead, int xfade, uint32_t seed,

                            int firstWords, int laterWords) {
  firstWords_ = std::max(1, firstWords); laterWords_ = std::max(1, laterWords);
  scales_[0] = noise; scales_[1] = length; scales_[2] = noisew;
  ctxWords_ = std::max(0, ctxWords); look_ = std::max(0, lookahead); xfade_ = std::max(0, xfade);
  finished_ = false; ids_.clear(); wordStarts_.clear(); committed_ = 0; committedFrames_ = 0; tail_.clear();
  std::mt19937 rng(seed); std::normal_distribution<float> nd;
  noiseDp_.resize(2 * NDP); for (auto& v : noiseDp_) v = nd(rng);
  noiseZ_.resize((size_t)C * NZ); for (auto& v : noiseZ_) v = nd(rng);
}

void WordStreamPiper::addWord(const std::vector<int>& ids) {
  if (finished_) throw std::runtime_error("Stream already finished");
  wordStarts_.push_back((int)ids_.size());
  ids_.insert(ids_.end(), ids.begin(), ids.end());
}

void WordStreamPiper::finish(const std::vector<int>& eosIds) {
  if (!finished_) { ids_.insert(ids_.end(), eosIds.begin(), eosIds.end()); finished_ = true; }
}

void WordStreamPiper::encodeWindow(int ctx, int n) {
  int* x = inIds_->writeMap<int>();
  std::fill(x, x + N_IDS, 0);
  std::copy(ids_.begin() + ctx, ids_.begin() + ctx + n, x);
  inLen_->writeMap<int>()[0] = n;
  std::copy(scales_, scales_ + 3, inScales_->writeMap<float>());
  float* dp = inNoiseDp_->writeMap<float>();
  for (int c = 0; c < 2; ++c)
    for (int i = 0; i < N_IDS; ++i) dp[c * N_IDS + i] = noiseDp_[c * NDP + (ctx + i) % NDP];
  // Window frame 0 is roughly the global frame where the context starts.
  long g0 = std::max(0L, committedFrames_ - 2L * DEC_CTX * ctxWords_);
  float* nz = inNoiseZ_->writeMap<float>();
  for (int c = 0; c < C; ++c)
    for (int f = 0; f < N_FRAMES; ++f) nz[c * N_FRAMES + f] = noiseZ_[(size_t)c * NZ + (g0 + f) % NZ];
  auto out = enc_->onForward({inIds_, inLen_, inScales_, inNoiseDp_, inNoiseZ_});
  if (out.size() < 3 || out[0] == nullptr) throw std::runtime_error("Encoder returned nothing");
  const float* z = out[0]->readMap<float>();
  const float* d = out[1]->readMap<float>();
  const float* f = out[2]->readMap<float>();
  if (!z || !d || !f) throw std::runtime_error("Encoder output read failed");
  z_.assign(z, z + (size_t)C * N_FRAMES);
  dur_.assign(d, d + N_IDS);
  validFrames_ = std::min(N_FRAMES, (int)f[0]);
}

std::vector<float> WordStreamPiper::decodeRange(int start, int end, int limit) {
  std::vector<float> pcm;
  pcm.reserve((size_t)(end - start) * HOP);
  for (int s = start; s < end; s += DEC_EMIT) {
    int e = std::min(end, s + DEC_EMIT);
    int a = std::max(0, s - DEC_CTX), b = std::min(limit, e + DEC_CTX);
    int dst = DEC_CTX - (s - a), len = b - a;
    float* w = decZ_->writeMap<float>();
    std::memset(w, 0, sizeof(float) * C * DEC_FRAMES);
    for (int c = 0; c < C; ++c) std::memcpy(w + c * DEC_FRAMES + dst, z_.data() + (size_t)c * N_FRAMES + a, len * sizeof(float));
    for (int i = 0; i < 4; ++i) {
      int sc = MASK_SCALES[i];
      float* m = decMask_[i]->writeMap<float>();
      std::memset(m, 0, sizeof(float) * DEC_FRAMES * sc);
      std::fill(m + dst * sc, m + (dst + len) * sc, 1.f);
    }
    auto out = dec_->onForward({decZ_, decMask_[0], decMask_[1], decMask_[2], decMask_[3]});
    if (out.empty() || out[0] == nullptr) throw std::runtime_error("Decoder returned nothing");
    const float* y = out[0]->readMap<float>();
    if (!y) throw std::runtime_error("Decoder output read failed");
    pcm.insert(pcm.end(), y + DEC_CTX * HOP, y + DEC_CTX * HOP + (size_t)(e - s) * HOP);
  }
  return pcm;
}

std::vector<float> WordStreamPiper::step() {
  ExecutorScope scope(exec_);
  info_.clear();
  const int total = (int)ids_.size();
  if (committed_ >= total) return {};

  // Left context: back ctxWords word starts, but keep it within a third of the window.
  int ctx = committed_;
  {
    int count = 0;
    for (int i = (int)wordStarts_.size() - 1; i >= 0 && count < ctxWords_; --i)
      if (wordStarts_[i] < committed_) { ctx = wordStarts_[i]; ++count; }
    ctx = std::max(ctx, committed_ - N_IDS / 3);
  }
  const int n = std::min(N_IDS, total - ctx);
  const bool windowReachesEnd = ctx + n == total;

  // Candidate commit points: word starts after the commit point with enough look-ahead.
  const bool windowFull = total - ctx >= N_IDS;
  auto candidates = [&](int minWords) {
    std::vector<int> v;
    for (int b : wordStarts_) {
      if (b <= committed_ || b >= ctx + n) continue;
      int after = 0, words = 0;
      for (int s : wordStarts_) { if (s > b) ++after; if (s >= committed_ && s < b) ++words; }
      if (words >= minWords && after >= look_) v.push_back(b);
    }
    return v;
  };
  // Wait until enough words can be committed together (fewer, larger windows),
  // unless the window is already full or the text has ended.
  std::vector<int> ok = candidates(finished_ ? 1 : (committed_ == 0 ? firstWords_ : laterWords_));
  if (ok.empty() && (windowFull || finished_)) ok = candidates(1);
  if (ok.empty() && !finished_ && !windowFull) return {};

  auto t0 = std::chrono::steady_clock::now();
  encodeWindow(ctx, n);
  double encMs = msSince(t0);

  std::vector<int> cum(n + 1, 0);
  for (int i = 0; i < n; ++i) cum[i + 1] = cum[i] + (int)dur_[i];
  const int limit = validFrames_;
  const int f0 = std::min(cum[committed_ - ctx], limit);

  // Choose the furthest target whose audio still has decoder right context in this window.
  int target = -1, f1 = 0;
  if (finished_ && windowReachesEnd && cum[n] <= N_FRAMES) {
    target = total; f1 = limit;
  } else {
    for (int i = (int)ok.size() - 1; i >= 0; --i) {
      int fe = cum[ok[i] - ctx];
      if (fe <= limit - DEC_CTX) { target = ok[i]; f1 = fe; break; }
    }
    if (target < 0 && (finished_ || (int)ids_.size() - ctx >= N_IDS)) {
      // A single stretch too long for one window: cut at the last ID that fits.
      for (int i = n - 1; i > committed_ - ctx; --i)
        if (cum[i] <= limit - DEC_CTX) { target = ctx + i; f1 = cum[i]; break; }
    }
  }
  if (target <= committed_ || f1 <= f0) return {};

  auto t1 = std::chrono::steady_clock::now();
  int ext = std::min(f1 + xfade_, limit);
  std::vector<float> out = decodeRange(f0, ext, limit);
  double decMs = msSince(t1);
  size_t body = (size_t)(f1 - f0) * HOP;
  std::vector<float> nextTail(out.begin() + std::min(body, out.size()), out.end());
  out.resize(std::min(body, out.size()));
  if (!tail_.empty()) {
    size_t k = std::min(tail_.size(), out.size());
    for (size_t i = 0; i < k; ++i) { float r = (float)i / (float)k; out[i] = tail_[i] * (1 - r) + out[i] * r; }
  }
  tail_ = std::move(nextTail);
  if (gpu_) rt_->updateCache();

  char buf[200];
  std::snprintf(buf, sizeof buf, "ids[%d:%d] commit %d->%d frames %d->%d of %d; enc %.1f ms dec %.1f ms; %.2f s",
                ctx, ctx + n, committed_, target, f0, f1, limit, encMs, decMs, out.size() / 22050.0);
  info_ = buf;
  committed_ = target;
  committedFrames_ += f1 - f0;
  return out;
}