forwardProjection method

Implementation

KittenGeneratorProjection forwardProjection({
  required MlxArray asr,
  required MlxArray f0Curve,
  required MlxArray noise,
  required MlxArray style,
  required MlxArray har,
}) {
  final style2d =
      style.ndim == 1 ? style.reshape([1, style.shape[0]]) : style;
  final f0In = mx.transposeAxes(f0Curve.expandDims(1), [0, 2, 1]);
  final nIn = mx.transposeAxes(noise.expandDims(1), [0, 2, 1]);
  final f0Down0 = f0Conv(f0In);
  final f0Down = mx.transposeAxes(f0Down0, [0, 2, 1]);
  final nDown0 = nConv(nIn);
  final nDown = mx.transposeAxes(nDown0, [0, 2, 1]);
  f0In.close();
  nIn.close();
  f0Down0.close();
  nDown0.close();

  final merged0 = mx.concatenate([asr, f0Down, nDown], axis: 1);
  var hidden = encode(merged0, style2d);
  merged0.close();
  final asrIn = mx.transposeAxes(asr, [0, 2, 1]);
  final asrRes0 = asrRes(asrIn);
  final asrSkip = mx.transposeAxes(asrRes0, [0, 2, 1]);
  asrIn.close();
  asrRes0.close();
  var useResidual = true;
  for (final block in decode) {
    if (useResidual) {
      final cat = mx.concatenate([hidden, asrSkip, f0Down, nDown], axis: 1);
      hidden.close();
      hidden = cat;
    }
    final next = block(hidden, style2d);
    hidden.close();
    hidden = next;
    if (block.upsample) {
      useResidual = false;
    }
  }
  asrSkip.close();
  f0Down.close();
  nDown.close();
  if (style.ndim == 1) {
    style2d.close();
  }
  final projection = generator.forwardWithHar(hidden, style, har);
  hidden.close();
  return projection;
}