7b92c9cAdd dataset generation scripts, ML fitting code and reference recordingsJeremy Magland 1"""A small DDSP-style generative model of the struck-sheet clips.
3 python3 ml/ddsp.py train --data DATASET_DIR --out RUN_DIR [--minutes 55] [--latent 8]
4 python3 ml/ddsp.py sample --model RUN_DIR/model.pt --out DIR [--n 8] [--seed 1]
6Encoder: a convolutional network on the cached log-frequency spectrogram
7(384 x 133) gives the mean and log-variance of an 8-dimensional latent z
8(a VAE, so that draws from the prior and interpolations are valid sounds).
10Decoder: z -> the controls of a parametric synthesizer:
11 - K partials with frequencies f_k (sorted, 30 Hz to 8 kHz), initial
12 amplitudes a_k, decay rates d_k, and a pitch-glide amount g_k, so that
13 partial k is a_k exp(-d_k t) sin(phase of f_k (1 + g_k exp(-t / tau)));
14 - a noise burst with a 48-band spectral shape, an amplitude, and a decay,
15 for the attack transient and the broadband part of the ring.
16Synthesis is plain arithmetic at 24 kHz (the simulations carry little energy
17above 6 kHz), so the same model can be run in JavaScript.
19Loss: multi-scale log-spectral distance between the synthesized and the
20original clip (FFT sizes 256 to 4096), plus a small KL term.
21"""
23import argparse
24import json
25import math
26import os
27import sys
28import time
29import wave
31import numpy as np
32import torch
33import torch.nn as nn
34import torch.nn.functional as Fn
35from scipy.signal import resample_poly
37sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
38from features import build_cache # noqa: E402
40torch.backends.cudnn.enabled = False
41SR = 24000
42DUR = 3.0
43N = int(SR * DUR)
44K = 128
45NOISE_BANDS = 48
46F_LO, F_HI = 30.0, 8000.0
49def read_wav24(path):
50 with wave.open(path) as w:
51 x = np.frombuffer(w.readframes(w.getnframes()), dtype="<i2").astype(np.float32) / 32768
52 y = resample_poly(x, 1, 2).astype(np.float32)
53 out = np.zeros(N, np.float32)
54 out[: min(N, len(y))] = y[:N]
55 return out
58def write_wav(path, y, sr=SR):
59 y = np.asarray(y, np.float32)
60 y = y / (np.abs(y).max() + 1e-9) * 0.95
61 with wave.open(path, "wb") as w:
62 w.setnchannels(1); w.setsampwidth(2); w.setframerate(sr)
63 w.writeframes((y * 32767).astype("<i2").tobytes())
66class Encoder(nn.Module):
67 def __init__(self, latent):
68 super().__init__()
69 ch, layers, c_in = (16, 32, 64, 96, 128), [], 1
70 for c in ch:
71 layers += [nn.Conv2d(c_in, c, 3, padding=1), nn.BatchNorm2d(c), nn.ReLU(inplace=True), nn.MaxPool2d(2)]
72 c_in = c
73 self.conv = nn.Sequential(*layers)
74 self.fc = nn.Sequential(nn.Linear(c_in * 12 * 4, 256), nn.ReLU(inplace=True), nn.Linear(256, 2 * latent))
76 def forward(self, x): # (B, 1, 384, 133)
77 h = self.conv(x) # (B, 128, 12, 4)
78 return self.fc(h.flatten(1)).chunk(2, -1)
81class Decoder(nn.Module):
82 def __init__(self, latent):
83 super().__init__()
84 self.mlp = nn.Sequential(nn.Linear(latent, 256), nn.LayerNorm(256), nn.ReLU(inplace=True), nn.Linear(256, 512), nn.LayerNorm(512), nn.ReLU(inplace=True))
85 self.freq = nn.Linear(512, K)
86 self.amp = nn.Linear(512, K)
87 self.decay = nn.Linear(512, K)
88 self.glide = nn.Linear(512, K)
89 self.noise_shape = nn.Linear(512, NOISE_BANDS)
90 self.noise_amp = nn.Linear(512, 1)
91 self.noise_decay = nn.Linear(512, 1)
92 self.glide_tau = nn.Linear(512, 1)
94 def forward(self, z):
95 h = self.mlp(z)
96 # Sorted frequencies: cumulative positive increments spanning the log range.
97 inc = Fn.softplus(self.freq(h)) + 1e-3
98 pos = torch.cumsum(inc, -1) / inc.sum(-1, keepdim=True) # (B, K) in (0, 1]
99 f = F_LO * (F_HI / F_LO) ** (pos * 0.98 + 0.005)
100 a = torch.sigmoid(self.amp(h) - 2.0) # amplitudes in (0, 1)
101 d = 0.3 + 40 * torch.sigmoid(self.decay(h) - 1.0) # 1/s
102 g = 0.15 * torch.tanh(self.glide(h)) # relative initial detuning
103 tau = 0.02 + 0.5 * torch.sigmoid(self.glide_tau(h)) # s
104 ns = torch.sigmoid(self.noise_shape(h) - 2.0)
105 na = torch.sigmoid(self.noise_amp(h) - 2.0)
106 nd = 2 + 60 * torch.sigmoid(self.noise_decay(h))
107 return dict(f=f, a=a, d=d, g=g, tau=tau, noise_shape=ns, noise_amp=na, noise_decay=nd)
110def synthesize(c, n=N, noise_seed=None):
111 """Renders a batch of control dicts to audio (B, n) at SR."""
112 B = c["f"].shape[0]
113 dev = c["f"].device
114 t = torch.arange(n, device=dev, dtype=torch.float32) / SR
115 # Partials: instantaneous frequency with a decaying glide; phase by cumulative sum.
116 f_inst = c["f"][:, :, None] * (1 + c["g"][:, :, None] * torch.exp(-t[None, None, :] / c["tau"][:, :, None])) # (B, K, n)
117 phase = 2 * math.pi * torch.cumsum(f_inst, -1) / SR
118 env = c["a"][:, :, None] * torch.exp(-c["d"][:, :, None] * t[None, None, :])
119 nyq = (c["f"] < SR / 2 - 200).float()[:, :, None]
120 y = (env * nyq * torch.sin(phase)).sum(1)
121 # Noise burst: white noise shaped in the STFT domain by a fixed spectral shape and a decaying envelope.
122 gen = torch.Generator(device=dev)
123 if noise_seed is not None:
124 gen.manual_seed(noise_seed)
125 noise = torch.randn(B, n, device=dev, generator=gen)
126 n_fft, hop = 512, 128
127 win = torch.hann_window(n_fft, device=dev)
128 S = torch.view_as_real(torch.stft(noise, n_fft, hop, window=win, return_complex=True)) # (B, 257, frames, 2)
129 shape = Fn.interpolate(c["noise_shape"][:, None, :], size=S.shape[1], mode="linear", align_corners=True)[:, 0, :] # (B, 257)
130 tf = torch.arange(S.shape[2], device=dev, dtype=torch.float32) * hop / SR
131 envn = c["noise_amp"] * torch.exp(-c["noise_decay"] * tf[None, :]) # (B, frames)
132 S = S * (shape[:, :, None] * envn[:, None, :])[..., None]
133 yn = torch.istft(torch.view_as_complex(S.contiguous()), n_fft, hop, window=win, length=n)
134 return y + 0.05 * yn
137def magnitude(x, n_fft):
138 """STFT magnitude without complex-tensor ops (the office PyTorch cannot JIT-compile those)."""
139 win = torch.hann_window(n_fft, device=x.device)
140 S = torch.view_as_real(torch.stft(x, n_fft, n_fft // 4, window=win, return_complex=True))
141 return torch.sqrt(S[..., 0] ** 2 + S[..., 1] ** 2 + 1e-12)
144def spectral_loss(y, x):
145 loss = 0.0
146 for n_fft in (256, 512, 1024, 2048, 4096):
147 Y = magnitude(y, n_fft)
148 X = magnitude(x, n_fft)
149 loss = loss + (Y - X).abs().mean() + 0.5 * (torch.log(Y + 1e-4) - torch.log(X + 1e-4)).abs().mean()
150 return loss / 5
153class Model(nn.Module):
154 def __init__(self, latent=8):
155 super().__init__()
156 self.latent = latent
157 self.enc = Encoder(latent)
158 self.dec = Decoder(latent)
160 def encode(self, spec):
161 mu, logvar = self.enc(spec)
162 return mu, logvar.clamp(-8, 4)
165def load_audio(data, ids):
166 cache = os.path.join(data, "ml", "audio24.npy")
167 if os.path.exists(cache):
168 a = np.load(cache, mmap_mode="r")
169 if a.shape[0] == len(ids):
170 return np.array(a)
171 a = np.stack([read_wav24(os.path.join(data, "clips", i + ".wav")) for i in ids])
172 np.save(cache, a)
173 return a
176def train(args):
177 os.makedirs(os.path.join(args.out, "samples"), exist_ok=True)
178 rows, feats = build_cache(args.data, os.path.join(args.data, "ml", "cache-logf"), kind="logf")
179 ids = [r["id"] for r in rows]
180 audio = load_audio(args.data, ids)
181 print(f"{len(ids)} clips, spectrograms {feats.shape}, audio {audio.shape}", flush=True)
182 dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
183 X = torch.from_numpy(feats)
184 A = torch.from_numpy(audio)
185 model = Model(args.latent).to(dev)
186 opt = torch.optim.Adam(model.parameters(), lr=args.lr)
187 rng = np.random.default_rng(0)
188 fixed = [0, 1234, 2500, 4321]
189 t0, step, last_sample = time.time(), 0, 0
190 log = open(os.path.join(args.out, "train.log"), "a")
191 while (time.time() - t0) / 60 < args.minutes:
192 perm = rng.permutation(len(ids))
193 for k in range(0, len(perm) - args.batch + 1, args.batch):
194 b = perm[k:k + args.batch]
195 spec = X[b].to(dev).float().unsqueeze(1)
196 x = A[b].to(dev)
197 mu, logvar = model.encode(spec)
198 z = mu + torch.randn_like(mu) * torch.exp(0.5 * logvar)
199 y = synthesize(model.dec(z))
200 rec = spectral_loss(y, x)
201 kl = -0.5 * torch.mean(1 + logvar - mu ** 2 - logvar.exp())
202 beta = args.beta * min(1.0, step / 2000)
203 loss = rec + beta * kl
204 opt.zero_grad(); loss.backward()
205 nn.utils.clip_grad_norm_(model.parameters(), 5.0)
206 opt.step()
207 step += 1
208 if step % 50 == 0:
209 msg = f"step {step} ({(time.time() - t0) / 60:.1f} min): rec {rec.item():.4f} kl {kl.item():.3f}"
210 print(msg, flush=True); log.write(msg + "\n"); log.flush()
211 if (time.time() - t0) / 60 - last_sample >= args.sample_every or (time.time() - t0) / 60 >= args.minutes:
212 last_sample = (time.time() - t0) / 60
213 torch.save({"state": model.state_dict(), "latent": args.latent, "step": step}, os.path.join(args.out, "model.pt"))
214 write_samples(model, X, ids, fixed, os.path.join(args.out, "samples", f"m{int(last_sample):03d}"), dev, step)
215 write_index(args.out)
216 if (time.time() - t0) / 60 >= args.minutes:
217 break
218 print("done", step, "steps")
222def write_samples(model, X, ids, fixed, out, dev, step, n_prior=6):
223 model.eval()
224 os.makedirs(out, exist_ok=True)
225 names = []
226 spec = X[fixed].to(dev).float().unsqueeze(1)
227 mu, _ = model.encode(spec)
228 y = synthesize(model.dec(mu), noise_seed=1).cpu().numpy()
229 for i, k in enumerate(fixed):
230 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"))
231 g = torch.Generator(device=dev); g.manual_seed(7)
232 zp = torch.randn(n_prior, model.latent, device=dev, generator=g)
233 y = synthesize(model.dec(zp), noise_seed=2).cpu().numpy()
234 for i in range(n_prior):
235 write_wav(os.path.join(out, f"prior-{i}.wav"), y[i]); names.append((f"prior sample {i}", f"prior-{i}.wav"))
236 steps = 5
237 zi = torch.stack([mu[0] + (mu[1] - mu[0]) * s / (steps - 1) for s in range(steps)])
238 y = synthesize(model.dec(zi), noise_seed=3).cpu().numpy()
239 for s in range(steps):
240 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"))
241 with open(os.path.join(out, "list.json"), "w") as f:
242 json.dump({"step": step, "items": names}, f)
243 model.train()
246def write_index(run):
247 sdir = os.path.join(run, "samples")
248 snaps = sorted(d for d in os.listdir(sdir) if os.path.isdir(os.path.join(sdir, d)))
249 blocks = []
250 for d in reversed(snaps):
251 try:
252 meta = json.load(open(os.path.join(sdir, d, "list.json")))
253 except Exception:
254 continue
255 rows = "".join(f"<tr><td class='l'>{label}</td><td class='l'><audio controls preload='none' src='samples/{d}/{fn}'></audio></td></tr>" for label, fn in meta["items"])
256 blocks.append(f"<h2>after {int(d[1:])} min, step {meta['step']}</h2><table>{rows}</table>")
257 html = ("<!doctype html><meta charset='utf-8'><title>DDSP sheet model samples</title><style>body{font:14px system-ui,sans-serif;margin:16px} table{border-collapse:collapse} td{padding:3px 8px;border-bottom:1px solid #ddd} td.l{text-align:left} audio{height:28px;width:260px} h2{font-size:15px;margin-top:18px}</style>"
258 "<h1>DDSP-style sheet model: samples during training</h1><p>Reconstructions of four fixed dataset clips (compare with clips/ID.wav in the gallery), six draws from the prior, and a five-step interpolation between two clips. Newest first.</p>" + "".join(blocks))
259 with open(os.path.join(run, "samples.html"), "w") as f:
260 f.write(html)
264def sample(args):
265 ck = torch.load(args.model, map_location="cpu", weights_only=False)
266 model = Model(ck["latent"]); model.load_state_dict(ck["state"]); model.eval()
267 os.makedirs(args.out, exist_ok=True)
268 g = torch.Generator(); g.manual_seed(args.seed)
269 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(",")]])
270 y = synthesize(model.dec(z), noise_seed=args.seed).numpy()
271 for i in range(y.shape[0]):
272 write_wav(os.path.join(args.out, f"sample-{i}.wav"), y[i])
273 print(f"sample-{i}.wav z = " + ", ".join(f"{v:.2f}" for v in z[i].tolist()))
276if __name__ == "__main__":
277 ap = argparse.ArgumentParser()
278 sub = ap.add_subparsers(dest="cmd", required=True)
279 t = sub.add_parser("train"); t.add_argument("--data", required=True); t.add_argument("--out", required=True)
280 t.add_argument("--minutes", type=float, default=55); t.add_argument("--latent", type=int, default=8); t.add_argument("--batch", type=int, default=16)
281 t.add_argument("--lr", type=float, default=1e-3); t.add_argument("--beta", type=float, default=1e-3); t.add_argument("--sample-every", type=float, default=10)
282 s = sub.add_parser("sample"); s.add_argument("--model", required=True); s.add_argument("--out", required=True)
283 s.add_argument("--n", type=int, default=8); s.add_argument("--seed", type=int, default=1); s.add_argument("--z", default=None, help="comma-separated latent vector for one clip")
284 a = ap.parse_args()
285 (train if a.cmd == "train" else sample)(a)