File size: 1,628 Bytes
63a4356 | 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 | import json, random, sys
from pathlib import Path
import torch
from kokoro import KModel
TEACHER, CORPUS, PACK, OUT, COUNT = sys.argv[1:6]
MAX_PHONEMES = 500
device = "cuda" if torch.cuda.is_available() else "cpu"
model = KModel(repo_id="hexgrad/Kokoro-82M", config=TEACHER + "/config.json", model=TEACHER + "/kokoro.pth").to(device).eval()
model.decoder = None
pack = torch.load(PACK, map_location=device, weights_only=True)
sentences = [json.loads(line)["phonemes"] for line in open(CORPUS, encoding="utf-8")]
rng = random.Random(7)
def paragraph():
parts = [rng.choice(sentences)]
while True:
candidate = rng.choice(sentences)
if len(" ".join(parts + [candidate])) > MAX_PHONEMES:
return " ".join(parts)
parts.append(candidate)
with open(OUT, "w") as sink, torch.no_grad():
for _ in range(int(COUNT)):
phonemes = paragraph()
ids = [0, *[model.vocab[p] for p in phonemes if p in model.vocab], 0]
tokens = torch.LongTensor([ids]).to(device)
lengths = torch.LongTensor([len(ids)]).to(device)
mask = torch.zeros(1, len(ids), dtype=torch.bool, device=device)
d_en = model.bert_encoder(model.bert(tokens, attention_mask=(~mask).int())).transpose(-1, -2)
d = model.predictor.text_encoder(d_en, pack[len(phonemes) - 1][:, 128:], lengths, mask)
x, _ = model.predictor.lstm(d)
dur = torch.round(torch.sigmoid(model.predictor.duration_proj(x)).sum(-1)).clamp(min=1).long().reshape(-1).tolist()
sink.write(json.dumps({"ids": ids, "dur": dur, "lang": "pt"}) + "\n")
print("CONTEXT_DUMP_DONE", flush=True)
|