| import re, json |
|
|
| CONTRACTIONS = { |
| "I'M": "I AM", "I'LL": "I WILL", "I'VE": "I HAVE", |
| "DON'T": "DO NOT", "DOESN'T": "DOES NOT", "CAN'T": "CANNOT", |
| "WON'T": "WILL NOT", "IT'S": "IT IS", "HE'S": "HE IS", |
| "SHE'S": "SHE IS", "THEY'RE": "THEY ARE", "WE'RE": "WE ARE", |
| "YOU'RE": "YOU ARE", "THAT'S": "THAT IS", "LET'S": "LET US", |
| "WASN'T": "WAS NOT", "WEREN'T": "WERE NOT", |
| "HAVEN'T": "HAVE NOT", "HASN'T": "HAS NOT", |
| "WOULDN'T": "WOULD NOT", "COULDN'T": "COULD NOT", |
| } |
|
|
| def load_ke_model(model_dir): |
| """Load saved Orpheus Kenyan English model for inference.""" |
| from unsloth import FastLanguageModel |
| import torch |
|
|
| model, tokenizer = FastLanguageModel.from_pretrained( |
| model_name = model_dir, |
| max_seq_length = 2048, |
| dtype = None, |
| load_in_4bit = True, |
| ) |
| FastLanguageModel.for_inference(model) |
|
|
| with open(f"{model_dir}/ke_lexicon.json") as f: |
| lexicon = json.load(f) |
|
|
| return model, tokenizer, lexicon |
|
|
|
|
| def apply_g2p(text, lexicon): |
| """Convert text to Kenyan English phonemes using saved lexicon.""" |
| words = text.split() |
| expanded = [] |
| for w in words: |
| clean = re.sub(r"^[^\w\']+|[^\w\']+$", "", w).upper() |
| if clean in CONTRACTIONS: |
| expanded.extend(CONTRACTIONS[clean].split()) |
| else: |
| expanded.append(w) |
| parts = [] |
| for token in expanded: |
| word = re.sub(r"^[^\w\']+|[^\w\']+$", "", token).upper() |
| if not word: |
| continue |
| parts.append(lexicon.get(word, word)) |
| return " || ".join(parts) |
|
|
|
|
| def synthesise(text, model, tokenizer, lexicon, snac_model, |
| temperature=0.6, top_p=0.95): |
| """Full TTS pipeline: text -> Kenyan phonemes -> audio.""" |
| import torch, numpy as np |
|
|
| START_OF_HUMAN = 128259 |
| END_OF_TEXT = 128009 |
| END_OF_HUMAN = 128260 |
| START_OF_SPEECH = 128257 |
| END_OF_SPEECH = 128258 |
|
|
| phoneme_text = apply_g2p(text, lexicon) |
| print(f" G2P: {phoneme_text}") |
|
|
| input_ids = tokenizer.encode(phoneme_text, return_tensors="pt") |
| prompt = torch.cat([ |
| torch.tensor([[START_OF_HUMAN]]), |
| input_ids, |
| torch.tensor([[END_OF_TEXT, END_OF_HUMAN]]), |
| ], dim=1).to("cuda") |
|
|
| with torch.no_grad(): |
| gen_ids = model.generate( |
| input_ids = prompt, |
| attention_mask = torch.ones_like(prompt), |
| max_new_tokens = 1200, |
| do_sample = True, |
| temperature = temperature, |
| top_p = top_p, |
| repetition_penalty = 1.1, |
| eos_token_id = END_OF_SPEECH, |
| use_cache = True, |
| ) |
|
|
| speech_start = (gen_ids == START_OF_SPEECH).nonzero(as_tuple=True) |
| if len(speech_start[1]) == 0: |
| print(" WARNING: no speech tokens generated") |
| return np.zeros(24000) |
|
|
| cropped = gen_ids[:, speech_start[1][-1].item() + 1:] |
| row = cropped[0][cropped[0] != END_OF_SPEECH] |
| trimmed = row[:(len(row) // 7) * 7] |
| code_list = [t.item() - 128266 for t in trimmed] |
|
|
| def redistribute(codes): |
| l1, l2, l3 = [], [], [] |
| for i in range(len(codes) // 7): |
| l1.append(max(0, min(4095, codes[7*i]))) |
| l2.append(max(0, min(4095, codes[7*i+1] - 4096))) |
| l3.append(max(0, min(4095, codes[7*i+2] - 2*4096))) |
| l3.append(max(0, min(4095, codes[7*i+3] - 3*4096))) |
| l2.append(max(0, min(4095, codes[7*i+4] - 4*4096))) |
| l3.append(max(0, min(4095, codes[7*i+5] - 5*4096))) |
| l3.append(max(0, min(4095, codes[7*i+6] - 6*4096))) |
| return ( |
| torch.tensor(l1).unsqueeze(0), |
| torch.tensor(l2).unsqueeze(0), |
| torch.tensor(l3).unsqueeze(0), |
| ) |
|
|
| snac_model.to("cpu") |
| t1, t2, t3 = redistribute(code_list) |
| with torch.inference_mode(): |
| audio_hat = snac_model.decode([t1, t2, t3]) |
|
|
| return audio_hat.squeeze().cpu().numpy() |
|
|
|
|
| |
| |
| |
| |
| |
| |
| |
|
|
