Remove obsolete package file train_hifigan_oracle_v1.py

2026-06-17 23:54:46 +00:00
parent 7d0adbcad9
commit 0bb3fe1dda
1 changed files with 0 additions and 830 deletions
@@ -1,830 +0,0 @@
 from __future__ import annotations
 import argparse
 import json
 import math
 import random
 import time
 from dataclasses import asdict, dataclass
 from pathlib import Path
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import torchaudio
 from torch.nn.utils import remove_weight_norm, spectral_norm, weight_norm
 from torch.utils.data import DataLoader, Dataset
@dataclass(frozen=True)
 class HifiGanConfig:
    variant: str
    sample_rate: int = 24000
    n_fft: int = 1024
    hop_size: int = 256
    win_size: int = 1024
    num_mels: int = 80
    fmin: float = 0.0
    fmax: float = 12000.0
    resblock: str = "1"
    upsample_rates: tuple[int, ...] = (8, 8, 2, 2)
    upsample_kernel_sizes: tuple[int, ...] = (16, 16, 4, 4)
    upsample_initial_channel: int = 128
    resblock_kernel_sizes: tuple[int, ...] = (3, 7, 11)
    resblock_dilation_sizes: tuple[tuple[int, ...], ...] = ((1, 3, 5), (1, 3, 5), (1, 3, 5))
    activation: str = "lrelu"
    conditioning_channels: int = 0
 def make_config(variant: str) -> HifiGanConfig:
    if variant == "v2":
        return HifiGanConfig(variant="v2")
    if variant == "v2plus":
        return HifiGanConfig(variant="v2plus", upsample_initial_channel=160)
    if variant == "v2wide":
        return HifiGanConfig(variant="v2wide", upsample_initial_channel=176)
    if variant == "snake_v2mid":
        return HifiGanConfig(variant="snake_v2mid", upsample_initial_channel=144, activation="snake")
    if variant == "snake_v2balanced":
        return HifiGanConfig(variant="snake_v2balanced", upsample_initial_channel=160, activation="snake")
    if variant == "source_snake_v2balanced":
        return HifiGanConfig(
            variant="source_snake_v2balanced",
            upsample_initial_channel=160,
            activation="snake",
            conditioning_channels=5,
        )
    if variant == "v3":
        return HifiGanConfig(
            variant="v3",
            resblock="2",
            upsample_rates=(8, 8, 4),
            upsample_kernel_sizes=(16, 16, 8),
            upsample_initial_channel=256,
            resblock_kernel_sizes=(3, 5, 7),
            resblock_dilation_sizes=((1, 2), (2, 6), (3, 12)),
        )
    raise ValueError(f"Unknown variant: {variant}")
 def get_padding(kernel_size: int, dilation: int = 1) -> int:
    return int((kernel_size * dilation - dilation) / 2)
 class SnakeActivation(nn.Module):
    def __init__(self, channels: int):
        super().__init__()
        self.log_alpha = nn.Parameter(torch.zeros(1, channels, 1))
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        alpha = self.log_alpha.exp().clamp(1e-4, 100.0)
        return x + torch.sin(alpha * x).pow(2) / alpha
 def make_activation(channels: int, activation: str) -> nn.Module:
    if activation == "snake":
        return SnakeActivation(channels)
    return nn.LeakyReLU(0.1)
 class ResBlock1(nn.Module):
    def __init__(self, channels: int, kernel_size: int, dilations: tuple[int, ...], activation: str = "lrelu"):
        super().__init__()
        self.convs1 = nn.ModuleList(
            [
                weight_norm(
                    nn.Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=d,
                        padding=get_padding(kernel_size, d),
                    )
                )
                for d in dilations
            ]
        )
        self.convs2 = nn.ModuleList(
            [
                weight_norm(
                    nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1))
                )
                for _ in dilations
            ]
        )
        self.acts1 = nn.ModuleList([make_activation(channels, activation) for _ in dilations])
        self.acts2 = nn.ModuleList([make_activation(channels, activation) for _ in dilations])
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for c1, c2, a1, a2 in zip(self.convs1, self.convs2, self.acts1, self.acts2):
            y = a1(x)
            y = c1(y)
            y = a2(y)
            y = c2(y)
            x = x + y
        return x
    def remove_weight_norm(self) -> None:
        for layer in list(self.convs1) + list(self.convs2):
            remove_weight_norm(layer)
 class ResBlock2(nn.Module):
    def __init__(self, channels: int, kernel_size: int, dilations: tuple[int, ...], activation: str = "lrelu"):
        super().__init__()
        self.convs = nn.ModuleList(
            [
                weight_norm(
                    nn.Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=d,
                        padding=get_padding(kernel_size, d),
                    )
                )
                for d in dilations
            ]
        )
        self.acts = nn.ModuleList([make_activation(channels, activation) for _ in dilations])
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for conv, act in zip(self.convs, self.acts):
            y = act(x)
            y = conv(y)
            x = x + y
        return x
    def remove_weight_norm(self) -> None:
        for layer in self.convs:
            remove_weight_norm(layer)
 class HifiGanGenerator(nn.Module):
    def __init__(self, cfg: HifiGanConfig):
        super().__init__()
        self.cfg = cfg
        self.num_kernels = len(cfg.resblock_kernel_sizes)
        self.num_upsamples = len(cfg.upsample_rates)
        self.conv_pre = weight_norm(
            nn.Conv1d(cfg.num_mels + cfg.conditioning_channels, cfg.upsample_initial_channel, 7, 1, padding=3)
        )
        self.ups = nn.ModuleList()
        self.up_acts = nn.ModuleList()
        self.resblocks = nn.ModuleList()
        resblock_cls = ResBlock1 if cfg.resblock == "1" else ResBlock2
        for i, (rate, kernel) in enumerate(zip(cfg.upsample_rates, cfg.upsample_kernel_sizes)):
            in_ch = cfg.upsample_initial_channel // (2**i)
            out_ch = cfg.upsample_initial_channel // (2 ** (i + 1))
            self.up_acts.append(make_activation(in_ch, cfg.activation))
            self.ups.append(
                weight_norm(
                    nn.ConvTranspose1d(
                        in_ch,
                        out_ch,
                        kernel,
                        rate,
                        padding=(kernel - rate) // 2,
                    )
                )
            )
            for k, d in zip(cfg.resblock_kernel_sizes, cfg.resblock_dilation_sizes):
                self.resblocks.append(resblock_cls(out_ch, k, d, cfg.activation))
        final_ch = cfg.upsample_initial_channel // (2 ** len(cfg.upsample_rates))
        self.post_act = make_activation(final_ch, cfg.activation)
        self.conv_post = weight_norm(nn.Conv1d(final_ch, 1, 7, 1, padding=3))
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv_pre(x)
        for i, up in enumerate(self.ups):
            x = self.up_acts[i](x)
            x = up(x)
            xs = 0.0
            for j in range(self.num_kernels):
                xs = xs + self.resblocks[i * self.num_kernels + j](x)
            x = xs / self.num_kernels
        x = self.post_act(x)
        x = self.conv_post(x)
        return torch.tanh(x)
    def remove_weight_norm(self) -> None:
        remove_weight_norm(self.conv_pre)
        for up in self.ups:
            remove_weight_norm(up)
        for block in self.resblocks:
            block.remove_weight_norm()
        remove_weight_norm(self.conv_post)
 def extract_source_features(
    wav: torch.Tensor,
    cfg: HifiGanConfig,
    frames: int,
    dropout: float = 0.0,
    noise: float = 0.0,
 ) -> torch.Tensor:
    """Return low-rate F0/voicing features for source-conditioned generators."""
    pitch = torchaudio.functional.detect_pitch_frequency(
        wav.detach().cpu(),
        sample_rate=cfg.sample_rate,
        frame_time=cfg.hop_size / cfg.sample_rate,
        win_length=30,
    ).to(wav.device)
    if pitch.ndim == 1:
        pitch = pitch.unsqueeze(0)
    if pitch.shape[-1] < frames:
        pitch = F.pad(pitch, (0, frames - pitch.shape[-1]), value=0.0)
    pitch = pitch[..., :frames]
    voiced = ((pitch >= 55.0) & (pitch <= 420.0)).float()
    pitch = pitch.clamp(55.0, 420.0)
    log_f0 = ((torch.log(pitch) - math.log(140.0)) / 0.45).clamp(-3.0, 3.0) * voiced
    if noise > 0.0:
        log_f0 = (log_f0 + torch.randn_like(log_f0) * noise * voiced).clamp(-3.0, 3.0)
    jump = F.pad((log_f0[..., 1:] - log_f0[..., :-1]).abs(), (1, 0))
    confidence = torch.exp(-1.5 * jump) * voiced
    reconstructed_f0 = torch.exp(log_f0 * 0.45 + math.log(140.0))
    phase = torch.cumsum(2.0 * math.pi * reconstructed_f0 * (cfg.hop_size / cfg.sample_rate), dim=-1)
    source = torch.stack(
        [log_f0, voiced, confidence, torch.sin(phase) * confidence, torch.cos(phase) * confidence],
        dim=1,
    )
    if dropout > 0.0:
        # Drop the complete source sketch for some examples so inference remains
        # stable when predicted F0 confidence is poor.
        keep = (torch.rand(source.shape[0], 1, 1, device=source.device) >= dropout).to(source.dtype)
        source = source * keep
    return source
 class DiscriminatorP(nn.Module):
    def __init__(self, period: int):
        super().__init__()
        self.period = period
        self.convs = nn.ModuleList(
            [
                weight_norm(nn.Conv2d(1, 32, (5, 1), (3, 1), padding=(2, 0))),
                weight_norm(nn.Conv2d(32, 128, (5, 1), (3, 1), padding=(2, 0))),
                weight_norm(nn.Conv2d(128, 512, (5, 1), (3, 1), padding=(2, 0))),
                weight_norm(nn.Conv2d(512, 1024, (5, 1), (3, 1), padding=(2, 0))),
                weight_norm(nn.Conv2d(1024, 1024, (5, 1), 1, padding=(2, 0))),
            ]
        )
        self.conv_post = weight_norm(nn.Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]:
        fmap = []
        b, c, t = x.shape
        if t % self.period != 0:
            x = F.pad(x, (0, self.period - (t % self.period)), mode="reflect")
            t = x.shape[-1]
        x = x.view(b, c, t // self.period, self.period)
        for conv in self.convs:
            x = F.leaky_relu(conv(x), 0.1)
            fmap.append(x)
        x = self.conv_post(x)
        fmap.append(x)
        return torch.flatten(x, 1, -1), fmap
 class MultiPeriodDiscriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.discriminators = nn.ModuleList([DiscriminatorP(p) for p in (2, 3, 5, 7, 11)])
    def forward(self, y: torch.Tensor, y_hat: torch.Tensor):
        y_d_rs, y_d_gs, fmap_rs, fmap_gs = [], [], [], []
        for d in self.discriminators:
            y_d_r, fmap_r = d(y)
            y_d_g, fmap_g = d(y_hat)
            y_d_rs.append(y_d_r)
            y_d_gs.append(y_d_g)
            fmap_rs.append(fmap_r)
            fmap_gs.append(fmap_g)
        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
 class DiscriminatorS(nn.Module):
    def __init__(self, use_spectral_norm: bool = False):
        super().__init__()
        norm = spectral_norm if use_spectral_norm else weight_norm
        self.convs = nn.ModuleList(
            [
                norm(nn.Conv1d(1, 128, 15, 1, padding=7)),
                norm(nn.Conv1d(128, 128, 41, 2, groups=4, padding=20)),
                norm(nn.Conv1d(128, 256, 41, 2, groups=16, padding=20)),
                norm(nn.Conv1d(256, 512, 41, 4, groups=16, padding=20)),
                norm(nn.Conv1d(512, 1024, 41, 4, groups=16, padding=20)),
                norm(nn.Conv1d(1024, 1024, 41, 1, groups=16, padding=20)),
                norm(nn.Conv1d(1024, 1024, 5, 1, padding=2)),
            ]
        )
        self.conv_post = norm(nn.Conv1d(1024, 1, 3, 1, padding=1))
    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]:
        fmap = []
        for conv in self.convs:
            x = F.leaky_relu(conv(x), 0.1)
            fmap.append(x)
        x = self.conv_post(x)
        fmap.append(x)
        return torch.flatten(x, 1, -1), fmap
 class MultiScaleDiscriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.discriminators = nn.ModuleList([DiscriminatorS(True), DiscriminatorS(), DiscriminatorS()])
        self.meanpools = nn.ModuleList([nn.AvgPool1d(4, 2, padding=2), nn.AvgPool1d(4, 2, padding=2)])
    def forward(self, y: torch.Tensor, y_hat: torch.Tensor):
        y_d_rs, y_d_gs, fmap_rs, fmap_gs = [], [], [], []
        for i, d in enumerate(self.discriminators):
            if i:
                y = self.meanpools[i - 1](y)
                y_hat = self.meanpools[i - 1](y_hat)
            y_d_r, fmap_r = d(y)
            y_d_g, fmap_g = d(y_hat)
            y_d_rs.append(y_d_r)
            y_d_gs.append(y_d_g)
            fmap_rs.append(fmap_r)
            fmap_gs.append(fmap_g)
        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
 class SpectrogramDiscriminator(nn.Module):
    def __init__(self):
        super().__init__()
        channels = (32, 64, 128, 128)
        layers: list[nn.Module] = []
        in_ch = 1
        for out_ch, stride in zip(channels, ((1, 2), (2, 2), (2, 2), (2, 1))):
            layers.append(weight_norm(nn.Conv2d(in_ch, out_ch, (5, 5), stride=stride, padding=(2, 2))))
            in_ch = out_ch
        self.convs = nn.ModuleList(layers)
        self.conv_post = weight_norm(nn.Conv2d(in_ch, 1, (3, 3), padding=(1, 1)))
    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]:
        fmap = []
        for conv in self.convs:
            x = F.leaky_relu(conv(x), 0.1)
            fmap.append(x)
        x = self.conv_post(x)
        fmap.append(x)
        return torch.flatten(x, 1, -1), fmap
 class MultiResolutionSpectrogramDiscriminator(nn.Module):
    def __init__(self, fft_sizes: tuple[int, ...] = (256, 512, 1024), hop_sizes: tuple[int, ...] = (64, 128, 256), win_lengths: tuple[int, ...] = (256, 512, 1024)):
        super().__init__()
        self.fft_sizes = fft_sizes
        self.hop_sizes = hop_sizes
        self.win_lengths = win_lengths
        self.discriminators = nn.ModuleList([SpectrogramDiscriminator() for _ in fft_sizes])
    def _features(self, wav: torch.Tensor, fft: int, hop: int, win_len: int) -> torch.Tensor:
        wav = wav.squeeze(1)
        window = torch.hann_window(win_len, device=wav.device)
        spec = torch.stft(wav, n_fft=fft, hop_length=hop, win_length=win_len, window=window, return_complex=True)
        mag = torch.log(spec.abs().clamp_min(1e-5))
        mean = mag.mean(dim=(1, 2), keepdim=True)
        std = mag.std(dim=(1, 2), keepdim=True).clamp_min(1e-4)
        return ((mag - mean) / std).unsqueeze(1)
    def forward(self, y: torch.Tensor, y_hat: torch.Tensor):
        y_d_rs, y_d_gs, fmap_rs, fmap_gs = [], [], [], []
        for disc, fft, hop, win_len in zip(self.discriminators, self.fft_sizes, self.hop_sizes, self.win_lengths):
            y_feat = self._features(y, fft, hop, win_len)
            y_hat_feat = self._features(y_hat, fft, hop, win_len)
            y_d_r, fmap_r = disc(y_feat)
            y_d_g, fmap_g = disc(y_hat_feat)
            y_d_rs.append(y_d_r)
            y_d_gs.append(y_d_g)
            fmap_rs.append(fmap_r)
            fmap_gs.append(fmap_g)
        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
 class MelFrontend(nn.Module):
    def __init__(self, cfg: HifiGanConfig):
        super().__init__()
        self.mel = torchaudio.transforms.MelSpectrogram(
            sample_rate=cfg.sample_rate,
            n_fft=cfg.n_fft,
            win_length=cfg.win_size,
            hop_length=cfg.hop_size,
            f_min=cfg.fmin,
            f_max=cfg.fmax,
            n_mels=cfg.num_mels,
            power=1.0,
            center=True,
            norm="slaney",
            mel_scale="slaney",
        )
    def forward(self, wav: torch.Tensor) -> torch.Tensor:
        return torch.log(torch.clamp(self.mel(wav), min=1e-5))
 def load_rows(path: Path, max_rows: int, min_seconds: float, max_seconds: float) -> list[dict]:
    rows = []
    with path.open("r", encoding="utf-8-sig") as f:
        for line in f:
            if not line.strip():
                continue
            row = json.loads(line)
            audio = Path(str(row.get("target_audio") or ""))
            text = str(row.get("target_text") or "").strip()
            dur = float(row.get("target_duration_s") or 0.0)
            if audio.is_file() and text and min_seconds <= (dur or 4.0) <= max_seconds:
                rows.append({"audio": str(audio), "text": text, "duration": dur})
                if max_rows > 0 and len(rows) >= max_rows:
                    break
    if not rows:
        raise RuntimeError(f"No rows loaded from {path}")
    return rows
 def load_audio(path: str, sample_rate: int) -> torch.Tensor:
    wav, sr = torchaudio.load(path)
    if wav.shape[0] > 1:
        wav = wav.mean(dim=0, keepdim=True)
    if sr != sample_rate:
        wav = torchaudio.functional.resample(wav, sr, sample_rate)
    wav = wav.squeeze(0)
    return wav.clamp(-1, 1)
 class AudioDataset(Dataset):
    def __init__(self, rows: list[dict], cfg: HifiGanConfig, segment_size: int, seed: int):
        self.rows = rows
        self.cfg = cfg
        self.segment_size = segment_size
        self.rng = random.Random(seed)
    def __len__(self) -> int:
        return len(self.rows)
    def __getitem__(self, idx: int) -> torch.Tensor:
        wav = load_audio(self.rows[idx]["audio"], self.cfg.sample_rate)
        if wav.numel() >= self.segment_size:
            start = self.rng.randint(0, wav.numel() - self.segment_size)
            return wav[start : start + self.segment_size]
        return F.pad(wav, (0, self.segment_size - wav.numel()))
 def feature_loss(fmap_r, fmap_g) -> torch.Tensor:
    loss = 0.0
    for dr, dg in zip(fmap_r, fmap_g):
        for rl, gl in zip(dr, dg):
            loss = loss + F.l1_loss(rl.detach(), gl)
    return loss * 2
 def discriminator_loss(disc_real_outputs, disc_generated_outputs) -> torch.Tensor:
    loss = 0.0
    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
        loss = loss + torch.mean((1 - dr) ** 2) + torch.mean(dg**2)
    return loss
 def generator_loss(disc_outputs) -> torch.Tensor:
    loss = 0.0
    for dg in disc_outputs:
        loss = loss + torch.mean((1 - dg) ** 2)
    return loss
 def stft_mag_loss(y_hat: torch.Tensor, y: torch.Tensor, fft_sizes: tuple[int, ...], hop_sizes: tuple[int, ...], win_lengths: tuple[int, ...]) -> torch.Tensor:
    # Multi-resolution spectral loss catches buzz/shimmer that can hide behind
    # mel loss, especially for a small generator near convergence.
    y_hat = y_hat.squeeze(1)
    y = y.squeeze(1)
    total = torch.zeros((), device=y.device)
    for fft, hop, win_len in zip(fft_sizes, hop_sizes, win_lengths):
        window = torch.hann_window(win_len, device=y.device)
        pred = torch.stft(y_hat, n_fft=fft, hop_length=hop, win_length=win_len, window=window, return_complex=True)
        target = torch.stft(y, n_fft=fft, hop_length=hop, win_length=win_len, window=window, return_complex=True)
        pred_mag = pred.abs().clamp_min(1e-7)
        target_mag = target.abs().clamp_min(1e-7)
        sc = torch.linalg.vector_norm(target_mag - pred_mag) / torch.linalg.vector_norm(target_mag).clamp_min(1e-7)
        log_mag = F.l1_loss(torch.log(pred_mag), torch.log(target_mag))
        total = total + sc + log_mag
    return total / max(1, len(fft_sizes))
 def count_parameters(module: nn.Module) -> int:
    return sum(p.numel() for p in module.parameters())
 def jsonable_args(args: argparse.Namespace) -> dict:
    return {k: str(v) if isinstance(v, Path) else v for k, v in vars(args).items()}
 def save_checkpoint(
    path: Path,
    generator: nn.Module,
    mpd: nn.Module,
    msd: nn.Module,
    optim_g,
    optim_d,
    cfg: HifiGanConfig,
    step: int,
    args,
    mrsd: nn.Module | None = None,
 ) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    payload = {
            "generator": generator.state_dict(),
            "mpd": mpd.state_dict(),
            "msd": msd.state_dict(),
            "optim_g": optim_g.state_dict(),
            "optim_d": optim_d.state_dict(),
            "config": asdict(cfg),
            "step": step,
            "args": jsonable_args(args),
            "generator_params": count_parameters(generator),
        }
    if mrsd is not None:
        payload["mrsd"] = mrsd.state_dict()
    torch.save(payload, tmp)
    tmp.replace(path)
 def checkpoint_step(path: Path) -> int:
    stem = path.stem
    tail = stem.rsplit("-", 1)[-1]
    return int(tail) if tail.isdigit() else -1
 def prune_checkpoints(out_dir: Path, variant: str, keep: int) -> None:
    if keep <= 0:
        return
    numbered = [p for p in out_dir.glob(f"hifigan-{variant}-*.pt") if checkpoint_step(p) >= 0]
    numbered.sort(key=checkpoint_step, reverse=True)
    for old in numbered[keep:]:
        old.unlink(missing_ok=True)
 def latest_checkpoint(out_dir: Path) -> Path | None:
    numbered = [p for p in out_dir.glob("hifigan-*-*.pt") if checkpoint_step(p) >= 0]
    if numbered:
        return max(numbered, key=checkpoint_step)
    ckpts = sorted(out_dir.glob("hifigan-*-latest.pt"), key=lambda p: p.stat().st_mtime, reverse=True)
    return ckpts[0] if ckpts else None
 def partial_load_state(module: nn.Module, state: dict[str, torch.Tensor]) -> tuple[int, int]:
    current = module.state_dict()
    patched: dict[str, torch.Tensor] = {}
    copied = 0
    skipped = 0
    for name, target in current.items():
        source = state.get(name)
        if source is None:
            skipped += 1
            continue
        if source.shape == target.shape:
            patched[name] = source
            copied += 1
            continue
        if source.ndim != target.ndim:
            skipped += 1
            continue
        value = target.clone()
        slices = tuple(slice(0, min(a, b)) for a, b in zip(target.shape, source.shape))
        value[slices] = source[slices].to(value.device, value.dtype)
        patched[name] = value
        copied += 1
    module.load_state_dict(patched, strict=False)
    return copied, skipped
 def train(args: argparse.Namespace) -> None:
    torch.backends.cudnn.benchmark = True
    cfg = make_config(args.variant)
    device = torch.device(args.device)
    rows = load_rows(args.train_jsonl, args.max_rows, args.min_seconds, args.max_seconds)
    rng = random.Random(args.seed)
    rng.shuffle(rows)
    dataset = AudioDataset(rows, cfg, args.segment_size, args.seed)
    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=args.num_workers)
    mel_frontend = MelFrontend(cfg).to(device)
    generator = HifiGanGenerator(cfg).to(device)
    mpd = MultiPeriodDiscriminator().to(device)
    msd = MultiScaleDiscriminator().to(device)
    mrsd = MultiResolutionSpectrogramDiscriminator().to(device) if args.spec_disc_weight > 0.0 else None
    optim_g = torch.optim.AdamW(generator.parameters(), lr=args.lr, betas=(0.8, 0.99))
    disc_params = list(mpd.parameters()) + list(msd.parameters())
    if mrsd is not None:
        disc_params += list(mrsd.parameters())
    optim_d = torch.optim.AdamW(disc_params, lr=args.lr, betas=(0.8, 0.99))
    start_step = 0
    if args.init_checkpoint and not args.resume:
        ckpt = torch.load(args.init_checkpoint, map_location=device, weights_only=False)
        if args.partial_init:
            copied, skipped = partial_load_state(generator, ckpt["generator"])
            print(f"Partially initialized generator from {args.init_checkpoint}: copied={copied} skipped={skipped}")
        else:
            generator.load_state_dict(ckpt["generator"])
        if "mpd" in ckpt and "msd" in ckpt:
            mpd.load_state_dict(ckpt["mpd"])
            msd.load_state_dict(ckpt["msd"])
        if mrsd is not None and "mrsd" in ckpt:
            mrsd.load_state_dict(ckpt["mrsd"])
        can_load_disc_optim = mrsd is None or "mrsd" in ckpt
        if not args.partial_init and not args.reset_optim and "optim_g" in ckpt:
            optim_g.load_state_dict(ckpt["optim_g"])
        if not args.partial_init and not args.reset_optim and can_load_disc_optim and "optim_d" in ckpt:
            optim_d.load_state_dict(ckpt["optim_d"])
        for group in optim_g.param_groups:
            group["lr"] = args.lr
        for group in optim_d.param_groups:
            group["lr"] = args.lr
        start_step = int(ckpt.get("step") or 0)
        print(f"Initialized {args.init_checkpoint} at step {start_step}; lr={args.lr:g}")
    if args.resume:
        ckpt_path = latest_checkpoint(args.out_dir)
        if ckpt_path:
            ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
            generator.load_state_dict(ckpt["generator"])
            mpd.load_state_dict(ckpt["mpd"])
            msd.load_state_dict(ckpt["msd"])
            if mrsd is not None and "mrsd" in ckpt:
                mrsd.load_state_dict(ckpt["mrsd"])
            optim_g.load_state_dict(ckpt["optim_g"])
            optim_d.load_state_dict(ckpt["optim_d"])
            for group in optim_g.param_groups:
                group["lr"] = args.lr
            for group in optim_d.param_groups:
                group["lr"] = args.lr
            start_step = int(ckpt.get("step") or 0)
            print(f"Resumed {ckpt_path} at step {start_step}; lr={args.lr:g}")
    args.out_dir.mkdir(parents=True, exist_ok=True)
    prune_checkpoints(args.out_dir, args.variant, args.keep_checkpoints)
    (args.out_dir / "config.json").write_text(
        json.dumps(
            {
                "config": asdict(cfg),
                "args": jsonable_args(args),
                "rows": len(rows),
                "generator_params": count_parameters(generator),
                "mpd_params": count_parameters(mpd),
                "msd_params": count_parameters(msd),
                "mrsd_params": count_parameters(mrsd) if mrsd is not None else 0,
            },
            indent=2,
        ),
        encoding="utf-8",
    )
    print(f"Variant: {args.variant}")
    print(f"Rows: {len(rows)}")
    print(f"Generator params: {count_parameters(generator):,} ({count_parameters(generator)/1_000_000:.3f}M)")
    print(f"MPD params: {count_parameters(mpd):,} MSD params: {count_parameters(msd):,} (training only)")
    if mrsd is not None:
        print(f"MRSD params: {count_parameters(mrsd):,} (training only)")
    if args.steps == 0:
        return
    step = start_step
    started = time.time()
    try:
        while step < args.steps:
            for wav in loader:
                step += 1
                y = wav.unsqueeze(1).to(device)
                with torch.no_grad():
                    mel = mel_frontend(wav.to(device))
                    if cfg.conditioning_channels:
                        source = extract_source_features(
                            wav.to(device),
                            cfg,
                            mel.shape[-1],
                            dropout=args.source_dropout,
                            noise=args.source_noise,
                        )
                        generator_input = torch.cat([mel, source], dim=1)
                    else:
                        generator_input = mel
                y_hat = generator(generator_input)
                common = min(y.shape[-1], y_hat.shape[-1])
                y = y[..., :common]
                y_hat = y_hat[..., :common]
                y_mel = mel_frontend(y.squeeze(1))
                y_hat_mel = mel_frontend(y_hat.squeeze(1))
                optim_d.zero_grad(set_to_none=True)
                y_df_hat_r, y_df_hat_g, _, _ = mpd(y, y_hat.detach())
                y_ds_hat_r, y_ds_hat_g, _, _ = msd(y, y_hat.detach())
                loss_disc = discriminator_loss(y_df_hat_r, y_df_hat_g) + discriminator_loss(y_ds_hat_r, y_ds_hat_g)
                loss_spec_disc = torch.zeros((), device=device)
                if mrsd is not None:
                    y_dm_hat_r, y_dm_hat_g, _, _ = mrsd(y, y_hat.detach())
                    loss_spec_disc = discriminator_loss(y_dm_hat_r, y_dm_hat_g)
                    loss_disc = loss_disc + args.spec_disc_weight * loss_spec_disc
                loss_disc.backward()
                torch.nn.utils.clip_grad_norm_(disc_params, args.grad_clip)
                optim_d.step()
                optim_g.zero_grad(set_to_none=True)
                mel_loss = F.l1_loss(y_mel, y_hat_mel) * args.mel_weight
                y_df_hat_r, y_df_hat_g, fmap_f_r, fmap_f_g = mpd(y, y_hat)
                y_ds_hat_r, y_ds_hat_g, fmap_s_r, fmap_s_g = msd(y, y_hat)
                loss_fm = feature_loss(fmap_f_r, fmap_f_g) + feature_loss(fmap_s_r, fmap_s_g)
                loss_gen = generator_loss(y_df_hat_g) + generator_loss(y_ds_hat_g)
                loss_spec_gen = torch.zeros((), device=device)
                loss_spec_fm = torch.zeros((), device=device)
                if mrsd is not None:
                    y_dm_hat_r, y_dm_hat_g, fmap_m_r, fmap_m_g = mrsd(y, y_hat)
                    loss_spec_gen = generator_loss(y_dm_hat_g)
                    loss_spec_fm = feature_loss(fmap_m_r, fmap_m_g)
                wav_l1 = F.l1_loss(y_hat, y) * args.wav_weight
                stft_loss = torch.zeros((), device=device)
                if args.stft_weight > 0.0:
                    stft_loss = stft_mag_loss(y_hat, y, (512, 1024, 2048), (128, 256, 512), (512, 1024, 2048)) * args.stft_weight
                loss_g = (
                    mel_loss
                    + args.fm_weight * loss_fm
                    + args.adv_weight * loss_gen
                    + wav_l1
                    + stft_loss
                    + args.spec_disc_weight * loss_spec_gen
                    + args.spec_fm_weight * loss_spec_fm
                )
                loss_g.backward()
                grad_g = torch.nn.utils.clip_grad_norm_(generator.parameters(), args.grad_clip)
                optim_g.step()
                if step == 1 or step % args.log_interval == 0:
                    elapsed = max(time.time() - started, 1e-6)
                    speed = (step - start_step) / elapsed
                    eta = (args.steps - step) / max(speed, 1e-6)
                    print(
                        f"step={step}/{args.steps} g={loss_g.item():.4f} d={loss_disc.item():.4f} "
                        f"mel={mel_loss.item():.4f} fm={loss_fm.item():.4f} adv={loss_gen.item():.4f} "
                        f"wav={wav_l1.item():.4f} stft={stft_loss.item():.4f} "
                        f"sd={loss_spec_disc.item():.4f} sfm={loss_spec_fm.item():.4f} sadv={loss_spec_gen.item():.4f} "
                        f"grad={float(grad_g):.3f} speed={speed:.3f} step/s eta={eta/60:.1f}m",
                        flush=True,
                    )
                if step % args.save_interval == 0 or step >= args.steps:
                    prune_checkpoints(args.out_dir, args.variant, max(args.keep_checkpoints - 1, 0))
                    save_checkpoint(args.out_dir / f"hifigan-{args.variant}-{step}.pt", generator, mpd, msd, optim_g, optim_d, cfg, step, args, mrsd)
                    save_checkpoint(args.out_dir / f"hifigan-{args.variant}-latest.pt", generator, mpd, msd, optim_g, optim_d, cfg, step, args, mrsd)
                if step >= args.steps:
                    break
    except KeyboardInterrupt:
        if step > start_step:
            save_checkpoint(args.out_dir / f"hifigan-{args.variant}-interrupt-{step}.pt", generator, mpd, msd, optim_g, optim_d, cfg, step, args, mrsd)
            save_checkpoint(args.out_dir / f"hifigan-{args.variant}-latest.pt", generator, mpd, msd, optim_g, optim_d, cfg, step, args, mrsd)
            print(f"Interrupted. Saved checkpoint at step {step}.", flush=True)
        raise
    save_checkpoint(args.out_dir / f"hifigan-{args.variant}-final.pt", generator, mpd, msd, optim_g, optim_d, cfg, step, args, mrsd)
    print(f"Done. {args.out_dir}")
 def main() -> None:
    ap = argparse.ArgumentParser(description="Train exact-ish HiFi-GAN V2/V3 oracle vocoders on corrected Mark audio.")
    ap.add_argument("--train-jsonl", type=Path, required=True)
    ap.add_argument("--out-dir", type=Path, required=True)
    ap.add_argument(
        "--variant",
        choices=["v2", "v2plus", "v2wide", "snake_v2mid", "snake_v2balanced", "source_snake_v2balanced", "v3"],
        required=True,
    )
    ap.add_argument("--steps", type=int, default=5000)
    ap.add_argument("--max-rows", type=int, default=0)
    ap.add_argument("--min-seconds", type=float, default=1.0)
    ap.add_argument("--max-seconds", type=float, default=12.0)
    ap.add_argument("--segment-size", type=int, default=8192)
    ap.add_argument("--batch-size", type=int, default=8)
    ap.add_argument("--num-workers", type=int, default=0)
    ap.add_argument("--lr", type=float, default=2.0e-4)
    ap.add_argument("--mel-weight", type=float, default=45.0)
    ap.add_argument("--wav-weight", type=float, default=1.0)
    ap.add_argument("--fm-weight", type=float, default=1.0)
    ap.add_argument("--adv-weight", type=float, default=1.0)
    ap.add_argument("--stft-weight", type=float, default=0.0)
    ap.add_argument("--spec-disc-weight", type=float, default=0.0, help="Training-only multi-resolution spectrogram adversarial weight.")
    ap.add_argument("--spec-fm-weight", type=float, default=0.0, help="Training-only spectrogram discriminator feature-matching weight.")
    ap.add_argument("--source-dropout", type=float, default=0.0, help="Probability of dropping source conditioning per training example.")
    ap.add_argument("--source-noise", type=float, default=0.0, help="Stddev of normalized log-F0 corruption for source conditioning.")
    ap.add_argument("--grad-clip", type=float, default=1000.0)
    ap.add_argument("--log-interval", type=int, default=50)
    ap.add_argument("--save-interval", type=int, default=1000)
    ap.add_argument("--keep-checkpoints", type=int, default=12)
    ap.add_argument("--seed", type=int, default=1234)
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--resume", action="store_true")
    ap.add_argument("--init-checkpoint", type=Path)
    ap.add_argument("--partial-init", action="store_true", help="Slice-copy compatible generator weights from init-checkpoint into a resized generator.")
    ap.add_argument("--reset-optim", action="store_true", help="When initializing from a checkpoint, load model/discriminators but start fresh optimizers.")
    args = ap.parse_args()
    train(args)
 if __name__ == "__main__":
    main()