data.py

"""Synthetic audio: harmonic + formant-modulated signals that resemble voiced
speech enough to make codec reconstruction quality visually obvious."""

import math
import torch

SR = 24000


def _harmonic(f0, dur, sr=SR, n_harm=6):
    t = torch.arange(int(dur * sr)) / sr
    y = torch.zeros_like(t)
    for h in range(1, n_harm + 1):
        amp = 1.0 / h
        y = y + amp * torch.sin(2 * math.pi * f0 * h * t + 0.5 * h)
    return y / y.abs().max().clamp(min=1e-6)


def _formant_mod(y, sr=SR):
    # Slow AM + slow formant-like band shifting via LFOs
    t = torch.arange(y.shape[-1]) / sr
    am = 0.6 + 0.4 * torch.sin(2 * math.pi * 4.0 * t)
    fm = torch.sin(2 * math.pi * 0.8 * t) * 0.3
    return y * am * (1.0 + fm)


def make_clip(dur=1.0, sr=SR):
    """Random pitched clip that sweeps in pitch and amplitude."""
    f0 = float(torch.empty(1).uniform_(80, 280))
    dur = float(dur)
    base = _harmonic(f0, dur, sr=sr)
    # pitch glide
    t = torch.arange(base.shape[-1]) / sr
    glide = 1.0 + 0.15 * torch.sin(2 * math.pi * 0.5 * t)
    idx = (torch.cumsum(glide, dim=0) - glide[0]).clamp(0, base.shape[-1] - 1).long()
    base = base[idx]
    base = _formant_mod(base, sr=sr)
    # add a touch of pink-ish noise
    n = torch.randn_like(base) * 0.03
    y = (base + n)
    y = y / y.abs().max().clamp(min=1e-6) * 0.9
    return y


def batch(bs=8, dur=1.0, sr=SR, device='cpu'):
    clips = torch.stack([make_clip(dur=dur, sr=sr) for _ in range(bs)])
    return clips.unsqueeze(1).to(device)  # [B, 1, T]