"""Stage 1: distill Nabra-82M into the 5.25M Kokoro student. The student is trained on the TEACHER's audio using the TEACHER's own durations, so the alignment is correct by construction. That is the whole reason for distilling: kitten-arabic established that with mismatched durations this architecture tops out around WER 0.85 and 62,000 further steps change nothing, because free-running inference beats teacher-forcing -- the durations, not the capacity, were the ceiling. Adversarial loss is NOT optional. The same project measured reconstruction loss falling from 6.445 to 2.996 while the audio degraded to pure static: multi-resolution STFT and mel distance constrain magnitude but not phase, so the optimizer can win on the objective and lose the waveform. MPD + MSD from StyleTTS2 are what make the objective mean something. Block names mirror KModel (bert/bert_encoder/predictor/text_encoder/decoder) so a finished student loads straight into KPipeline with the teacher's G2P and voicepack. """ from __future__ import annotations import _env # noqa: F401 import argparse, contextlib, json, math, os, random, time from pathlib import Path import numpy as np, torch, torch.nn as nn, torch.nn.functional as F from torch.utils.data import Dataset, DataLoader from transformers import AlbertConfig from kokoro.istftnet import Decoder from kokoro.modules import ProsodyPredictor, TextEncoder, CustomAlbert from style_encoder import StyleEncoder from styletts2_gan.discriminators import (MultiPeriodDiscriminator, MultiResSpecDiscriminator, WavLMDiscriminator) from styletts2_gan.losses_styletts2 import (GeneratorLoss, DiscriminatorLoss, MultiResolutionSTFTLoss, WavLMLoss) class MelL1(nn.Module): """L1 on log-mel. The existing multi-resolution STFT loss is spectral convergence on LINEAR-frequency bins, so every bin counts equally in Hz; mel spacing instead weights by where hearing resolves detail. Measured motivation: against the teacher the student runs +0.025 energy in 0-1 kHz and -0.025 in 1-4 kHz -- the consonant/F2-F3 band -- with a 12.28 dB mel L1 gap, while intelligibility is already fine (WER 0.064, median 0.000). So the deficit is timbre in a specific band, which is exactly what a mel-domain term penalises and a linear-bin one under-weights.""" def __init__(self, sr=24000, n_fft=1024, hop=256, n_mels=80): super().__init__() import torchaudio self.mel = torchaudio.transforms.MelSpectrogram( sample_rate=sr, n_fft=n_fft, hop_length=hop, n_mels=n_mels, power=1.0) def forward(self, y_hat, y): a = torch.log(self.mel(y_hat) + 1e-5) b = torch.log(self.mel(y) + 1e-5) return F.l1_loss(a, b) R1R2_EVERY = 8 SPF = 600 # samples per alignment frame, MEASURED in dump_teacher.py; --sr rewrites it AUDIO_SR = 0 # set from --audio-sr when the files differ from the model's rate SPACE = 16 FRAME_HZ = 40 # the alignment grid, held constant across sample rates so a duration means # the same thing whether the decoder emits 600 or 400 samples per frame NB = "/home/ahmed-wasfy/Documents/projects/kikiri-tts/training/nabra" def build_student(cfg_path): c = json.load(open(cfg_path)) m = nn.Module() m.bert = CustomAlbert(AlbertConfig(vocab_size=c["n_token"], **c["plbert"])) predictor_hidden = c.get("predictor_hidden", c["hidden_dim"]) m.bert_encoder = nn.Linear(m.bert.config.hidden_size, predictor_hidden) m.predictor = ProsodyPredictor(style_dim=c["style_dim"], d_hid=predictor_hidden, nlayers=c["n_layer"], max_dur=c["max_dur"], dropout=c["dropout"]) m.text_encoder = TextEncoder(channels=c["hidden_dim"], kernel_size=c["text_encoder_kernel_size"], depth=c["n_layer"], n_symbols=c["n_token"]) decoder_dim_in = c.get("decoder_dim_in", c["hidden_dim"]) if "decoder_dim_in" in c: m.feature_proj = nn.Sequential(nn.Conv1d(c["hidden_dim"], decoder_dim_in, 1), nn.SiLU(), nn.Conv1d(decoder_dim_in, decoder_dim_in, 1)) m.decoder = Decoder(dim_in=decoder_dim_in, style_dim=c["style_dim"], dim_out=c["n_mels"], disable_complex=True, **c["istftnet"]) return m, c def clip_pitch(audio, floor=60, ceiling=800): import parselmouth rate, hop = SPF * FRAME_HZ, SPF // 2 if len(audio) < rate // 10: return np.zeros(len(audio) // hop, dtype=np.float32) pitch = parselmouth.Sound(audio.astype(np.float64), rate).to_pitch_ac(time_step=hop / rate, pitch_floor=floor, pitch_ceiling=ceiling) return np.nan_to_num([pitch.get_value_at_time((i + 0.5) * hop / rate) for i in range(len(audio) // hop)]).astype(np.float32) class DistDataset(Dataset): """Teacher audio from the packed i16 memmap, or -- when a row carries "wav" -- Jenny's real recording with the DTW-transferred durations from align_real.py. For the real branch the clip is cut to sum(dur)*SPF samples, because that is exactly the span the durations describe; anything past it has no label.""" def __init__(self, rows, bin_path, crop, pitch=False): self.rows, self.bin_path, self.crop, self.pitch = rows, bin_path, crop, pitch self._mm = {} def _mem(self, path=None): # a row may name its own bin file, so a bilingual index can span two teacher dumps path = path or self.bin_path if path not in self._mm: self._mm[path] = np.memmap(path, dtype=np.int16, mode="r") return self._mm[path] def __len__(self): return len(self.rows) def __getitem__(self, i): r = self.rows[i] if "wav" in r: import soundfile as sf # "start" lets one long recording supply many training segments without being cut # into files on disk; v9 clips run to 123 s and are segmented at aligner-found pauses. need = int(sum(r["dur"])) * SPF off = int(r.get("start", 0)) * SPF if AUDIO_SR and AUDIO_SR != SPF * FRAME_HZ: import soxr nspf = AUDIO_SR // FRAME_HZ # frames are 40 Hz at either rate a, _ = sf.read(r["wav"], dtype="float32", start=int(r.get("start", 0)) * nspf, frames=int(sum(r["dur"])) * nspf) if a.ndim > 1: a = a.mean(axis=1) a = soxr.resample(a, AUDIO_SR, SPF * FRAME_HZ).astype(np.float32) else: a, _ = sf.read(r["wav"], dtype="float32", start=off, frames=need) if a.ndim > 1: a = a.mean(axis=1) if len(a) < need: a = np.pad(a, (0, need - len(a))) a = a[:need] elif "offset" not in r: a = np.zeros(0, dtype=np.float32) else: a = np.asarray(self._mem(r.get("bin"))[r["offset"]: r["offset"] + r["samples"]], dtype=np.float32) / 32767.0 return {"ids": r["ids"], "dur": r["dur"], "audio": a, "lang": r.get("lang", "en"), **({"f0": clip_pitch(a)} if self.pitch else {})} def collate(batch): B = len(batch) L = max(len(b["ids"]) for b in batch) ids = torch.zeros(B, L, dtype=torch.long) dur = torch.zeros(B, L, dtype=torch.float) lens = torch.tensor([len(b["ids"]) for b in batch], dtype=torch.long) for i, b in enumerate(batch): ids[i, : len(b["ids"])] = torch.tensor(b["ids"]) dur[i, : len(b["dur"])] = torch.tensor(b["dur"], dtype=torch.float) return {"ids": ids, "dur": dur, "lens": lens, "lang": [b.get("lang", "en") for b in batch], "audio": [torch.from_numpy(b["audio"]) for b in batch], **({"f0": [torch.from_numpy(b["f0"]) for b in batch]} if "f0" in batch[0] else {})} def build_aln(dur, lens, F_max, device): """[B,L] integer durations -> [B,L,F] hard monotonic alignment.""" B, L = dur.shape aln = torch.zeros(B, L, F_max, device=device) for b in range(B): pos = 0 for j in range(int(lens[b])): d = int(dur[b, j]) if d <= 0 or pos >= F_max: continue aln[b, j, pos: min(pos + d, F_max)] = 1.0 pos += d return aln def main(): ap = argparse.ArgumentParser() ap.add_argument("--steps", type=int, default=60000) ap.add_argument("--batch", type=int, default=16) ap.add_argument("--crop", type=int, default=40) # frames; 40 * 600 = 1.0 s ap.add_argument("--lr", type=float, default=2e-4) ap.add_argument("--lr-disc", type=float, default=2e-4) ap.add_argument("--lr-decay", action="store_true", help="cosine-decay the generator LR to 10%% of peak. The 15.34M student " "diverged at a CONSTANT 2e-4: its duration L1 stepped from 0.33 to 0.76 " "between 20.5k and 21k and never recovered, leaving predictions " "uncorrelated with the true durations (r=+0.019 against the 7.48M's " "+0.381). The acoustic path kept improving throughout, so nothing in the " "headline losses showed it.") ap.add_argument("--dur-weight", type=float, default=5.0) ap.add_argument("--stft-weight", type=float, default=5.0) ap.add_argument("--disc-warmup", type=int, default=500) ap.add_argument("--slm-weight", type=float, default=0.0, help="WavLM feature-matching loss. Every other term here is MAGNITUDE-domain " "(log-STFT, mel, silence RMS), so nothing but the GAN constrains phase " "-- and our GAN is MPD (waveform) plus MultiResSpec (magnitude, " "phase-blind). Measured: the student matches the teacher on band " "energy, F0 variation, harmonic ratio, modulation spectrum and " "frame-rate buzz, yet sounds robotic. SLM losses are what StyleTTS2 " "credits for naturalness, and this repo vendored them unused.") ap.add_argument("--slm-adv-weight", type=float, default=0.0, help="SLM adversarial term on top of feature matching") ap.add_argument("--mel-weight", type=float, default=0.0, help="weight on the log-mel L1 term (see MelL1)") ap.add_argument("--sil-weight", type=float, default=0.0, help="penalise output energy where the TEACHER is silent. Kokoro-82M emits " "true digital silence in pauses (frame RMS 2e-5); the 5.42M student " "floors at 5.4e-3, 270x higher, which is audible as a constant hiss. " "Multi-resolution STFT is computed on log-magnitude, where -45 dB vs " "-90 dB is a small loss difference and a large perceptual one, so " "nothing in the existing objective punishes it.") ap.add_argument("--sil-thresh", type=float, default=1e-3, help="frame RMS below which the teacher counts as silent") ap.add_argument("--init-from", default=None, help="load weights from a finished run and start a fresh schedule; --resume " "cannot do this because it restores a scheduler that is already done") ap.add_argument("--r1r2-weight", type=float, default=0.0) ap.add_argument("--index", default="dist/index.jsonl") ap.add_argument("--bin", default="dist/audio.i16.bin") ap.add_argument("--config", default="config_student.json") ap.add_argument("--out", default="runs/distill") ap.add_argument("--holdout", type=int, default=200) ap.add_argument("--eval-every", type=int, default=1000) ap.add_argument("--ckpt-every", type=int, default=2000) ap.add_argument("--sr", type=int, default=24000, help="output sample rate. 16000 for corpora that are natively 16 kHz: " "upsampling them to 24 kHz invents nothing above 8 kHz and misreports " "the bandwidth. Sets samples-per-frame to sr/40, and the mel, silence " "and WavLM losses with it.") ap.add_argument("--audio-sr", type=int, default=0, help="the wav files' own rate, when it differs from --sr. The audio is then " "resampled on read. Durations are in 40 Hz FRAMES, which mean the same " "thing at any rate, so an alignment built at one rate is valid at the " "other and does not need redoing.") ap.add_argument("--init-loose", action="store_true", help="load only the tensors whose shapes match. Needed when the decoder is " "rebuilt for another sample rate: the 16 kHz upsample stack differs from " "the 24 kHz one in exactly two tensors, and the other 99.6%% transfers.") ap.add_argument("--bilingual", action="store_true", help="rows carry a 'lang' field; condition English rows on af_heart and " "Arabic rows on af_msa so one student learns both, with language " "selectable at inference by choosing the pack") ap.add_argument("--mode", choices=["joint", "text", "decoder"], default="joint") ap.add_argument("--teacher", default=None) ap.add_argument("--feat-weight", type=float, default=5.0) ap.add_argument("--prefix-prob", type=float, default=0.0) ap.add_argument("--lookahead-words", type=int, default=2) ap.add_argument("--resume", action="store_true") ap.add_argument("--style-lr", type=float, default=0.0) ap.add_argument("--style-labels", default="") ap.add_argument("--probe", action="store_true") ap.add_argument("--freeze-text", action="store_true") ap.add_argument("--pitch-weight", type=float, default=0.0) ap.add_argument("--mark-crop-prob", type=float, default=0.0) ap.add_argument("--mark-ids", default="26,27,28") a = ap.parse_args() dev = "cuda" if torch.cuda.is_available() else "cpu" torch.manual_seed(42) global SPF, AUDIO_SR SPF = a.sr // FRAME_HZ AUDIO_SR = a.audio_sr if AUDIO_SR and AUDIO_SR != a.sr: print(f"[*] reading {AUDIO_SR} Hz audio, resampling to {a.sr} Hz on the fly", flush=True) if a.sr != 24000: print(f"[*] {a.sr} Hz: {SPF} samples per frame at {FRAME_HZ} Hz", flush=True) out = Path(a.out); out.mkdir(parents=True, exist_ok=True) rows = [json.loads(l) for l in open(a.index, encoding="utf-8")] rows = [r for r in rows if sum(r["dur"]) >= a.crop + 2] random.Random(1234).shuffle(rows) train_rows, eval_rows = rows[: -a.holdout], rows[-a.holdout:] print(f"[*] train {len(train_rows):,} | held out {len(eval_rows):,}", flush=True) model, cfg = build_student(a.config) model.to(dev) n = sum(p.numel() for p in model.parameters()) print(f"[*] student {n/1e6:.2f}M", flush=True) pack = torch.load(os.environ["STUDENT_PACK"], map_location="cpu", weights_only=True) PACKS = {"ar": pack} for name, path in json.loads(os.environ.get("STUDENT_PACK_MAP", "{}")).items(): PACKS[name] = torch.load(path, map_location="cpu", weights_only=True) if a.bilingual: import _env # noqa: F401 from kokoro import KPipeline PACKS["en"] = KPipeline(lang_code="a", repo_id="hexgrad/Kokoro-82M", model=False).load_voice("af_heart").cpu() if os.path.exists("pack_cs.pt"): # the code-switch corpus is a THIRD voice (239.5 Hz vs 292.6 Arabic / 197.5 English), # so it gets its own style vector instead of overwriting one of theirs PACKS["cs"] = torch.load("pack_cs.pt", map_location="cpu", weights_only=True) print(f"[*] bilingual packs: {sorted(PACKS)}", flush=True) teacher = None if a.mode != "joint": from kokoro import KModel teacher = KModel(repo_id="hexgrad/Kokoro-82M", config=a.teacher + "/config.json", model=a.teacher + "/kokoro.pth").to(dev).eval() teacher.decoder = None mpd, msd = MultiPeriodDiscriminator().to(dev), MultiResSpecDiscriminator().to(dev) gen_loss = GeneratorLoss(mpd, msd).to(dev) disc_loss = DiscriminatorLoss(mpd, msd).to(dev) stft_loss = MultiResolutionSTFTLoss().to(dev) mel_loss = MelL1(sr=a.sr).to(dev) wl = wd = opt_wd = None if a.slm_weight > 0 or a.slm_adv_weight > 0: # WavLM-base-plus emits 13 hidden states x 768 dims; the discriminator's first conv takes # slm_hidden * slm_layers = 9984 channels. Passing (1024, 512, 256) builds a 524288-channel # conv and fails on the first real batch. wd = WavLMDiscriminator(slm_hidden=768, slm_layers=13, initial_channel=64).to(dev) wl = WavLMLoss("microsoft/wavlm-base-plus", wd, a.sr, 16000).to(dev) for p_ in wl.wavlm.parameters(): p_.requires_grad = False wl.wavlm.eval() opt_wd = torch.optim.AdamW(wd.parameters(), lr=a.lr_disc, betas=(0.8, 0.99)) print(f"[*] SLM: WavLM feature matching (w={a.slm_weight}) " f"adversarial (w={a.slm_adv_weight})", flush=True) opt_g = torch.optim.AdamW((model.decoder if a.mode == "decoder" or a.freeze_text else model).parameters(), lr=a.lr, betas=(0.8, 0.99), weight_decay=1e-2) styler = StyleEncoder(sr=a.sr).to(dev) if a.style_lr > 0 and a.mode == "joint" else None if styler is not None and not a.freeze_text: opt_g.add_param_group({"params": styler.parameters(), "lr": a.style_lr}) labels = [label for label in a.style_labels.split(",") if label] label_style = torch.nn.Embedding(len(labels) + 1, 256, padding_idx=0).to(dev) if labels and a.mode == "joint" else None if label_style is not None: torch.nn.init.zeros_(label_style.weight) if label_style is not None and not a.freeze_text: opt_g.add_param_group({"params": label_style.parameters(), "lr": a.style_lr or a.lr}) sched_g = (torch.optim.lr_scheduler.CosineAnnealingLR(opt_g, T_max=a.steps, eta_min=a.lr * 0.1) if a.lr_decay else None) opt_d = torch.optim.AdamW(list(mpd.parameters()) + list(msd.parameters()), lr=a.lr_disc, betas=(0.8, 0.99), weight_decay=1e-2) mark_ids = torch.tensor([int(mark) for mark in a.mark_ids.split(",")], device=dev) dl = DataLoader(DistDataset(train_rows, a.bin, a.crop, pitch=a.pitch_weight > 0), batch_size=a.batch, shuffle=True, collate_fn=collate, num_workers=int(os.environ.get("WORKERS", 4)), drop_last=True, pin_memory=True) def style_rows(lens, langs): return torch.cat([PACKS.get(lg, pack)[max(min(int(l) - 3, pack.shape[0] - 1), 0)] for l, lg in zip(lens, langs)]).to(dev) def prefix_cuts(ids, lens): cuts, keeps = lens.clone(), lens.clone() for i in range(ids.shape[0]): spaces = (ids[i, :lens[i]] == SPACE).nonzero().flatten().tolist() if random.random() >= a.prefix_prob or len(spaces) <= a.lookahead_words + 1: continue k = random.randint(a.lookahead_words + 1, len(spaces) - 1) cuts[i], keeps[i] = spaces[k], spaces[k - a.lookahead_words] return cuts, keeps def prepare(b): ids = b["ids"].to(dev); lens = b["lens"].to(dev); durt = b["dur"].to(dev) B, L = ids.shape tmask = torch.arange(L, device=dev)[None].expand(B, -1) + 1 > lens[:, None] langs = b.get("lang", ["ar"] * B) ref = style_rows(lens, langs) if styler is not None: ref = ref + styler(b["audio"]) if label_style is not None: ref = ref + label_style(torch.tensor([labels.index(lg) + 1 if lg in labels else 0 for lg in langs], device=dev)) frames = durt.sum(dim=1).long() aln = build_aln(durt, lens, int(frames.max()), dev) return ids, lens, durt, tmask, ref[:, 128:], ref[:, :128], frames, aln def text_side(net, ids, lens, tmask, s_pred, aln): d_en = net.bert_encoder(net.bert(ids, attention_mask=(~tmask).int())).transpose(-1, -2) d = net.predictor.text_encoder(d_en, s_pred, lens, tmask) F0, N = net.predictor.F0Ntrain(d.transpose(-1, -2) @ aln, s_pred) feats = net.text_encoder(ids, lens, tmask) if hasattr(net, "feature_proj"): feats = net.feature_proj(feats) return d, F0, N, feats def duration_loss(d, durt, ignored): x, _ = model.predictor.lstm(d) dur_pred = torch.sigmoid(model.predictor.duration_proj(x)).sum(dim=-1) return (F.l1_loss(dur_pred, durt, reduction="none") * (~ignored)).sum() / (~ignored).sum() def text_losses(b): ids, lens, durt, tmask, s_pred, _, frames, aln = prepare(b) with torch.no_grad(): _, t_F0, t_N, t_feats = text_side(teacher, ids, lens, tmask, s_pred, aln) cuts, keeps = prefix_cuts(ids, lens) L = int(cuts.max()) positions = torch.arange(L, device=dev)[None] inside, scored = positions < cuts[:, None], positions < keeps[:, None] ids, durt = ids[:, :L] * inside, durt[:, :L] * inside T = int(durt.sum(1).max()) aln = aln[:, :L, :T] * inside[:, :, None] d, F0, N, feats = text_side(model, ids, cuts, ~inside, style_rows(cuts, b.get("lang", ["ar"] * len(cuts)))[:, 128:], aln) scored_frames = (torch.arange(2 * T, device=dev)[None] < 2 * (durt * scored).sum(1)[:, None]).float() masked = lambda loss, mask: (loss * mask).sum() / mask.sum() l_f0 = masked(F.smooth_l1_loss(F0, t_F0[:, :2 * T], reduction="none"), scored_frames) / 10 l_n = masked(F.smooth_l1_loss(N, t_N[:, :2 * T], reduction="none"), scored_frames) l_feat = masked(F.l1_loss(feats, t_feats[:, :, :L], reduction="none").mean(1), scored.float()) return duration_loss(d, durt, ~scored), l_f0, l_n, l_feat def run_batch(b, probe=False): with torch.no_grad() if a.freeze_text else contextlib.nullcontext(): ids, lens, durt, tmask, s_pred, s_dec, frames, aln = prepare(b) B = ids.shape[0] if a.mode == "decoder" or a.freeze_text: with torch.no_grad(): _, F0, N, feats = text_side(teacher if teacher is not None else model, ids, lens, tmask, s_pred, aln) l_dur = torch.zeros((), device=dev) else: d, F0, N, feats = text_side(model, ids, lens, tmask, s_pred, aln) l_dur = duration_loss(d, durt, tmask) asr = feats @ aln up = F0.shape[-1] / asr.shape[-1] if probe: print(f" asr {tuple(asr.shape)} F0 {tuple(F0.shape)} up={up:.2f} " f"frames={frames.tolist()[:4]}", flush=True) assert abs(up - round(up)) < 1e-6, f"F0/asr ratio not integral: {up}" up = round(up) # Random crop: the decoder is convolutional, so a window is enough and lets the batch # be rectangular. Audio, asr and F0/N are cropped on the SAME frame window. C = a.crop asr_c = torch.empty(B, asr.shape[1], C, device=dev) F0_c = torch.empty(B, C * up, device=dev); N_c = torch.empty(B, C * up, device=dev) y = torch.empty(B, 1, C * SPF, device=dev) pitch_c = torch.zeros(B, C * up, device=dev) for i in range(B): hi = max(int(frames[i]) - C, 0) marks = torch.isin(ids[i, :lens[i]], mark_ids).nonzero().flatten().tolist() if random.random() < a.mark_crop_prob else [] if marks: mark = random.choice(marks) centre = int(durt[i, :mark].sum() + durt[i, mark] / 2) f0 = min(max(centre - C // 2 + random.randint(-C // 4, C // 4), 0), hi) else: f0 = random.randint(0, hi) if a.pitch_weight > 0: clip = b["f0"][i][f0 * up:(f0 + C) * up] pitch_c[i, :clip.shape[0]] = clip.to(dev) asr_c[i] = asr[i, :, f0:f0 + C] F0_c[i] = F0[i, f0 * up:(f0 + C) * up] N_c[i] = N[i, f0 * up:(f0 + C) * up] w = b["audio"][i] seg = w[f0 * SPF:(f0 + C) * SPF] if seg.shape[0] < C * SPF: seg = F.pad(seg, (0, C * SPF - seg.shape[0])) y[i, 0] = seg.to(dev) l_pitch = torch.zeros((), device=dev) if a.pitch_weight > 0: model.decoder.generator.teacher_f0 = pitch_c y_hat = model.decoder(asr_c, F0_c, N_c, s_dec) if a.pitch_weight > 0: l_pitch = model.decoder.generator.pitch_loss(pitch_c) model.decoder.generator.teacher_f0 = None if y_hat.dim() == 2: y_hat = y_hat.unsqueeze(1) y_hat = y_hat[..., : y.shape[-1]] if probe: print(f" y {tuple(y.shape)} y_hat {tuple(y_hat.shape)} pitch {l_pitch.item():.3f} voiced {(pitch_c > 0).float().mean().item():.2f}", flush=True) return y, y_hat, l_dur, l_pitch if a.probe and a.mode != "text": b = next(iter(dl)) y, y_hat, l_dur, _ = run_batch(b, probe=True) print(f"[+] probe OK dur_L1 {l_dur.item():.3f} " f"stft {stft_loss(y_hat.squeeze(1), y.squeeze(1)).item():.3f}", flush=True) return step, best = 0, math.inf if a.init_from: c = torch.load(a.init_from, map_location="cpu", weights_only=False) src = c["model"] if "model" in c else c # KModel checkpoints are stored NESTED as {block: {name: tensor}}; the trainer's own # last.pt is flat. Flatten so --init-from accepts either, instead of silently matching # nothing and training from scratch. if src and all(isinstance(v, dict) for v in src.values()): src = {f"{b}.{k}": t for b, d in src.items() for k, t in d.items()} # Checkpoints saved under DataParallel/DDP carry a "module." prefix inside each block # (Nabra-82M does). Left in place it matches nothing and the model trains from scratch. if any(".module." in k for k in src): src = {k.replace(".module.", ".", 1): v for k, v in src.items()} if a.init_loose: tgt = model.state_dict() fit = {k: v for k, v in src.items() if k in tgt and tgt[k].shape == v.shape} skipped = [k for k in tgt if k not in fit] model.load_state_dict(fit, strict=False) n_fit = sum(tgt[k].numel() for k in fit) n_all = sum(v.numel() for v in tgt.values()) print(f"[*] loose init: {n_fit/1e6:.2f}M of {n_all/1e6:.2f}M " f"({100*n_fit/n_all:.1f}%) loaded, {len(skipped)} tensors left random", flush=True) for k in skipped: print(f" random: {k}", flush=True) else: model.load_state_dict(src) if "mpd" in c: mpd.load_state_dict(c["mpd"]); msd.load_state_dict(c["msd"]) if styler is not None and "styler" in c: styler.load_state_dict(c["styler"]) if label_style is not None and "label_style" in c: label_style.load_state_dict(c["label_style"]) print(f"[*] init weights (+discriminators) from {a.init_from}", flush=True) ck = out / "last.pt" if a.resume and ck.exists(): c = torch.load(ck, map_location="cpu", weights_only=False) model.load_state_dict(c["model"]); mpd.load_state_dict(c["mpd"]) msd.load_state_dict(c["msd"]); opt_g.load_state_dict(c["opt_g"]) opt_d.load_state_dict(c["opt_d"]); step = c["step"] if styler is not None and "styler" in c: styler.load_state_dict(c["styler"]) if label_style is not None and "label_style" in c: label_style.load_state_dict(c["label_style"]) if sched_g is not None: if "sched_g" in c: sched_g.load_state_dict(c["sched_g"]) else: # checkpoint predates the scheduler for _ in range(step): sched_g.step() print(f"[*] resumed at step {step}", flush=True) if a.mode == "text": model.train() step, acc, t0 = 0, {}, time.time() while step < a.steps: for b in dl: if step >= a.steps: break l_dur, l_f0, l_n, l_feat = text_losses(b) opt_g.zero_grad(set_to_none=True) (a.dur_weight * l_dur + l_f0 + l_n + a.feat_weight * l_feat).backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0) opt_g.step() if sched_g is not None: sched_g.step() step += 1 for k, v in (("dur", l_dur), ("f0", l_f0), ("n", l_n), ("feat", l_feat)): acc[k] = acc.get(k, 0.0) + float(v) if step % 100 == 0 or a.probe: print(f" step {step}/{a.steps} " + " ".join(f"{k} {v/(1 if a.probe else 100):.4f}" for k, v in acc.items()) + f" {step * a.batch / (time.time() - t0):.0f} clip/s", flush=True) acc = {} if a.probe and step >= 3: return if step % a.ckpt_every == 0 or step == a.steps: torch.save(model.state_dict(), out / "student.pth") print(f"[+] text side done at step {step}", flush=True) return t0 = time.time(); acc = {} model.train() if a.freeze_text: for part in (model.bert, model.bert_encoder, model.predictor, model.text_encoder, getattr(model, "feature_proj", None), styler, label_style): if part is not None: part.eval() while step < a.steps: for b in dl: if step >= a.steps: break y, y_hat, l_dur, l_pitch = run_batch(b) opt_d.zero_grad(set_to_none=True) d_l = disc_loss(y, y_hat.detach()) l_r1r2 = torch.zeros((), device=dev) if a.r1r2_weight > 0 and step % R1R2_EVERY == 0: real_wave = y.detach().requires_grad_(True) fake_wave = y_hat.detach().requires_grad_(True) for disc in (mpd, msd): real_out, fake_out = disc(real_wave, fake_wave)[:2] grads = torch.autograd.grad(sum(o.mean() for o in real_out + fake_out), [real_wave, fake_wave], create_graph=True) l_r1r2 = l_r1r2 + sum(g.pow(2).sum(dim=(1, 2)).mean() for g in grads) d_l = d_l + R1R2_EVERY * a.r1r2_weight / 2 * l_r1r2 d_l.backward(); torch.nn.utils.clip_grad_norm_( list(mpd.parameters()) + list(msd.parameters()), 10.0) opt_d.step() opt_g.zero_grad(set_to_none=True) l_stft = stft_loss(y_hat.squeeze(1), y.squeeze(1)) l_mel = (mel_loss(y_hat.squeeze(1), y.squeeze(1)) if a.mel_weight > 0 else torch.zeros((), device=dev)) l_slm = torch.zeros((), device=dev) if wl is not None: if a.slm_weight > 0: l_slm = wl(y.squeeze(1), y_hat.squeeze(1)) if a.slm_adv_weight > 0 and step >= a.disc_warmup: l_slm = l_slm + a.slm_adv_weight / max(a.slm_weight, 1e-9) * \ wl.generator(y_hat.squeeze(1)) l_sil = torch.zeros((), device=dev) if a.sil_weight > 0: H = a.sr // 100 # 10 ms n = y.shape[-1] // H yr = y[..., : n * H].reshape(y.shape[0], n, H) hr = y_hat[..., : n * H].reshape(y_hat.shape[0], n, H) yrms = (yr.pow(2).mean(-1) + 1e-12).sqrt() tsil = (yrms < a.sil_thresh).float() if tsil.sum() > 0: # RMS, not mean-square: the floor is ~5e-3, so squared it is ~3e-5 and # would need a weight of ~10,000 to matter beside an stft term near 1.9. # In RMS the weight reads directly as "how many dB do I care about silence". # # Hinged at the TARGET's own floor. Kokoro emits true digital silence # (1e-6), so against the teacher this is the original loss to 4 decimals. # Jenny is a real microphone: her quiet frames sit at 2.75e-4, and driving # the student below its own target's room tone is not a fix, it is a # different artefact (pumping between speech and dead air). hrms = (hr.pow(2).mean(-1) + 1e-12).sqrt() l_sil = (F.relu(hrms - yrms) * tsil).sum() / tsil.sum() # The discriminator is noise until it can tell real from fake; connecting it to the # generator immediately injects that noise into the weights. l_g = gen_loss(y, y_hat) if step >= a.disc_warmup else torch.zeros((), device=dev) loss = a.stft_weight * l_stft + a.dur_weight * l_dur + l_g \ + a.sil_weight * l_sil + a.mel_weight * l_mel \ + a.slm_weight * l_slm + a.pitch_weight * l_pitch loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0) if styler is not None: torch.nn.utils.clip_grad_norm_(styler.parameters(), 5.0) if label_style is not None: torch.nn.utils.clip_grad_norm_(label_style.parameters(), 5.0) opt_g.step() if sched_g is not None: sched_g.step() step += 1 for k, v in (("stft", l_stft), ("dur", l_dur), ("gen", l_g), ("disc", d_l), ("sil", l_sil), ("mel", l_mel), ("slm", l_slm), ("r1r2", l_r1r2), ("pitch", l_pitch)): acc[k] = acc.get(k, 0.0) + float(v) if step % 100 == 0: r = step * a.batch / (time.time() - t0) print(f" step {step}/{a.steps} " + " ".join(f"{k} {v/100:.3f}" for k, v in acc.items()) + f" {r:.0f} clip/s", flush=True) acc = {} if a.ckpt_every and step % a.ckpt_every == 0: tmp = out / "last.pt.tmp" torch.save({"model": model.state_dict(), "mpd": mpd.state_dict(), "msd": msd.state_dict(), "opt_g": opt_g.state_dict(), "opt_d": opt_d.state_dict(), "step": step, "config": cfg, **({"sched_g": sched_g.state_dict()} if sched_g else {}), **({"styler": styler.state_dict()} if styler is not None else {}), **({"label_style": label_style.state_dict()} if label_style is not None else {})}, tmp) os.replace(tmp, ck) torch.save(model.state_dict(), out / "student.pth") if styler is not None: torch.save(styler.state_dict(), out / "styler.pth") if label_style is not None: torch.save({label: label_style.weight[i + 1].detach().cpu() for i, label in enumerate(labels)}, out / "label_style.pt") if step % 4000 == 0: torch.save(model.state_dict(), out / f"student_{step}.pth") print(f"[+] done at step {step}", flush=True) if __name__ == "__main__": main()