Skip to content

feat(tts): add a Core ML Kokoro variant #1384

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
msluszniak wants to merge 2 commits into
base: rne-rewrite
Choose a base branch
from
+178 −14
Open

feat(tts): add a Core ML Kokoro variant #1384

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
fa26094
feat(tts): accept a padded-token Kokoro contract
msluszniak Aug 25, 2026
6c8c654
feat(tts): add the Polish and German Core ML Kokoro variants
msluszniak Aug 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 29 additions & 2 deletions apps/speech/app/kokoro-text-to-speech/index.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import React, { useCallback, useEffect, useRef, useState } from 'react';
import { View, Text, StyleSheet, ScrollView, TextInput } from 'react-native';
import { View, Text, StyleSheet, ScrollView, TextInput, Platform } from 'react-native';
import {
useTextToSpeech,
models,
Expand Down Expand Up @@ -28,6 +28,20 @@ const LANGUAGE_OPTIONS = [

type KokoroLanguage = (typeof LANGUAGE_OPTIONS)[number]['value'];

// Core ML is exported for the standard weights only, and only runs on iOS.
const BACKEND_OPTIONS = [
{ label: 'XNNPACK', value: 'XNNPACK_FP32' as const },
{ label: 'Core ML', value: 'COREML_FP32' as const },
];

type KokoroBackend = (typeof BACKEND_OPTIONS)[number]['value'];

const kokoroModel = (language: KokoroLanguage, backend: KokoroBackend) => {
const entry = models.textToSpeech.KOKORO[language];
const model = backend in entry ? entry[backend as keyof typeof entry] : entry.XNNPACK_FP32;
return model as KokoroTtsModel<string>;
};

// cspell:disable
const SAMPLE_TEXTS: Record<KokoroLanguage, string> = {
EN_US:
Expand Down Expand Up @@ -65,7 +79,9 @@ function KokoroContent() {
const [runError, setRunError] = useState<string | null>(null);
const [totalDuration, setTotalDuration] = useState<number | null>(null);

const model = models.textToSpeech.KOKORO[language].XNNPACK_FP32 as KokoroTtsModel<string>;
const [backend, setBackend] = useState<KokoroBackend>('XNNPACK_FP32');

const model = kokoroModel(language, backend);
const voiceNames = Object.keys(model.voices);
const [voice, setVoice] = useState(voiceNames[0]!);

Expand All @@ -77,6 +93,9 @@ function KokoroContent() {
useEffect(() => {
setVoice(Object.keys(models.textToSpeech.KOKORO[language].XNNPACK_FP32.voices)[0]!);
setText(SAMPLE_TEXTS[language]);
if (!('COREML_FP32' in models.textToSpeech.KOKORO[language])) {
setBackend('XNNPACK_FP32');
}
}, [language]);

const getAudioContext = useCallback(async () => {
Expand Down Expand Up @@ -198,6 +217,14 @@ function KokoroContent() {
selectedValue={language}
onValueChange={setLanguage}
/>
{Platform.OS === 'ios' && 'COREML_FP32' in models.textToSpeech.KOKORO[language] && (
<ModelPicker
label="Backend"
options={BACKEND_OPTIONS.map((b) => ({ ...b, disabled: isBusy }))}
selectedValue={backend}
onValueChange={setBackend}
/>
)}
<ModelPicker
label="Voice"
options={voiceNames.map((name) => ({ label: name, value: name, disabled: isBusy }))}
Expand Down
73 changes: 67 additions & 6 deletions packages/react-native-executorch/src/extensions/speech/tasks/kokoroTextToSpeech.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -147,7 +147,7 @@ export async function createKokoroTextToSpeech<K extends PropertyKey>(

try {
const predictorSpec = validateSpec(models.durationPredictor.schema, {
default: method(
dynamic: method(
'forward',
[
i64(1, Dyn('T')), // tokens
Expand All @@ -168,10 +168,25 @@ export async function createKokoroTextToSpeech<K extends PropertyKey>(
),
]
),
// Core ML cannot express the token axis dynamically, so its export fixes
// it and every chunk is padded up to that length.
padded: method(
'forward',
[
i64(1, 'T'), // tokens
bool(1, 'T'), // textMask
f32(1, VOICE_REF_HALF_SIZE), // voiceRef
f32(1), // speed
],
[
i64('T'), // predictedDurations
f32(1, 'T', DURATION_FEATURE_DIM), // durationFeatures
]
),
});

const synthesizerSpec = validateSpec(models.synthesizer.schema, {
default: method(
dynamic: method(
'forward',
[
i64(1, Dyn('T')), // tokens
Expand All @@ -194,15 +209,61 @@ export async function createKokoroTextToSpeech<K extends PropertyKey>(
),
]
),
padded: method(
'forward',
[
i64(1, 'T'), // tokens
bool(1, 'T'), // textMask
i64(Dyn('D')), // indices
f32(1, 'T', DURATION_FEATURE_DIM), // durationFeatures
f32(1, VOICE_REF_SIZE), // voiceRef
],
[f32(1, 1, Dyn('AUDIO_LEN'))], // audio
[
constr.linear(
{ paramSide: 'output', tensorIdx: 0, dimIdx: 2 },
{ paramSide: 'input', tensorIdx: 2, dimIdx: 0 },
TICKS_PER_DURATION
),
]
),
});

const [predictorTokens] = predictorSpec.dims.range('T');
const [synthesizerTokens, durations] = synthesizerSpec.dims.range('T', 'D');
if (predictorSpec.variant !== synthesizerSpec.variant) {
throw RnExecuTorchError(
'SCHEMA_MISMATCH',
`The duration predictor and the synthesizer declare different token axes ` +
`('${String(predictorSpec.variant)}' vs '${String(synthesizerSpec.variant)}').`
);
}

const minTokens = Math.max(predictorTokens.min, synthesizerTokens.min);
const maxTokens = predictorTokens.max;
const [durations] = synthesizerSpec.dims.range('D');
const maxDurationTicks = durations.max;

let minTokens: number;
let maxTokens: number;

if (predictorSpec.variant === 'padded') {
// A padded model takes one token count and one only, so both bounds
// collapse onto it and every chunk is padded up to it.
const predictorTokens = predictorSpec.dim('T', 'constant');
const synthesizerTokens = synthesizerSpec.dim('T', 'constant');
if (predictorTokens !== synthesizerTokens) {
throw RnExecuTorchError(
'SCHEMA_MISMATCH',
`The duration predictor and the synthesizer are padded to different token ` +
`counts (${predictorTokens} and ${synthesizerTokens}).`
);
}
minTokens = predictorTokens;
maxTokens = predictorTokens;
} else {
const [predictorTokens] = predictorSpec.dims.range('T');
const [synthesizerTokens] = synthesizerSpec.dims.range('T');
minTokens = Math.max(predictorTokens.min, synthesizerTokens.min);
maxTokens = predictorTokens.max;
}

const phonemizer = await wrapAsync(createPhonemizer, runtime)(config.phonemizer);
allocated.push(phonemizer);

Expand Down
88 changes: 82 additions & 6 deletions packages/react-native-executorch/src/models.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -880,14 +880,23 @@ const SUPERTONIC_3_MLX_FP32: SupertonicTtsModel<SupertonicDefaultVoiceName> = {
const KOKORO_ROOT = `${BASE_URL}-kokoro/${NEXT_VERSION_TAG}`;
const KOKORO_PHONEMIZER_ROOT = `${KOKORO_ROOT}/phonemizer`;

const kokoroModelPaths = (variant: 'std' | 'pl' | 'de', dir: string) => ({
durationPredictor: `${KOKORO_ROOT}/xnnpack/${dir}/duration_predictor_${variant}_xnnpack_fp32.pte`,
synthesizer: `${KOKORO_ROOT}/xnnpack/${dir}/synthesizer_${variant}_xnnpack_fp32.pte`,
const kokoroModelPaths = (
backend: 'xnnpack' | 'coreml',
variant: 'std' | 'pl' | 'de',
dir: string
) => ({
durationPredictor: `${KOKORO_ROOT}/${backend}/${dir}/duration_predictor_${variant}_${backend}_fp32.pte`,
synthesizer: `${KOKORO_ROOT}/${backend}/${dir}/synthesizer_${variant}_${backend}_fp32.pte`,
});

const KOKORO_STANDARD_PATHS = kokoroModelPaths('std', 'standard');
const KOKORO_POLISH_PATHS = kokoroModelPaths('pl', 'polish');
const KOKORO_GERMAN_PATHS = kokoroModelPaths('de', 'german');
const KOKORO_STANDARD_PATHS = kokoroModelPaths('xnnpack', 'std', 'standard');
const KOKORO_POLISH_PATHS = kokoroModelPaths('xnnpack', 'pl', 'polish');
const KOKORO_GERMAN_PATHS = kokoroModelPaths('xnnpack', 'de', 'german');

// Core ML is exported for the standard (multi-language) weights only.
const KOKORO_STANDARD_COREML_PATHS = kokoroModelPaths('coreml', 'std', 'standard');
const KOKORO_POLISH_COREML_PATHS = kokoroModelPaths('coreml', 'pl', 'polish');
const KOKORO_GERMAN_COREML_PATHS = kokoroModelPaths('coreml', 'de', 'german');

const kokoroVoices = <const N extends string>(names: readonly N[]) =>
names.reduce(
Expand Down Expand Up @@ -955,6 +964,55 @@ const KOKORO_HI_XNNPACK_FP32: KokoroTtsModel<'hf_alpha' | 'hm_omega' | 'hm_psi'>
phonemizer: kokoroNeuralPhonemizer('hi'),
voices: kokoroVoices(['hf_alpha', 'hm_omega', 'hm_psi']),
};

// Core ML counterparts of the standard-weight presets. They pad the token axis
// to the models' fixed length, so only the languages served by the standard
// weights have a Core ML variant.
const KOKORO_EN_US_COREML_FP32: KokoroTtsModel<
'af_heart' | 'af_river' | 'af_sarah' | 'am_adam' | 'am_michael' | 'am_santa'
> = {
name: 'kokoro',
modelPaths: KOKORO_STANDARD_COREML_PATHS,
phonemizer: kokoroEnglishPhonemizer('en-us'),
voices: kokoroVoices(['af_heart', 'af_river', 'af_sarah', 'am_adam', 'am_michael', 'am_santa']),
};
const KOKORO_EN_GB_COREML_FP32: KokoroTtsModel<'bf_emma' | 'bm_daniel'> = {
name: 'kokoro',
modelPaths: KOKORO_STANDARD_COREML_PATHS,
phonemizer: kokoroEnglishPhonemizer('en-gb'),
voices: kokoroVoices(['bf_emma', 'bm_daniel']),
};
const KOKORO_ES_COREML_FP32: KokoroTtsModel<'ef_dora' | 'em_alex'> = {
name: 'kokoro',
modelPaths: KOKORO_STANDARD_COREML_PATHS,
phonemizer: kokoroNeuralPhonemizer('es'),
voices: kokoroVoices(['ef_dora', 'em_alex']),
};
const KOKORO_FR_COREML_FP32: KokoroTtsModel<'ff_siwis'> = {
name: 'kokoro',
modelPaths: KOKORO_STANDARD_COREML_PATHS,
phonemizer: kokoroNeuralPhonemizer('fr'),
voices: kokoroVoices(['ff_siwis']),
};
const KOKORO_IT_COREML_FP32: KokoroTtsModel<'if_sara' | 'im_nicola'> = {
name: 'kokoro',
modelPaths: KOKORO_STANDARD_COREML_PATHS,
phonemizer: kokoroNeuralPhonemizer('it'),
voices: kokoroVoices(['if_sara', 'im_nicola']),
};
const KOKORO_PT_COREML_FP32: KokoroTtsModel<'pf_dora' | 'pm_santa'> = {
name: 'kokoro',
modelPaths: KOKORO_STANDARD_COREML_PATHS,
phonemizer: kokoroNeuralPhonemizer('pt'),
voices: kokoroVoices(['pf_dora', 'pm_santa']),
};
const KOKORO_HI_COREML_FP32: KokoroTtsModel<'hf_alpha' | 'hm_omega' | 'hm_psi'> = {
name: 'kokoro',
modelPaths: KOKORO_STANDARD_COREML_PATHS,
phonemizer: kokoroNeuralPhonemizer('hi'),
voices: kokoroVoices(['hf_alpha', 'hm_omega', 'hm_psi']),
};

const KOKORO_PL_XNNPACK_FP32: KokoroTtsModel<'pm_mateusz'> = {
name: 'kokoro',
modelPaths: KOKORO_POLISH_PATHS,
Expand All @@ -968,6 +1026,15 @@ const KOKORO_DE_XNNPACK_FP32: KokoroTtsModel<'df_anna'> = {
voices: kokoroVoices(['df_anna']),
};

const KOKORO_PL_COREML_FP32: KokoroTtsModel<'pm_mateusz'> = {
...KOKORO_PL_XNNPACK_FP32,
modelPaths: KOKORO_POLISH_COREML_PATHS,
};
const KOKORO_DE_COREML_FP32: KokoroTtsModel<'df_anna'> = {
...KOKORO_DE_XNNPACK_FP32,
modelPaths: KOKORO_GERMAN_COREML_PATHS,
};

// =============================================================================
// Privacy Filter
// =============================================================================
Expand Down Expand Up @@ -2197,38 +2264,47 @@ export const models = {
EN_US: {
...KOKORO_EN_US_XNNPACK_FP32,
XNNPACK_FP32: KOKORO_EN_US_XNNPACK_FP32,
COREML_FP32: KOKORO_EN_US_COREML_FP32,
},
EN_GB: {
...KOKORO_EN_GB_XNNPACK_FP32,
XNNPACK_FP32: KOKORO_EN_GB_XNNPACK_FP32,
COREML_FP32: KOKORO_EN_GB_COREML_FP32,
},
ES: {
...KOKORO_ES_XNNPACK_FP32,
XNNPACK_FP32: KOKORO_ES_XNNPACK_FP32,
COREML_FP32: KOKORO_ES_COREML_FP32,
},
FR: {
...KOKORO_FR_XNNPACK_FP32,
XNNPACK_FP32: KOKORO_FR_XNNPACK_FP32,
COREML_FP32: KOKORO_FR_COREML_FP32,
},
IT: {
...KOKORO_IT_XNNPACK_FP32,
XNNPACK_FP32: KOKORO_IT_XNNPACK_FP32,
COREML_FP32: KOKORO_IT_COREML_FP32,
},
PT: {
...KOKORO_PT_XNNPACK_FP32,
XNNPACK_FP32: KOKORO_PT_XNNPACK_FP32,
COREML_FP32: KOKORO_PT_COREML_FP32,
},
HI: {
...KOKORO_HI_XNNPACK_FP32,
XNNPACK_FP32: KOKORO_HI_XNNPACK_FP32,
COREML_FP32: KOKORO_HI_COREML_FP32,
},
PL: {
...KOKORO_PL_XNNPACK_FP32,
XNNPACK_FP32: KOKORO_PL_XNNPACK_FP32,
COREML_FP32: KOKORO_PL_COREML_FP32,
},
DE: {
...KOKORO_DE_XNNPACK_FP32,
XNNPACK_FP32: KOKORO_DE_XNNPACK_FP32,
COREML_FP32: KOKORO_DE_COREML_FP32,
},
},
},
Expand Down