7b92c9cAdd dataset generation scripts, ML fitting code and reference recordingsJeremy Magland 1"""Second generative model: an oscillator bank on a fixed log-frequency grid.
3 python3 ml/ddsp2.py train --data DATASET_DIR --out RUN_DIR [--minutes 20] [--latent 8]
4 python3 ml/ddsp2.py sample --model RUN_DIR/model.pt --out DIR [--n 8] [--seed 1] [--z v1,...]
6Instead of letting the decoder choose partial frequencies (which the spectral
7loss cannot steer, see ddsp.py), the synthesizer has one oscillator at each of
8M = 96 bins per octave from 30 Hz to 8 kHz (about 770), and the decoder emits
9for each an amplitude, a decay rate, and a frequency offset within its bin.
10A partial is then an active bin, its exact frequency the bin plus the offset,
11and the amplitudes have well-behaved gradients. The decoder is a 1D
12transposed-convolutional network over the frequency axis, so that patterns of
13lines can shift along it easily (a change of the sheet's frequency scale is a
14translation on this axis). A shaped noise burst covers the attack. The loss
15adds the smoothed log-frequency "line picture" of ml/linesim.py, computed in
16torch, to the multi-scale spectral loss.
17"""
19import argparse
20import json
21import math
22import os
23import sys
24import time
26import numpy as np
27import torch
28import torch.nn as nn
29import torch.nn.functional as Fn
31sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
32from features import build_cache # noqa: E402
33from ddsp import SR, DUR, N, NOISE_BANDS, Encoder, magnitude, spectral_loss, load_audio, write_wav, write_index # noqa: E402
35torch.backends.cudnn.enabled = False
36F_LO, F_HI = 30.0, 8000.0
37BPO = 96
38M = int(math.log2(F_HI / F_LO) * BPO) # oscillators
39GRID = F_LO * 2 ** (np.arange(M) / BPO)
40HALF_BIN = 2 ** (0.5 / BPO) - 1 # relative half width of a bin
43class Decoder(nn.Module):
44 def __init__(self, latent, base=25, ch=64):
45 super().__init__()
46 self.base = base
47 self.fc = nn.Sequential(nn.Linear(latent, 256), nn.ReLU(inplace=True), nn.Linear(256, ch * base), nn.ReLU(inplace=True))
48 self.ch = ch
49 ups = []
50 c = ch
51 for k in range(5): # base * 32 = 800 >= M
52 ups += [nn.ConvTranspose1d(c + 1, ch, 4, stride=2, padding=1), nn.GroupNorm(8, ch), nn.ReLU(inplace=True)]
53 c = ch
54 self.ups = nn.ModuleList(ups)
55 self.head = nn.Conv1d(ch + 1, 3, 5, padding=2)
56 self.noise = nn.Sequential(nn.Linear(latent, 128), nn.ReLU(inplace=True), nn.Linear(128, NOISE_BANDS + 2))
58 def forward(self, z):
59 h = self.fc(z).view(-1, self.ch, self.base)
60 i = 0
61 while i < len(self.ups):
62 coord = torch.linspace(0, 1, h.shape[-1], device=h.device).expand(h.shape[0], 1, -1)
63 h = self.ups[i](torch.cat([h, coord], 1)); h = self.ups[i + 1](h); h = self.ups[i + 2](h)
64 i += 3
65 coord = torch.linspace(0, 1, h.shape[-1], device=h.device).expand(h.shape[0], 1, -1)
66 out = self.head(torch.cat([h, coord], 1))[:, :, :M] # (B, 3, M)
67 amp = torch.sigmoid(out[:, 0] - 3.0)
68 decay = 0.3 + 40 * torch.sigmoid(out[:, 1] - 1.0)
69 offset = torch.tanh(out[:, 2]) * HALF_BIN
70 nz = self.noise(z)
71 return dict(amp=amp, decay=decay, offset=offset, noise_shape=torch.sigmoid(nz[:, :NOISE_BANDS] - 2.0),
72 noise_amp=torch.sigmoid(nz[:, NOISE_BANDS] - 2.0), noise_decay=2 + 60 * torch.sigmoid(nz[:, NOISE_BANDS + 1]))
75_GRID_T = None
78def synthesize(c, n=N, noise_seed=None):
79 global _GRID_T
80 B = c["amp"].shape[0]
81 dev = c["amp"].device
82 if _GRID_T is None or _GRID_T.device != dev:
83 _GRID_T = torch.tensor(GRID, dtype=torch.float32, device=dev)
84 t = torch.arange(n, device=dev, dtype=torch.float32) / SR
85 f = _GRID_T[None, :] * (1 + c["offset"]) # (B, M)
86 usable = (f < SR / 2 - 200).float()
87 phase = 2 * math.pi * f[:, :, None] * t[None, None, :]
88 env = (c["amp"] * usable)[:, :, None] * torch.exp(-c["decay"][:, :, None] * t[None, None, :])
89 y = (env * torch.sin(phase)).sum(1)
90 gen = torch.Generator(device=dev)
91 if noise_seed is not None:
92 gen.manual_seed(noise_seed)
93 noise = torch.randn(B, n, device=dev, generator=gen)
94 n_fft, hop = 512, 128
95 win = torch.hann_window(n_fft, device=dev)
96 S = torch.view_as_real(torch.stft(noise, n_fft, hop, window=win, return_complex=True))
97 shape = Fn.interpolate(c["noise_shape"][:, None, :], size=S.shape[1], mode="linear", align_corners=True)[:, 0, :]
98 tf = torch.arange(S.shape[2], device=dev, dtype=torch.float32) * hop / SR
99 envn = c["noise_amp"][:, None] * torch.exp(-c["noise_decay"][:, None] * tf[None, :])
100 S = S * (shape[:, :, None] * envn[:, None, :])[..., None]
101 yn = torch.istft(torch.view_as_complex(S.contiguous()), n_fft, hop, window=win, length=n)
102 return y + 0.05 * yn
105# ---------- line-picture loss (linesim.py in torch) ----------
107_PIC = {}
110def picture_torch(x, tol=0.02, n_fft=4096, hop=1024, slices=((0, 0.3), (0.3, 1.0), (1.0, 2.0), (2.0, 3.0))):
111 dev = x.device
112 key = (tol, str(dev))
113 if key not in _PIC:
114 f_lin = np.arange(n_fft // 2 + 1) * SR / n_fft
115 centers = F_LO * 2 ** (np.arange(int(math.log2(5000 / F_LO) * 48)) / 48)
116 W = np.exp(-0.5 * ((np.log(np.maximum(f_lin[None, :], 1e-3) / centers[:, None])) / tol) ** 2)
117 W[:, 0] = 0
118 W /= np.maximum(W.sum(1, keepdims=True), 1e-12)
119 _PIC[key] = torch.tensor(W, dtype=torch.float32, device=dev)
120 W = _PIC[key]
121 P = magnitude(x, n_fft) ** 2 # (B, bins, frames), hop n_fft // 4
122 fps = SR / (n_fft // 4)
123 outs = []
124 for a, b in slices:
125 sl = P[:, :, int(a * fps): max(int(a * fps) + 1, int(b * fps))].mean(-1)
126 outs.append(W @ sl.T) # (n_log, B)
127 db = 10 * torch.log10(torch.stack(outs, 0) + 1e-12) # (K, n_log, B)
128 db = db - db.amax((0, 1), keepdim=True)
129 return torch.clamp(db, min=-70)
132def picture_loss(y, x):
133 return (picture_torch(y) - picture_torch(x)).abs().mean() / 10
136class Model(nn.Module):
137 def __init__(self, latent=8):
138 super().__init__()
139 self.latent = latent
140 self.enc = Encoder(latent)
141 self.dec = Decoder(latent)
143 def encode(self, spec):
144 mu, logvar = self.enc(spec)
145 return mu, logvar.clamp(-8, 4)
148def train(args):
149 os.makedirs(os.path.join(args.out, "samples"), exist_ok=True)
150 rows, feats = build_cache(args.data, os.path.join(args.data, "ml", "cache-logf"), kind="logf")
151 ids = [r["id"] for r in rows]
152 audio = load_audio(args.data, ids)
153 print(f"{len(ids)} clips, {M} oscillators, audio {audio.shape}", flush=True)
154 dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
155 X = torch.from_numpy(feats); A = torch.from_numpy(audio)
156 model = Model(args.latent).to(dev)
157 print(f"{sum(p.numel() for p in model.parameters()) / 1e6:.2f} M parameters", flush=True)
158 opt = torch.optim.Adam(model.parameters(), lr=args.lr)
159 rng = np.random.default_rng(0)
160 fixed = [0, 1234, 2500, 4321]
161 t0, step, last_sample = time.time(), 0, 0
162 log = open(os.path.join(args.out, "train.log"), "a")
163 done = False
164 while not done:
165 perm = rng.permutation(len(ids))
166 for k in range(0, len(perm) - args.batch + 1, args.batch):
167 b = perm[k:k + args.batch]
168 spec = X[b].to(dev).float().unsqueeze(1)
169 x = A[b].to(dev)
170 mu, logvar = model.encode(spec)
171 z = mu + torch.randn_like(mu) * torch.exp(0.5 * logvar)
172 y = synthesize(model.dec(z))
173 rec = spectral_loss(y, x)
174 pic = picture_loss(y, x)
175 kl = -0.5 * torch.mean(1 + logvar - mu ** 2 - logvar.exp())
176 beta = args.beta * min(1.0, step / 3000)
177 loss = rec + pic + beta * torch.clamp(kl, min=args.free_bits)
178 opt.zero_grad(); loss.backward()
179 nn.utils.clip_grad_norm_(model.parameters(), 5.0)
180 opt.step()
181 step += 1
182 minutes = (time.time() - t0) / 60
183 if step % 50 == 0:
184 msg = f"step {step} ({minutes:.1f} min): rec {rec.item():.4f} pic {pic.item():.4f} ({10 * pic.item():.2f} dB) kl {kl.item():.3f}"
185 print(msg, flush=True); log.write(msg + "\n"); log.flush()
186 if minutes - last_sample >= args.sample_every or minutes >= args.minutes:
187 last_sample = minutes
188 torch.save({"state": model.state_dict(), "latent": args.latent, "step": step, "kind": "bank"}, os.path.join(args.out, "model.pt"))
189 write_samples(model, X, ids, fixed, os.path.join(args.out, "samples", f"m{int(minutes):03d}"), dev, step)
190 write_index(args.out)
191 if minutes >= args.minutes:
192 done = True
193 break
194 print("done", step, "steps")
198def write_samples(model, X, ids, fixed, out, dev, step, n_prior=6):
199 model.eval()
200 os.makedirs(out, exist_ok=True)
201 names = []
202 spec = X[fixed].to(dev).float().unsqueeze(1)
203 mu, _ = model.encode(spec)
204 y = synthesize(model.dec(mu), noise_seed=1).cpu().numpy()
205 for i, k in enumerate(fixed):
206 write_wav(os.path.join(out, f"recon-{ids[k]}.wav"), y[i]); names.append((f"reconstruction of {ids[k]}", f"recon-{ids[k]}.wav"))
207 g = torch.Generator(device=dev); g.manual_seed(7)
208 zp = torch.randn(n_prior, model.latent, device=dev, generator=g)
209 y = synthesize(model.dec(zp), noise_seed=2).cpu().numpy()
210 for i in range(n_prior):
211 write_wav(os.path.join(out, f"prior-{i}.wav"), y[i]); names.append((f"prior sample {i}", f"prior-{i}.wav"))
212 steps = 5
213 zi = torch.stack([mu[0] + (mu[1] - mu[0]) * s / (steps - 1) for s in range(steps)])
214 y = synthesize(model.dec(zi), noise_seed=3).cpu().numpy()
215 for s in range(steps):
216 write_wav(os.path.join(out, f"interp-{s}.wav"), y[s]); names.append((f"interpolation {ids[fixed[0]]} -> {ids[fixed[1]]}, step {s}", f"interp-{s}.wav"))
217 with open(os.path.join(out, "list.json"), "w") as f:
218 json.dump({"step": step, "items": names}, f)
219 model.train()
223def sample(args):
224 ck = torch.load(args.model, map_location="cpu", weights_only=False)
225 model = Model(ck["latent"]); model.load_state_dict(ck["state"]); model.eval()
226 os.makedirs(args.out, exist_ok=True)
227 g = torch.Generator(); g.manual_seed(args.seed)
228 z = torch.randn(args.n, ck["latent"], generator=g) if args.z is None else torch.tensor([[float(v) for v in args.z.split(",")]])
229 y = synthesize(model.dec(z), noise_seed=args.seed).numpy()
230 for i in range(y.shape[0]):
231 write_wav(os.path.join(args.out, f"sample-{i}.wav"), y[i])
232 print(f"sample-{i}.wav z = " + ", ".join(f"{v:.2f}" for v in z[i].tolist()))
235if __name__ == "__main__":
236 ap = argparse.ArgumentParser()
237 sub = ap.add_subparsers(dest="cmd", required=True)
238 t = sub.add_parser("train"); t.add_argument("--data", required=True); t.add_argument("--out", required=True)
239 t.add_argument("--minutes", type=float, default=20); t.add_argument("--latent", type=int, default=8); t.add_argument("--batch", type=int, default=16)
240 t.add_argument("--lr", type=float, default=5e-4); t.add_argument("--beta", type=float, default=2e-3); t.add_argument("--free-bits", type=float, default=0.5)
241 t.add_argument("--sample-every", type=float, default=5)
242 s = sub.add_parser("sample"); s.add_argument("--model", required=True); s.add_argument("--out", required=True)
243 s.add_argument("--n", type=int, default=8); s.add_argument("--seed", type=int, default=1); s.add_argument("--z", default=None)
244 a = ap.parse_args()
245 (train if a.cmd == "train" else sample)(a)