Buckets:
| """Sigma VO bake-off, SNAC-family Indian voices: Veena (agastya, vinaya), Maya1 (Indian male, designed), Svara (English male). | |
| Every line: N sampled takes -> Parakeet + Moonshine transcription -> best take kept. Writes to /data (private) and /pub.""" | |
| import os, sys, json, time, tarfile, traceback | |
| import numpy as np, torch, soundfile as sf | |
| sys.path.insert(0, '/data/code'); import eval_lib as E | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| from snac import SNAC | |
| N = int(os.environ.get('TAKES', '3')) | |
| ONLY = [x for x in os.environ.get('ENGINES', '').split(',') if x] | |
| LINES = json.load(open(os.environ.get('LINES', '/data/code/fix_lines.json'))) | |
| TAG = os.environ.get('TAG', 'fix') | |
| OUT = '/data/out'; PUB = '/pub' | |
| os.makedirs(OUT, exist_ok=True) | |
| SR = 24000 | |
| snac = SNAC.from_pretrained('hubertsiuzdak/snac_24khz').eval().to('cuda') | |
| BASE = 128266 | |
| def decode(tokens): | |
| tok = [t for t in tokens if BASE <= t < BASE + 7 * 4096] | |
| n = len(tok) // 7 | |
| if n < 4: return None | |
| tok = tok[:n * 7] | |
| l1, l2, l3 = [], [], [] | |
| for i in range(n): | |
| s = [tok[i * 7 + j] - BASE - j * 4096 for j in range(7)] | |
| if any(x < 0 or x > 4095 for x in s): s = [(tok[i * 7 + j] - BASE) % 4096 for j in range(7)] | |
| l1.append(s[0]); l2 += [s[1], s[4]]; l3 += [s[2], s[3], s[5], s[6]] | |
| codes = [torch.tensor(x, dtype=torch.long, device='cuda')[None] for x in (l1, l2, l3)] | |
| with torch.inference_mode(): a = snac.decode(codes) | |
| return a.squeeze().float().clamp(-1, 1).cpu().numpy() | |
| def until_eos(seq, eos): | |
| out = [] | |
| for t in seq: | |
| if t in eos: break | |
| out.append(t) | |
| return out | |
| def gen(model, ids, n, eos, **kw): | |
| ids = torch.tensor([ids], device='cuda') | |
| with torch.inference_mode(): | |
| o = model.generate(ids, attention_mask=torch.ones_like(ids), do_sample=True, num_return_sequences=n, eos_token_id=eos, pad_token_id=eos[0], **kw) | |
| return [until_eos(o[i, ids.shape[1]:].tolist(), set(eos)) for i in range(n)] | |
| def load(name): | |
| tok = AutoTokenizer.from_pretrained(name, trust_remote_code=True) | |
| m = AutoModelForCausalLM.from_pretrained(name, torch_dtype=torch.bfloat16, device_map='cuda', trust_remote_code=True).eval() | |
| return tok, m | |
| def veena(spk): | |
| def run(): | |
| tok, m = load('maya-research/Veena') | |
| def f(text, n): | |
| ids = [128259] + tok.encode(f'<spk_{spk}> {text}', add_special_tokens=False) + [128260, 128261, 128257] | |
| return [decode(s) for s in gen(m, ids, n, [128258, 128262], max_new_tokens=900, temperature=0.4, top_p=0.9, repetition_penalty=1.05)] | |
| return f, m | |
| return run | |
| MAYA_DESC = ('Realistic male voice in the 40s age with a indian accent. Low pitch, deep timbre, slow pacing, ' | |
| 'neutral tone delivery at med intensity, corporate domain, explainer_video_voice role, formal delivery') | |
| def maya1(): | |
| tok, m = load('maya-research/maya1') | |
| soh, eoh, soa, sos, eot = (tok.decode([i]) for i in (128259, 128260, 128261, 128257, 128009)) | |
| def f(text, n): | |
| prompt = soh + tok.bos_token + f'<description="{MAYA_DESC}"> {text}' + eot + eoh + soa + sos | |
| ids = tok(prompt).input_ids | |
| outs = gen(m, ids, n, [128258], max_new_tokens=1400, min_new_tokens=28, temperature=0.4, top_p=0.9, repetition_penalty=1.1) | |
| res = [] | |
| for s in outs: | |
| a = decode(s); res.append(a[2048:] if a is not None and len(a) > 4096 else a) | |
| return res | |
| return f, m | |
| def svara(): | |
| tok, m = load('kenpath/svara-tts-v1') | |
| def f(text, n): | |
| ids = [128259] + tok(f'English (Male): {text} <formal>').input_ids + [128009, 128260] | |
| outs = gen(m, ids, n, [128258], max_new_tokens=1500, temperature=0.7, top_p=0.8, repetition_penalty=1.1) | |
| res = [] | |
| for s in outs: | |
| if 128257 in s: s = s[len(s) - s[::-1].index(128257):] | |
| res.append(decode(s)) | |
| return res | |
| return f, m | |
| ENGINES = {'veena_agastya': veena('agastya'), 'veena_vinaya': veena('vinaya'), 'maya1_indian': maya1, 'svara_male': svara} | |
| summary = {} | |
| for name, builder in ENGINES.items(): | |
| if ONLY and name not in ONLY: continue | |
| t0 = time.time(); d = f'{OUT}/{name}_{TAG}'; os.makedirs(f'{d}/best', exist_ok=True) | |
| try: | |
| f, model = builder() | |
| except Exception as e: | |
| print('LOAD FAIL', name, e); traceback.print_exc(); summary[name] = {'error': str(e)}; continue | |
| rep = {} | |
| for L in LINES: | |
| k, text, ref, mx = L['id'], L['tts'], L['ref'], L['max'] | |
| takes = [] | |
| variants = text if isinstance(text, list) else [text] | |
| for attempt in range(int(os.environ.get('ROUNDS', '2'))): | |
| for vi, vt in enumerate(variants): | |
| try: auds = f(vt, N) | |
| except Exception as e: print('GEN FAIL', name, k, e); auds = [] | |
| for a in auds: | |
| if a is None or len(a) < SR * 0.3: continue | |
| a = E.trim(a, SR) | |
| ev = E.evaluate(a, SR, ref); ev['audio'] = a; ev['variant'] = vt; takes.append(ev) | |
| print(f" take {k} v{vi} P={ev['wer_p']:.2f} M={ev['wer_m']:.2f} dur={ev['dur']:.2f} | {ev['parakeet']} | {ev['moonshine']}", flush=True) | |
| if sum(1 for t in takes if t['wer_p'] == 0 and t['wer_m'] == 0 and t['dur'] <= mx) >= int(os.environ.get('NEED', '2')): break | |
| if not takes: rep[k] = {'error': 'no audio'}; continue | |
| takes.sort(key=lambda t: (t['wer_p'] + t['wer_m'], t['dur'] > mx, -t['indian'])) | |
| b = takes[0] | |
| sf.write(f'{d}/best/{k}.flac', b['audio'], SR) | |
| rep[k] = {kk: vv for kk, vv in b.items() if kk != 'audio'} | |
| rep[k]['n_takes'] = len(takes); rep[k]['n_pass'] = sum(1 for t in takes if t['wer_p'] == 0 and t['wer_m'] == 0) | |
| rep[k]['max'] = mx | |
| for i, t in enumerate(sorted([t for t in takes if t['wer_p'] == 0 and t['wer_m'] == 0 and t['dur'] <= mx], key=lambda t: -t['indian'])[:3]): | |
| os.makedirs(f'{d}/alt', exist_ok=True); sf.write(f'{d}/alt/{k}_{i}.flac', t['audio'], SR) | |
| print(f"{name} {k} best P={b['wer_p']:.2f} M={b['wer_m']:.2f} ind={b['indian']:.2f} {b['accent']} dur={b['dur']:.2f}/{mx} | {b['parakeet']} | {b['moonshine']}", flush=True) | |
| del model; torch.cuda.empty_cache() | |
| ok = [v for v in rep.values() if 'error' not in v] | |
| both = sum(1 for v in ok if v['wer_p'] == 0 and v['wer_m'] == 0) | |
| fits = sum(1 for v in ok if v['dur'] <= v['max']) | |
| allaud = np.concatenate([sf.read(f'{d}/best/{k}.flac')[0] for k in rep if 'error' not in rep[k]] + [np.zeros(SR)]) | |
| s = {'both_exact': both, 'lines': len(LINES), 'fit': fits, 'mean_indian': float(np.mean([v['indian'] for v in ok])), | |
| 'f0': E.f0med(E.to16(allaud[:SR * 20], SR)), 'secs': round(time.time() - t0)} | |
| summary[name] = s; print('SUMMARY', name, s, flush=True) | |
| json.dump(rep, open(f'{d}/report.json', 'w'), indent=1) | |
| os.makedirs(PUB, exist_ok=True) | |
| with tarfile.open(f'{PUB}/{name}_{TAG}.tar.gz', 'w:gz') as tf: | |
| tf.add(f'{d}/best', arcname=f'{name}_{TAG}/best'); tf.add(f'{d}/report.json', arcname=f'{name}_{TAG}/report.json') | |
| if os.path.isdir(f'{d}/alt'): tf.add(f'{d}/alt', arcname=f'{name}_{TAG}/alt') | |
| json.dump(summary, open(f'{OUT}/summary_{TAG}.json', 'w'), indent=1) | |
| json.dump(summary, open(f'{PUB}/summary_{TAG}.json', 'w'), indent=1) | |
| print('ALL DONE', json.dumps(summary)) | |
Xet Storage Details
- Size:
- 7.35 kB
- Xet hash:
- 4a8f70ea4b72cd4a5eb4259788e3bb275020798d4491e989466c3a5ae0008bd1
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.