import torchfrom snac import SNACfrom unsloth import FastLanguageModelimport soundfile as sf # Load model + adaptermodel, tokenizer = FastLanguageModel.from_pretrained( model_name="MosesJoshuaCoker/Novax-Krio-TTS-v3", max_seq_length=4096, dtype=None, load_in_4bit=True,)FastLanguageModel.for_inference(model) # Load SNAC decodersnac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").eval().cuda() # Constants (fixed by Orpheus)END_OF_TEXT, START_OF_SPEECH, END_OF_SPEECH = 128009, 128257, 128258START_OF_HUMAN, END_OF_HUMAN, START_OF_AI = 128259, 128260, 128261AUDIO_CODE_OFFSET, CODEBOOK = 128266, 4096 def synthesize(text, temperature=0.6, top_p=0.95, max_tokens=2400): """Generate Krio speech from text.""" ids = tokenizer(text, return_tensors="pt").input_ids start = torch.tensor([[START_OF_HUMAN]], dtype=torch.long) end = torch.tensor([[END_OF_TEXT, END_OF_HUMAN, START_OF_AI, START_OF_SPEECH]], dtype=torch.long) input_ids = torch.cat([start, ids, end], dim=1).cuda() gen = model.generate( input_ids=input_ids, attention_mask=torch.ones_like(input_ids), max_new_tokens=max_tokens, do_sample=True, temperature=temperature, top_p=top_p, eos_token_id=END_OF_SPEECH, pad_token_id=tokenizer.pad_token_id, ) seq = gen[0] sos = (seq == START_OF_SPEECH).nonzero() seq = seq[sos[-1].item() + 1:] if len(sos) else seq raw = [t.item() - AUDIO_CODE_OFFSET for t in seq if t.item() >= AUDIO_CODE_OFFSET] raw = raw[: (len(raw) // 7) * 7] cap = lambda v: max(0, min(CODEBOOK - 1, v)) l1, l2, l3 = [], [], [] for i in range(len(raw) // 7): l1.append(cap(raw[7*i])) l2.append(cap(raw[7*i+1] - 1*CODEBOOK)) l3.append(cap(raw[7*i+2] - 2*CODEBOOK)) l3.append(cap(raw[7*i+3] - 3*CODEBOOK)) l2.append(cap(raw[7*i+4] - 4*CODEBOOK)) l3.append(cap(raw[7*i+5] - 5*CODEBOOK)) l3.append(cap(raw[7*i+6] - 6*CODEBOOK)) layers = [torch.tensor(l1).unsqueeze(0).cuda(), torch.tensor(l2).unsqueeze(0).cuda(), torch.tensor(l3).unsqueeze(0).cuda()] with torch.inference_mode(): wav = snac.decode(layers).squeeze().cpu().numpy() return wav # Exampletext = "kushɛ, aw di bɔdi? wi de tray fɔ mek krio tɛkst-tu-spich voice."wav = synthesize(text)sf.write("krio_sample.wav", wav, 24000)