diff --git a/apps/speech/app/kokoro-text-to-speech/index.tsx b/apps/speech/app/kokoro-text-to-speech/index.tsx index 121bfa702d..18fd602174 100644 --- a/apps/speech/app/kokoro-text-to-speech/index.tsx +++ b/apps/speech/app/kokoro-text-to-speech/index.tsx @@ -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, @@ -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; +}; + // cspell:disable const SAMPLE_TEXTS: Record = { EN_US: @@ -65,7 +79,9 @@ function KokoroContent() { const [runError, setRunError] = useState(null); const [totalDuration, setTotalDuration] = useState(null); - const model = models.textToSpeech.KOKORO[language].XNNPACK_FP32 as KokoroTtsModel; + const [backend, setBackend] = useState('XNNPACK_FP32'); + + const model = kokoroModel(language, backend); const voiceNames = Object.keys(model.voices); const [voice, setVoice] = useState(voiceNames[0]!); @@ -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 () => { @@ -198,6 +217,14 @@ function KokoroContent() { selectedValue={language} onValueChange={setLanguage} /> + {Platform.OS === 'ios' && 'COREML_FP32' in models.textToSpeech.KOKORO[language] && ( + ({ ...b, disabled: isBusy }))} + selectedValue={backend} + onValueChange={setBackend} + /> + )} ({ label: name, value: name, disabled: isBusy }))} diff --git a/packages/react-native-executorch/src/extensions/speech/tasks/kokoroTextToSpeech.ts b/packages/react-native-executorch/src/extensions/speech/tasks/kokoroTextToSpeech.ts index ee3d0459d7..d11783f507 100644 --- a/packages/react-native-executorch/src/extensions/speech/tasks/kokoroTextToSpeech.ts +++ b/packages/react-native-executorch/src/extensions/speech/tasks/kokoroTextToSpeech.ts @@ -147,7 +147,7 @@ export async function createKokoroTextToSpeech( try { const predictorSpec = validateSpec(models.durationPredictor.schema, { - default: method( + dynamic: method( 'forward', [ i64(1, Dyn('T')), // tokens @@ -168,10 +168,25 @@ export async function createKokoroTextToSpeech( ), ] ), + // 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 @@ -194,15 +209,61 @@ export async function createKokoroTextToSpeech( ), ] ), + 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); diff --git a/packages/react-native-executorch/src/models.ts b/packages/react-native-executorch/src/models.ts index b4db68689c..1645859b15 100644 --- a/packages/react-native-executorch/src/models.ts +++ b/packages/react-native-executorch/src/models.ts @@ -880,14 +880,23 @@ const SUPERTONIC_3_MLX_FP32: SupertonicTtsModel = { 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 = (names: readonly N[]) => names.reduce( @@ -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, @@ -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 // ============================================================================= @@ -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, }, }, },