File size: 4,923 Bytes
ac573f1 | 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 | """Streaming TTS with Flash-DPO (Dynamic Accent, CosyVoice3 engine): the tagged reply is cut into chunks
(a short first chunk for latency, then sentence by sentence) and each chunk is synthesised as soon as the previous is done."""
import os, sys, re, time
import numpy as np
REL = os.environ.get("DA_REL", "/home/dev/fi/reltest")
sys.path.insert(0, REL)
os.environ.setdefault("CV3_DIR", "/home/dev/models/Fun-CosyVoice3-0.5B-2512")
os.environ.setdefault("DA_PROMPTS", REL + "/prompts")
TAG = re.compile(r"<(eng|ibo|yor|hau|pcm)>(.*?)</\1>", re.S)
VOICE = {"anchor_male": "anchor", "waxal_ibo_5_female": "waxal_ibo_5", "waxal_yor_5_female": "waxal_yor_5"}
def chunk_tagged(tagged, first_max_words=6):
"""-> list of tagged chunk strings. Splits at sentence ends (. ! ?), keeps tags, first chunk short (also cut at a comma/greeting)."""
units, cur = [], []
for l, t in TAG.findall(tagged):
parts = re.split(r"(?<=[.!?])\s+", t.strip())
for k, s in enumerate(parts):
if not s: continue
cur.append((l, s))
if re.search(r"[.!?]$", s):
units.append(cur); cur = []
if cur: units.append(cur)
if not units: return []
first = units[0]
n = sum(len(s.split()) for _, s in first)
if n > first_max_words: # fast first chunk: cut the first sentence at a comma / tag boundary / conjunction
if len(first) > 1 and len(first[0][1].split()) <= first_max_words:
units[0] = [first[0]]; units.insert(1, first[1:])
else:
l, s = first[0]
m = re.search(r"[,;:]\s", s)
if not (m and 1 <= len(s[:m.start()].split()) <= first_max_words):
m = None
ws = s.split()
for k in range(3, min(first_max_words, len(ws) - 3)):
if ws[k].lower() in ("and", "but", "so", "because", "then", "or"):
m = re.search(r"\s" + re.escape(ws[k]) + r"\s", s); break
if m:
cut = m.end() if s[m.start()] in ",;:" else m.start()
units[0] = [(l, s[:cut].strip())]; units.insert(1, [(l, s[cut:].strip())] + first[1:])
# hard caps: first chunk <= 10 words, later chunks <= 22 words (split at a comma if possible, else at a word boundary)
def split_unit(u, cap):
out, cur, n = [], [], 0
for l, s in u:
ws = s.split()
while n + len(ws) > cap:
k = cap - n
sub = " ".join(ws[:k]); c = sub.rfind(",")
if c > len(sub) // 3: k = len(sub[:c + 1].split())
if k > 0: cur.append((l, " ".join(ws[:k])))
out.append(cur); cur, n, ws = [], 0, ws[k:]
if ws: cur.append((l, " ".join(ws))); n += len(ws)
if cur: out.append(cur)
return out
capped = []
for i, u in enumerate(units):
capped += split_unit(u, 10 if i == 0 else 22)
units = [u for u in capped if u]
return [" ".join(f"<{l}>{s}</{l}>" for l, s in u) for u in units]
class TTS:
def __init__(self, voice="anchor_male", lm=None, diac=None, tail_pad=4):
import flash_lm
from frontend import FrontEnd
lm = lm or f"{REL}/checkpoints/flash/flash_dpo/best.pt"
diac = diac or f"{REL}/checkpoints/flash/diacritics_byt5"
self.S = flash_lm.Synth(lm)
self.fe = FrontEnd(diac if os.path.isdir(diac) else None)
self.sr = self.S.cosy.sample_rate
self.voice = VOICE.get(voice, voice); self.tail_pad = tail_pad
self.synth("<eng>Warm up.</eng>") # compile/cache voice prompt
def synth(self, tagged, seed=0, voice=None):
import torch
v = VOICE.get(voice, voice) if voice else self.voice
t0 = time.perf_counter()
# diacritic restoration only for ibo/yor spans written without any tone marks (the LLM usually writes them)
need = any(l in ("ibo", "yor") and not re.search(r"[\u0300-\u036f\u1ea0-\u1ef9ịụọẹṣńḿ]", __import__("unicodedata").normalize("NFD", s) + s) for l, s in TAG.findall(tagged))
t = self.fe(tagged, "codeswitch", "eng") if need else tagged
t1 = time.perf_counter()
toks = list(self.S.tokens(t, v, "normal", seed)[0]); torch.cuda.synchronize()
t2 = time.perf_counter()
if self.tail_pad: toks += [1] * self.tail_pad
w = np.asarray(self.S.wav(toks, v), dtype=np.float32)
t3 = time.perf_counter()
self.last = dict(fe=round(t1 - t0, 3), lm=round(t2 - t1, 3), flow_voc=round(t3 - t2, 3), n_tok=len(toks), audio_s=round(len(w) / self.sr, 2))
return w
def stream(self, tagged, seed=0, voice=None):
"""yields (chunk_text, wav, synth_seconds)"""
for i, c in enumerate(chunk_tagged(tagged)):
t = time.perf_counter(); w = self.synth(c, seed + 101 * i, voice)
yield c, w, time.perf_counter() - t
|