jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(self, in_dim, out_dim, kernel_size, stride=1, padding=0, 
        bias=True, weight_norm=True, scale=False):
        """Intializes a Conv2d augmented with weight normalization.

        (See torch.nn.utils.weight_norm for detail.)

        Args:
            in_dim: number of input channels.
            out_dim: number of output channels.
            kernel_size: size of convolving kernel.
            stride: stride of convolution.
            padding: zero-padding added to both sides of input.
            bias: True if include learnable bias parameters, False otherwise.
            weight_norm: True if apply weight normalization, False otherwise.
            scale: True if include magnitude parameters, False otherwise.
        """
        super(WeightNormConv2d, self).__init__()

        if weight_norm:
            self.conv = nn.utils.parametrizations.weight_norm(
                nn.Conv2d(in_dim, out_dim, kernel_size, 
                    stride=stride, padding=padding, bias=bias))
            if not scale:
                self.conv.weight_g.data = torch.ones_like(self.conv.weight_g.data)
                self.conv.weight_g.requires_grad = False    # freeze scaling
        else:
            self.conv = nn.Conv2d(in_dim, out_dim, kernel_size, 
                stride=stride, padding=padding, bias=bias)
//  [python]
def __init__(self, n_feats, act=nn.ReLU(True)):
        super(MY, self).__init__()

        self.act = activation('lrelu', neg_slope=0.05)
        def wn(x): return torch.nn.utils.weight_norm(x)
        self.srb1 = SRBW1(n_feats)
        self.srb2 = SRBW1(n_feats)
        self.distilled_channels1 = n_feats // 2
        self.remaining_channels1 = n_feats // 2
        self.distilled_channels2 = n_feats // 2
        self.remaining_channels2 = n_feats // 2
        self.rb1 = SRBW(n_feats // 2, wn=wn, act=act)
        self.rb2 = SRBW(n_feats // 2, wn=wn, act=act)
        self.A1_coffconv = CoffConv(n_feats)
        self.B1_coffconv = CoffConv(n_feats)
        self.A2_coffconv = CoffConv(n_feats)
        self.B2_coffconv = CoffConv(n_feats)
        self.conv_distilled1 = nn.Conv2d(
            n_feats // 2, n_feats, kernel_size=3, stride=1, padding=1, bias=False)
        self.conv_distilled2 = nn.Conv2d(
            n_feats // 2, n_feats, kernel_size=3, stride=1, padding=1, bias=False)
        self.sigmoid1 = nn.Sigmoid()
        self.sigmoid2 = nn.Sigmoid()
        self.sigmoid3 = nn.Sigmoid()
        self.scale_x1 = Scale(1)
        self.scale_x2 = Scale(1)
        self.srb3 = SRBW1(n_feats)
        self.srb4 = SRBW1(n_feats)
        self.fuse1 = SRBW2(n_feats * 2)
        self.fuse2 = nn.Conv2d(
            2 * n_feats, n_feats, kernel_size=1, stride=1, padding=0, bias=False, dilation=1)
        self.SAlayer = sa_layer(n_feats)
//  [python]
def __init__(self, n_feats, act=nn.ReLU(True)):
        super(MY, self).__init__()

        self.act = activation('lrelu', neg_slope=0.05)
        def wn(x): return torch.nn.utils.parametrizations.weight_norm(x)
        self.srb1 = SRBW1(n_feats)
        self.srb2 = SRBW1(n_feats)
        self.distilled_channels1 = n_feats // 2
        self.remaining_channels1 = n_feats // 2
        self.distilled_channels2 = n_feats // 2
        self.remaining_channels2 = n_feats // 2
        self.rb1 = SRBW(n_feats // 2, wn=wn, act=act)
        self.rb2 = SRBW(n_feats // 2, wn=wn, act=act)
        self.A1_coffconv = CoffConv(n_feats)
        self.B1_coffconv = CoffConv(n_feats)
        self.A2_coffconv = CoffConv(n_feats)
        self.B2_coffconv = CoffConv(n_feats)
        self.conv_distilled1 = nn.Conv2d(
            n_feats // 2, n_feats, kernel_size=3, stride=1, padding=1, bias=False)
        self.conv_distilled2 = nn.Conv2d(
            n_feats // 2, n_feats, kernel_size=3, stride=1, padding=1, bias=False)
        self.sigmoid1 = nn.Sigmoid()
        self.sigmoid2 = nn.Sigmoid()
        self.sigmoid3 = nn.Sigmoid()
        self.scale_x1 = Scale(1)
        self.scale_x2 = Scale(1)
        self.srb3 = SRBW1(n_feats)
        self.srb4 = SRBW1(n_feats)
        self.fuse1 = SRBW2(n_feats * 2)
        self.fuse2 = nn.Conv2d(
            2 * n_feats, n_feats, kernel_size=1, stride=1, padding=0, bias=False, dilation=1)
        self.SAlayer = sa_layer(n_feats)
//  [python]
def __init__(self, in_channels):
        super(Generator, self).__init__()

        self.generator = nn.Sequential(
            nn.ReflectionPad1d(3),
            nn.utils.weight_norm(nn.Conv1d(in_channels, 512, kernel_size=7)),
            nn.LeakyReLU(0.2, True),
            UpsampleNet(512, 256, 8),
            ResStack(256),
            nn.LeakyReLU(0.2, True),
            UpsampleNet(256, 128, 6),
            ResStack(128),
            nn.LeakyReLU(0.2, True),
            UpsampleNet(128, 64, 5),
            ResStack(64),
            nn.LeakyReLU(0.2, True),
            nn.ReflectionPad1d(3),
            nn.utils.weight_norm(nn.Conv1d(64, 1, kernel_size=7)),
            nn.Tanh(),
        )

        self.num_params()
//  [python]
def __init__(self, in_channels):
        super(Generator, self).__init__()

        self.generator = nn.Sequential(
            nn.ReflectionPad1d(3),
            nn.utils.parametrizations.weight_norm(nn.Conv1d(in_channels, 512, kernel_size=7)),
            nn.LeakyReLU(0.2, True),
            UpsampleNet(512, 256, 8),
            ResStack(256),
            nn.LeakyReLU(0.2, True),
            UpsampleNet(256, 128, 6),
            ResStack(128),
            nn.LeakyReLU(0.2, True),
            UpsampleNet(128, 64, 5),
            ResStack(64),
            nn.LeakyReLU(0.2, True),
            nn.ReflectionPad1d(3),
            nn.utils.parametrizations.weight_norm(nn.Conv1d(64, 1, kernel_size=7)),
            nn.Tanh(),
        )

        self.num_params()
//  [python]
def __init__(
        self,
        d_model,
        d_hidden,
        n_layers,
        learn_h0=False,
        rnn_type='gru',
        weight_norm=False,
        skip_connections=True,
        dropout=0.0,
    ):
        super().__init__()

        self.d_model = d_model
        self.d_hidden = d_hidden
        self.n_layers = n_layers
        self.learn_h0 = learn_h0
        self.weight_norm = weight_norm
        self.skip_connections = skip_connections
        self.rnn_type = rnn_type

        self.rnn = StackedRNN(
            d_model=d_model,
            d_hidden=d_hidden,
            n_layers=n_layers,
            rnn_type=rnn_type,
            skip_connections=skip_connections,
            weight_norm=weight_norm,
            dropout=dropout,
            output_linear=False,
        )

        self.lin1 = torch.nn.Linear(d_hidden, d_hidden)
        self.lin2 = torch.nn.Linear(d_hidden, d_hidden)
        self.lin3 = torch.nn.Linear(d_hidden, d_hidden)

        if weight_norm:
            self.lin1 = torch.nn.utils.weight_norm(self.lin1)
            self.lin2 = torch.nn.utils.weight_norm(self.lin2)
            self.lin3 = torch.nn.utils.weight_norm(self.lin3)
//  [python]
def __init__(
        self,
        d_model,
        d_hidden,
        n_layers,
        learn_h0=False,
        rnn_type='gru',
        weight_norm=False,
        skip_connections=True,
        dropout=0.0,
    ):
        super().__init__()

        self.d_model = d_model
        self.d_hidden = d_hidden
        self.n_layers = n_layers
        self.learn_h0 = learn_h0
        self.weight_norm = weight_norm
        self.skip_connections = skip_connections
        self.rnn_type = rnn_type

        self.rnn = StackedRNN(
            d_model=d_model,
            d_hidden=d_hidden,
            n_layers=n_layers,
            rnn_type=rnn_type,
            skip_connections=skip_connections,
            weight_norm=weight_norm,
            dropout=dropout,
            output_linear=False,
        )

        self.lin1 = torch.nn.Linear(d_hidden, d_hidden)
        self.lin2 = torch.nn.Linear(d_hidden, d_hidden)
        self.lin3 = torch.nn.Linear(d_hidden, d_hidden)

        if weight_norm:
            self.lin1 = torch.nn.utils.parametrizations.weight_norm(self.lin1)
            self.lin2 = torch.nn.utils.parametrizations.weight_norm(self.lin2)
            self.lin3 = torch.nn.utils.parametrizations.weight_norm(self.lin3)
//  [python]
def __init__(self, channel, kernel_size=3, base=3, nums=4):
        super(ResStack, self).__init__()

        self.layers = nn.ModuleList([
            nn.Sequential(
                nn.LeakyReLU(),
                nn.utils.weight_norm(nn.Conv1d(channel, channel,
                                               kernel_size=kernel_size, dilation=base ** i, padding=base ** i)),
                nn.LeakyReLU(),
                nn.utils.weight_norm(nn.Conv1d(channel, channel,
                                               kernel_size=kernel_size, dilation=1, padding=1)),
            )
            for i in range(nums)
        ])
//  [python]
def __init__(self, channel, kernel_size=3, base=3, nums=4):
        super(ResStack, self).__init__()

        self.layers = nn.ModuleList([
            nn.Sequential(
                nn.LeakyReLU(),
                nn.utils.parametrizations.weight_norm(nn.Conv1d(channel, channel,
                                               kernel_size=kernel_size, dilation=base ** i, padding=base ** i)),
                nn.LeakyReLU(),
                nn.utils.parametrizations.weight_norm(nn.Conv1d(channel, channel,
                                               kernel_size=kernel_size, dilation=1, padding=1)),
            )
            for i in range(nums)
        ])
//  [python]
def LinearActivation(
    d_input,
    d_output,
    bias=True,
    zero_bias_init=False,
    transposed=False,
    initializer=None,
    activation=None,
    activate=False,  # Apply activation as part of this module
    weight_norm=False,
    **kwargs,
):
    """Returns a linear nn.Module with control over axes order, initialization, and activation"""

    # Construct core module
    linear_cls = TransposedLinear if transposed else nn.Linear
    if activation == "glu":
        d_output *= 2
    linear = linear_cls(d_input, d_output, bias=bias, **kwargs)

    # Initialize weight
    if initializer is not None:
        get_initializer(initializer, activation)(linear.weight)

    # Initialize bias
    if bias and zero_bias_init:
        nn.init.zeros_(linear.bias)

    # Weight norm
    if weight_norm:
        linear = nn.utils.weight_norm(linear)

    if activate and activation is not None:
        activation = Activation(activation, dim=-2 if transposed else -1)
        linear = nn.Sequential(linear, activation)
    return linear
//  [python]
def LinearActivation(
    d_input,
    d_output,
    bias=True,
    zero_bias_init=False,
    transposed=False,
    initializer=None,
    activation=None,
    activate=False,  # Apply activation as part of this module
    weight_norm=False,
    **kwargs,
):
    """Returns a linear nn.Module with control over axes order, initialization, and activation"""

    # Construct core module
    linear_cls = TransposedLinear if transposed else nn.Linear
    if activation == "glu":
        d_output *= 2
    linear = linear_cls(d_input, d_output, bias=bias, **kwargs)

    # Initialize weight
    if initializer is not None:
        get_initializer(initializer, activation)(linear.weight)

    # Initialize bias
    if bias and zero_bias_init:
        nn.init.zeros_(linear.bias)

    # Weight norm
    if weight_norm:
        linear = nn.utils.parametrizations.weight_norm(linear)

    if activate and activation is not None:
        activation = Activation(activation, dim=-2 if transposed else -1)
        linear = nn.Sequential(linear, activation)
    return linear
//  [python]
def __init__(self,
                 d_feature,
                 mode,
                 d_in,
                 d_out,
                 d_hidden,
                 n_layers,
                 weight_norm=True,
                 multires_view=0,
                 squeeze_out=True):
        super().__init__()

        self.mode = mode
        self.squeeze_out = squeeze_out
        dims = [d_in + d_feature] + [d_hidden
                                     for _ in range(n_layers)] + [d_out]

        self.embedview_fn = None
        if multires_view > 0:
            embedview_fn, input_ch = get_embedder(multires_view)
            self.embedview_fn = embedview_fn
            dims[0] += (input_ch - 3)

        self.num_layers = len(dims)

        for layer in range(0, self.num_layers - 1):
            out_dim = dims[layer + 1]
            lin = nn.Linear(dims[layer], out_dim)

            if weight_norm:
                lin = nn.utils.weight_norm(lin)

            setattr(self, 'lin' + str(layer), lin)

        self.relu = nn.ReLU()
//  [python]
def __init__(self,
                 d_feature,
                 mode,
                 d_in,
                 d_out,
                 d_hidden,
                 n_layers,
                 weight_norm=True,
                 multires_view=0,
                 squeeze_out=True):
        super().__init__()

        self.mode = mode
        self.squeeze_out = squeeze_out
        dims = [d_in + d_feature] + [d_hidden
                                     for _ in range(n_layers)] + [d_out]

        self.embedview_fn = None
        if multires_view > 0:
            embedview_fn, input_ch = get_embedder(multires_view)
            self.embedview_fn = embedview_fn
            dims[0] += (input_ch - 3)

        self.num_layers = len(dims)

        for layer in range(0, self.num_layers - 1):
            out_dim = dims[layer + 1]
            lin = nn.Linear(dims[layer], out_dim)

            if weight_norm:
                lin = nn.utils.parametrizations.weight_norm(lin)

            setattr(self, 'lin' + str(layer), lin)

        self.relu = nn.ReLU()
//  [python]
def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, n_speakers=0, spk_channels=0,
                 p_dropout=0):
        super(WN, self).__init__()
        assert (kernel_size % 2 == 1)
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size,
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        self.n_speakers = n_speakers
        self.spk_channels = spk_channels
        self.p_dropout = p_dropout

        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.drop = nn.Dropout(p_dropout)

        if n_speakers > 0:
            cond_layer = torch.nn.Conv1d(spk_channels, 2 * hidden_channels * n_layers, 1)
            self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')

        for i in range(n_layers):
            dilation = dilation_rate ** i
            padding = int((kernel_size * dilation - dilation) / 2)
            in_layer = torch.nn.Conv1d(hidden_channels, 2 * hidden_channels, kernel_size,
                                       dilation=dilation, padding=padding)
            in_layer = torch.nn.utils.weight_norm(in_layer, name='weight')
            self.in_layers.append(in_layer)

            # last one is not necessary
            if i < n_layers - 1:
                res_skip_channels = 2 * hidden_channels
            else:
                res_skip_channels = hidden_channels

            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
            res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight')
            self.res_skip_layers.append(res_skip_layer)
//  [python]
def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, n_speakers=0, spk_channels=0,
                 p_dropout=0):
        super(WN, self).__init__()
        assert (kernel_size % 2 == 1)
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size,
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        self.n_speakers = n_speakers
        self.spk_channels = spk_channels
        self.p_dropout = p_dropout

        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.drop = nn.Dropout(p_dropout)

        if n_speakers > 0:
            cond_layer = torch.nn.Conv1d(spk_channels, 2 * hidden_channels * n_layers, 1)
            self.cond_layer = torch.nn.utils.parametrizations.weight_norm(cond_layer, name='weight')

        for i in range(n_layers):
            dilation = dilation_rate ** i
            padding = int((kernel_size * dilation - dilation) / 2)
            in_layer = torch.nn.Conv1d(hidden_channels, 2 * hidden_channels, kernel_size,
                                       dilation=dilation, padding=padding)
            in_layer = torch.nn.utils.parametrizations.weight_norm(in_layer, name='weight')
            self.in_layers.append(in_layer)

            # last one is not necessary
            if i < n_layers - 1:
                res_skip_channels = 2 * hidden_channels
            else:
                res_skip_channels = hidden_channels

            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
            res_skip_layer = torch.nn.utils.parametrizations.weight_norm(res_skip_layer, name='weight')
            self.res_skip_layers.append(res_skip_layer)
//  [python]
def __init__(self, n_in_channels, n_mel_channels, n_layers, n_channels, kernel_size):
        super(WN, self).__init__()
        assert kernel_size % 2 == 1
        assert n_channels % 2 == 0
        self.n_layers = n_layers
        self.n_channels = n_channels
        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.cond_layers = torch.nn.ModuleList()

        start = torch.nn.Conv1d(n_in_channels, n_channels, 1)
        start = torch.nn.utils.weight_norm(start, name="weight")
        self.start = start

        # Initializing last layer to 0 makes the affine coupling layers
        # do nothing at first.  This helps with training stability
        end = torch.nn.Conv1d(n_channels, 2 * n_in_channels, 1)
        end.weight.data.zero_()
        end.bias.data.zero_()
        self.end = end

        for i in range(n_layers):
            dilation = 2**i
            padding = int((kernel_size * dilation - dilation) / 2)
            in_layer = torch.nn.Conv1d(n_channels, 2 * n_channels, kernel_size, dilation=dilation, padding=padding)
            in_layer = torch.nn.utils.weight_norm(in_layer, name="weight")
            self.in_layers.append(in_layer)

            cond_layer = torch.nn.Conv1d(n_mel_channels, 2 * n_channels, 1)
            cond_layer = torch.nn.utils.weight_norm(cond_layer, name="weight")
            self.cond_layers.append(cond_layer)

            # last one is not necessary
            if i < n_layers - 1:
                res_skip_channels = 2 * n_channels
            else:
                res_skip_channels = n_channels
            res_skip_layer = torch.nn.Conv1d(n_channels, res_skip_channels, 1)
            res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name="weight")
            self.res_skip_layers.append(res_skip_layer)
//  [python]
def __init__(self, n_in_channels, n_mel_channels, n_layers, n_channels, kernel_size):
        super(WN, self).__init__()
        assert kernel_size % 2 == 1
        assert n_channels % 2 == 0
        self.n_layers = n_layers
        self.n_channels = n_channels
        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.cond_layers = torch.nn.ModuleList()

        start = torch.nn.Conv1d(n_in_channels, n_channels, 1)
        start = torch.nn.utils.parametrizations.weight_norm(start, name="weight")
        self.start = start

        # Initializing last layer to 0 makes the affine coupling layers
        # do nothing at first.  This helps with training stability
        end = torch.nn.Conv1d(n_channels, 2 * n_in_channels, 1)
        end.weight.data.zero_()
        end.bias.data.zero_()
        self.end = end

        for i in range(n_layers):
            dilation = 2**i
            padding = int((kernel_size * dilation - dilation) / 2)
            in_layer = torch.nn.Conv1d(n_channels, 2 * n_channels, kernel_size, dilation=dilation, padding=padding)
            in_layer = torch.nn.utils.parametrizations.weight_norm(in_layer, name="weight")
            self.in_layers.append(in_layer)

            cond_layer = torch.nn.Conv1d(n_mel_channels, 2 * n_channels, 1)
            cond_layer = torch.nn.utils.parametrizations.weight_norm(cond_layer, name="weight")
            self.cond_layers.append(cond_layer)

            # last one is not necessary
            if i < n_layers - 1:
                res_skip_channels = 2 * n_channels
            else:
                res_skip_channels = n_channels
            res_skip_layer = torch.nn.Conv1d(n_channels, res_skip_channels, 1)
            res_skip_layer = torch.nn.utils.parametrizations.weight_norm(res_skip_layer, name="weight")
            self.res_skip_layers.append(res_skip_layer)
//  [python]
def apply_norm(net, weight_norm_type):
    for m in net.modules():
        if isinstance(m, nn.Conv2d):
            if weight_norm_type.lower() == 'spectral_norm':
                tutils.spectral_norm(m)
            elif weight_norm_type.lower() == 'weight_norm':
                tutils.weight_norm(m)
            else:
                pass
//  [python]
def __init__(self, layer_dims, skip_connection=[], activ=None, use_layernorm=False, use_weightnorm=False,
                 geometric_init=False, out_bias=0., invert=False):
        """Initialize a multi-layer perceptron with skip connection.
        Args:
            layer_dims: A list of integers representing the number of channels in each layer.
            skip_connection: A list of integers representing the index of layers to add skip connection.
        """
        super().__init__()
        self.skip_connection = skip_connection
        self.use_layernorm = use_layernorm
        self.linears = torch.nn.ModuleList()
        if use_layernorm:
            self.layer_norm = torch.nn.ModuleList()
        # Hidden layers
        layer_dim_pairs = list(zip(layer_dims[:-1], layer_dims[1:]))
        for li, (k_in, k_out) in enumerate(layer_dim_pairs):
            if li in self.skip_connection:
                k_in += layer_dims[0]
            linear = torch.nn.Linear(k_in, k_out)
            if geometric_init:
                self._geometric_init(linear, k_in, k_out, first=(li == 0),
                                     skip_dim=(layer_dims[0] if li in self.skip_connection else 0))
            if use_weightnorm:
                linear = torch.nn.utils.weight_norm(linear)
            self.linears.append(linear)
            if use_layernorm and li != len(layer_dim_pairs) - 1:
                self.layer_norm.append(torch.nn.LayerNorm(k_out))
            if li == len(layer_dim_pairs) - 1:
                self.linears[-1].bias.data.fill_(0.0)
        # SDF prediction layer
        self.linear_sdf = torch.nn.Linear(k_in, 1)
        if geometric_init:
            self._geometric_init_sdf(self.linear_sdf, k_in, out_bias=out_bias, invert=invert)
        self.activ = activ or torch_F.relu_
//  [python]
def __init__(self, layer_dims, skip_connection=[], activ=None, use_layernorm=False, use_weightnorm=False,
                 geometric_init=False, out_bias=0., invert=False):
        """Initialize a multi-layer perceptron with skip connection.
        Args:
            layer_dims: A list of integers representing the number of channels in each layer.
            skip_connection: A list of integers representing the index of layers to add skip connection.
        """
        super().__init__()
        self.skip_connection = skip_connection
        self.use_layernorm = use_layernorm
        self.linears = torch.nn.ModuleList()
        if use_layernorm:
            self.layer_norm = torch.nn.ModuleList()
        # Hidden layers
        layer_dim_pairs = list(zip(layer_dims[:-1], layer_dims[1:]))
        for li, (k_in, k_out) in enumerate(layer_dim_pairs):
            if li in self.skip_connection:
                k_in += layer_dims[0]
            linear = torch.nn.Linear(k_in, k_out)
            if geometric_init:
                self._geometric_init(linear, k_in, k_out, first=(li == 0),
                                     skip_dim=(layer_dims[0] if li in self.skip_connection else 0))
            if use_weightnorm:
                linear = torch.nn.utils.parametrizations.weight_norm(linear)
            self.linears.append(linear)
            if use_layernorm and li != len(layer_dim_pairs) - 1:
                self.layer_norm.append(torch.nn.LayerNorm(k_out))
            if li == len(layer_dim_pairs) - 1:
                self.linears[-1].bias.data.fill_(0.0)
        # SDF prediction layer
        self.linear_sdf = torch.nn.Linear(k_in, 1)
        if geometric_init:
            self._geometric_init_sdf(self.linear_sdf, k_in, out_bias=out_bias, invert=invert)
        self.activ = activ or torch_F.relu_
//  [python]
def apply_weight_norm(self):
        """Apply weight normalization module from all of the layers."""

        def _apply_weight_norm(m):
            if isinstance(m, torch.nn.Conv1d):
                torch.nn.utils.weight_norm(m)
                logging.debug(f"Weight norm is applied to {m}.")

        self.apply(_apply_weight_norm)
//  [python]
def apply_weight_norm(self):
        """Apply weight normalization module from all of the layers."""

        def _apply_weight_norm(m):
            if isinstance(m, torch.nn.Conv1d):
                torch.nn.utils.parametrizations.weight_norm(m)
                logging.debug(f"Weight norm is applied to {m}.")

        self.apply(_apply_weight_norm)
//  [python]
def __init__(self, query_size, key_size, value_size=None, mode='bahdanau', normalize=False, dropout=0, batch_first=False, weight_norm=False, bias=True, query_transform=True, output_transform=True, output_nonlinearity='tanh', output_size=None):
        super(AttentionLayer, self).__init__()
        assert mode == 'bahdanau' or mode == 'dot_prod'
        value_size = value_size or key_size
        self.mode = mode
        self.query_size = query_size
        self.key_size = key_size
        self.value_size = value_size
        self.normalize = normalize
        wn_func = wn if weight_norm else lambda x: x
        if mode == 'bahdanau':
            self.linear_att = nn.Linear(key_size, 1, bias=bias)
            if normalize:
                self.linear_att = nn.utils.weight_norm(self.linear_att)
        elif normalize:
            self.scale = nn.Parameter(torch.Tensor([1]))
        if output_transform:
            output_size = output_size or query_size
            self.linear_out = wn_func(nn.Linear(query_size + value_size, output_size, bias=bias))
            self.output_size = output_size
        else:
            self.output_size = value_size
        if query_transform:
            self.linear_q = wn_func(nn.Linear(query_size, key_size, bias=bias))
        self.dropout = nn.Dropout(dropout)
        self.batch_first = batch_first
        self.output_nonlinearity = output_nonlinearity
        self.mask = None
//  [python]
def __init__(self, query_size, key_size, value_size=None, mode='bahdanau', normalize=False, dropout=0, batch_first=False, weight_norm=False, bias=True, query_transform=True, output_transform=True, output_nonlinearity='tanh', output_size=None):
        super(AttentionLayer, self).__init__()
        assert mode == 'bahdanau' or mode == 'dot_prod'
        value_size = value_size or key_size
        self.mode = mode
        self.query_size = query_size
        self.key_size = key_size
        self.value_size = value_size
        self.normalize = normalize
        wn_func = wn if weight_norm else lambda x: x
        if mode == 'bahdanau':
            self.linear_att = nn.Linear(key_size, 1, bias=bias)
            if normalize:
                self.linear_att = nn.utils.parametrizations.weight_norm(self.linear_att)
        elif normalize:
            self.scale = nn.Parameter(torch.Tensor([1]))
        if output_transform:
            output_size = output_size or query_size
            self.linear_out = wn_func(nn.Linear(query_size + value_size, output_size, bias=bias))
            self.output_size = output_size
        else:
            self.output_size = value_size
        if query_transform:
            self.linear_q = wn_func(nn.Linear(query_size, key_size, bias=bias))
        self.dropout = nn.Dropout(dropout)
        self.batch_first = batch_first
        self.output_nonlinearity = output_nonlinearity
        self.mask = None
//  [python]
def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0,
                 p_dropout=0, share_cond_layers=False):
        super(WN, self).__init__()
        assert (kernel_size % 2 == 1)
        assert (hidden_channels % 2 == 0)
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        self.gin_channels = gin_channels
        self.p_dropout = p_dropout
        self.share_cond_layers = share_cond_layers

        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        
        self.drop = nn.Dropout(p_dropout)

        self.use_adapters = hparams.get("use_adapters", False)
        if self.use_adapters:
            self.adapter_layers = torch.nn.ModuleList()

        if gin_channels != 0 and not share_cond_layers:
            cond_layer = torch.nn.Conv1d(gin_channels, 2 * hidden_channels * n_layers, 1)
            self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')

        for i in range(n_layers):
            dilation = dilation_rate ** i
            padding = int((kernel_size * dilation - dilation) / 2)
            in_layer = torch.nn.Conv1d(hidden_channels, 2 * hidden_channels, kernel_size,
                                       dilation=dilation, padding=padding)
            in_layer = torch.nn.utils.weight_norm(in_layer, name='weight')
            self.in_layers.append(in_layer)

            # last one is not necessary
            if i < n_layers - 1:
                res_skip_channels = 2 * hidden_channels
            else:
                res_skip_channels = hidden_channels

            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
            res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight')
            self.res_skip_layers.append(res_skip_layer)

            if self.use_adapters:
                adapter_layer = MlpAdapter(in_out_dim=res_skip_channels, hid_dim=res_skip_channels//4)
                self.adapter_layers.append(adapter_layer)
//  [python]
def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0,
                 p_dropout=0, share_cond_layers=False):
        super(WN, self).__init__()
        assert (kernel_size % 2 == 1)
        assert (hidden_channels % 2 == 0)
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        self.gin_channels = gin_channels
        self.p_dropout = p_dropout
        self.share_cond_layers = share_cond_layers

        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        
        self.drop = nn.Dropout(p_dropout)

        self.use_adapters = hparams.get("use_adapters", False)
        if self.use_adapters:
            self.adapter_layers = torch.nn.ModuleList()

        if gin_channels != 0 and not share_cond_layers:
            cond_layer = torch.nn.Conv1d(gin_channels, 2 * hidden_channels * n_layers, 1)
            self.cond_layer = torch.nn.utils.parametrizations.weight_norm(cond_layer, name='weight')

        for i in range(n_layers):
            dilation = dilation_rate ** i
            padding = int((kernel_size * dilation - dilation) / 2)
            in_layer = torch.nn.Conv1d(hidden_channels, 2 * hidden_channels, kernel_size,
                                       dilation=dilation, padding=padding)
            in_layer = torch.nn.utils.parametrizations.weight_norm(in_layer, name='weight')
            self.in_layers.append(in_layer)

            # last one is not necessary
            if i < n_layers - 1:
                res_skip_channels = 2 * hidden_channels
            else:
                res_skip_channels = hidden_channels

            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
            res_skip_layer = torch.nn.utils.parametrizations.weight_norm(res_skip_layer, name='weight')
            self.res_skip_layers.append(res_skip_layer)

            if self.use_adapters:
                adapter_layer = MlpAdapter(in_out_dim=res_skip_channels, hid_dim=res_skip_channels//4)
                self.adapter_layers.append(adapter_layer)
//  [python]
def get_dino_head_weights(pretrain_path):

    """
    :param pretrain_path: Path to full DINO pretrained checkpoint as in https://github.com/facebookresearch/dino
     'full_ckpt'
    :return: weights only for the projection head
    """

    all_weights = torch.load(pretrain_path)

    head_state_dict = {}
    for k, v in all_weights['teacher'].items():
        if 'head' in k and 'last_layer' not in k:
            head_state_dict[k] = v

    head_state_dict = strip_state_dict(head_state_dict, strip_key='head.')

    # Deal with weight norm
    weight_norm_state_dict = {}
    for k, v in all_weights['teacher'].items():
        if 'last_layer' in k:
            weight_norm_state_dict[k.split('.')[2]] = v

    linear_shape = weight_norm_state_dict['weight'].shape
    dummy_linear = torch.nn.Linear(in_features=linear_shape[1], out_features=linear_shape[0], bias=False)
    dummy_linear.load_state_dict(weight_norm_state_dict)
    dummy_linear = torch.nn.utils.weight_norm(dummy_linear)

    for k, v in dummy_linear.state_dict().items():

        head_state_dict['last_layer.' + k] = v

    return head_state_dict
//  [python]
def get_dino_head_weights(pretrain_path):

    """
    :param pretrain_path: Path to full DINO pretrained checkpoint as in https://github.com/facebookresearch/dino
     'full_ckpt'
    :return: weights only for the projection head
    """

    all_weights = torch.load(pretrain_path)

    head_state_dict = {}
    for k, v in all_weights['teacher'].items():
        if 'head' in k and 'last_layer' not in k:
            head_state_dict[k] = v

    head_state_dict = strip_state_dict(head_state_dict, strip_key='head.')

    # Deal with weight norm
    weight_norm_state_dict = {}
    for k, v in all_weights['teacher'].items():
        if 'last_layer' in k:
            weight_norm_state_dict[k.split('.')[2]] = v

    linear_shape = weight_norm_state_dict['weight'].shape
    dummy_linear = torch.nn.Linear(in_features=linear_shape[1], out_features=linear_shape[0], bias=False)
    dummy_linear.load_state_dict(weight_norm_state_dict)
    dummy_linear = torch.nn.utils.parametrizations.weight_norm(dummy_linear)

    for k, v in dummy_linear.state_dict().items():

        head_state_dict['last_layer.' + k] = v

    return head_state_dict
//  [python]
def __init__(self,
                 in_channels,
                 hidden_channels,
                 kernel_size,
                 dilation_rate,
                 n_blocks,
                 n_layers,
                 p_dropout=0.,
                 n_split=4,
                 n_sqz=2,
                 sigmoid_scale=False,
                 gin_channels=0,
                 inv_conv_type='near',
                 share_cond_layers=False,
                 share_wn_layers=0,
                 ):
        super().__init__()
        """
        Note that regularization likes weight decay can leads to Nan error!
        """

        self.in_channels = in_channels
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size
        self.dilation_rate = dilation_rate
        self.n_blocks = n_blocks
        self.n_layers = n_layers
        self.p_dropout = p_dropout
        self.n_split = n_split
        self.n_sqz = n_sqz
        self.sigmoid_scale = sigmoid_scale
        self.gin_channels = gin_channels
        self.share_cond_layers = share_cond_layers
        if gin_channels != 0 and share_cond_layers:
            cond_layer = torch.nn.Conv1d(gin_channels * n_sqz, 2 * hidden_channels * n_layers, 1)
            self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')
        wn = None
        self.flows = nn.ModuleList()
        for b in range(n_blocks):
            self.flows.append(ActNorm(channels=in_channels * n_sqz))
            if inv_conv_type == 'near':
                self.flows.append(InvConvNear(channels=in_channels * n_sqz, n_split=n_split, n_sqz=n_sqz))
            if inv_conv_type == 'invconv':
                self.flows.append(InvConv(channels=in_channels * n_sqz))
            if share_wn_layers > 0:
                if b % share_wn_layers == 0:
                    wn = WN(hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels * n_sqz,
                            p_dropout, share_cond_layers)
            self.flows.append(
                CouplingBlock(
                    in_channels * n_sqz,
                    hidden_channels,
                    kernel_size=kernel_size,
                    dilation_rate=dilation_rate,
                    n_layers=n_layers,
                    gin_channels=gin_channels * n_sqz,
                    p_dropout=p_dropout,
                    sigmoid_scale=sigmoid_scale,
                    share_cond_layers=share_cond_layers,
                    wn=wn
                ))
//  [python]
def __init__(self,
                 in_channels,
                 hidden_channels,
                 kernel_size,
                 dilation_rate,
                 n_blocks,
                 n_layers,
                 p_dropout=0.,
                 n_split=4,
                 n_sqz=2,
                 sigmoid_scale=False,
                 gin_channels=0,
                 inv_conv_type='near',
                 share_cond_layers=False,
                 share_wn_layers=0,
                 ):
        super().__init__()
        """
        Note that regularization likes weight decay can leads to Nan error!
        """

        self.in_channels = in_channels
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size
        self.dilation_rate = dilation_rate
        self.n_blocks = n_blocks
        self.n_layers = n_layers
        self.p_dropout = p_dropout
        self.n_split = n_split
        self.n_sqz = n_sqz
        self.sigmoid_scale = sigmoid_scale
        self.gin_channels = gin_channels
        self.share_cond_layers = share_cond_layers
        if gin_channels != 0 and share_cond_layers:
            cond_layer = torch.nn.Conv1d(gin_channels * n_sqz, 2 * hidden_channels * n_layers, 1)
            self.cond_layer = torch.nn.utils.parametrizations.weight_norm(cond_layer, name='weight')
        wn = None
        self.flows = nn.ModuleList()
        for b in range(n_blocks):
            self.flows.append(ActNorm(channels=in_channels * n_sqz))
            if inv_conv_type == 'near':
                self.flows.append(InvConvNear(channels=in_channels * n_sqz, n_split=n_split, n_sqz=n_sqz))
            if inv_conv_type == 'invconv':
                self.flows.append(InvConv(channels=in_channels * n_sqz))
            if share_wn_layers > 0:
                if b % share_wn_layers == 0:
                    wn = WN(hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels * n_sqz,
                            p_dropout, share_cond_layers)
            self.flows.append(
                CouplingBlock(
                    in_channels * n_sqz,
                    hidden_channels,
                    kernel_size=kernel_size,
                    dilation_rate=dilation_rate,
                    n_layers=n_layers,
                    gin_channels=gin_channels * n_sqz,
                    p_dropout=p_dropout,
                    sigmoid_scale=sigmoid_scale,
                    share_cond_layers=share_cond_layers,
                    wn=wn
                ))
//  [python]
def update_model(old_model):
    if not _check_model_old_version(old_model):
        return old_model
    new_model = copy.deepcopy(old_model)
    for idx in range(0, len(new_model.WN)):
        wavenet = new_model.WN[idx]
        wavenet.res_skip_layers = torch.nn.ModuleList()
        n_channels = wavenet.n_channels
        n_layers = wavenet.n_layers
        for i in range(0, n_layers):
            if i < n_layers - 1:
                res_skip_channels = 2*n_channels
            else:
                res_skip_channels = n_channels
            res_skip_layer = torch.nn.Conv1d(n_channels, res_skip_channels, 1)
            skip_layer = torch.nn.utils.remove_weight_norm(wavenet.skip_layers[i])
            if i < n_layers - 1:
                res_layer = torch.nn.utils.remove_weight_norm(wavenet.res_layers[i])
                res_skip_layer.weight = torch.nn.Parameter(torch.cat([res_layer.weight, skip_layer.weight]))
                res_skip_layer.bias = torch.nn.Parameter(torch.cat([res_layer.bias, skip_layer.bias]))
            else:
                res_skip_layer.weight = torch.nn.Parameter(skip_layer.weight)
                res_skip_layer.bias = torch.nn.Parameter(skip_layer.bias)
            res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight')
            wavenet.res_skip_layers.append(res_skip_layer)
        del wavenet.res_layers
        del wavenet.skip_layers
    return new_model
//  [python]
def update_model(old_model):
    if not _check_model_old_version(old_model):
        return old_model
    new_model = copy.deepcopy(old_model)
    for idx in range(0, len(new_model.WN)):
        wavenet = new_model.WN[idx]
        wavenet.res_skip_layers = torch.nn.ModuleList()
        n_channels = wavenet.n_channels
        n_layers = wavenet.n_layers
        for i in range(0, n_layers):
            if i < n_layers - 1:
                res_skip_channels = 2*n_channels
            else:
                res_skip_channels = n_channels
            res_skip_layer = torch.nn.Conv1d(n_channels, res_skip_channels, 1)
            skip_layer = torch.nn.utils.remove_weight_norm(wavenet.skip_layers[i])
            if i < n_layers - 1:
                res_layer = torch.nn.utils.remove_weight_norm(wavenet.res_layers[i])
                res_skip_layer.weight = torch.nn.Parameter(torch.cat([res_layer.weight, skip_layer.weight]))
                res_skip_layer.bias = torch.nn.Parameter(torch.cat([res_layer.bias, skip_layer.bias]))
            else:
                res_skip_layer.weight = torch.nn.Parameter(skip_layer.weight)
                res_skip_layer.bias = torch.nn.Parameter(skip_layer.bias)
            res_skip_layer = torch.nn.utils.parametrizations.weight_norm(res_skip_layer, name='weight')
            wavenet.res_skip_layers.append(res_skip_layer)
        del wavenet.res_layers
        del wavenet.skip_layers
    return new_model
//  [python]
def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0):
    super(WN, self).__init__()
    assert(kernel_size % 2 == 1)
    self.hidden_channels =hidden_channels
    self.kernel_size = kernel_size,
    self.dilation_rate = dilation_rate
    self.n_layers = n_layers
    self.gin_channels = gin_channels
    self.p_dropout = p_dropout

    self.in_layers = torch.nn.ModuleList()
    self.res_skip_layers = torch.nn.ModuleList()
    self.drop = nn.Dropout(p_dropout)

    if gin_channels != 0:
      cond_layer = torch.nn.Conv1d(gin_channels, 2*hidden_channels*n_layers, 1).to(device)
      self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')

    for i in range(n_layers):
      dilation = dilation_rate ** i
      padding = int((kernel_size * dilation - dilation) / 2)
      in_layer = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, kernel_size,
                                 dilation=dilation, padding=padding).to(device)
      in_layer = torch.nn.utils.weight_norm(in_layer, name='weight')
      self.in_layers.append(in_layer)

      # last one is not necessary
      if i < n_layers - 1:
        res_skip_channels = 2 * hidden_channels
      else:
        res_skip_channels = hidden_channels

      res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1).to(device)
      res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight')
      self.res_skip_layers.append(res_skip_layer)
//  [python]
def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0):
    super(WN, self).__init__()
    assert(kernel_size % 2 == 1)
    self.hidden_channels =hidden_channels
    self.kernel_size = kernel_size,
    self.dilation_rate = dilation_rate
    self.n_layers = n_layers
    self.gin_channels = gin_channels
    self.p_dropout = p_dropout

    self.in_layers = torch.nn.ModuleList()
    self.res_skip_layers = torch.nn.ModuleList()
    self.drop = nn.Dropout(p_dropout)

    if gin_channels != 0:
      cond_layer = torch.nn.Conv1d(gin_channels, 2*hidden_channels*n_layers, 1).to(device)
      self.cond_layer = torch.nn.utils.parametrizations.weight_norm(cond_layer, name='weight')

    for i in range(n_layers):
      dilation = dilation_rate ** i
      padding = int((kernel_size * dilation - dilation) / 2)
      in_layer = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, kernel_size,
                                 dilation=dilation, padding=padding).to(device)
      in_layer = torch.nn.utils.parametrizations.weight_norm(in_layer, name='weight')
      self.in_layers.append(in_layer)

      # last one is not necessary
      if i < n_layers - 1:
        res_skip_channels = 2 * hidden_channels
      else:
        res_skip_channels = hidden_channels

      res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1).to(device)
      res_skip_layer = torch.nn.utils.parametrizations.weight_norm(res_skip_layer, name='weight')
      self.res_skip_layers.append(res_skip_layer)
//  [python]
def __init__(self, n_in_channels, n_mel_channels, d_i, n_layers, n_channels, kernel_size=3):
        super().__init__()
        assert(kernel_size % 2 == 1)
        assert(n_channels % 2 == 0)
        self.n_layers = n_layers
        self.n_channels = n_channels
        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()

        start = torch.nn.Conv1d(n_in_channels, n_channels, 1)
        start = torch.nn.utils.weight_norm(start, name='weight')
        self.start = start

        # Initializing last layer to 0 makes the affine coupling layers
        # do nothing at first.  This helps with training stability
        end = torch.nn.Conv1d(n_channels, 2 * n_in_channels, 1)
        end.weight.data.zero_()
        end.bias.data.zero_()
        self.end = end

        cond_layer = torch.nn.Conv1d(n_mel_channels, 2*n_channels*n_layers, 1)
        self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')

        for i in range(n_layers):
            dilation = d_i ** i
            padding = int((kernel_size*dilation - dilation)/2)
            in_layer = torch.nn.Conv1d(n_channels, 2*n_channels, kernel_size,
                                       dilation=dilation, padding=padding)
            in_layer = torch.nn.utils.weight_norm(in_layer, name='weight')
            self.in_layers.append(in_layer)

            # last one is not necessary
            if i < n_layers - 1:
                res_skip_channels = 2*n_channels
            else:
                res_skip_channels = n_channels
            res_skip_layer = torch.nn.Conv1d(n_channels, res_skip_channels, 1)
            res_skip_layer = torch.nn.utils.weight_norm(
                res_skip_layer, name='weight')
            self.res_skip_layers.append(res_skip_layer)