File size: 26,554 Bytes
e9b69d2
 
 
 
 
 
 
76f722d
e9b69d2
7c600ad
8ab445f
e9b69d2
 
 
d989475
c8ab947
 
 
d989475
 
c8ab947
d989475
 
 
 
 
 
 
c8ab947
d989475
c8ab947
 
e9b69d2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
61044da
22f4889
 
dd3c146
61044da
24f33ad
dd3c146
 
 
 
 
 
 
 
 
f3d4c6b
61044da
 
dd3c146
f3d4c6b
61044da
dd3c146
5895ac2
dd3c146
5895ac2
8c17d76
61044da
e9b69d2
 
9e3a2ff
 
 
 
 
a7e7cf3
7c600ad
 
 
 
 
dd3c146
40576c0
61044da
0662719
 
343970c
 
7c600ad
e9b69d2
 
 
 
7c600ad
e9b69d2
 
7c600ad
e9b69d2
7c600ad
 
e9b69d2
 
 
 
 
24f33ad
 
dd3b355
f3d4c6b
343970c
85cde2d
dd3b355
e9b69d2
7c600ad
d989475
e9b69d2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7c600ad
 
 
 
7b74eba
e9b69d2
480e1dc
c27ca74
c70c0b7
c27ca74
 
9b9ab2a
c27ca74
c70c0b7
 
e9b69d2
d989475
e9b69d2
 
 
 
 
 
 
 
 
 
 
 
 
 
d989475
 
e9b69d2
 
 
d989475
 
e9b69d2
d989475
 
 
e9b69d2
7c600ad
 
e9b69d2
 
7c600ad
 
 
d989475
 
7c600ad
d989475
 
 
 
 
6766b80
123602a
9733762
123602a
 
 
64fb7dd
480b27e
6766b80
7c600ad
6766b80
 
 
 
 
7c600ad
7b74eba
d989475
7b74eba
 
7c600ad
 
 
2acb6ca
e9b69d2
2acb6ca
 
e9b69d2
2acb6ca
 
e9b69d2
2acb6ca
 
e9b69d2
2acb6ca
e9b69d2
 
9b9ab2a
 
 
7b74eba
 
 
 
 
 
 
 
 
7c600ad
e9b69d2
 
7b74eba
 
 
 
 
 
 
 
 
 
 
e9b69d2
 
 
 
7c600ad
 
e9b69d2
 
7c600ad
e9b69d2
 
9b9ab2a
d989475
 
7c600ad
 
 
e9b69d2
82642e5
e9b69d2
 
82642e5
e9b69d2
82642e5
e9b69d2
 
82642e5
e9b69d2
82642e5
123602a
 
e9b69d2
7c600ad
d989475
 
7c600ad
 
9b9ab2a
 
7c600ad
 
 
d989475
 
 
 
 
7c600ad
 
 
 
 
 
 
 
d989475
 
7c600ad
d989475
7c600ad
 
d989475
 
 
 
7c600ad
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9733762
7c600ad
 
d959cb1
 
 
7c600ad
 
 
d989475
 
 
 
7c600ad
 
76f722d
7c600ad
 
 
 
 
2acb6ca
7c600ad
2acb6ca
 
7c600ad
2acb6ca
 
7c600ad
2acb6ca
d959cb1
 
7c600ad
 
 
 
 
 
 
 
 
 
76f722d
 
 
 
 
 
 
 
 
 
7c600ad
 
 
76f722d
7b74eba
76f722d
7b74eba
 
 
 
 
 
 
 
 
7c600ad
d989475
7c600ad
 
 
 
 
 
 
 
82642e5
7c600ad
 
7b74eba
7c600ad
7b74eba
7c600ad
2acb6ca
7b74eba
2acb6ca
7b74eba
d959cb1
d989475
7c600ad
 
 
 
d989475
7c600ad
8fe5320
 
7b74eba
8fe5320
7b74eba
 
 
 
 
 
 
 
 
 
 
61044da
8fe5320
c8be4ae
 
 
8fe5320
7b74eba
4b9083a
 
 
 
 
 
 
 
 
 
 
 
480b27e
4b9083a
 
 
 
 
343970c
4b9083a
 
8fe5320
4b9083a
7b74eba
 
 
9b9ab2a
 
 
7b74eba
 
 
 
 
 
 
 
 
 
 
 
 
 
9b9ab2a
 
7c600ad
 
7b74eba
 
7c600ad
 
e9b69d2
1a61201
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
from huggingface_hub import snapshot_download
from katsu import Katsu
from models import build_model
import gradio as gr
import numpy as np
import os
import phonemizer
import pypdf
import random
import re
import spaces
import torch
import yaml

CUDA_AVAILABLE = torch.cuda.is_available()

snapshot = snapshot_download(repo_id='hexgrad/kokoro', allow_patterns=['*.pt', '*.pth', '*.yml'], use_auth_token=os.environ['TOKEN'])
config = yaml.safe_load(open(os.path.join(snapshot, 'config.yml')))

models = {device: build_model(config['model_params'], device) for device in ['cpu'] + (['cuda'] if CUDA_AVAILABLE else [])}
for key, state_dict in torch.load(os.path.join(snapshot, 'net.pth'), map_location='cpu', weights_only=True)['net'].items():
    for device in models:
        assert key in models[device], key
        try:
            models[device][key].load_state_dict(state_dict)
        except:
            state_dict = {k[7:]: v for k, v in state_dict.items()}
            models[device][key].load_state_dict(state_dict, strict=False)

PARAM_COUNT = sum(p.numel() for value in models['cpu'].values() for p in value.parameters())
assert PARAM_COUNT < 82_000_000, PARAM_COUNT

random_texts = {}
for lang in ['en', 'ja']:
    with open(f'{lang}.txt', 'r') as r:
        random_texts[lang] = [line.strip() for line in r]

def get_random_text(voice):
    if voice[0] == 'j':
        lang = 'ja'
    else:
        lang = 'en'
    return random.choice(random_texts[lang])

def parens_to_angles(s):
    return s.replace('(', '«').replace(')', '»')

def split_num(num):
    num = num.group()
    if '.' in num:
        # Decimal
        a, b = num.split('.')
        return ' point '.join([a, ' '.join(b)])
    elif ':' in num:
        # Time
        h, m = [int(n) for n in num.split(':')]
        if m == 0:
            return f"{h} o'clock"
        elif m < 10:
            return f'{h} oh {m}'
        return f'{h} {m}'
    # Year
    year = int(num[:4])
    if year < 1100 or year % 1000 < 10:
        return num
    left, right = num[:2], int(num[2:4])
    s = 's' if num.endswith('s') else ''
    if 100 <= year % 1000 <= 999:
        if right == 0:
            return f'{left} hundred{s}'
        elif right < 10:
            return f'{left} oh {right}{s}'
    return f'{left} {right}{s}'

def normalize(text):
    # TODO: Custom text normalization rules?
    text = re.sub(r'\bD[Rr]\.(?= [A-Z])', 'Doctor', text)
    text = re.sub(r'\b(?:Mr\.|MR\.(?= [A-Z]))', 'Mister', text)
    text = re.sub(r'\b(?:Ms\.|MS\.(?= [A-Z]))', 'Miss', text)
    text = re.sub(r'\b(?:Mrs\.|MRS\.(?= [A-Z]))', 'Mrs', text)
    text = re.sub(r'\betc\.(?! [A-Z])', 'etc', text)
    text = re.sub(r'\b([Yy])eah\b', r"\1e'a", text)
    text = text.replace(chr(8216), "'").replace(chr(8217), "'")
    text = text.replace(chr(8220), '"').replace(chr(8221), '"')
    text = re.sub(r'[^\S \n]', ' ', text)
    text = re.sub(r'  +', ' ', text)
    text = re.sub(r'(?<=\n) +(?=\n)', '', text)
    text = re.sub(r'\d*\.\d+|\b\d{4}s?\b|(?<!:)\b(?:[1-9]|1[0-2]):[0-5]\d\b(?!:)', split_num, text)
    text = re.sub(r'(?<=\d),(?=\d)', '', text)
    text = re.sub(r'(?<=\d)-(?=\d)', ' to ', text) # TODO: could be minus
    text = re.sub(r'(?<=\d)S', ' S', text)
    text = re.sub(r"(?<=[A-Z])'?s", lambda m: m.group().upper(), text)
    text = re.sub(r'(?:[A-Za-z]\.){2,} [a-z]', lambda m: m.group().replace('.', '-'), text)
    text = re.sub(r'(?i)(?<=[A-Z])\.(?=[A-Z])', '-', text)
    return parens_to_angles(text).strip()

phonemizers = dict(
    a=phonemizer.backend.EspeakBackend(language='en-us', preserve_punctuation=True, with_stress=True),
    b=phonemizer.backend.EspeakBackend(language='en-gb', preserve_punctuation=True, with_stress=True),
    j=Katsu(),
)

def phonemize(text, voice, norm=True):
    lang = voice[0]
    if norm:
        text = normalize(text)
    ps = phonemizers[lang].phonemize([text])
    ps = ps[0] if ps else ''
    # TODO: Custom phonemization rules?
    ps = parens_to_angles(ps)
    # https://en.wiktionary.org/wiki/kokoro#English
    if lang in 'ab':
        ps = ps.replace('kəkˈoːɹoʊ', 'kˈoʊkəɹoʊ').replace('kəkˈɔːɹəʊ', 'kˈəʊkəɹəʊ')
        ps = ps.replace('ʲ', 'j').replace('r', 'ɹ').replace('x', 'k').replace('ɬ', 'l')
        ps = re.sub(r'(?<=[a-zɹː])(?=hˈʌndɹɪd)', ' ', ps)
        ps = re.sub(r' z(?=[;:,.!?¡¿—…"«»“” ]|$)', 'z', ps)
        if lang == 'a':
            ps = re.sub(r'(?<=nˈaɪn)ti(?!ː)', 'di', ps)
    ps = ''.join(filter(lambda p: p in VOCAB, ps))
    if lang == 'j' and any(p in 'ABCDEFGHIJKLMNOPQRSTUVWXYZ' for p in ps):
        gr.Warning('Japanese tokenizer does not handle English letters')
    return ps.strip()

def length_to_mask(lengths):
    mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
    mask = torch.gt(mask+1, lengths.unsqueeze(1))
    return mask

def get_vocab():
    _pad = "$"
    _punctuation = ';:,.!?¡¿—…"«»“” '
    _letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
    _letters_ipa = "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"
    symbols = [_pad] + list(_punctuation) + list(_letters) + list(_letters_ipa)
    dicts = {}
    for i in range(len((symbols))):
        dicts[symbols[i]] = i
    return dicts

VOCAB = get_vocab()

def tokenize(ps):
    return [i for i in map(VOCAB.get, ps) if i is not None]

# ⭐ voices are stable, 🧪 voices are unstable
CHOICES = {
'🇺🇸 🚺 American Female ⭐': 'af',
'🇺🇸 🚺 Bella': 'af_bella',
'🇺🇸 🚺 Nicole': 'af_nicole',
'🇺🇸 🚺 Sarah': 'af_sarah',
'🇺🇸 🚺 Sky 🧪': 'af_sky',
'🇺🇸 🚹 Adam 🧪': 'am_adam',
'🇺🇸 🚹 Michael': 'am_michael',
'🇬🇧 🚹 Lewis 🧪': 'bm_lewis',
'🇯🇵 🚺 Japanese Female': 'jf_0',
}
VOICES = {device: {k: torch.load(os.path.join(snapshot, 'voicepacks', f'{k}.pt'), weights_only=True).to(device) for k in CHOICES.values()} for device in models}

np_log_99 = np.log(99)
def s_curve(p):
    if p <= 0:
        return 0
    elif p >= 1:
        return 1
    s = 1 / (1 + np.exp((1-p*2)*np_log_99))
    s = (s-0.01) * 50/49
    return s

SAMPLE_RATE = 24000

@torch.no_grad()
def forward(tokens, voice, speed, device='cpu'):
    ref_s = VOICES[device][voice][len(tokens)]
    tokens = torch.LongTensor([[0, *tokens, 0]]).to(device)
    input_lengths = torch.LongTensor([tokens.shape[-1]]).to(device)
    text_mask = length_to_mask(input_lengths).to(device)
    bert_dur = models[device].bert(tokens, attention_mask=(~text_mask).int())
    d_en = models[device].bert_encoder(bert_dur).transpose(-1, -2)
    s = ref_s[:, 128:]
    d = models[device].predictor.text_encoder(d_en, s, input_lengths, text_mask)
    x, _ = models[device].predictor.lstm(d)
    duration = models[device].predictor.duration_proj(x)
    duration = torch.sigmoid(duration).sum(axis=-1) / speed
    pred_dur = torch.round(duration).clamp(min=1).long()
    pred_aln_trg = torch.zeros(input_lengths, pred_dur.sum().item())
    c_frame = 0
    for i in range(pred_aln_trg.size(0)):
        pred_aln_trg[i, c_frame:c_frame + pred_dur[0,i].item()] = 1
        c_frame += pred_dur[0,i].item()
    en = d.transpose(-1, -2) @ pred_aln_trg.unsqueeze(0).to(device)
    F0_pred, N_pred = models[device].predictor.F0Ntrain(en, s)
    t_en = models[device].text_encoder(tokens, input_lengths, text_mask)
    asr = t_en @ pred_aln_trg.unsqueeze(0).to(device)
    return models[device].decoder(asr, F0_pred, N_pred, ref_s[:, :128]).squeeze().cpu().numpy()

@spaces.GPU(duration=10)
def forward_gpu(tokens, voice, speed):
    return forward(tokens, voice, speed, device='cuda')

# Must be backwards compatible with https://huggingface.co./spaces/Pendrokar/TTS-Spaces-Arena
def generate(text, voice, ps, speed, _reduce_noise, opening_cut, closing_cut, ease_in, ease_out, _pad_before, _pad_after, use_gpu):
    return _generate(text, voice, ps, speed, opening_cut, closing_cut, ease_in, ease_out, use_gpu)

def _generate(text, voice, ps, speed, opening_cut, closing_cut, ease_in, ease_out, use_gpu):
    if voice not in VOICES['cpu']:
        voice = 'af'
    ps = ps or phonemize(text, voice)
    tokens = tokenize(ps)
    if not tokens:
        return (None, '')
    elif len(tokens) > 510:
        tokens = tokens[:510]
    ps = ''.join(next(k for k, v in VOCAB.items() if i == v) for i in tokens)
    try:
        if not use_gpu or (use_gpu == 'auto' and len(ps) < 100):
            out = forward(tokens, voice, speed)
        else:
            out = forward_gpu(tokens, voice, speed)
    except gr.exceptions.Error as e:
        raise gr.Error(e)
        return (None, '')
    opening_cut = int(opening_cut / speed)
    if opening_cut > 0:
        out = out[opening_cut:]
    closing_cut = int(closing_cut / speed)
    if closing_cut > 0:
        out = out[:-closing_cut]
    ease_in = min(int(ease_in / speed), len(out)//2)
    for i in range(ease_in):
        out[i] *= s_curve(i / ease_in)
    ease_out = min(int(ease_out / speed), len(out)//2)
    for i in range(ease_out):
        out[-i-1] *= s_curve(i / ease_out)
    return ((SAMPLE_RATE, out), ps)

def toggle_autoplay(autoplay):
    return gr.Audio(interactive=False, label='Output Audio', autoplay=autoplay)

USE_GPU_CHOICES = [('Auto', 'auto'), ('CPU', False), ('ZeroGPU', True)]
USE_GPU_INFOS = {
    'auto': 'Use CPU or GPU, whichever is faster',
    False: 'CPU is ~faster <100 tokens',
    True: 'ZeroGPU is ~faster >100 tokens',
}
def change_use_gpu(value):
    return gr.Dropdown(USE_GPU_CHOICES, value=value, label='Hardware', info=USE_GPU_INFOS[value], interactive=CUDA_AVAILABLE)

with gr.Blocks() as basic_tts:
    with gr.Row():
        with gr.Column():
            text = gr.Textbox(label='Input Text', info='Generate speech for one segment of text using Kokoro, a TTS model with 80 million parameters')
            with gr.Row():
                voice = gr.Dropdown(list(CHOICES.items()), value='af', label='Voice', info='⭐ voices are stable, 🧪 voices are unstable')
                use_gpu = gr.Dropdown(
                    USE_GPU_CHOICES,
                    value='auto' if CUDA_AVAILABLE else False,
                    label='Hardware',
                    info=USE_GPU_INFOS['auto' if CUDA_AVAILABLE else False],
                    interactive=CUDA_AVAILABLE
                )
                use_gpu.change(fn=change_use_gpu, inputs=[use_gpu], outputs=[use_gpu])
            with gr.Row():
                random_btn = gr.Button('Random Text', variant='secondary')
                generate_btn = gr.Button('Generate', variant='primary')
            random_btn.click(get_random_text, inputs=[voice], outputs=[text])
            with gr.Accordion('Input Tokens', open=False):
                in_ps = gr.Textbox(show_label=False, info='Override the input text with custom phonemes. Leave this blank to automatically tokenize the input text instead.')
                with gr.Row():
                    clear_btn = gr.ClearButton(in_ps)
                    phonemize_btn = gr.Button('Tokenize Input Text', variant='primary')
            phonemize_btn.click(phonemize, inputs=[text, voice], outputs=[in_ps])
        with gr.Column():
            audio = gr.Audio(interactive=False, label='Output Audio', autoplay=True)
            autoplay = gr.Checkbox(value=True, label='Autoplay')
            autoplay.change(toggle_autoplay, inputs=[autoplay], outputs=[audio])
            with gr.Accordion('Output Tokens', open=True):
                out_ps = gr.Textbox(interactive=False, show_label=False, info='Tokens used to generate the audio, up to 510 allowed. Same as input tokens if supplied, excluding unknowns.')
    with gr.Accordion('Audio Settings', open=False):
        with gr.Row():
            speed = gr.Slider(minimum=0.5, maximum=2, value=1, step=0.1, label='⚡️ Speed', info='Adjust the speed of the audio; the settings below are auto-scaled by speed')
        with gr.Row():
            with gr.Column():
                opening_cut = gr.Slider(minimum=0, maximum=24000, value=4000, step=1000, label='✂️ Opening Cut', info='Cut samples from the start')
            with gr.Column():
                closing_cut = gr.Slider(minimum=0, maximum=24000, value=2000, step=1000, label='🎬 Closing Cut', info='Cut samples from the end')
        with gr.Row():
            with gr.Column():
                ease_in = gr.Slider(minimum=0, maximum=24000, value=3000, step=1000, label='🎢 Ease In', info='Ease in samples, after opening cut')
            with gr.Column():
                ease_out = gr.Slider(minimum=0, maximum=24000, value=1000, step=1000, label='🛝 Ease Out', info='Ease out samples, before closing cut')
    text.submit(_generate, inputs=[text, voice, in_ps, speed, opening_cut, closing_cut, ease_in, ease_out, use_gpu], outputs=[audio, out_ps])
    generate_btn.click(_generate, inputs=[text, voice, in_ps, speed, opening_cut, closing_cut, ease_in, ease_out, use_gpu], outputs=[audio, out_ps])

@torch.no_grad()
def lf_forward(token_lists, voice, speed, device='cpu'):
    voicepack = VOICES[device][voice]
    outs = []
    for tokens in token_lists:
        ref_s = voicepack[len(tokens)]
        s = ref_s[:, 128:]
        tokens = torch.LongTensor([[0, *tokens, 0]]).to(device)
        input_lengths = torch.LongTensor([tokens.shape[-1]]).to(device)
        text_mask = length_to_mask(input_lengths).to(device)
        bert_dur = models[device].bert(tokens, attention_mask=(~text_mask).int())
        d_en = models[device].bert_encoder(bert_dur).transpose(-1, -2)
        d = models[device].predictor.text_encoder(d_en, s, input_lengths, text_mask)
        x, _ = models[device].predictor.lstm(d)
        duration = models[device].predictor.duration_proj(x)
        duration = torch.sigmoid(duration).sum(axis=-1) / speed
        pred_dur = torch.round(duration).clamp(min=1).long()
        pred_aln_trg = torch.zeros(input_lengths, pred_dur.sum().item())
        c_frame = 0
        for i in range(pred_aln_trg.size(0)):
            pred_aln_trg[i, c_frame:c_frame + pred_dur[0,i].item()] = 1
            c_frame += pred_dur[0,i].item()
        en = d.transpose(-1, -2) @ pred_aln_trg.unsqueeze(0).to(device)
        F0_pred, N_pred = models[device].predictor.F0Ntrain(en, s)
        t_en = models[device].text_encoder(tokens, input_lengths, text_mask)
        asr = t_en @ pred_aln_trg.unsqueeze(0).to(device)
        outs.append(models[device].decoder(asr, F0_pred, N_pred, ref_s[:, :128]).squeeze().cpu().numpy())
    return outs

@spaces.GPU
def lf_forward_gpu(token_lists, voice, speed):
    return lf_forward(token_lists, voice, speed, device='cuda')

def resplit_strings(arr):
    # Handle edge cases
    if not arr:
        return '', ''
    if len(arr) == 1:
        return arr[0], ''
    # Try each possible split point
    min_diff = float('inf')
    best_split = 0
    # Calculate lengths when joined with spaces
    lengths = [len(s) for s in arr]
    spaces = len(arr) - 1  # Total spaces needed
    # Try each split point
    left_len = 0
    right_len = sum(lengths) + spaces
    for i in range(1, len(arr)):
        # Add current word and space to left side
        left_len += lengths[i-1] + (1 if i > 1 else 0)
        # Remove current word and space from right side
        right_len -= lengths[i-1] + 1
        diff = abs(left_len - right_len)
        if diff < min_diff:
            min_diff = diff
            best_split = i
    # Join the strings with the best split point
    return ' '.join(arr[:best_split]), ' '.join(arr[best_split:])

def recursive_split(text, voice):
    if not text:
        return []
    tokens = phonemize(text, voice, norm=False)
    if len(tokens) < 511:
        return [(text, tokens, len(tokens))] if tokens else []
    if ' ' not in text:
        return []
    for punctuation in ['!.?…', ':;', ',—']:
        splits = re.split(f'(?:(?<=[{punctuation}])|(?<=[{punctuation}]["\'»])|(?<=[{punctuation}]["\'»]["\'»])) ', text)
        if len(splits) > 1:
            break
        else:
            splits = None
    splits = splits or text.split(' ')
    a, b = resplit_strings(splits)
    return recursive_split(a, voice) + recursive_split(b, voice)

def segment_and_tokenize(text, voice, skip_square_brackets=True, newline_split=2):
    if skip_square_brackets:
        text = re.sub(r'\[.*?\]', '', text)
    texts = [t.strip() for t in re.split('\n{'+str(newline_split)+',}', normalize(text))] if newline_split > 0 else [normalize(text)]
    segments = [row for t in texts for row in recursive_split(t, voice)]
    return [(i, *row) for i, row in enumerate(segments)]

def lf_generate(segments, voice, speed, opening_cut, closing_cut, ease_in, ease_out, pad_between, use_gpu):
    token_lists = list(map(tokenize, segments['Tokens']))
    wavs = []
    opening_cut = int(opening_cut / speed)
    closing_cut = int(closing_cut / speed)
    pad_between = int(pad_between / speed)
    batch_size = 100
    for i in range(0, len(token_lists), batch_size):
        try:
            if use_gpu:
                outs = lf_forward_gpu(token_lists[i:i+batch_size], voice, speed)
            else:
                outs = lf_forward(token_lists[i:i+batch_size], voice, speed)
        except gr.exceptions.Error as e:
            if wavs:
                gr.Warning(str(e))
            else:
                raise gr.Error(e)
            break
        for out in outs:
            if opening_cut > 0:
                out = out[opening_cut:]
            if closing_cut > 0:
                out = out[:-closing_cut]
            ease_in = min(int(ease_in / speed), len(out)//2)
            for i in range(ease_in):
                out[i] *= s_curve(i / ease_in)
            ease_out = min(int(ease_out / speed), len(out)//2)
            for i in range(ease_out):
                out[-i-1] *= s_curve(i / ease_out)
            if wavs and pad_between > 0:
                wavs.append(np.zeros(pad_between))
            wavs.append(out)
    return (SAMPLE_RATE, np.concatenate(wavs)) if wavs else None

def did_change_segments(segments):
    x = len(segments) if segments['Length'].any() else 0
    return [
        gr.Button('Tokenize', variant='secondary' if x else 'primary'),
        gr.Button(f'Generate x{x}', variant='primary' if x else 'secondary', interactive=x > 0),
    ]

def extract_text(file):
    if file.endswith('.pdf'):
        with open(file, 'rb') as rb:
            pdf_reader = pypdf.PdfReader(rb)
            return '\n'.join([page.extract_text() for page in pdf_reader.pages])
    elif file.endswith('.txt'):
        with open(file, 'r') as r:
            return '\n'.join([line for line in r])
    return None

with gr.Blocks() as lf_tts:
    with gr.Row():
        with gr.Column():
            file_input = gr.File(file_types=['.pdf', '.txt'], label='Input File: pdf or txt')
            text = gr.Textbox(label='Input Text', info='Generate speech in batches of 100 text segments and automatically join them together')
            file_input.upload(fn=extract_text, inputs=[file_input], outputs=[text])
            with gr.Row():
                voice = gr.Dropdown(list(CHOICES.items()), value='af', label='Voice', info='⭐ voices are stable, 🧪 voices are unstable')
                use_gpu = gr.Dropdown(
                    [('ZeroGPU', True), ('CPU', False)],
                    value=CUDA_AVAILABLE,
                    label='Hardware',
                    info='GPU is >10x faster but has a usage quota',
                    interactive=CUDA_AVAILABLE
                )
            with gr.Accordion('Text Settings', open=False):
                skip_square_brackets = gr.Checkbox(True, label='Skip [Square Brackets]', info='Recommended for academic papers, Wikipedia articles, or texts with citations')
                newline_split = gr.Number(2, label='Newline Split', info='Split the input text on this many newlines. Affects how the text is segmented.', precision=0, minimum=0)
            with gr.Row():
                segment_btn = gr.Button('Tokenize', variant='primary')
                generate_btn = gr.Button('Generate x0', variant='secondary', interactive=False)
        with gr.Column():
            audio = gr.Audio(interactive=False, label='Output Audio')
            with gr.Accordion('Audio Settings', open=False):
                with gr.Row():
                    speed = gr.Slider(minimum=0.5, maximum=2, value=1, step=0.1, label='⚡️ Speed', info='Adjust the speed of the audio; the settings below are auto-scaled by speed')
                with gr.Row():
                    with gr.Column():
                        opening_cut = gr.Slider(minimum=0, maximum=24000, value=4000, step=1000, label='✂️ Opening Cut', info='Cut samples from the start')
                    with gr.Column():
                        closing_cut = gr.Slider(minimum=0, maximum=24000, value=2000, step=1000, label='🎬 Closing Cut', info='Cut samples from the end')
                with gr.Row():
                    with gr.Column():
                        ease_in = gr.Slider(minimum=0, maximum=24000, value=3000, step=1000, label='🎢 Ease In', info='Ease in samples, after opening cut')
                    with gr.Column():
                        ease_out = gr.Slider(minimum=0, maximum=24000, value=1000, step=1000, label='🛝 Ease Out', info='Ease out samples, before closing cut')
                with gr.Row():
                    pad_between = gr.Slider(minimum=0, maximum=24000, value=10000, step=1000, label='🔇 Pad Between', info='How many samples of silence to insert between segments')
    with gr.Row():
        segments = gr.Dataframe(headers=['#', 'Text', 'Tokens', 'Length'], row_count=(1, 'dynamic'), col_count=(4, 'fixed'), label='Segments', interactive=False, wrap=True)
        segments.change(fn=did_change_segments, inputs=[segments], outputs=[segment_btn, generate_btn])
    segment_btn.click(segment_and_tokenize, inputs=[text, voice, skip_square_brackets, newline_split], outputs=[segments])
    generate_btn.click(lf_generate, inputs=[segments, voice, speed, opening_cut, closing_cut, ease_in, ease_out, pad_between, use_gpu], outputs=[audio])

with gr.Blocks() as about:
    gr.Markdown("""
Kokoro is a frontier TTS model for its size. It has 80 million parameters,<sup>[1]</sup> uses a lean StyleTTS 2 architecture,<sup>[2]</sup> and was trained on high-quality data. The weights are currently private, but a free public demo is hosted here, at `https://hf.co/spaces/hexgrad/Kokoro-TTS`. The Community tab is open for feature requests, bug reports, etc. For other inquiries, contact @rzvzn on Discord.

### FAQ
#### Will this be open sourced?
There currently isn't a release date scheduled for the weights. The inference code in this space is MIT licensed. The architecture was already published by Li et al, with MIT licensed code and pretrained weights.<sup>[2]</sup>

#### What is an unstable voice?
An unstable voice is more likely to stumble or produce unnatural artifacts, especially on short or strange texts.

#### How can CPU be faster than ZeroGPU?
The CPU is a dedicated resource for this Space, while the ZeroGPU pool is shared and dynamically allocated across all of HF. The ZeroGPU queue/allocator system inevitably adds latency to each request.<br/>
For Basic TTS under ~100 tokens or characters, only a few seconds of audio need to be generated, so the actual compute is not that heavy. In these short bursts, the dedicated CPU can often compute the result faster than the total time it takes to: enter the ZeroGPU queue, wait to get allocated, and have a GPU compute and deliver the result.<br/>
ZeroGPU catches up beyond 100 tokens and especially closer to the ~500 token context window. Long-Form mode processes batches of 100 segments at a time, so the GPU should outspeed the CPU by 1-2 orders of magnitude.

### Compute
The model was trained on 1x A100-class 80GB instances rented from [Vast.ai](https://cloud.vast.ai/?ref_id=79907).<sup>[3]</sup><br/>
Vast was chosen over other compute providers due to its competitive on-demand hourly rates.<br/>
The average hourly cost for the 1x A100-class 80GB VRAM instances used for training was below $1/hr — around half the quoted rates from other providers.

### Gradio API
This Space can be used via API. The following code block can be copied and run in one Google Colab cell.
```
# 1. Install the Gradio Python client
!pip install -q gradio_client

# 2. Initialize the client
from gradio_client import Client
client = Client('hexgrad/Kokoro-TTS')

# 3. Call the generate endpoint, which returns a pair: an audio path and a string of output phonemes
audio_path, out_ps = client.predict(
    text="How could I know? It's an unanswerable question. Like asking an unborn child if they'll lead a good life. They haven't even been born.",
    voice='af',
    api_name='/generate'
)

# 4. Display the audio and print the output phonemes
from IPython.display import display, Audio
display(Audio(audio_path, autoplay=True))
print(out_ps)
```
Note that this Space and the underlying Kokoro model are both under development and subject to change. Reliability is not guaranteed. Hugging Face and/or Gradio might enforce their own rate limits.

### Model Version History
| Version | Date | Val mel / dur / f0 Losses |
| ------- | ---- | ------------------------- |
| v0.19 | 2024 Nov 22 | 0.261 / 0.627 / 1.897 |
| v0.16 | 2024 Nov 15 | 0.263 / 0.646 / 1.934 |
| v0.14 | 2024 Nov 12 | 0.262 / 0.642 / 1.889 |

### Licenses
Inference code: MIT<br/>
espeak-ng dependency: GPL-3.0<sup>[4]</sup><br/>
Random English texts: Unknown<sup>[5]</sup><br/>
Random Japanese texts: CC0 public domain<sup>[6]</sup>

### References
1. Kokoro parameter count | https://hf.co/spaces/hexgrad/Kokoro-TTS/blob/main/app.py#L31
2. StyleTTS 2 | https://github.com/yl4579/StyleTTS2
3. Vast.ai referral link | https://cloud.vast.ai/?ref_id=79907
4. eSpeak NG | https://github.com/espeak-ng/espeak-ng
5. Quotable Data | https://github.com/quotable-io/data/blob/master/data/quotes.json
6. Common Voice Japanese sentences | https://github.com/common-voice/common-voice/tree/main/server/data/ja
""")

with gr.Blocks() as app:
    gr.TabbedInterface(
        [basic_tts, lf_tts, about],
        ['🔥 Basic TTS', '📖 Long-Form', 'ℹ️ About'],
    )

if __name__ == '__main__':
    app.queue(api_open=True).launch()