@scrypted/libav 1.0.169 → 1.0.171

package/opus/dnn/torch/osce/models/fd_discriminator.py

@@ -1,974 +0,0 @@
-"""
-/* Copyright (c) 2023 Amazon
-   Written by Jan Buethe */
-/*
-   Redistribution and use in source and binary forms, with or without
-   modification, are permitted provided that the following conditions
-   are met:
-
-   - Redistributions of source code must retain the above copyright
-   notice, this list of conditions and the following disclaimer.
-
-   - Redistributions in binary form must reproduce the above copyright
-   notice, this list of conditions and the following disclaimer in the
-   documentation and/or other materials provided with the distribution.
-
-   THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
-   ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
-   LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
-   A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER
-   OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
-   EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
-   PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
-   PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
-   LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
-   NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
-   SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-*/
-"""
-
-import math as m
-import copy
-
-import torch
-import torch.nn.functional as F
-from torch import nn
-from torch.nn.utils import weight_norm, spectral_norm
-import torchaudio
-
-from utils.spec import gen_filterbank
-
-# auxiliary functions
-
-def remove_all_weight_norms(module):
-    for m in module.modules():
-        if hasattr(m, 'weight_v'):
-            nn.utils.remove_weight_norm(m)
-
-
-def create_smoothing_kernel(h, w, gamma=1.5):
-
-    ch = h / 2 - 0.5
-    cw = w / 2 - 0.5
-
-    sh = gamma * ch
-    sw = gamma * cw
-
-    vx = ((torch.arange(h) - ch) / sh) ** 2
-    vy = ((torch.arange(w) - cw) / sw) ** 2
-    vals = vx.view(-1, 1) + vy.view(1, -1)
-    kernel = torch.exp(- vals)
-    kernel = kernel / kernel.sum()
-
-    return kernel
-
-
-def create_kernel(h, w, sh, sw):
-    # proto kernel gives disjoint partition of 1
-    proto_kernel = torch.ones((sh, sw))
-
-    # create smoothing kernel eta
-    h_eta, w_eta = h - sh + 1, w - sw + 1
-    assert h_eta > 0 and w_eta > 0
-    eta = create_smoothing_kernel(h_eta, w_eta).view(1, 1, h_eta, w_eta)
-
-    kernel0 = F.pad(proto_kernel, [w_eta - 1, w_eta - 1, h_eta - 1, h_eta - 1]).unsqueeze(0).unsqueeze(0)
-    kernel = F.conv2d(kernel0, eta)
-
-    return kernel
-
-# positional embeddings
-class FrequencyPositionalEmbedding(nn.Module):
-    def __init__(self):
-
-        super().__init__()
-
-    def forward(self, x):
-
-        N = x.size(2)
-        args = torch.arange(0, N, dtype=x.dtype, device=x.device) * torch.pi * 2 / N
-        cos = torch.cos(args).reshape(1, 1, -1, 1)
-        sin = torch.sin(args).reshape(1, 1, -1, 1)
-        zeros = torch.zeros_like(x[:, 0:1, :, :])
-
-        y = torch.cat((x, zeros + sin, zeros + cos), dim=1)
-
-        return y
-
-
-class PositionalEmbedding2D(nn.Module):
-    def __init__(self, d=5):
-
-        super().__init__()
-
-        self.d = d
-
-    def forward(self, x):
-
-        N = x.size(2)
-        M = x.size(3)
-
-        h_args = torch.arange(0, N, dtype=x.dtype, device=x.device).reshape(1, 1, -1, 1)
-        w_args = torch.arange(0, M, dtype=x.dtype, device=x.device).reshape(1, 1, 1, -1)
-        coeffs = (10000 ** (-2 * torch.arange(0, self.d, dtype=x.dtype, device=x.device) / self.d)).reshape(1, -1, 1, 1)
-
-        h_sin = torch.sin(coeffs * h_args)
-        h_cos = torch.sin(coeffs * h_args)
-        w_sin = torch.sin(coeffs * w_args)
-        w_cos = torch.sin(coeffs * w_args)
-
-        zeros = torch.zeros_like(x[:, 0:1, :, :])
-
-        y = torch.cat((x, zeros + h_sin, zeros + h_cos, zeros + w_sin, zeros + w_cos), dim=1)
-
-        return y
-
-
-# spectral discriminator base class
-class SpecDiscriminatorBase(nn.Module):
-    RECEPTIVE_FIELD_MAX_WIDTH=10000
-    def __init__(self,
-                 layers,
-                 resolution,
-                 fs=16000,
-                 freq_roi=[50, 7000],
-                 noise_gain=1e-3,
-                 fmap_start_index=0
-                 ):
-        super().__init__()
-
-
-        self.layers = nn.ModuleList(layers)
-        self.resolution = resolution
-        self.fs = fs
-        self.noise_gain = noise_gain
-        self.fmap_start_index = fmap_start_index
-
-        if fmap_start_index >= len(layers):
-            raise ValueError(f'fmap_start_index is larger than number of layers')
-
-        # filter bank for noise shaping
-        n_fft = resolution[0]
-
-        self.filterbank = nn.Parameter(
-            gen_filterbank(n_fft // 2, fs, keep_size=True),
-            requires_grad=False
-        )
-
-        # roi bins
-        f_step = fs / n_fft
-        self.start_bin = int(m.ceil(freq_roi[0] / f_step - 0.01))
-        self.stop_bin = min(int(m.floor(freq_roi[1] / f_step + 0.01)), n_fft//2 + 1)
-
-        self.init_weights()
-
-        # determine receptive field size, offsets and strides
-
-        hw = 1000
-        while True:
-            x = torch.zeros((1, hw, hw))
-            with torch.no_grad():
-                y = self.run_layer_stack(x)[-1]
-
-            pos0 = [y.size(-2) // 2, y.size(-1) // 2]
-            pos1 = [t + 1 for t in pos0]
-
-            hs0, ws0 = self._receptive_field((hw, hw), pos0)
-            hs1, ws1 = self._receptive_field((hw, hw), pos1)
-
-            h0 = hs0[1] - hs0[0] + 1
-            h1 = hs1[1] - hs1[0] + 1
-            w0 = ws0[1] - ws0[0] + 1
-            w1 = ws1[1] - ws1[0] + 1
-
-            if h0 != h1 or w0 != w1:
-                hw = 2 * hw
-            else:
-
-                # strides
-                sh = hs1[0] - hs0[0]
-                sw = ws1[0] - ws0[0]
-
-                if sh == 0 or sw == 0: continue
-
-                # offsets
-                oh = hs0[0] - sh * pos0[0]
-                ow = ws0[0] - sw * pos0[1]
-
-                # overlap factor
-                overlap = w0 / sw + h0 / sh
-
-                #print(f"{w0=} {h0=} {sw=} {sh=} {overlap=}")
-                self.receptive_field_params = {'width': [sw, ow, w0], 'height': [sh, oh, h0], 'overlap': overlap}
-
-                break
-
-            if hw > self.RECEPTIVE_FIELD_MAX_WIDTH:
-                print("warning: exceeded max size while trying to determine receptive field")
-
-        # create transposed convolutional kernel
-        #self.tconv_kernel = nn.Parameter(create_kernel(h0, w0, sw, sw), requires_grad=False)
-
-    def run_layer_stack(self, spec):
-
-        output = []
-
-        x = spec.unsqueeze(1)
-
-        for layer in self.layers:
-            x = layer(x)
-            output.append(x)
-
-        return output
-
-    def forward(self, x):
-        """ returns array with feature maps and final score at index -1 """
-
-        output = []
-
-        x = self.spectrogram(x)
-
-        output = self.run_layer_stack(x)
-
-        return output[self.fmap_start_index:]
-
-    def receptive_field(self, output_pos):
-
-        if self.receptive_field_params is not None:
-            s, o, h = self.receptive_field_params['height']
-            h_min = output_pos[0] * s + o + self.start_bin
-            h_max = h_min + h
-            h_min = max(h_min, self.start_bin)
-            h_max = min(h_max, self.stop_bin)
-
-            s, o, w = self.receptive_field_params['width']
-            w_min = output_pos[1] * s + o
-            w_max = w_min + w
-
-            return (h_min, h_max), (w_min, w_max)
-
-        else:
-            return None, None
-
-
-    def _receptive_field(self, input_dims, output_pos):
-        """ determines receptive field probabilistically via autograd (slow) """
-
-        x = torch.randn((1,) + input_dims, requires_grad=True)
-
-        # run input through layers
-        y = self.run_layer_stack(x)[-1]
-        b, c, h, w = y.shape
-
-        if output_pos[0] >= h or output_pos[1] >= w:
-            raise ValueError("position out of range")
-
-        mask = torch.zeros((b, c, h, w))
-        mask[0, 0, output_pos[0], output_pos[1]] = 1
-
-        (mask * y).sum().backward()
-
-        hs, ws = torch.nonzero(x.grad[0], as_tuple=True)
-
-        h_min, h_max = hs.min().item(), hs.max().item()
-        w_min, w_max = ws.min().item(), ws.max().item()
-
-        return [h_min, h_max], [w_min, w_max]
-
-
-
-    def init_weights(self):
-
-        for m in self.modules():
-            if isinstance(m, nn.Conv1d) or isinstance(m, nn.ConvTranspose1d) or isinstance(m, nn.Linear) or isinstance(m, nn.Embedding):
-                nn.init.orthogonal_(m.weight.data)
-
-
-    def spectrogram(self, x):
-        n_fft, hop_length, win_length = self.resolution
-        x = x.squeeze(1)
-        window = getattr(torch, 'hann_window')(win_length).to(x.device)
-
-        x = torch.stft(x, n_fft=n_fft, hop_length=hop_length, win_length=win_length,\
-                       window=window, return_complex=True) #[B, F, T]
-        x = torch.abs(x)
-
-        # noise floor following spectral envelope
-        smoothed_x = torch.matmul(self.filterbank, x)
-        noise = torch.randn_like(x) * smoothed_x * self.noise_gain
-        x = x + noise
-
-        # frequency ROI
-        x = x[:, self.start_bin : self.stop_bin + 1, ...]
-
-        return torchaudio.functional.amplitude_to_DB(x,db_multiplier=0.0, multiplier=20,amin=1e-05,top_db=80)#torch.sqrt(x)
-
-    def grad_map(self, x):
-        self.zero_grad()
-
-        n_fft, hop_length, win_length = self.resolution
-
-        window = getattr(torch, 'hann_window')(win_length).to(x.device)
-
-        y = torch.stft(x.squeeze(1), n_fft=n_fft, hop_length=hop_length, win_length=win_length,
-                       window=window, return_complex=True) #[B, F, T]
-        y = torch.abs(y)
-
-        specgram  = torchaudio.functional.amplitude_to_DB(y,db_multiplier=0.0, multiplier=20,amin=1e-05,top_db=80)
-
-        specgram.requires_grad = True
-        specgram.retain_grad()
-
-        if specgram.grad is not None:
-            specgram.grad.zero_()
-
-        y = specgram[:, self.start_bin : self.stop_bin + 1, ...]
-
-        scores = self.run_layer_stack(y)[-1]
-
-        loss = torch.mean((1 - scores) ** 2)
-        loss.backward()
-
-        return specgram.data[0], torch.abs(specgram.grad)[0]
-
-    def relevance_map(self, x):
-
-        n_fft, hop_length, win_length = self.resolution
-        y = x.view(-1)
-        window = getattr(torch, 'hann_window')(win_length).to(x.device)
-
-        y = torch.stft(y, n_fft=n_fft, hop_length=hop_length, win_length=win_length,\
-                       window=window, return_complex=True) #[B, F, T]
-        y = torch.abs(y)
-
-        specgram  = torchaudio.functional.amplitude_to_DB(y,db_multiplier=0.0, multiplier=20,amin=1e-05,top_db=80)
-
-
-        scores = self.forward(x)[-1]
-
-        sh, _, h = self.receptive_field_params['height']
-        sw, _, w = self.receptive_field_params['width']
-        kernel = create_kernel(h, w, sh, sw).float().to(scores.device)
-        with torch.no_grad():
-            pad_w = (w + sw - 1) // sw
-            pad_h = (h + sh - 1) // sh
-            padded_scores = F.pad(scores, (pad_w, pad_w, pad_h, pad_h), mode='replicate')
-            # CAVE: padding should be derived from offsets
-            rv = F.conv_transpose2d(padded_scores, kernel, bias=None, stride=(sh, sw), padding=(h//2, w//2))
-            rv = rv[..., pad_h * sh : - pad_h * sh,  pad_w * sw : -pad_w * sw]
-
-            relevance = torch.zeros_like(specgram)
-            relevance[..., self.start_bin : self.start_bin + rv.size(-2), : rv.size(-1)] = rv
-
-
-        return specgram, relevance
-
-
-    def lrp(self, x, eps=1e-9, label='both', threshold=0.5, low=None, high=None, verbose=False):
-        """ layer-wise relevance propagation (https://git.tu-berlin.de/gmontavon/lrp-tutorial) """
-
-        # ToDo: this code is highly unsafe as it assumes that layers are nn.Sequential with suitable activations
-
-        def newconv2d(layer,g):
-
-            new_layer = nn.Conv2d(layer.in_channels,
-                                  layer.out_channels,
-                                  layer.kernel_size,
-                                  stride=layer.stride,
-                                  padding=layer.padding,
-                                  dilation=layer.dilation,
-                                  groups=layer.groups)
-
-            try: new_layer.weight = nn.Parameter(g(layer.weight.data.clone()))
-            except AttributeError: pass
-
-            try: new_layer.bias   = nn.Parameter(g(layer.bias.data.clone()))
-            except AttributeError: pass
-
-            return new_layer
-
-        bounds = {
-            64: [-85.82449722290039, 2.1755014657974243],
-            128: [-84.49211349487305, 3.5078893899917607],
-            256: [-80.33127822875977, 7.6687201976776125],
-            512: [-73.79328079223633, 14.20672025680542],
-            1024: [-67.59239501953125, 20.40760498046875],
-            2048: [-62.31902580261231, 25.680974197387698],
-        }
-
-        nfft = self.resolution[0]
-        if low is None: low = bounds[nfft][0]
-        if high is None: high = bounds[nfft][1]
-
-        remove_all_weight_norms(self)
-
-        for p in self.parameters():
-            if p.grad is not None:
-                p.grad.zero_()
-
-        num_layers = len(self.layers)
-        X = self.spectrogram(x). detach()
-
-
-        # forward pass
-        A = [X.unsqueeze(1)] + [None] * len(self.layers)
-
-        for i in range(num_layers - 1):
-            A[i + 1] = self.layers[i](A[i])
-
-        # initial relevance is last layer without activation
-        r = A[-2]
-        last_layer_rs = [r]
-        layer = self.layers[-1]
-        for sublayer in list(layer)[:-1]:
-            r = sublayer(r)
-            last_layer_rs.append(r)
-
-
-        mask = torch.zeros_like(r)
-        mask.requires_grad_(False)
-        if verbose:
-            print(r.min(), r.max())
-        if label in {'both', 'fake'}:
-            mask[r < -threshold] = 1
-        if label in {'both', 'real'}:
-            mask[r > threshold] = 1
-        r = r * mask
-
-        # backward pass
-        R = [None] * num_layers + [r]
-
-        for l in range(1, num_layers)[::-1]:
-            A[l] = (A[l]).data.requires_grad_(True)
-
-            layer = nn.Sequential(*(list(self.layers[l])[:-1]))
-            z = layer(A[l]) + eps
-            s = (R[l+1] / z).data
-            (z*s).sum().backward()
-            c = A[l].grad
-            R[l] = (A[l] * c).data
-
-        # first layer
-        A[0] = (A[0].data).requires_grad_(True)
-
-        Xl = (torch.zeros_like(A[0].data) + low).requires_grad_(True)
-        Xh = (torch.zeros_like(A[0].data) + high).requires_grad_(True)
-
-        if len(list(self.layers)) > 2:
-            # unsafe way to check for embedding layer
-            embed = list(self.layers[0])[0]
-            conv  = list(self.layers[0])[1]
-
-            layer = nn.Sequential(embed, conv)
-            layerl = nn.Sequential(embed, newconv2d(conv, lambda p: p.clamp(min=0)))
-            layerh = nn.Sequential(embed, newconv2d(conv, lambda p: p.clamp(max=0)))
-
-        else:
-            layer = list(self.layers[0])[0]
-            layerl = newconv2d(layer, lambda p: p.clamp(min=0))
-            layerh = newconv2d(layer, lambda p: p.clamp(max=0))
-
-
-        z = layer(A[0])
-        z -= layerl(Xl) + layerh(Xh)
-        s = (R[1] / z).data
-        (z * s).sum().backward()
-        c, cp, cm = A[0].grad, Xl.grad, Xh.grad
-
-        R[0] = (A[0] * c + Xl * cp + Xh * cm)
-        #R[0] = (A[0] * c).data
-
-        return X, R[0].mean(dim=1)
-
-
-
-
-
-
-
-
-
-
-def create_3x3_conv_plan(num_layers : int,
-                         f_stretch : int,
-                         f_down : int,
-                         t_stretch : int,
-                         t_down : int
-                         ):
-
-
-    """ creates a stride, dilation, padding plan for a 2d conv network
-
-    Args:
-        num_layers (int): number of layers
-        f_stretch (int): log_2 of stretching factor along frequency axis
-        f_down (int): log_2 of downsampling factor along frequency axis
-        t_stretch (int): log_2 of stretching factor along time axis
-        t_down (int): log_2 of downsampling factor along time axis
-
-    Returns:
-        list(list(tuple)): list containing entries [(stride_t, stride_f), (dilation_t, dilation_f), (padding_t, padding_f)]
-    """
-
-    assert num_layers > 0 and t_stretch >= 0 and t_down >= 0 and f_stretch >= 0 and f_down >= 0
-    assert f_stretch < num_layers and t_stretch < num_layers
-
-    def process_dimension(n_layers, stretch, down):
-
-        stack_layers = n_layers - 1
-
-        stride_layers = min(min(down, stretch) , stack_layers)
-        dilation_layers = max(min(stack_layers - stride_layers - 1, stretch - stride_layers), 0)
-        final_stride = 2 ** (max(down - stride_layers, 0))
-
-        final_dilation = 1
-        if stride_layers < stack_layers and stretch - stride_layers - dilation_layers > 0:
-                final_dilation = 2
-
-        strides, dilations, paddings = [], [], []
-        processed_layers = 0
-        current_dilation = 1
-
-        for _ in range(stride_layers):
-            # increase receptive field and downsample via stride = 2
-            strides.append(2)
-            dilations.append(1)
-            paddings.append(1)
-            processed_layers += 1
-
-        if processed_layers < stack_layers:
-            strides.append(1)
-            dilations.append(1)
-            paddings.append(1)
-            processed_layers += 1
-
-        for _ in range(dilation_layers):
-            # increase receptive field via dilation = 2
-            strides.append(1)
-            current_dilation *= 2
-            dilations.append(current_dilation)
-            paddings.append(current_dilation)
-            processed_layers += 1
-
-        while processed_layers < n_layers - 1:
-            # fill up with std layers
-            strides.append(1)
-            dilations.append(current_dilation)
-            paddings.append(current_dilation)
-            processed_layers += 1
-
-        # final layer
-        strides.append(final_stride)
-        current_dilation * final_dilation
-        dilations.append(current_dilation)
-        paddings.append(current_dilation)
-        processed_layers += 1
-
-        assert processed_layers == n_layers
-
-        return strides, dilations, paddings
-
-    t_strides, t_dilations, t_paddings = process_dimension(num_layers, t_stretch, t_down)
-    f_strides, f_dilations, f_paddings = process_dimension(num_layers, f_stretch, f_down)
-
-    plan = []
-
-    for i in range(num_layers):
-        plan.append([
-            (f_strides[i], t_strides[i]),
-            (f_dilations[i], t_dilations[i]),
-            (f_paddings[i], t_paddings[i]),
-            ])
-
-    return plan
-
-
-class DiscriminatorExperimental(SpecDiscriminatorBase):
-
-    def __init__(self,
-                 resolution,
-                 fs=16000,
-                 freq_roi=[50, 7400],
-                 noise_gain=0,
-                 num_channels=16,
-                 max_channels=512,
-                 num_layers=5,
-                 use_spectral_norm=False):
-
-        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
-
-        self.num_channels = num_channels
-        self.num_channels_max = max_channels
-        self.num_layers = num_layers
-
-        layers = []
-        stride = (2, 1)
-        padding= (1, 1)
-        in_channels = 1 + 2
-        out_channels = self.num_channels
-        for _ in range(self.num_layers):
-            layers.append(
-                nn.Sequential(
-                    FrequencyPositionalEmbedding(),
-                    norm_f(nn.Conv2d(in_channels, out_channels, (3, 3), stride=stride, padding=padding)),
-                    nn.ReLU(inplace=True)
-                )
-            )
-            in_channels = out_channels + 2
-            out_channels = min(2 * out_channels, self.num_channels_max)
-
-        layers.append(
-            nn.Sequential(
-                FrequencyPositionalEmbedding(),
-                norm_f(nn.Conv2d(in_channels, 1, (3, 3), padding=padding)),
-                nn.Sigmoid()
-            )
-        )
-
-        super().__init__(layers=layers, resolution=resolution, fs=fs, freq_roi=freq_roi, noise_gain=noise_gain)
-
-        # bias biases
-        bias_val = 0.1
-        with torch.no_grad():
-            for name, weight in self.named_parameters():
-                if 'bias' in name:
-                    weight = weight + bias_val
-
-
-configs = {
-    'f_down': {
-        'stretch' : {
-            64 : (0, 0),
-            128: (1, 0),
-            256: (2, 0),
-            512: (3, 0),
-            1024: (4, 0),
-            2048: (5, 0)
-        },
-        'down' : {
-            64 : (0, 0),
-            128: (1, 0),
-            256: (2, 0),
-            512: (3, 0),
-            1024: (4, 0),
-            2048: (5, 0)
-        }
-    },
-    'ft_down': {
-        'stretch' : {
-            64 : (0, 4),
-            128: (1, 3),
-            256: (2, 2),
-            512: (3, 1),
-            1024: (4, 0),
-            2048: (5, 0)
-        },
-        'down' : {
-            64 : (0, 4),
-            128: (1, 3),
-            256: (2, 2),
-            512: (3, 1),
-            1024: (4, 0),
-            2048: (5, 0)
-        }
-    },
-    'dilated': {
-        'stretch' : {
-            64 : (0, 4),
-            128: (1, 3),
-            256: (2, 2),
-            512: (3, 1),
-            1024: (4, 0),
-            2048: (5, 0)
-        },
-        'down' : {
-            64 : (0, 0),
-            128: (0, 0),
-            256: (0, 0),
-            512: (0, 0),
-            1024: (0, 0),
-            2048: (0, 0)
-        }
-    },
-    'mixed': {
-        'stretch' : {
-            64 : (0, 4),
-            128: (1, 3),
-            256: (2, 2),
-            512: (3, 1),
-            1024: (4, 0),
-            2048: (5, 0)
-        },
-        'down' : {
-            64 : (0, 0),
-            128: (1, 0),
-            256: (2, 0),
-            512: (3, 0),
-            1024: (4, 0),
-            2048: (5, 0)
-        }
-    },
-}
-
-
-class DiscriminatorMagFree(SpecDiscriminatorBase):
-
-    def __init__(self,
-                 resolution,
-                 fs=16000,
-                 freq_roi=[50, 7400],
-                 noise_gain=0,
-                 num_channels=16,
-                 max_channels=256,
-                 num_layers=5,
-                 use_spectral_norm=False,
-                 design=None):
-
-        if design is None:
-            raise ValueError('error: arch required in DiscriminatorMagFree')
-
-        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
-
-        stretch = configs[design]['stretch'][resolution[0]]
-        down = configs[design]['down'][resolution[0]]
-
-        self.num_channels = num_channels
-        self.num_channels_max = max_channels
-        self.num_layers = num_layers
-        self.stretch = stretch
-        self.down = down
-
-        layers = []
-        plan = create_3x3_conv_plan(num_layers + 1, stretch[0], down[0], stretch[1], down[1])
-        in_channels = 1 + 2
-        out_channels = self.num_channels
-        for i in range(self.num_layers):
-            layers.append(
-                nn.Sequential(
-                    FrequencyPositionalEmbedding(),
-                    norm_f(nn.Conv2d(in_channels, out_channels, (3, 3), stride=plan[i][0], dilation=plan[i][1], padding=plan[i][2])),
-                    nn.ReLU(inplace=True)
-                )
-            )
-            in_channels = out_channels + 2
-            # product over strides
-            channel_factor = plan[i][0][0] * plan[i][0][1]
-            out_channels = min(channel_factor * out_channels, self.num_channels_max)
-
-        layers.append(
-            nn.Sequential(
-                FrequencyPositionalEmbedding(),
-                norm_f(nn.Conv2d(in_channels, 1, (3, 3), stride=plan[-1][0], dilation=plan[-1][1], padding=plan[-1][2])),
-                nn.Sigmoid()
-            )
-        )
-
-
-
-        # for layer in layers:
-        #     print(layer)
-
-        # print("end\n\n")
-
-        super().__init__(layers=layers, resolution=resolution, fs=fs, freq_roi=freq_roi, noise_gain=noise_gain)
-
-        # bias biases
-        bias_val = 0.1
-        with torch.no_grad():
-            for name, weight in self.named_parameters():
-                if 'bias' in name:
-                    weight = weight + bias_val
-
-class DiscriminatorMagFreqPosition(SpecDiscriminatorBase):
-
-    def __init__(self,
-                 resolution,
-                 fs=16000,
-                 freq_roi=[50, 7400],
-                 noise_gain=0,
-                 num_channels=16,
-                 max_channels=512,
-                 num_layers=5,
-                 use_spectral_norm=False):
-
-        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
-
-        self.num_channels = num_channels
-        self.num_channels_max = max_channels
-        self.num_layers = num_layers
-
-        layers = []
-        stride = (2, 1)
-        padding= (1, 1)
-        in_channels = 1 + 2
-        out_channels = self.num_channels
-        for _ in range(self.num_layers):
-            layers.append(
-                nn.Sequential(
-                    FrequencyPositionalEmbedding(),
-                    norm_f(nn.Conv2d(in_channels, out_channels, (3, 3), stride=stride, padding=padding)),
-                    nn.LeakyReLU(0.2, inplace=True)
-                )
-            )
-            in_channels = out_channels + 2
-            out_channels = min(2 * out_channels, self.num_channels_max)
-
-        layers.append(
-            nn.Sequential(
-                FrequencyPositionalEmbedding(),
-                norm_f(nn.Conv2d(in_channels, 1, (3, 3), padding=padding))
-            )
-        )
-
-        super().__init__(layers=layers, resolution=resolution, fs=fs, freq_roi=freq_roi, noise_gain=noise_gain)
-
-
-
-class DiscriminatorMag2dPositional(SpecDiscriminatorBase):
-
-    def __init__(self,
-                 resolution,
-                 fs=16000,
-                 freq_roi=[50, 7400],
-                 noise_gain=0,
-                 num_channels=16,
-                 max_channels=512,
-                 num_layers=5,
-                 d=5,
-                 use_spectral_norm=False):
-
-        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
-        self.resolution = resolution
-        self.num_channels = num_channels
-        self.num_channels_max = max_channels
-        self.num_layers = num_layers
-        self.d = d
-        embedding_dim = 4 * d
-
-
-        layers = []
-        stride = (2, 2)
-        padding= (1, 1)
-        in_channels = 1 + embedding_dim
-        out_channels = self.num_channels
-        for _ in range(self.num_layers):
-            layers.append(
-                nn.Sequential(
-                    PositionalEmbedding2D(d),
-                    norm_f(nn.Conv2d(in_channels, out_channels, (3, 3), stride=stride, padding=padding)),
-                    nn.LeakyReLU(0.2, inplace=True)
-                )
-            )
-            in_channels = out_channels + embedding_dim
-            out_channels = min(2 * out_channels, self.num_channels_max)
-
-
-        layers.append(
-            nn.Sequential(
-                PositionalEmbedding2D(),
-                norm_f(nn.Conv2d(in_channels, 1, (3, 3), padding=padding))
-            )
-        )
-
-        super().__init__(layers=layers, resolution=resolution, fs=fs, freq_roi=freq_roi, noise_gain=noise_gain)
-
-
-
-class DiscriminatorMag(SpecDiscriminatorBase):
-    def __init__(self,
-                 resolution,
-                 fs=16000,
-                 freq_roi=[50, 7400],
-                 noise_gain=0,
-                 num_channels=32,
-                 num_layers=5,
-                 use_spectral_norm=False):
-
-        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
-
-        self.num_channels = num_channels
-        self.num_layers = num_layers
-
-        layers = []
-        stride = (1, 1)
-        padding= (1, 1)
-        in_channels = 1
-        out_channels = self.num_channels
-        for _ in range(self.num_layers):
-            layers.append(
-                nn.Sequential(
-                    norm_f(nn.Conv2d(in_channels, out_channels, (3, 3), stride=stride, padding=padding)),
-                    nn.LeakyReLU(0.2, inplace=True)
-                )
-            )
-            in_channels = out_channels
-
-        layers.append(norm_f(nn.Conv2d(in_channels, 1, (3, 3), padding=padding)))
-
-        super().__init__(layers=layers, resolution=resolution, fs=fs, freq_roi=freq_roi, noise_gain=noise_gain)
-
-
-discriminators = {
-    'mag': DiscriminatorMag,
-    'freqpos': DiscriminatorMagFreqPosition,
-    '2dpos': DiscriminatorMag2dPositional,
-    'experimental': DiscriminatorExperimental,
-    'free': DiscriminatorMagFree
-}
-
-class TFDMultiResolutionDiscriminator(torch.nn.Module):
-    def __init__(self,
-                 fft_sizes_16k=[64, 128, 256, 512, 1024, 2048],
-                 architecture='mag',
-                 fs=16000,
-                 freq_roi=[50, 7400],
-                 noise_gain=0,
-                 use_spectral_norm=False,
-                 **kwargs):
-
-        super().__init__()
-
-
-        fft_sizes = [int(round(fft_size_16k * fs / 16000)) for fft_size_16k in fft_sizes_16k]
-
-        resolutions = [[n_fft, n_fft // 4, n_fft] for n_fft in fft_sizes]
-
-
-        Disc = discriminators[architecture]
-
-        discs = [Disc(resolutions[i], fs=fs, freq_roi=freq_roi, noise_gain=noise_gain, use_spectral_norm=use_spectral_norm, **kwargs) for i in range(len(resolutions))]
-
-        self.discriminators = nn.ModuleList(discs)
-
-    def forward(self, y):
-        outputs = []
-
-        for  disc in self.discriminators:
-            outputs.append(disc(y))
-
-        return outputs
-
-
-class FWGAN_disc_wrapper(nn.Module):
-    def __init__(self, disc):
-        super().__init__()
-
-        self.disc = disc
-
-    def forward(self, y, y_hat):
-
-        out_real = self.disc(y)
-        out_fake = self.disc(y_hat)
-
-        y_d_rs = []
-        y_d_gs = []
-        fmap_rs = []
-        fmap_gs = []
-
-        for y_real, y_fake in zip(out_real, out_fake):
-            y_d_rs.append(y_real[-1])
-            y_d_gs.append(y_fake[-1])
-            fmap_rs.append(y_real[:-1])
-            fmap_gs.append(y_fake[:-1])
-
-        return y_d_rs, y_d_gs, fmap_rs, fmap_gs