prosatts-work / scripts /pttrain14 /train_student.py
Checkpoints, banco de risos e scripts
63a4356 verified about 14 hours ago
35.8 kB
"""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()