run method

Implementation

KittenFrontResult run(
  MlxArray inputIds,
  MlxArray refS, {
  double speed = 1.0,
}) {
  if (inputIds.shape[0] != 1 || refS.shape[0] != 1) {
    throw ArgumentError('KittenFrontRunner currently expects batch size 1.');
  }

  final textMask = MlxArray.fromBoolList(
    List<bool>.filled(inputIds.shape[1], false),
    shape: [1, inputIds.shape[1]],
  );
  final attentionMask = MlxArray.ones([
    1,
    inputIds.shape[1],
  ], dtype: MlxDType.MLX_INT32);

  final bertOutput = bert(inputIds, attentionMask: attentionMask);
  attentionMask.close();

  final dEn = mx.transposeAxes(bertEncoder(bertOutput.sequenceOutput), [
    0,
    2,
    1,
  ]);
  bertOutput.sequenceOutput.close();
  bertOutput.pooledOutput.close();

  final prosodyStyle = refS.slice(start: [0, 128], stop: [1, 256]);
  final decoderStyle = refS.slice(start: [0, 0], stop: [1, 128]);

  final d = predictor.textEncoder(dEn, prosodyStyle, textMask);
  dEn.close();
  final x = predictor.lstm(d).output;
  final durationLogits = predictor.durationProj(x);
  x.close();
  final speedScalar = scalar(speed);
  final duration = durationLogits.sigmoid().sum(axis: -1) / speedScalar;
  speedScalar.close();
  durationLogits.close();
  final rounded = duration.round().reshape([inputIds.shape[1]]);
  final predDurValues = rounded.toList().cast<double>();
  final predDurList = predDurValues
      .map((value) => value < 1.0 ? 1 : value.toInt())
      .toList(growable: false);
  final predDur = MlxArray.fromInt32List(
    predDurList,
    shape: [inputIds.shape[1]],
  );
  rounded.close();
  duration.close();

  final alignment = _buildAlignment(predDurList, inputIds.shape[1]);
  final transposedD = mx.transposeAxes(d, [0, 2, 1]);
  final en = mx.matmul(transposedD, alignment);
  transposedD.close();
  final prosody = predictor.f0Ntrain(en, prosodyStyle);
  en.close();

  final encodedText = textEncoder(inputIds, textMask);
  textMask.close();
  final asr = mx.matmul(encodedText, alignment);
  encodedText.close();
  alignment.close();
  d.close();
  prosodyStyle.close();
  return KittenFrontResult(
    asr: asr,
    f0Pred: prosody.f0,
    nPred: prosody.noise,
    style: decoderStyle,
    predDur: predDur,
  );
}