Kotlin code to play the audio generated by the Kokoro model using LiteRT

KokoroAudio.kt

import android.media.AudioAttributes
import android.media.AudioFormat
import android.media.AudioTrack
import java.nio.ByteBuffer
import java.nio.ByteOrder

private const val SAMPLE_RATE = 22050
private const val PCM_WRITE_BUFFER_SIZE = 4096

fun playKokoroAudio(audioData: FloatArray) {
    val pcmBytes = audioData.toPcm16Bytes()
    val track = createAudioTrack()

    try {
        track.play()

        var offset = 0
        while (offset < pcmBytes.size) {
            val written = track.write(
                pcmBytes,
                offset,
                pcmBytes.size - offset
            )

            if (written <= 0) break
            offset += written
        }
    } finally {
        track.release()
    }
}

private fun FloatArray.toPcm16Bytes(): ByteArray {
    val byteBuffer = ByteBuffer
        .allocate(size * 2)
        .order(ByteOrder.LITTLE_ENDIAN)

    for (sample in this) {
        val clamped = sample.coerceIn(-1f, 1f)
        val pcmValue = (clamped * Short.MAX_VALUE).toInt().toShort()
        byteBuffer.putShort(pcmValue)
    }

    return byteBuffer.array()
}

private fun createAudioTrack(): AudioTrack {
    val minBufferSize = AudioTrack.getMinBufferSize(
        SAMPLE_RATE,
        AudioFormat.CHANNEL_OUT_MONO,
        AudioFormat.ENCODING_PCM_16BIT
    )

    val audioFormat = AudioFormat.Builder()
        .setEncoding(AudioFormat.ENCODING_PCM_16BIT)
        .setSampleRate(SAMPLE_RATE)
        .setChannelMask(AudioFormat.CHANNEL_OUT_MONO)
        .build()

    return AudioTrack.Builder()
        .setAudioAttributes(
            AudioAttributes.Builder()
                .setUsage(AudioAttributes.USAGE_MEDIA)
                .setContentType(AudioAttributes.CONTENT_TYPE_SPEECH)
                .build()
        )
        .setAudioFormat(audioFormat)
        .setBufferSizeInBytes(
            minBufferSize.coerceAtLeast(PCM_WRITE_BUFFER_SIZE)
        )
        .setTransferMode(AudioTrack.MODE_STREAM)
        .build()
}