fakesound 0.0.0__py3-none-any.whl

fakesound/data/dataset.py

@@ -0,0 +1,58 @@
+from collections.abc import Callable
+import numpy as np
+import glob, os
+import torch
+from scipy.io import wavfile
+from numpy.lib.stride_tricks import sliding_window_view
+from torch.utils.data import Dataset
+
+class AudioDataSet(Dataset[torch.Tensor]):
+    def __init__(self, folder_path:str, chunk_samples:int,
+                       hop_samples:int, transform:Callable = None,
+                       reset = False, pad_start:int = 0):
+        self.sample_rate = None # set when 1st loading audio
+        self.transform = transform
+        files = glob.glob(os.path.join(folder_path, "*.wav"))
+        arrays = []
+        self.reset_flags = []
+        # if reset is used, all files should be equal length
+        self.num_chunks = []
+        for f in files:
+            x = self.load_audio(f)
+            x = np.pad(x, (pad_start, 0))
+            chunks = sliding_window_view(x, chunk_samples)[::hop_samples,:]
+            self.num_chunks.append(len(chunks))
+            file_reset =  [False] *  self.num_chunks[-1]
+            if reset:
+                file_reset[0] = True
+            arrays.append(chunks)
+            self.reset_flags = self.reset_flags + file_reset
+        self.data = np.concatenate(arrays, 0)
+
+
+    def __len__(self):
+        return len(self.data)
+
+    # reset assumes no shuffling
+    def __getitem__(self, index:int):
+        sample = self.data[index]
+        reset = self.reset_flags[index]
+        sample = torch.from_numpy(sample).unsqueeze(0).float()
+        transformed = sample
+        if self.transform is not None:
+            transformed = self.transform(sample)
+        return transformed, sample, reset
+
+    def load_audio(self, file_path:str):
+        samplerate, data = wavfile.read(file_path)
+        if self.sample_rate is None:
+            self.sample_rate = samplerate
+        if samplerate != self.sample_rate:
+            raise Exception("Different sample rates in dataset")
+        if data.dtype == np.int16:
+            data = data / 2**15
+        if data.dtype == np.int32:
+            data = data / 2**31
+        if len(data.shape) > 1:
+            data = data[:,0]
+        return data

fakesound/data/mel_bands.py

@@ -0,0 +1,34 @@
+import numpy as np
+import torch
+
+def mel2hz(mel):
+  return 700. * (np.exp(mel / 1127.) - 1.0)
+
+def hz2mel(hz):
+    return 1127. * np.log(1.0 + (hz / 700.0))
+
+def linear2mel(bin_freqs:np.ndarray, n_mels:int, fmin:float, fmax:float):
+    return get_mel_matrix(bin_freqs, n_mels, fmin, fmax)
+
+def mel2linear(bin_freqs:np.ndarray, n_mels:int, fmin:float, fmax:float):
+    M = get_mel_matrix(bin_freqs, n_mels, fmin, fmax)
+    Mt = M.T
+    S = np.sum(np.dot(M, Mt), 0)
+    d = np.divide(1.0, S, out=np.zeros_like(S), where = S!=0)
+    return np.dot(Mt, np.diag(d))
+
+def get_mel_matrix(bin_freqs:np.ndarray, n_mels:int, fmin:float, fmax:float):
+    # based on https://librosa.org/doc/main/_modules/librosa/filters.html#mel
+    N = bin_freqs.shape[0]
+    M = np.zeros((n_mels,N))
+    mels =  np.linspace(hz2mel(fmin), hz2mel(fmax), n_mels+2)
+    mel_freqs = mel2hz(mels)
+    fdiff = np.diff(mel_freqs)
+    ramps = np.subtract.outer(mel_freqs, bin_freqs)
+    for i in range(n_mels):
+        rising = -ramps[i, :] / fdiff[i]
+        falling = ramps[i + 2, :] / fdiff[i + 1]
+        M[i, :] = np.maximum(0, np.minimum(rising, falling))
+    norm = 2.0 / (mel_freqs[2 : n_mels + 2] - mel_freqs[:n_mels])
+    M = M * norm[:, None]
+    return M

fakesound/data/pqmf.py

@@ -0,0 +1,21 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+from scipy import signal
+
+def get_pqmf_filter( n_bands:int = 16, atten:float=100, analysis = True):
+    cutoff = 1  / (2 * n_bands)
+    n_taps, beta = signal.kaiserord(atten, cutoff)
+    h = signal.firwin(n_taps + 1, cutoff,
+        window=('kaiser', beta), scale = False )
+    k = torch.arange(n_bands).reshape(-1, 1)
+    alpha = (2 * k + 1) * (torch.pi /
+            (2 * n_bands)) * (torch.arange( n_taps + 1)-(( n_taps - 1) / 2))
+    phase = (-1)**k * np.pi / 4
+    if analysis:
+        f = 2 * torch.from_numpy(h) * torch.cos(alpha + phase)
+        f = f[:, None, :].float()
+    else:
+        f = 2 * torch.from_numpy(h) * torch.cos(alpha - phase)
+        f = f[None, :, :].float()
+    return f

fakesound/data/transforms.py

@@ -0,0 +1,168 @@
+import numpy as np
+import torch
+from torch import nn
+import torch.nn.functional as F
+from .mel_bands import linear2mel
+from .pqmf import get_pqmf_filter
+
+class muLawEncoder(nn.Module):
+    def __init__(self, n_levels:int = 256):
+        super().__init__()
+        self.num_levels:int = n_levels
+        bins:torch.Tensor = torch.linspace(-1, 1, n_levels)
+        mu:torch.Tensor = torch.Tensor([n_levels - 1])
+        self.register_buffer('bins', bins)
+        self.register_buffer('mu', mu)
+
+    def forward(self, x:torch.Tensor) -> torch.Tensor:
+        x = x.clone()
+        compressed = torch.sign(x) * torch.log1p(self.mu * torch.abs(x)) / torch.log1p(self.mu)
+        return  torch.bucketize(compressed, self.bins)
+
+class muLawDecoder(nn.Module):
+    def __init__(self, n_levels:int = 256):
+        super().__init__()
+        mu:torch.Tensor = torch.Tensor([n_levels - 1])
+        self.register_buffer('mu', mu)
+
+    def forward(self, x:torch.Tensor) -> torch.Tensor:
+        x = x / (1 + self.mu) * 2 - 1
+        return torch.sign(x) * (torch.pow(self.mu + 1, torch.abs(x)) - 1) / self.mu
+
+class Wave2LogMagIF(nn.Module):
+    def __init__(self, fft_size:int, hop_size:int):
+        super().__init__()
+        self.fft_size = fft_size
+        self.hop_size = hop_size
+        self.epsilon = 1e-7
+        window = torch.hann_window(fft_size)
+        self.register_buffer('window', window)
+
+    def forward(self, x:torch.Tensor, remove_nyquist:bool = True) -> tuple[torch.Tensor, torch.Tensor]:
+         x = x.squeeze(1)
+         spec = torch.stft(x,
+             n_fft = self.fft_size,
+             hop_length = self.hop_size,
+             win_length = self.fft_size,
+             window = self.window,
+             center = True,
+             pad_mode = 'reflect',
+             normalized = False,
+             onesided =  True,
+             return_complex = True).squeeze(0)
+         logmag = torch.log(torch.abs(spec).clip(min = self.epsilon))
+         phase = torch.angle(spec)
+         dphase = torch.diff(phase)
+         dphase = torch.where(dphase >  np.pi, dphase - 2 * np.pi, dphase)
+         dphase = torch.where(dphase < -np.pi, dphase + 2 * np.pi, dphase)
+         dphase = torch.hstack((phase[:,0, None], dphase))
+         if remove_nyquist:
+             return (logmag[:-1,:], dphase[:-1,:])
+         else:
+             return (logmag, dphase)
+
+class LogMagIF2Wave(nn.Module):
+        def __init__(self, fft_size:int, hop_size:int):
+            super().__init__()
+            self.fft_size = fft_size
+            self.hop_size = hop_size
+            window = torch.hann_window(fft_size)
+            self.register_buffer('window', window)
+
+        def forward(self, logmag:torch.Tensor, instfreq: torch.Tensor, add_nyquist:bool = True) -> torch.Tensor:
+            mag = torch.exp(logmag)
+            phase = torch.cumsum(instfreq, 1)
+            spec = torch.polar(mag, phase)
+            if add_nyquist:
+                new_spec = torch.empty(
+                    logmag.shape[0]+1, logmag.shape[1] ,
+                    dtype = spec.dtype).to(logmag.device)
+                new_spec[:logmag.shape[0], :] = spec
+                spec = new_spec
+            audio = torch.istft(spec,
+                n_fft = self.fft_size,
+                hop_length = self.hop_size,
+                win_length = self.fft_size,
+                window = self.window,
+                center = True,
+                normalized = False,
+                onesided =  True,
+                return_complex = False
+                )
+            return audio
+
+
+class Wave2LogMelIF(nn.Module):
+    def __init__(self, sr:float, fft_size:int, hop_size:int,
+                        n_mels:int, fmin:float, fmax:float,
+                        pad_amount:int = None, return_phase:bool = True):
+        super().__init__()
+        self.fft_size = fft_size
+        self.hop_size = hop_size
+        self.n_mels = n_mels
+        self.return_phase = return_phase
+        self.epsilon = 1e-7
+        if pad_amount is None:
+            self.pad_amount = fft_size //2
+        else:
+            self.pad_amount = pad_amount
+        dF = sr/fft_size
+        n_bins = fft_size / 2 + 1
+        ff = np.arange(0, dF * n_bins, dF)
+        mel_matrix =  torch.from_numpy(linear2mel(ff, n_mels, fmin, fmax)).float()
+        window = torch.hann_window(fft_size)
+        self.register_buffer('mel_matrix', mel_matrix)
+        self.register_buffer('window', window)
+
+    def forward(self, x:torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+         x = torch.nn.functional.pad(
+              x.unsqueeze(1), (self.pad_amount, self.pad_amount),
+              mode='reflect'
+          )
+         x = x.squeeze(1)
+         spec = torch.stft(x,
+             n_fft = self.fft_size,
+             hop_length = self.hop_size,
+             win_length = self.fft_size,
+             window = self.window,
+             center= False,
+             pad_mode = 'reflect',
+             normalized = False,
+             onesided =  True,
+             return_complex = True).squeeze(0)
+         mel_mag = torch.matmul(self.mel_matrix, torch.abs(spec))
+         mel_logmag =  torch.log(mel_mag.clip(min = self.epsilon))
+         if not self.return_phase: return (mel_logmag, torch.empty(0))
+         mel_phase = torch.matmul(self.mel_matrix, torch.angle(spec))
+         dphase = torch.diff(mel_phase)
+         dphase = torch.where(dphase >  np.pi, dphase - 2 * np.pi, dphase)
+         dphase = torch.where(dphase < -np.pi, dphase + 2 * np.pi, dphase)
+         dphase = torch.hstack((mel_phase[:,0, None], dphase))
+         return (mel_logmag, dphase)
+
+class Wave2PQMF(nn.Module):
+    def __init__(self, n_bands:int = 16, atten:float = 100):
+        super().__init__()
+        h = get_pqmf_filter(n_bands, atten)
+        self.n_taps = h.shape[-1]-1
+        self.n_bands = n_bands
+        self.register_buffer("h", h)
+    def forward(self, x:torch.Tensor)-> torch.Tensor:
+        return F.conv1d(x, self.h, padding = self.n_taps // 2, stride=self.n_bands)
+
+class PQMF2Wave(nn.Module):
+    def __init__(self, n_bands:int = 16, atten:float = 100):
+        super().__init__()
+        h = get_pqmf_filter(n_bands, atten, False)
+        self.n_taps = h.shape[-1]-1
+        self.n_bands = n_bands
+        self.register_buffer("h", h)
+        updown_filter = torch.zeros((n_bands, n_bands, n_bands)).float()
+        updown_filter[range(n_bands), range(n_bands), 0] = 1
+        self.register_buffer("updown_filter", updown_filter)
+        self.pad_fn = torch.nn.ConstantPad1d(self.n_taps // 2, 0.0)
+
+    def forward(self, x:torch.Tensor)-> torch.Tensor:
+        x = F.conv_transpose1d(x, self.updown_filter * self.n_bands,stride=self.n_bands)
+        x = F.conv1d(x, self.h, padding = self.n_taps // 2)
+        return x

fakesound/gansynth/dataset.py

@@ -0,0 +1,46 @@
+import numpy as np
+import glob, os
+import torch
+from fakesound.data.dataset import AudioDataSet
+from fakesound.data.transforms import Wave2LogMagIF
+
+# Lazy dataset for NSynth training set including pitch
+
+class NSynthDataset(AudioDataSet):
+    def __init__(self, folder_path:str, sample_rate:int,
+                       fft_size:int, hop_size:int, pad_start:int):
+        self.sample_rate = None
+        self.fft_size = fft_size
+        self.n_mels = fft_size // 2
+        self.pad_start = pad_start
+        self.transform = Wave2LogMagIF(fft_size, hop_size)
+
+        self.files = glob.glob(os.path.join(folder_path, "*.wav"))[:16]
+
+        self.max_mag = np.log(fft_size//4)
+        self.min_mag = np.log(self.transform.epsilon)
+        self.max_phase = np.pi
+        self.min_phase = -np.pi
+    def __len__(self):
+        return len(self.files)
+
+    def normalise_mag(self, x:torch.Tensor)-> torch.Tensor:
+        return 2 * (x - self.min_mag) / (self.max_mag - self.min_mag) -1
+    def normalise_phase(self, x:torch.Tensor)-> torch.Tensor:
+        return 2 * (x - self.min_phase) / (self.max_phase - self.min_phase) -1
+    def denormalise_mag(self, x:torch.Tensor)-> torch.Tensor:
+        return (x + 1)* (self.max_mag - self.min_mag) / 2 + self.min_mag
+    def denormalise_phase(self, x:torch.Tensor)-> torch.Tensor:
+        return (x + 1)* (self.max_phase - self.min_phase) / 2 + self.min_phase
+
+    # reset assumes no shuffling
+    def __getitem__(self, index:int):
+        filename = self.files[index]
+        pitch = int(filename.split("_")[-1].split("-")[1])
+        sample = self.load_audio(filename)
+        sample = np.pad(sample, (self.pad_start, 0))
+        sample = torch.from_numpy(sample).unsqueeze(0).float()
+        logmag, phase = self.transform(sample)
+        logmag = self.normalise_mag(logmag)
+        phase = self.normalise_phase(phase)
+        return torch.stack((logmag, phase)), pitch

fakesound/gansynth/models.py

@@ -0,0 +1,89 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+class PixelNorm(nn.Module):
+    def forward(self, x):
+        n_chan = x.shape[1]
+        return x  / (1/ n_chan * (x ** 2).sum(dim = 1, keepdims = True)).sqrt()
+
+class Generator(nn.Module):
+    def __init__(self, input_size:int = 317):
+        super().__init__()
+        self.input_size = input_size
+        self.input_block = nn.Sequential(
+                nn.Conv2d(input_size, 256, (16, 2)),
+                PixelNorm(), nn.LeakyReLU(),
+                nn.Conv2d(256, 256, 3, padding="same"),
+                PixelNorm(), nn.LeakyReLU())
+        self.upsample2x2 =  lambda c:  nn.ConvTranspose2d(c, c, (2,2),(2,2))
+        self.conv_block = lambda in_c, out_c: nn.Sequential(
+                nn.Conv2d(in_c, out_c, 3, padding="same"),
+                PixelNorm(), nn.LeakyReLU(),
+                nn.Conv2d(out_c, out_c, 3, padding="same"),
+                PixelNorm(), nn.LeakyReLU())
+        self.output = nn.Sequential(
+            nn.Conv2d(32, 2, 1),
+            nn.Tanh()
+        )
+
+        self.network = nn.Sequential(
+            self.input_block,
+            self.upsample2x2(256),
+            self.conv_block(256, 256),
+            self.upsample2x2(256),
+            self.conv_block(256, 256),
+            self.upsample2x2(256),
+            self.conv_block(256, 256),
+            self.upsample2x2(256),
+            self.conv_block(256, 128),
+            self.upsample2x2(128),
+            self.conv_block(128, 64),
+            self.upsample2x2(64),
+            self.conv_block(64, 32),
+            self.output
+        )
+    def forward(self, x:torch.Tensor)-> torch.Tensor:
+        x = F.pad(x,(1, 1, 15, 15))
+        return self.network(x)
+
+class Discriminator(nn.Module):
+    def __init__(self, num_classes:int):
+        super().__init__()
+        self.input = nn.Conv2d(2, 32, 1)
+        self.downsample2x2 =  lambda c:  nn.Conv2d(c, c, (2,2),(2,2))
+        self.conv_block = lambda in_c, out_c: nn.Sequential(
+                nn.Conv2d(in_c, out_c, 3, padding = "same"),
+                nn.LeakyReLU(),
+                nn.Conv2d(out_c, out_c, 3, padding = "same"),
+                nn.LeakyReLU()
+        )
+        self.last_conv = nn.Sequential(
+            nn.Conv2d(256, 256, 3, padding="same"),
+            nn.LeakyReLU(),
+            nn.Conv2d(256, 256, (16,2), padding="valid"),
+            nn.LeakyReLU()
+        )
+        self.disc_out = nn.Linear(256, 1)
+        self.classifier_out = nn.Linear(256, num_classes)
+
+        self.network = nn.Sequential(
+                    self.input,
+                    self.conv_block(32, 32),
+                    self.downsample2x2(32),
+                    self.conv_block(32, 64),
+                    self.downsample2x2(64),
+                    self.conv_block(64, 128),
+                    self.downsample2x2(128),
+                    self.conv_block(128, 256),
+                    self.downsample2x2(256),
+                    self.conv_block(256, 256),
+                    self.downsample2x2(256),
+                    self.conv_block(256, 256),
+                    self.downsample2x2(256),
+                    self.conv_block(256, 256),
+                    self.last_conv
+                )
+    def forward(self, x:torch.Tensor)-> torch.Tensor:
+        x =  self.network(x).flatten(1)
+        return  self.disc_out(x),  self.classifier_out(x)

fakesound/gansynth/utils.py

@@ -0,0 +1,31 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+gp_beta = 10
+pitch_beta = 10
+crossentropy = nn.CrossEntropyLoss()
+
+# https://github.com/EmilienDupont/wgan-gp/
+def gradient_penalty(real_data, fake_data, batch_size, disc, device):
+    alpha = torch.rand(batch_size, 1, 1, 1)
+    alpha = alpha.expand_as(real_data).to(device)
+    interpolated = alpha * real_data.detach() + (1 - alpha) * fake_data.detach()
+    interpolated = torch.autograd.Variable(interpolated, requires_grad=True)
+    prob, _ = disc(interpolated)
+    gradients = torch.autograd.grad(outputs = prob,
+                           inputs = interpolated,
+                           grad_outputs = torch.ones_like(prob),
+                           create_graph = True, retain_graph = True)[0]
+
+    gradients = gradients.view(batch_size, -1)
+    gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()
+    return gradient_penalty
+
+def create_fake_data(generator, config):
+    fake_pitch = torch.randint( 0, 61, (config.batch_size,)).unsqueeze(1).to(config.device)
+    fake_pitch_vector = torch.zeros(config.batch_size, 61).to(config.device).scatter_(1, fake_pitch, 1).unsqueeze(2).unsqueeze(3).to(config.device)
+    noise = torch.randn(config.batch_size, config.latent_size, 1, 1).to(config.device)
+    gen_input =  torch.cat((noise, fake_pitch_vector), 1)
+    fake_mel = generator(gen_input)
+    return fake_pitch, fake_mel

fakesound/hifigan/blocks.py

@@ -0,0 +1,57 @@
+# based on https://github.com/jik876/hifi-gan
+# only v1 implemented
+
+import torch
+from torch import nn
+from torch.nn.utils import remove_weight_norm
+from torch.nn.utils.parametrizations import weight_norm
+
+
+class ResidualBlock(torch.nn.Module):
+    def __init__(self, channels:int, kernel_size:int = 3, dilations:list[int] = [1, 3, 5]):
+        super().__init__()
+        self.conv = lambda d: weight_norm(
+            nn.Conv1d(channels, channels, kernel_size, 1,
+                dilation = d, padding = (kernel_size * d - d) // 2)
+        )
+        self.convs = nn.Sequential()
+        for d in dilations:
+            self.convs.extend(nn.Sequential(*[
+                nn.LeakyReLU(0.1), self.conv(d), nn.LeakyReLU(0.1), self.conv(1)
+            ]))
+    def forward(self, x:torch.Tensor)->torch.Tensor:
+        return x + self.convs(x)
+
+    def remove_weight_norm(self):
+        for l in self.convs:
+            if type(l) == nn.Conv1d:
+                remove_weight_norm(l)
+
+
+class UpsamplingBlock(torch.nn.Module):
+    def __init__(self, in_channels:int, out_channels:int, kernel_size:int, stride:int,
+                 res_channels:int, res_kernels:list[int], res_dilations:list[int]):
+        assert(len(res_kernels) == len(res_dilations))
+        super().__init__()
+        self.num_res_blocks = len(res_kernels)
+        self.relu = nn.LeakyReLU(0.1)
+        self.conv_transpose = nn.ConvTranspose1d(
+                            in_channels,out_channels, kernel_size, stride,
+                            (kernel_size - stride ) // 2)
+        self.res_blocks = nn.ModuleList(
+                [ResidualBlock(res_channels, res_kernels[i], res_dilations)
+                    for i in range(self.num_res_blocks)]
+        )
+
+    def forward(self, x:torch.Tensor)->torch.Tensor:
+        x = self.conv_transpose(self.relu(x))
+        xs = 0
+        for res_block in self.res_blocks:
+             xs += res_block(x)
+        x = xs / self.num_res_blocks
+        return x
+
+    def remove_weight_norm(self):
+        remove_weight_norm(self.conv_transpose)
+        for res_block in self.res_blocks:
+            res_block.remove_weight_norm()

fakesound/hifigan/loss_funcs.py

@@ -0,0 +1,29 @@
+import torch
+
+def feature_loss(fmap_real:list, fmap_fake:list):
+    loss = 0
+    for disc_real, disc_fake  in zip(fmap_real, fmap_fake):
+        for layer_real, layer_fake in zip(disc_real, disc_fake):
+            loss += torch.mean(torch.abs(layer_real - layer_fake))
+    return loss * 2
+
+def discriminator_loss(disc_real:list, disc_fake:list):
+    loss = 0
+    real_losses = []
+    fake_losses = []
+    for subdisc_real, subdisc_fake in zip(disc_real, disc_fake):
+        real_loss = torch.mean((1 - subdisc_real) ** 2)
+        fake_loss = torch.mean(subdisc_fake ** 2)
+        loss += (real_loss + fake_loss)
+        real_losses.append(real_loss.item())
+        fake_losses.append(fake_loss.item())
+    return loss, real_losses, fake_losses
+
+def generator_loss(disc_outputs:list):
+    loss = 0
+    generator_losses = []
+    for subdisc_fake in disc_outputs:
+        l = torch.mean((1 - subdisc_fake) ** 2)
+        generator_losses.append(l)
+        loss += l
+    return loss, generator_losses

fakesound/hifigan/models.py

@@ -0,0 +1,168 @@
+# based on https://github.com/jik876/hifi-gan
+# only v1 implemented
+#
+import torch
+from .blocks import ResidualBlock
+from .blocks import UpsamplingBlock
+
+from torch import nn
+import torch.nn.functional as F
+from torch.nn.utils.parametrizations import weight_norm, spectral_norm
+from torch.nn.utils import remove_weight_norm
+
+
+class Generator(torch.nn.Module):
+    def __init__(self,
+        n_mels,
+        upsample_kernels,
+        upsample_strides,
+        upsample_start,
+        residual_kernels,
+        residual_dilations
+    ):
+        super().__init__()
+        assert(len(upsample_kernels) == len(upsample_strides))
+        self.num_upsamples = len(upsample_kernels)
+        self.num_res_blocks = len(residual_kernels)
+        self.relu = nn.LeakyReLU(0.1)
+        self.tanh = nn.Tanh()
+        self.initial_conv = weight_norm(nn.Conv1d(n_mels, upsample_start, 7, 1, 3))
+        self.upsample = nn.Sequential(*[UpsamplingBlock(
+            upsample_start // (2 ** i),
+            upsample_start // (2 ** (i + 1)),
+            upsample_kernels[i],
+            upsample_strides[i],
+            upsample_start // (2 ** (i + 1)),
+            residual_kernels,
+            residual_dilations
+        ) for i in range(self.num_upsamples)])
+        self.final_conv = weight_norm(nn.Conv1d(
+            upsample_start//(2 ** (self.num_upsamples)), 1, 7, 1, 3)
+        )
+
+    def forward(self, x:torch.Tensor)->torch.Tensor:
+        x = self.initial_conv(x)
+        x = self.upsample(x)
+        x = self.relu(x)
+        x = self.final_conv(x)
+        x = self.tanh(x)
+        return x
+
+    def remove_weight_norm(self):
+        remove_weight_norm(self.initial_conv)
+        for u in self.upsample.modules():
+            u.remove_weight_norm()
+        remove_weight_norm(self.final_conv)
+
+class PeriodDiscriminator(torch.nn.Module):
+    def __init__(self, period:int, kernel_size:int = 5, stride:int = 3):
+        super().__init__()
+        self.period = period
+        self.conv = lambda in_c, out_c: weight_norm(
+            nn.Conv2d(in_c, out_c,
+                (kernel_size, 1), (stride, 1),
+                ((kernel_size -1)//2, 0))
+        )
+        self.convs = nn.ModuleList([
+            self.conv(1, 32),
+            self.conv(32, 128),
+            self.conv(128, 512),
+            self.conv(512, 1024),
+            weight_norm(nn.Conv2d(1024, 1024,
+                (kernel_size, 1), 1, (2, 0)))
+        ])
+        self.final_conv = weight_norm(
+            nn.Conv2d(1024, 1, (3, 1), 1, (1, 0))
+        )
+
+    def forward(self, x:torch.Tensor)->torch.Tensor:
+        fmap = []
+        n_batches, n_chans, n_frames = x.shape
+        if n_frames % self.period != 0:
+            n_pad = self.period - (n_frames % self.period)
+            x = F.pad(x, (0, n_pad), "reflect")
+            n_frames += n_pad
+        x = x.view(n_batches, n_chans, n_frames // self.period, self.period)
+
+        for conv in self.convs:
+            x = F.leaky_relu(conv(x), 0.1)
+            fmap.append(x)
+        x = self.final_conv(x)
+        fmap.append(x)
+        return x.flatten(1, -1), fmap
+
+class MultiPeriodDiscriminator(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.sub_discriminators = nn.ModuleList([
+            PeriodDiscriminator(n) for n in [2, 3, 5, 7, 11]
+        ])
+
+    def forward(self, y:torch.Tensor, y_hat:torch.Tensor):
+        y_real = []
+        y_fake = []
+        fmap_real = []
+        fmap_fake = []
+        for d in self.sub_discriminators:
+            y_r, fmap_r = d(y)
+            y_real.append(y_r)
+            fmap_real.append(fmap_r)
+            y_f, fmap_f = d(y_hat)
+            y_fake.append(y_f)
+            fmap_fake.append(fmap_f)
+        return y_real, y_fake, fmap_real, fmap_fake
+
+class ScaleDiscriminator(torch.nn.Module):
+    def __init__(self, use_spectral_norm:bool = False):
+        super().__init__()
+        norm = spectral_norm if not use_spectral_norm else weight_norm
+        self.convs = nn.ModuleList([
+             norm(nn.Conv1d(1, 128, 15, 1, 7)),
+             norm(nn.Conv1d(128, 128, 41, 2, 20, 1, 4)),
+             norm(nn.Conv1d(128, 256, 41, 2, 20, 1, 16)),
+             norm(nn.Conv1d(256, 512, 41, 4, 20, 1, 16)),
+             norm(nn.Conv1d(512, 1024, 41, 4, 20, 1, 16)),
+             norm(nn.Conv1d(1024, 1024, 41, 1, 20, 1, 16)),
+             norm(nn.Conv1d(1024, 1024, 5, 1, 2))
+        ])
+        self.final_conv = norm(nn.Conv1d(1024, 1, 3, 1, 1))
+
+    def forward(self, x:torch.Tensor)->torch.Tensor:
+        fmap = []
+        for conv in self.convs:
+            x = F.leaky_relu(conv(x), 0.1)
+            fmap.append(x)
+        x = self.final_conv(x)
+        fmap.append(x)
+        return x.flatten(1, -1), fmap
+
+
+class MultiScaleDiscriminator(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.sub_discriminators = nn.ModuleList([
+            ScaleDiscriminator(use_spectral_norm=True),
+            ScaleDiscriminator(),
+            ScaleDiscriminator(),
+        ])
+        self.avg_pools = nn.ModuleList([
+            nn.AvgPool1d(4, 2, 2),
+            nn.AvgPool1d(4, 2, 2)
+        ])
+
+    def forward(self, y:torch.Tensor, y_hat:torch.Tensor)->tuple[torch.Tensor,torch.Tensor,torch.Tensor,torch.Tensor]:
+        y_real = []
+        y_fake = []
+        fmap_real = []
+        fmap_fake = []
+        for i, d in enumerate(self.sub_discriminators):
+            if i != 0:
+                y = self.avg_pools[i-1](y)
+                y_hat = self.avg_pools[i-1](y_hat)
+            y_r, fmap_r = d(y)
+            y_real.append(y_r)
+            fmap_real.append(fmap_r)
+            y_f, fmap_f = d(y_hat)
+            y_fake.append(y_f)
+            fmap_fake.append(fmap_f)
+        return y_real, y_fake, fmap_real, fmap_fake

fakesound/rave/core_funcs.py

@@ -0,0 +1,46 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+import torch.fft as fft
+
+# From
+# https://github.com/acids-ircam/RAVE/blob/master/rave/core.py
+
+def mod_sigmoid(x):
+    return 2 * torch.sigmoid(x)**2.3 + 1e-7
+
+def get_augmented_latent_size(latent_size: int, noise_augmentation: int):
+    return latent_size + noise_augmentation
+
+def amp_to_impulse_response(amp, target_size):
+    """
+    transforms frequency amps to ir on the last dimension
+    """
+    amp = torch.stack([amp, torch.zeros_like(amp)], -1)
+    amp = torch.view_as_complex(amp)
+    amp = fft.irfft(amp)
+
+    filter_size = amp.shape[-1]
+
+    amp = torch.roll(amp, filter_size // 2, -1)
+    win = torch.hann_window(filter_size, dtype=amp.dtype, device=amp.device)
+
+    amp = amp * win
+
+    amp = F.pad(
+        amp,
+        (0, int(target_size) - int(filter_size)),
+    )
+    amp = torch.roll(amp, -filter_size // 2, -1)
+
+    return amp
+
+def fft_convolve(signal, kernel):
+    """
+    convolves signal by kernel on the last dimension
+    """
+    signal = F.pad(signal, (0, signal.shape[-1]))
+    kernel = F.pad(kernel, (kernel.shape[-1], 0))
+    output = fft.irfft(fft.rfft(signal) * fft.rfft(kernel))
+    output = output[..., output.shape[-1] // 2:]
+    return output

fakesound/rave/discriminator.py

@@ -0,0 +1,47 @@
+import torch
+from torch import nn
+from torch.nn.utils.parametrizations import weight_norm
+from torch.nn import Conv1d, Sequential
+
+class Discriminator(nn.Module):
+    def __init__(self, in_channels:int, out_channels:int, n_layers:int, stride:int):
+        super().__init__()
+        layers = []
+        for i in range(n_layers):
+            layers.append(
+                Sequential(
+                    weight_norm(
+                        Conv1d(in_channels, out_channels,
+                        15, stride, padding = 7)
+                    ),
+                    nn.LeakyReLU(0.2)
+                )
+            )
+            in_channels = out_channels
+            out_channels = out_channels * 2
+        layers.append(Conv1d(in_channels, 1, 1))
+        self.layers = Sequential(*layers)
+
+    def forward(self, x:torch.Tensor)->torch.Tensor:
+        features = []
+        for layer in self.layers:
+            x = layer(x)
+            # the original RAVE appends the outputs of convs
+            # here we follow hifigan and use the output of the LeakyReLU
+            features.append(x)
+        return features
+
+class MultiScaleDiscriminator(nn.Module):
+
+    def __init__(self, in_channels:int, out_channels:int, n_scales:int = 3):
+        super().__init__()
+        self.n_scales = n_scales
+        self.layers = nn.ModuleList([
+            Discriminator(in_channels, out_channels, 4, 4) for _ in range(n_scales)
+        ])
+    def forward(self, x:torch.Tensor)->list[torch.Tensor]:
+        features = []
+        for layer in self.layers:
+            features.append(layer(x))
+            x = nn.functional.avg_pool1d(x, 2)
+        return features

fakesound/rave/encoder.py

@@ -0,0 +1,28 @@
+import torch
+from torch import nn
+from torch.nn import Conv1d, Sequential
+import torch.nn.functional as F
+
+class Encoder(nn.Module):
+    def __init__(self, n_bands:int, capacity:int,
+                  latent_dim:int, ratios:list[int], n_out:int):
+        super().__init__()
+        self.input = nn.Conv1d(n_bands, capacity, 7 , padding = 3)
+        self.block = lambda n_in, n_out, ratio: Sequential(*[
+            nn.BatchNorm1d(n_in),
+            nn.LeakyReLU(.2),
+            Conv1d(n_in, n_out, 2 * ratio + 1,
+                   padding = ratio, stride = ratio)
+        ])
+        self.output = Conv1d(
+            2 ** len(ratios) * capacity, latent_dim * n_out, 5,
+            padding = 2, groups = n_out,
+        )
+        layers = [self.input]
+        for i, r in enumerate(ratios):
+            layers += self.block( 2**i * capacity, 2**(i + 1) * capacity, r)
+        layers += [nn.LeakyReLU(.2), self.output]
+        self.net = Sequential(*layers)
+
+    def forward(self, x:torch.Tensor)->torch.Tensor:
+        return self.net(x)

fakesound/rave/generator.py

@@ -0,0 +1,102 @@
+import torch
+from torch import nn
+from torch.nn import Conv1d, ConvTranspose1d, Sequential
+import torch.nn.functional as F
+from torch.nn.utils.parametrizations import weight_norm
+import numpy as np
+from .core_funcs import mod_sigmoid, amp_to_impulse_response, fft_convolve
+
+class ResidualLayer(nn.Module):
+    def __init__(self, dim:int, kernel_size:int, dilations:list[int]):
+        super().__init__()
+        layers = []
+        for d in dilations:
+            layers+=[nn.LeakyReLU(0.2),
+                weight_norm(Conv1d(dim, dim, kernel_size,
+                            padding = d, dilation = d))]
+        self.net = Sequential(*layers)
+
+    def forward(self, x:torch.Tensor)->torch.Tensor:
+        return x + self.net(x)
+
+class ResidualStack(nn.Module):
+    def __init__(self,
+                 dim,
+                 kernel_size,
+                 dilations_list):
+        super().__init__()
+        self.net = Sequential(*[
+                    ResidualLayer(
+                        dim,
+                        kernel_size,
+                        dilations
+                    ) for dilations in dilations_list]
+                )
+    def forward(self, x:torch.Tensor)->torch.Tensor:
+        return self.net(x)
+
+class NoiseGenerator(nn.Module):
+    def __init__(self, in_size:int, n_bands:int, ratios:list[int], noise_bands:int):
+        super().__init__()
+        net = []
+        channels = [in_size] * len(ratios) + [n_bands * noise_bands]
+        for i, r in enumerate(ratios):
+            net.append(
+                Conv1d(channels[i], channels[i + 1], 3, stride=r))
+            if i != len(ratios) - 1:
+                net.append(nn.LeakyReLU(.2))
+        self.net = Sequential(*net)
+        self.n_bands = n_bands
+        self.register_buffer(
+            "target_size",
+            torch.tensor(np.prod(ratios)).long(),
+        )
+
+    def forward(self, x:torch.Tensor)->torch.Tensor:
+        amp = mod_sigmoid(self.net(x) - 5)
+        amp = amp.permute(0, 2, 1)
+        amp = amp.reshape(amp.shape[0], amp.shape[1], self.n_bands, -1)
+        ir = amp_to_impulse_response(amp, self.target_size)
+        noise = torch.rand_like(ir) * 2 - 1
+        noise = fft_convolve(noise, ir).permute(0, 2, 1, 3)
+        noise = noise.reshape(noise.shape[0], noise.shape[1], -1)
+        return noise
+
+class Generator(nn.Module):
+    def __init__(self, latent_dim:int, capacity:int,
+                  n_bands:int, ratios:list[int]):
+        super().__init__()
+        self.input = weight_norm(
+            Conv1d(
+                latent_dim, 2 ** len(ratios) * capacity,
+                7, padding = 3
+            ))
+        self.block = lambda n_in, n_out, ratio: Sequential(*[
+            nn.LeakyReLU(.2),
+            ConvTranspose1d(n_in, n_out, 2 * ratio, ratio, ratio // 2),
+            ResidualStack(n_out,[3], [[1, 1], [3, 1], [5, 1]])
+        ])
+        self.wave_gen =  weight_norm(
+            Conv1d(capacity, n_bands, 7, padding = 3)
+        )
+        self.loud_gen = weight_norm(
+            Conv1d(capacity, 1, 3, padding = 1)
+        )
+
+        self.noise_gen = NoiseGenerator(capacity, n_bands, [4, 4, 4], 5)
+
+        layers = [self.input]
+        for i, r in enumerate(ratios):
+            layers += self.block( 2**(len(ratios) - i) * capacity, 2**(len(ratios) - i - 1) * capacity, r)
+        self.net = Sequential(*layers)
+        self.branches = [self.wave_gen, self.loud_gen, self.noise_gen]
+        self.add_noise = False
+    def forward(self, x:torch.Tensor)->torch.Tensor:
+            x = self.net(x)
+            waveform = self.wave_gen(x)
+            loudness = self.loud_gen(x)
+            noise = self.noise_gen(x)
+            loudness = loudness.reshape(x.shape[0], 1, -1)
+            waveform = torch.tanh(waveform) * mod_sigmoid(loudness)
+            if self.add_noise: waveform += noise
+            return waveform

fakesound/rave/model.py

@@ -0,0 +1,31 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from .encoder import Encoder
+from .generator import Generator
+
+class Rave(nn.Module):
+    def __init__(self, n_bands:int, capacity:int,
+                  latent_dim:int, ratios:list[int], n_out:int, beta:float = 1.):
+        super().__init__()
+        self.encoder = Encoder(n_bands, capacity, latent_dim, ratios, n_out)
+        self.generator = Generator(latent_dim, capacity, n_bands, ratios)
+        self.beta = beta
+
+    def encode(self, x:torch.Tensor)->tuple[torch.Tensor, torch.Tensor]:
+        z = self.encoder(x)
+        mean, scale = z.chunk(2, 1)
+        std = F.softplus(scale) + 1e-4
+        return mean, std
+
+    def decode(self, x:torch.Tensor)->torch.Tensor:
+        return self.generator(x)
+
+    def reparametrize(self, mean:torch.Tensor, std:torch.Tensor)->torch.Tensor:
+        z = mean + torch.randn_like(mean) * std
+        return z
+
+    def forward(self, x:torch.Tensor)->torch.Tensor:
+        mean, logvar = self.encode(x)
+        z = self.reparametrize(mean, logvar)
+        return self.decode(z)

fakesound/samplernn/blocks.py

@@ -0,0 +1,65 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+class SampleLevelModule(nn.Module):
+     def __init__(self, frame_size:int, embedding_size:int, dim:int, num_levels:int):
+        nn.Module.__init__(self)
+        self.num_levels = num_levels
+        self.embedding = nn.Embedding(
+            num_levels,
+            embedding_size
+        )
+        self.input = nn.Conv1d(num_levels, dim, frame_size, bias = False)
+        self.hidden = nn.Sequential(
+                nn.Conv1d(dim, dim, 1),
+                nn.ReLU(),
+                nn.Conv1d(dim, dim,1),
+                nn.ReLU()
+        )
+        self.output = nn.Conv1d(dim, num_levels, 1)
+
+     def forward(self, x:torch.Tensor, cond:torch.Tensor):
+        x = self.embedding(x)
+        # the following steps are implemented using Conv1d,
+        # so here per-sample vector becomes channels
+        x = x.permute(0, 2, 1)
+        cond = cond.permute(0, 2, 1)
+        x = self.input(x)
+        x = x + cond
+        x = self.hidden(x)
+        x = self.output(x)
+        # back to samples * levels
+        x = x.permute(0, 2, 1)
+        return x
+
+class FrameLevelModule(nn.Module):
+    def __init__(self, frame_size:int, num_steps:int,
+                 upsampling_factor:int, num_layers:int, dim:int):
+       nn.Module.__init__(self)
+       self.frame_size = frame_size
+       self.state = None
+       self.num_layers = num_layers
+       self.num_steps = num_steps
+       self.upsampling_factor = upsampling_factor
+       self.dim = dim
+       self.input = nn.Linear(
+           frame_size,
+           dim,
+           bias = False
+       )
+       self.rnn = nn.GRU(dim, dim, num_layers,
+                         bias = True, batch_first = True)
+       self.upsampling = nn.ConvTranspose1d(dim, dim, upsampling_factor, upsampling_factor, bias = False)
+
+    def forward(self, x:torch.tensor, reset:bool, cond:torch.Tensor)->torch.Tensor:
+        x = self.input(x) + (cond if cond is not None else 0)
+        # reset applies to whole batch
+        if reset or self.state == None:
+            batch_size = x.shape[0]
+            self.state = torch.zeros(self.num_layers, batch_size, self.dim)
+        x, h = self.rnn(x, self.state)
+        self.state = h.detach()
+        # ConvTranspose1D upsamples on last dimension
+        x = self.upsampling(x.permute(0, 2, 1)).permute(0, 2, 1)
+        return x

fakesound/samplernn/model.py

@@ -0,0 +1,75 @@
+from collections.abc import Callable
+import torch
+from torch import nn
+from torch.nn import functional as F
+from .blocks import FrameLevelModule, SampleLevelModule
+
+class SampleRNN(nn.Module):
+    def __init__(self, seq_len:int, sizes:list[int], dim:int,
+                 rnn_layers:int, num_levels:int):
+        nn.Module.__init__(self)
+        self.dim = dim
+        self.num_levels = num_levels
+        self.top_frame = sizes[-1]
+        self.bottom_frame = sizes[0]
+        self.seq_len = seq_len
+
+        self.rnns = nn.ModuleList()
+        steps_below = self.seq_len
+        for frame_size in sizes: # from bottom to top rnn
+            num_steps = self.seq_len // frame_size
+            self.rnns.append(FrameLevelModule(frame_size,
+                                 num_steps,
+                                 steps_below// num_steps,
+                                 rnn_layers,
+                                 dim))
+            steps_below = num_steps
+        self.mlp = SampleLevelModule(self.bottom_frame,
+            self.num_levels, dim, self.num_levels)
+
+    @property
+    def lookback(self):
+        return self.top_frame
+
+    def reset(self):
+        for rnn in self.rnns: rnn.state = None
+
+    def forward(self, x:torch.Tensor, original:torch.Tensor, reset = True)->torch.Tensor:
+        assert(x.shape[1] == self.seq_len+self.top_frame-1)
+        batch_size = x.shape[0]
+        cond = None
+        for rnn in reversed(self.rnns):
+            start = self.top_frame - rnn.frame_size
+            end = -rnn.frame_size + 1
+            rnn_input = original[:, start:end]
+            rnn_input = rnn_input.contiguous().view(batch_size, -1, rnn.frame_size)
+            cond = rnn(rnn_input, reset, cond)
+        out = self.mlp(x[:,self.top_frame - self.bottom_frame:], cond)
+        return out
+
+    def generate(self, batch_size:int, seq_len:int, decoder:Callable, temp:float = 1.0)->torch.Tensor:
+        device = next(self.parameters()).device
+        input_size = seq_len + self.lookback
+        samples = torch.full((batch_size, input_size), self.num_levels// 2, device = device)
+        frame_outs = [None for _ in self.rnns]
+        cond = None
+        self.reset()
+        for i in range(self.lookback, self.lookback + seq_len):
+            for (tier, rnn) in reversed(list(enumerate(self.rnns))):
+                if i % rnn.frame_size != 0:
+                    continue
+                prev_samples = decoder(samples[:, i - rnn.frame_size : i])
+                if tier == len(self.rnns) - 1:
+                    cond = None
+                else:
+                    frame_index = (i // rnn.frame_size) % self.rnns[tier + 1].upsampling_factor
+                    cond = frame_outs[tier + 1][:, frame_index, :].unsqueeze(1)
+                rnn_input = prev_samples.contiguous().view(batch_size, -1, rnn.frame_size)
+                frame_outs[tier] = rnn(rnn_input, False, cond)
+            mlp_input = samples[:, i - self.bottom_frame : i]
+            cond = frame_outs[0][:, i % self.bottom_frame, :].unsqueeze(1)
+            out = self.mlp(mlp_input, cond)
+            out = out.softmax(2) / temp
+            samples[:, i] = out.squeeze(1).multinomial(1).squeeze(1)
+        self.reset() # allow generation of different batch size
+        return samples[:, self.lookback :]

fakesound/util/train_utils.py

@@ -0,0 +1,31 @@
+import sys, json, torch
+from collections import namedtuple
+
+def get_config(default_filename:str):
+    config_file = default_filename
+    if len(sys.argv) > 1: config_file = sys.argv[1]
+    f = open(config_file)
+    try:
+        data = json.load(f)
+    except:
+        print("Error parsing configuration file")
+        sys.exit()
+
+    return namedtuple('Config', data.keys())(**data)
+
+def save_model(model:torch.nn.Module,
+               optimizer:torch.optim.Optimizer,
+               step:int, path:str):
+        torch.save({
+            'model': model.state_dict(),
+            'optimizer': optimizer.state_dict(),
+            'step': step,
+        }, path)
+
+def load_model(path:str, model:torch.nn.Module,
+               optimizer:torch.optim.Optimizer):
+        checkpoint = torch.load(path)
+        model.load_state_dict(checkpoint['model'])
+        optimizer.load_state_dict(checkpoint['optimizer'])
+        step = checkpoint['step']
+        return step

fakesound/wavenet/blocks.py

@@ -0,0 +1,18 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from .layers import GatedDilatedConv1d
+
+class WavenetBlock (nn.Module):
+    def __init__(self, dilation_channels, residual_channels, skip_channels,
+                 dilations, kernel_size = 2, bias = False):
+        super().__init__()
+        self.layers = nn.ModuleList([
+            GatedDilatedConv1d(dilation_channels, residual_channels, skip_channels,
+                               d, kernel_size, bias)
+            for d in dilations])
+
+    def forward(self, x:torch.Tensor, skip)->torch.Tensor:
+        for l in self.layers:
+            x, skip = l(x, skip)
+        return x, skip

fakesound/wavenet/layers.py

@@ -0,0 +1,41 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+class CausalConv1d(nn.Module):
+    def __init__(self, in_channels:int, out_channels:int, kernel_size:int = 2 , bias:bool = False):
+        super().__init__()
+        self.pad = (1, -1) # enforce causality
+        self.conv = nn.Conv1d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            bias = bias
+        )
+    def forward(self, x:torch.Tensor, pad = True)->torch.Tensor:
+        if pad:
+            return self.conv(F.pad(x, self.pad))
+        else:
+            return self.conv(x)
+
+class GatedDilatedConv1d(nn.Module):
+# AKA Wavenet Layer, with residual and skip connections
+    def __init__(self, dilation_channels:int, residual_channels:int, skip_channels:int,
+                       dilation:int, kernel_size:int = 2, bias:bool = False):
+        super().__init__()
+        self.pad =  ((kernel_size - 1) * dilation, 0)
+        self.conv1d = lambda in_c, out_c: nn.Conv1d(in_c, out_c, kernel_size,
+                                                    bias = bias, dilation = dilation)
+        self.conv1d_1x1 = lambda in_c, out_c: nn.Conv1d(in_c, out_c, 1,
+                                                        bias = bias, dilation = dilation)
+        self.filter_conv = self.conv1d(residual_channels, dilation_channels)
+        self.gate_conv = self.conv1d(residual_channels, dilation_channels)
+        self.residual_conv = self.conv1d_1x1(dilation_channels, residual_channels)
+        self.skip_conv = self.conv1d_1x1(dilation_channels, skip_channels)
+        self.test_conv = nn.Conv1d(16, 32, kernel_size)
+
+    def forward(self, x:torch.Tensor, skip)->tuple[torch.Tensor,torch.Tensor]:
+        padded = F.pad(x, self.pad)
+        out = F.tanh(self.filter_conv(padded)) * F.sigmoid(self.gate_conv(padded))
+        skip_slice = out[:,:,-skip.shape[2]:]
+        skip += self.skip_conv(skip_slice)
+        return self.residual_conv(out) + x, skip

fakesound/wavenet/model.py

@@ -0,0 +1,50 @@
+import torch
+from torch import nn
+import numpy as np
+from .blocks import WavenetBlock
+from .layers import CausalConv1d
+from torch.distributions import Categorical
+
+class WaveNet (nn.Module):
+    def __init__(self, dilation_channels:int, residual_channels:int, skip_channels:int,
+                 num_blocks:int, layers_per_block:int, output_length:int,
+                 dilation_factor:int = 2, kernel_size:int = 2, bias:bool = False, num_classes:int = 256):
+        nn.Module.__init__(self)
+        self.output_length:int = output_length
+        self.num_classes = num_classes
+        self.skip_channels:int = skip_channels
+        block_dilations = [dilation_factor ** i for i in range(layers_per_block)]
+        self.receptive_field:int =  (kernel_size - 1) * num_blocks * sum(block_dilations) + 1
+        self.causal_conv:CausalConv1d = CausalConv1d(num_classes, residual_channels, kernel_size, bias)
+        self.blocks:nn.ModuleList = nn.ModuleList([WavenetBlock(dilation_channels, residual_channels, skip_channels,
+                                      block_dilations, kernel_size, bias)
+                                    for _ in range(num_blocks)])
+        self.end_seq = nn.Sequential(
+            nn.ReLU(),
+            nn.Conv1d(skip_channels, skip_channels, 1, bias = True),
+            nn.ReLU(),
+            nn.Conv1d(skip_channels, num_classes, 1, bias = True)
+        )
+
+    def forward(self, x:torch.Tensor, output_length:int = None):
+         if  output_length == None: output_length = self.output_length
+         skip = torch.zeros(x.shape[0], self.skip_channels, output_length)
+         x = self.causal_conv(x, self.training)
+         for b in self.blocks:
+             x, skip = b(x, skip)
+         return self.end_seq(skip)
+
+    def generate(self, batch_size, num_samples:int)->torch.Tensor:
+        input_size = self.receptive_field
+        device = next(self.parameters()).device
+        generated = torch.full((batch_size, input_size), self.num_classes // 2, device = device)
+        model_input = torch.zeros((batch_size, self.num_classes, input_size), device = device)
+
+        for i in range(num_samples):
+            model_input.zero_()
+            model_input = model_input.scatter_(1, generated[:, -input_size:].view(-1, 1, input_size), 1.)
+            out = self.forward(model_input, 1)[:, :, -1]
+            pdf = Categorical(logits = out)
+            out = pdf.sample()
+            generated = torch.cat([generated, out.unsqueeze(1)], 1)
+        return generated

fakesound-0.0.0.dist-info/METADATA

@@ -0,0 +1,6 @@
+Metadata-Version: 2.4
+Name: fakesound
+Version: 0.0.0
+Requires-Python: >=3.10
+License-File: LICENSE
+Dynamic: license-file

fakesound-0.0.0.dist-info/RECORD

@@ -0,0 +1,26 @@
+fakesound/data/dataset.py,sha256=0rch8V3gDoTqbSvNRFG9fF8PiT_S7A_wv8oHVf8_Hb8,2133
+fakesound/data/mel_bands.py,sha256=cUqA-vkhuaHo2oadzJn8vQufPXJpsSaGM32FaELOlCU,1231
+fakesound/data/pqmf.py,sha256=uT04kwkXd3f9XuFSSVik5M-NX2LK_0fQ4w2x7Ybfuj4,764
+fakesound/data/transforms.py,sha256=aXMYucQMiSfnn0GkDqryZBQY5enAitFyrT9dvJyE5IA,6759
+fakesound/gansynth/dataset.py,sha256=Qhdc0gUHmyS3BZEHL5_DqwDz3S-44LyzffhM_wyHINs,1913
+fakesound/gansynth/models.py,sha256=n9gGBmXxd9B44Cs-Ml5-WOtkR4dyTAIaA2uCIk47ySs,3379
+fakesound/gansynth/utils.py,sha256=a98jcKJKV7EGVdi5BPj-bgbrjOfqwETxAfmyE2CVILA,1411
+fakesound/hifigan/blocks.py,sha256=UtwYbsVDq2xYJrLXmuKW6-afVm8ej3-7c7pN-NVjgeg,2135
+fakesound/hifigan/loss_funcs.py,sha256=CtAaFVbRJY-7Z1jGLVgBAa13N5YZCgiOhr6Ps9SRFHE,996
+fakesound/hifigan/models.py,sha256=CYPCLqRjrgIkc35U3n52G3kFBcoHXBEU-eSf4Tl8Elk,5751
+fakesound/rave/core_funcs.py,sha256=y38SQjHdS3ahALxhJZDt04PgvPW8xpzbiYC16yczoQQ,1248
+fakesound/rave/discriminator.py,sha256=X07NfiWmbmnmlLeR2UTZ3ytRj70GZkOqnTwR4LDmjhc,1653
+fakesound/rave/encoder.py,sha256=agU_87F9eYwYjnlW6zdMT6DIIQP6dtbXak4FLIPT2h8,1038
+fakesound/rave/generator.py,sha256=Pi-c77UO75IBUnOuLChZSnGu3FKxaqpAgLZYPAMKiZo,3850
+fakesound/rave/model.py,sha256=FFv22KD2Zeeh68xbOzAG23hqTN-fxPVBwn97BT2Iqsg,1087
+fakesound/samplernn/blocks.py,sha256=FV3A6y_AwmSfOpkFB5Bi_h06icF-Bh7LRaPEV-v41mo,2369
+fakesound/samplernn/model.py,sha256=AoLw6qmvi8T_cZqXbpIdk8xJBCkHhdNaZMf6U8-kwAQ,3295
+fakesound/util/train_utils.py,sha256=NCGBhQRELdjMJ8W2DdJlw5S3NLUeMALcU9cvYZ8MRcg,960
+fakesound/wavenet/blocks.py,sha256=HXTsPCxjoZdio0XwJvS9iG-pTslS0KCuqvRopOwJ5ZA,655
+fakesound/wavenet/layers.py,sha256=cjp6vcGfcEK5EkVsZ4tCUt0UGTBFgHwRYorXBsXhgBQ,1959
+fakesound/wavenet/model.py,sha256=tz_5yRtBir3L2xXehLkz-45lmGciYROYFS-8ELwuTLQ,2501
+fakesound-0.0.0.dist-info/licenses/LICENSE,sha256=9kw0BDLGvrSEZPBlUO1f7JrMReitfq9e74aIlyixxQk,1062
+fakesound-0.0.0.dist-info/METADATA,sha256=kzoYZVW2xwRLBn94Dh0lZ_2EY49mGSQdPfRbs5Z5kNk,121
+fakesound-0.0.0.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+fakesound-0.0.0.dist-info/top_level.txt,sha256=jCW2KhTkYY_3ss1hBHdhf-cpa5x0_obFzV30fE-i-So,10
+fakesound-0.0.0.dist-info/RECORD,,

fakesound-0.0.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.2.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

fakesound-0.0.0.dist-info/licenses/LICENSE

@@ -0,0 +1,18 @@
+MIT License
+
+Copyright (c) 2026 groma
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
+associated documentation files (the "Software"), to deal in the Software without restriction, including
+without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the
+following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial
+portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT
+LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO
+EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
+IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
+USE OR OTHER DEALINGS IN THE SOFTWARE.

fakesound-0.0.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+fakesound