1
import { debounce_timeout } from '../../constants.js';
2
import { debounceAsync, splitRecursive } from '../../utils.js';
3
import { getPreviewString, saveTtsProviderSettings } from './index.js';
4

5
export class KokoroTtsProvider {
6
    constructor() {
7
        this.settings = {
8
            modelId: 'onnx-community/Kokoro-82M-v1.0-ONNX',
9
            dtype: 'q8',
10
            device: 'wasm',
11
            voiceMap: {},
12
            defaultVoice: 'af_heart',
13
            speakingRate: 1.0,
14
        };
15
        this.ready = false;
16
        this.voices = [
17
            'af_heart',
18
            'af_alloy',
19
            'af_aoede',
20
            'af_bella',
21
            'af_jessica',
22
            'af_kore',
23
            'af_nicole',
24
            'af_nova',
25
            'af_river',
26
            'af_sarah',
27
            'af_sky',
28
            'am_adam',
29
            'am_echo',
30
            'am_eric',
31
            'am_fenrir',
32
            'am_liam',
33
            'am_michael',
34
            'am_onyx',
35
            'am_puck',
36
            'am_santa',
37
            'bf_emma',
38
            'bf_isabella',
39
            'bm_george',
40
            'bm_lewis',
41
            'bf_alice',
42
            'bf_lily',
43
            'bm_daniel',
44
            'bm_fable',
45
        ];
46
        this.worker = null;
47
        this.separator = ' ... ... ... ';
48
        this.pendingRequests = new Map();
49
        this.nextRequestId = 1;
50

51
        // Update display values immediately but only reinitialize TTS after a delay
52
        this.initTtsDebounced = debounceAsync(this.initializeWorker.bind(this), debounce_timeout.relaxed);
53
    }
54

55
    /**
56
     * Perform any text processing before passing to TTS engine.
57
     * @param {string} text Input text
58
     * @returns {string} Processed text
59
     */
60
    processText(text) {
61
        // TILDE!
62
        text = text.replace(/~/g, '.');
63
        return text;
64
    }
65

66
    async loadSettings(settings) {
67
        if (settings.modelId !== undefined) this.settings.modelId = settings.modelId;
68
        if (settings.dtype !== undefined) this.settings.dtype = settings.dtype;
69
        if (settings.device !== undefined) this.settings.device = settings.device;
70
        if (settings.voiceMap !== undefined) this.settings.voiceMap = settings.voiceMap;
71
        if (settings.defaultVoice !== undefined) this.settings.defaultVoice = settings.defaultVoice;
72
        if (settings.speakingRate !== undefined) this.settings.speakingRate = settings.speakingRate;
73

74
        $('#kokoro_model_id').val(this.settings.modelId).on('input', this.onSettingsChange.bind(this));
75
        $('#kokoro_dtype').val(this.settings.dtype).on('change', this.onSettingsChange.bind(this));
76
        $('#kokoro_device').val(this.settings.device).on('change', this.onSettingsChange.bind(this));
77
        $('#kokoro_speaking_rate').val(this.settings.speakingRate).on('input', this.onSettingsChange.bind(this));
78
        $('#kokoro_speaking_rate_output').text(this.settings.speakingRate + 'x');
79
    }
80

81
    initializeWorker() {
82
        return new Promise((resolve, reject) => {
83
            try {
84
                // Terminate the existing worker if it exists
85
                if (this.worker) {
86
                    this.worker.terminate();
87
                    $('#kokoro_status_text').text('Initializing...').removeAttr('style');
88
                }
89

90
                // Create a new worker
91
                this.worker = new Worker(new URL('./kokoro-worker.js', import.meta.url), { type: 'module' });
92

93
                // Set up message handling
94
                this.worker.onmessage = this.handleWorkerMessage.bind(this);
95

96
                // Initialize the worker with the current settings
97
                this.worker.postMessage({
98
                    action: 'initialize',
99
                    data: {
100
                        modelId: this.settings.modelId,
101
                        dtype: this.settings.dtype,
102
                        device: this.settings.device,
103
                    },
104
                });
105

106
                // Create a promise that will resolve when initialization completes
107
                const initPromise = new Promise((initResolve, initReject) => {
108
                    const timeoutId = setTimeout(() => {
109
                        initReject(new Error('Worker initialization timed out'));
110
                    }, 600000); // 600 second timeout
111

112
                    this.pendingRequests.set('initialization', {
113
                        resolve: (result) => {
114
                            clearTimeout(timeoutId);
115
                            initResolve(result);
116
                        },
117
                        reject: (error) => {
118
                            clearTimeout(timeoutId);
119
                            initReject(error);
120
                        },
121
                    });
122
                });
123

124
                // Resolve the outer promise when initialization completes
125
                initPromise.then(success => {
126
                    this.ready = success;
127
                    this.updateStatusDisplay();
128
                    resolve(success);
129
                }).catch(error => {
130
                    console.error('Worker initialization failed:', error);
131
                    this.ready = false;
132
                    this.updateStatusDisplay();
133
                    reject(error);
134
                });
135
            } catch (error) {
136
                console.error('Failed to create worker:', error);
137
                this.ready = false;
138
                this.updateStatusDisplay();
139
                reject(error);
140
            }
141
        });
142
    }
143

144
    handleWorkerMessage(event) {
145
        const { action, success, ready, error, requestId, blobUrl } = event.data;
146

147
        switch (action) {
148
            case 'initialized': {
149
                const initRequest = this.pendingRequests.get('initialization');
150
                if (initRequest) {
151
                    if (success) {
152
                        initRequest.resolve(true);
153
                    } else {
154
                        initRequest.reject(new Error(error || 'Initialization failed'));
155
                    }
156
                    this.pendingRequests.delete('initialization');
157
                }
158
            } break;
159
            case 'generatedTts': {
160
                const request = this.pendingRequests.get(requestId);
161
                if (request) {
162
                    if (success) {
163
                        fetch(blobUrl).then(response => response.blob()).then(audioBlob => {
164
                            // Clean up the blob URL
165
                            URL.revokeObjectURL(blobUrl);
166

167
                            request.resolve(new Response(audioBlob, {
168
                                headers: {
169
                                    'Content-Type': 'audio/wav',
170
                                },
171
                            }));
172
                        }).catch(error => {
173
                            request.reject(new Error('Failed to fetch TTS audio blob: ' + error));
174
                        });
175
                    } else {
176
                        request.reject(new Error(error || 'TTS generation failed'));
177
                    }
178
                    this.pendingRequests.delete(requestId);
179
                }
180
            } break;
181
            case 'readyStatus':
182
                this.ready = ready;
183
                this.updateStatusDisplay();
184
                break;
185
        }
186
    }
187

188
    updateStatusDisplay() {
189
        const statusText = this.ready ? 'Ready' : 'Failed';
190
        const statusColor = this.ready ? 'green' : 'red';
191
        $('#kokoro_status_text').text(statusText).css('color', statusColor);
192
    }
193

194
    async checkReady() {
195
        if (!this.worker) {
196
            return await this.initializeWorker();
197
        }
198

199
        this.worker.postMessage({ action: 'checkReady' });
200
        return this.ready;
201
    }
202

203
    async onRefreshClick() {
204
        return await this.initializeWorker();
205
    }
206

207
    get settingsHtml() {
208
        return `
209
            <div class="kokoro_tts_settings">
210
                <label for="kokoro_model_id">Model ID:</label>
211
                <input id="kokoro_model_id" type="text" class="text_pole" value="${this.settings.modelId}" />
212

213
                <label for="kokoro_dtype">Data Type:</label>
214
                <select id="kokoro_dtype" class="text_pole">
215
                    <option value="q8" ${this.settings.dtype === 'q8' ? 'selected' : ''}>q8 (Recommended)</option>
216
                    <option value="fp32" ${this.settings.dtype === 'fp32' ? 'selected' : ''}>fp32 (High Precision)</option>
217
                    <option value="fp16" ${this.settings.dtype === 'fp16' ? 'selected' : ''}>fp16</option>
218
                    <option value="q4" ${this.settings.dtype === 'q4' ? 'selected' : ''}>q4 (Low Memory)</option>
219
                    <option value="q4f16" ${this.settings.dtype === 'q4f16' ? 'selected' : ''}>q4f16</option>
220
                </select>
221

222
                <label for="kokoro_device">Device:</label>
223
                <select id="kokoro_device" class="text_pole">
224
                    <option value="wasm" ${this.settings.device === 'wasm' ? 'selected' : ''}>WebAssembly (CPU)</option>
225
                    <option value="webgpu" ${this.settings.device === 'webgpu' ? 'selected' : ''}>WebGPU (GPU Acceleration)</option>
226
                </select>
227

228
                <label for="kokoro_speaking_rate">Speaking Rate: <span id="kokoro_speaking_rate_output">${this.settings.speakingRate}x</span></label>
229
                <input id="kokoro_speaking_rate" type="range" value="${this.settings.speakingRate}" min="0.5" max="2.0" step="0.1" />
230

231
                <hr>
232
                <div>
233
                    Status: <span id="kokoro_status_text">Initializing...</span>
234
                </div>
235
            </div>
236
        `;
237
    }
238

239
    async onSettingsChange() {
240
        this.settings.modelId = $('#kokoro_model_id').val().toString();
241
        this.settings.dtype = $('#kokoro_dtype').val().toString();
242
        this.settings.device = $('#kokoro_device').val().toString();
243
        this.settings.speakingRate = parseFloat($('#kokoro_speaking_rate').val().toString());
244

245
        // Update UI display
246
        $('#kokoro_speaking_rate_output').text(this.settings.speakingRate + 'x');
247

248
        // Reinitialize TTS engine with debounce
249
        this.initTtsDebounced();
250
        saveTtsProviderSettings();
251
    }
252

253
    async fetchTtsVoiceObjects() {
254
        if (!this.ready) {
255
            await this.checkReady();
256
        }
257
        return this.voices.map(voice => ({
258
            name: voice,
259
            voice_id: voice,
260
            preview_url: null,
261
            lang: voice.startsWith('b') ? 'en-GB' : 'en-US',
262
        }));
263
    }
264

265
    async previewTtsVoice(voiceId) {
266
        if (!this.ready) {
267
            await this.checkReady();
268
        }
269

270
        const voice = this.getVoice(voiceId);
271
        const previewText = getPreviewString(voice.lang);
272
        for await (const response of this.generateTts(previewText, voiceId)) {
273
            const audio = await response.blob();
274
            const url = URL.createObjectURL(audio);
275
            await new Promise(resolve => {
276
                const audioElement = new Audio();
277
                audioElement.src = url;
278
                audioElement.play();
279
                audioElement.onended = () => resolve();
280
            });
281
            URL.revokeObjectURL(url);
282
        }
283
    }
284

285
    getVoiceDisplayName(voiceId) {
286
        return voiceId;
287
    }
288

289
    getVoice(voiceName) {
290
        const defaultVoice = this.settings.defaultVoice || 'af_heart';
291
        const actualVoiceName = this.voices.includes(voiceName) ? voiceName : defaultVoice;
292
        return {
293
            name: actualVoiceName,
294
            voice_id: actualVoiceName,
295
            preview_url: null,
296
            lang: actualVoiceName.startsWith('b') ? 'en-GB' : 'en-US',
297
        };
298
    }
299

300
    /**
301
     * Generate TTS audio for the given text using the specified voice.
302
     * @param {string} text Text to generate
303
     * @param {string} voiceId Voice ID
304
     * @returns {AsyncGenerator<Response>} Audio response generator
305
     */
306
    async* generateTts(text, voiceId) {
307
        if (!this.ready || !this.worker) {
308
            console.log('TTS not ready, initializing...');
309
            await this.initializeWorker();
310
        }
311

312
        if (!this.ready || !this.worker) {
313
            throw new Error('Failed to initialize TTS engine');
314
        }
315

316
        if (text.trim().length === 0) {
317
            throw new Error('Empty text');
318
        }
319

320
        const voice = this.getVoice(voiceId);
321
        const requestId = this.nextRequestId++;
322

323
        const chunkSize = 400;
324
        const chunks = splitRecursive(text, chunkSize, ['\n\n', '\n', '.', '?', '!', ',', ' ', '']);
325

326
        for (const chunk of chunks) {
327
            yield await new Promise((resolve, reject) => {
328
                // Store the promise callbacks
329
                this.pendingRequests.set(requestId, { resolve, reject });
330

331
                // Send the request to the worker
332
                this.worker.postMessage({
333
                    action: 'generateTts',
334
                    data: {
335
                        text: chunk,
336
                        voice: voice.voice_id,
337
                        speakingRate: this.settings.speakingRate || 1.0,
338
                        requestId,
339
                    },
340
                });
341
            });
342
        }
343
    }
344

345
    dispose() {
346
        // Clean up the worker when the provider is disposed
347
        if (this.worker) {
348
            this.worker.terminate();
349
            this.worker = null;
350
        }
351
    }
352
}