| 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) |
|
|
