from dataclasses import dataclass from functools import partial from typing import Optional from torch import nn import vampnet.dsp.signal as sn from vampnet.dsp.signal import Signal from vampnet.mask import random_along_time from torch import Tensor import torch class MedianFilterAugment(nn.Module): def __init__(self, kernel_size: int, train_min: int = 1, train_max: int = 20, ): super().__init__() self.kernel_size = kernel_size self.train_min = train_min self.train_max = train_max def forward(self, x: Tensor) -> Tensor: if self.training: sizes = torch.randint( self.train_min, self.train_max, size=(x.shape[0],) ) else: sizes = self.kernel_size # print(f"median filter sizes: {sizes}") return sn.median_filter_1d(x, sizes) class RMS(nn.Module): def __init__(self, hop_length, window_length=2048, n_quantize=None, sample_rate=44100, median_filter_size: Optional[int] = None, train_median_filter_min=1, train_median_filter_max=15, ): super().__init__() self.hop_length = hop_length self.window_length = window_length self.n_quantize = n_quantize self.sample_rate = sample_rate self.mf = MedianFilterAugment( kernel_size=median_filter_size, train_min=train_median_filter_min, train_max=train_median_filter_max ) if median_filter_size is not None else None @property def dim(self): return 1 def extract(self, sig: Signal) -> Tensor: rmsd = sn.rms(sig, window_length=self.window_length, hop_length=self.hop_length, )[:, :, :-1] # TODO: cutting the last frame to match DAC tokens but why :'( nb, _, _ = rmsd.shape if self.n_quantize is not None: # standardize to 0-1 rmsd = (rmsd - rmsd.min()) / (rmsd.max() - rmsd.min()) # quantize to 128 steps rmsd = torch.round(rmsd * self.n_quantize) rmsd = rmsd / self.n_quantize if self.mf is not None: rmsd = self.mf(rmsd) return rmsd class HarmonicChroma(nn.Module): def __init__(self, hop_length: int, window_length: int = 4096, n_chroma: int = 48, sample_rate: int = 44100, top_n: int = 0 ): super().__init__() from torchaudio.prototype.transforms import ChromaScale self.hop_length = hop_length self.window_length = window_length self.n_chroma = n_chroma self.sample_rate = sample_rate self.top_n = top_n # HUGO: this representation, as is, # encodes timbre information in the chroma # which is not what we want!!! # would a median filter help perhaps? self.chroma = ChromaScale( sample_rate=self.sample_rate, n_freqs=self.window_length // 2 + 1, n_chroma=self.n_chroma, octwidth=5.0, ) @property def dim(self): return self.n_chroma def extract(self, sig: Signal) -> Tensor: from vampnet.dsp.hpss import hpss self.chroma.to(sig.wav.device) # spectrogram spec = sn.stft(sig, window_length=self.window_length, hop_length=self.hop_length ) # magnitude spec = torch.abs(spec) # hpss spec = hpss(spec, kernel_size=51, hard=True)[0] # chroma chroma = self.chroma(spec) # get the rms of this spec rms_d = sn.rms_from_spec( spec, window_length=self.window_length ) # convert the rms to db rms_d = 10 * torch.log10(rms_d + 1e-7) # make a mask based on the rms < -40 mask = torch.where(rms_d < -40, torch.zeros_like(rms_d), torch.ones_like(rms_d)) # remove anything below 80 (where the fuck did I get this number from?) chroma = torch.where(chroma < 100, torch.zeros_like(chroma), chroma) # Get top 2 values and indices along the -2 dimension if self.top_n: _, topk_indices = torch.topk(chroma, self.top_n, dim=-2) # Create a mask for the top 2 values topk_mask = torch.zeros_like(chroma).scatter_(-2, topk_indices, 1.0) # Retain only the top 2 values chroma = chroma * topk_mask # apply the mask chroma = chroma * mask.unsqueeze(-2) # Apply softmax along dim=-2 if self.top_n > 0: chroma = torch.nn.functional.softmax(chroma, dim=-2) # mask out any timesteps whose chroma have all equal values (all 0s before softmax) # TODO: i did this with chatgpt, there's gott a be a better way chroma_mean = chroma.mean(dim=-2, keepdim=True) chroma_diff = torch.abs(chroma - chroma_mean) equal_mask = torch.all(chroma_diff < 1e-6, dim=-2, keepdim=True) # Set chroma values to zero for timesteps with all equal values chroma = torch.where(equal_mask, torch.zeros_like(chroma), chroma) return chroma[:, 0, :, :-1] # mono only :( FIX ME! # TODO: try harmonic mel? CONTROLLERS = { "rms": RMS, "rmsq128": partial(RMS, n_quantize=128), "rmsq16": partial(RMS, n_quantize=16), "rms-median": partial(RMS, median_filter_size=5), "rmsq16-median": partial(RMS, n_quantize=16, median_filter_size=3), "hchroma": HarmonicChroma, "hchroma-12c-top2": partial(HarmonicChroma, n_chroma=12, top_n=2), # TODO: refactor me. If this works, this should just be named hchroma. "hchroma-36c-top3": partial(HarmonicChroma, n_chroma=36, top_n=3) # TODO: refactor me. If this works, this should just be named hchroma. } class Sketch2SoundController(nn.Module): def __init__( self, ctrl_keys: list[str], hop_length: str, sample_rate: int, ): super().__init__() assert all([k in CONTROLLERS for k in ctrl_keys]), f"got an unsupported control key in {ctrl_keys}!\n supported: {CONTROLLERS.keys()}" self.hop_length = hop_length self.ctrl_keys = ctrl_keys self.sample_rate = sample_rate self.controllers = { k: CONTROLLERS[k](hop_length=hop_length, sample_rate=sample_rate) for k in self.ctrl_keys } @property def ctrl_dims(self, ) -> dict[str, int]: return { k: controller.dim for k, controller in self.controllers.items() } def extract(self, sig: Signal) -> dict[str, Tensor]: ctrls = { k: controller.extract(sig) for k, controller in self.controllers.items() } return ctrls def random_mask(self, ctrls: dict[str, Tensor], r: float): masks = {} for k, ctrl in ctrls.items(): masks[k] = 1-random_along_time(ctrl, r) return masks def empty_mask(self, ctrls: dict[str, Tensor]): first_key = next(iter(ctrls)) mask = torch.zeros_like(ctrls[first_key]) return {k: mask for k in ctrls} def test_controller(): controller = Sketch2SoundController( ctrl_keys=["rms-median", "rms", "rmsq128"], hop_length=512, sample_rate=44100 ) controller.train() # sig = sn.read_from_file("assets/example.wav") # sig = sn.read_from_file("/Users/hugo/Downloads/DCS_SE_FullChoir_ScaleUpDown06_A2_DYN.wav") # sig = sn.excerpt('/Users/hugo/Downloads/(guitarra - hugo mix) bubararu - tambor negro.wav', offset=0, duration=10) sig = sn.read_from_file("assets/voice-prompt.wav") ctrls = controller.extract(sig) print(f"given sig of shape {sig.wav.shape}, extracted controls: {ctrls}") # print the whole thing # torch.set_printoptions(profile="full") # print(ctrls["hchroma"][0][0][:, 200:210]) # imshow the chroma import matplotlib.pyplot as plt # Define relative heights for the subplots fig, (ax1, ax2, ax3, ax4) = plt.subplots( 4, 1, sharex=True, ) # Display the spectrogram on the top ax1.imshow(sn.stft(sig, hop_length=512, window_length=2048).abs()[0][0].cpu().log().numpy(), aspect='auto', origin='lower') # display rms on the bottom ax2.plot(ctrls["rms-median"][0][0]) ax3.plot(ctrls["rms"][0][0]) ax4.plot(ctrls["rmsq128"][0][0]) plt.tight_layout() # Ensure proper spacing plt.savefig("img.png") if __name__ == "__main__": test_controller()