jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(self, input_size, hidden_size, bias=True):
        super().__init__()
        self.f = torch.nn.Sequential(
            torch.nn.utils.parametrizations.weight_norm(torch.nn.Linear(input_size, hidden_size)),
            torch.nn.Tanh(),
            torch.nn.utils.parametrizations.weight_norm(torch.nn.Linear(hidden_size, 1, bias=bias)),
        )
        if bias:
            self.f[-1].bias.data[:] = 5.0
//  [python]
def __init__(self, input_size, hidden_size, max_activation=10, bias=True):
        super().__init__()
        self.f = torch.nn.Sequential(
            torch.nn.utils.weight_norm(torch.nn.Linear(input_size, hidden_size)),
            torch.nn.Tanh(),
            torch.nn.utils.weight_norm(torch.nn.Linear(hidden_size, 1, bias=bias)),
            torch.nn.Tanh(),
        )
        self.bias = torch.nn.Parameter(torch.tensor(5.0))
        self.max_activation = max_activation
//  [python]
def __init__(self, input_size, hidden_size, max_activation=10, bias=True):
        super().__init__()
        self.f = torch.nn.Sequential(
            torch.nn.utils.parametrizations.weight_norm(torch.nn.Linear(input_size, hidden_size)),
            torch.nn.Tanh(),
            torch.nn.utils.parametrizations.weight_norm(torch.nn.Linear(hidden_size, 1, bias=bias)),
            torch.nn.Tanh(),
        )
        self.bias = torch.nn.Parameter(torch.tensor(5.0))
        self.max_activation = max_activation
//  [python]
def __init__(self, hidden_size, kernel_size, dilation_rate, n_layers, c_cond=0,
                 p_dropout=0, share_cond_layers=False, is_BTC=False, use_weightnorm=False):
        """
        If weightnorm is set to false, we can deepcopy the module, which we need to be able to do to perform SWA.
        Without weightnorm, the module will probably take a little longer to converge.
        """
        super(WN, self).__init__()
        assert (kernel_size % 2 == 1)
        assert (hidden_size % 2 == 0)
        self.is_BTC = is_BTC
        self.hidden_size = hidden_size
        self.kernel_size = kernel_size
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        self.gin_channels = c_cond
        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)

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

        for i in range(n_layers):
            dilation = dilation_rate ** i
            padding = int((kernel_size * dilation - dilation) / 2)
            in_layer = torch.nn.Conv1d(hidden_size, 2 * hidden_size, kernel_size,
                                       dilation=dilation, padding=padding)
            if use_weightnorm:
                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_size
            else:
                res_skip_channels = hidden_size

            res_skip_layer = torch.nn.Conv1d(hidden_size, res_skip_channels, 1)
            if use_weightnorm:
                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_size, kernel_size, dilation_rate, n_layers, c_cond=0,
                 p_dropout=0, share_cond_layers=False, is_BTC=False, use_weightnorm=False):
        """
        If weightnorm is set to false, we can deepcopy the module, which we need to be able to do to perform SWA.
        Without weightnorm, the module will probably take a little longer to converge.
        """
        super(WN, self).__init__()
        assert (kernel_size % 2 == 1)
        assert (hidden_size % 2 == 0)
        self.is_BTC = is_BTC
        self.hidden_size = hidden_size
        self.kernel_size = kernel_size
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        self.gin_channels = c_cond
        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)

        if c_cond != 0 and not share_cond_layers:
            cond_layer = torch.nn.Conv1d(c_cond, 2 * hidden_size * n_layers, 1)
            if use_weightnorm:
                self.cond_layer = torch.nn.utils.parametrizations.weight_norm(cond_layer)
            else:
                self.cond_layer = cond_layer

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

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

            res_skip_layer = torch.nn.Conv1d(hidden_size, res_skip_channels, 1)
            if use_weightnorm:
                res_skip_layer = torch.nn.utils.parametrizations.weight_norm(res_skip_layer)
            self.res_skip_layers.append(res_skip_layer)
//  [python]
def __init__(self, args):
        super(WDSR_B_Model, self).__init__()
        # hyper-params
        self.args = args
        scale = args['scale']
        n_resblocks = args['n_resblocks']
        n_feats = args['n_feats']
        kernel_size = 3
        act = nn.ReLU(True)
        # wn = lambda x: x
        wn = lambda x: torch.nn.utils.weight_norm(x)

        # define head module
        head = []
        head.append(
            wn(nn.Conv2d(3, n_feats, 3, padding=3//2)))

        # define body module
        body = []
        for i in range(n_resblocks):
            body.append(
                Block(n_feats, kernel_size, act=act, res_scale=args['res_scale'], wn=wn))

        # define tail module
        tail = []
        out_feats = scale*scale*3
        tail.append(
            wn(nn.Conv2d(n_feats, out_feats, 3, padding=3//2)))
        tail.append(nn.PixelShuffle(scale))

        skip = []
        skip.append(
            wn(nn.Conv2d(3, out_feats, 5, padding=5//2))
        )
        skip.append(nn.PixelShuffle(scale))

        # make object members
        self.head = nn.Sequential(*head)
        self.body = nn.Sequential(*body)
        self.tail = nn.Sequential(*tail)
        self.skip = nn.Sequential(*skip)
//  [python]
def __init__(self, args):
        super(WDSR_B_Model, self).__init__()
        # hyper-params
        self.args = args
        scale = args['scale']
        n_resblocks = args['n_resblocks']
        n_feats = args['n_feats']
        kernel_size = 3
        act = nn.ReLU(True)
        # wn = lambda x: x
        wn = lambda x: torch.nn.utils.parametrizations.weight_norm(x)

        # define head module
        head = []
        head.append(
            wn(nn.Conv2d(3, n_feats, 3, padding=3//2)))

        # define body module
        body = []
        for i in range(n_resblocks):
            body.append(
                Block(n_feats, kernel_size, act=act, res_scale=args['res_scale'], wn=wn))

        # define tail module
        tail = []
        out_feats = scale*scale*3
        tail.append(
            wn(nn.Conv2d(n_feats, out_feats, 3, padding=3//2)))
        tail.append(nn.PixelShuffle(scale))

        skip = []
        skip.append(
            wn(nn.Conv2d(3, out_feats, 5, padding=5//2))
        )
        skip.append(nn.PixelShuffle(scale))

        # make object members
        self.head = nn.Sequential(*head)
        self.body = nn.Sequential(*body)
        self.tail = nn.Sequential(*tail)
        self.skip = nn.Sequential(*skip)
//  [python]
def make_conv_pos(e, k, g):
    pos_conv = nn.Conv1d(e, e, kernel_size=k, padding=k // 2, groups=g)
    dropout = 0
    std = math.sqrt(4 * (1.0 - dropout) / (k * e))
    nn.init.normal_(pos_conv.weight, mean=0, std=std)
    nn.init.constant_(pos_conv.bias, 0)
    pos_conv = nn.utils.weight_norm(pos_conv, name='weight', dim=2)
    pos_conv = nn.Sequential(pos_conv, SamePad(k), nn.GELU())
    return pos_conv
//  [python]
def make_conv_pos(e, k, g):
    pos_conv = nn.Conv1d(e, e, kernel_size=k, padding=k // 2, groups=g)
    dropout = 0
    std = math.sqrt(4 * (1.0 - dropout) / (k * e))
    nn.init.normal_(pos_conv.weight, mean=0, std=std)
    nn.init.constant_(pos_conv.bias, 0)
    pos_conv = nn.utils.parametrizations.weight_norm(pos_conv, name='weight', dim=2)
    pos_conv = nn.Sequential(pos_conv, SamePad(k), nn.GELU())
    return pos_conv
//  [python]
def conv_layer(in_channels, out_channels, kernel_size, stride=1, dilation=1, groups=1):
    padding = int((kernel_size - 1) / 2) * dilation
    wn = lambda x: torch.nn.utils.weight_norm(x)
    return wn(nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding=padding, bias=True, dilation=dilation,
                     groups=groups))
//  [python]
def conv_layer(in_channels, out_channels, kernel_size, stride=1, dilation=1, groups=1):
    padding = int((kernel_size - 1) / 2) * dilation
    wn = lambda x: torch.nn.utils.parametrizations.weight_norm(x)
    return wn(nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding=padding, bias=True, dilation=dilation,
                     groups=groups))
//  [python]
def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
        super(ResBlock, self).__init__()
        self.convs1 = nn.ModuleList([weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0], padding=get_padding(kernel_size, dilation[0]))), weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1], padding=get_padding(kernel_size, dilation[1]))), weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2], padding=get_padding(kernel_size, dilation[2])))])
        self.convs1.apply(init_weights)
        self.convs2 = nn.ModuleList([weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1))), weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1))), weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)))])
        self.convs2.apply(init_weights)
//  [python]
def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
        super(ResBlock, self).__init__()
        self.convs1 = nn.ModuleList([parametrizations.weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0], padding=get_padding(kernel_size, dilation[0]))), parametrizations.weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1], padding=get_padding(kernel_size, dilation[1]))), parametrizations.weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2], padding=get_padding(kernel_size, dilation[2])))])
        self.convs1.apply(init_weights)
        self.convs2 = nn.ModuleList([parametrizations.weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1))), parametrizations.weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1))), parametrizations.weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)))])
        self.convs2.apply(init_weights)
//  [python]
def __init__(self, cfg):
        super(Generator, self).__init__()
        self.num_kernels = len(cfg['resblock_kernel_sizes'])
        self.num_upsamples = len(cfg['upsample_rates'])
        self.conv_pre = weight_norm(Conv1d(cfg.get('model_in_dim', 80), cfg['upsample_initial_channel'], 7, 1, padding=3))
        self.ups = nn.ModuleList()
        for i, (u, k) in enumerate(zip(cfg['upsample_rates'], cfg['upsample_kernel_sizes'])):
            self.ups.append(weight_norm(ConvTranspose1d(cfg['upsample_initial_channel'] // 2 ** i, cfg['upsample_initial_channel'] // 2 ** (i + 1), k, u, padding=(k - u) // 2)))
        self.resblocks = nn.ModuleList()
        for i in range(len(self.ups)):
            ch = cfg['upsample_initial_channel'] // 2 ** (i + 1)
            for k, d in zip(cfg['resblock_kernel_sizes'], cfg['resblock_dilation_sizes']):
                self.resblocks.append(ResBlock(ch, k, d))
        self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3))
        self.ups.apply(init_weights)
        self.conv_post.apply(init_weights)
//  [python]
def __init__(self, cfg):
        super(Generator, self).__init__()
        self.num_kernels = len(cfg['resblock_kernel_sizes'])
        self.num_upsamples = len(cfg['upsample_rates'])
        self.conv_pre = nn.utils.parametrizations.weight_norm(Conv1d(cfg.get('model_in_dim', 80), cfg['upsample_initial_channel'], 7, 1, padding=3))
        self.ups = nn.ModuleList()
        for i, (u, k) in enumerate(zip(cfg['upsample_rates'], cfg['upsample_kernel_sizes'])):
            self.ups.append(nn.utils.parametrizations.weight_norm(ConvTranspose1d(cfg['upsample_initial_channel'] // 2 ** i, cfg['upsample_initial_channel'] // 2 ** (i + 1), k, u, padding=(k - u) // 2)))
        self.resblocks = nn.ModuleList()
        for i in range(len(self.ups)):
            ch = cfg['upsample_initial_channel'] // 2 ** (i + 1)
            for k, d in zip(cfg['resblock_kernel_sizes'], cfg['resblock_dilation_sizes']):
                self.resblocks.append(ResBlock(ch, k, d))
        self.conv_post = nn.utils.parametrizations.weight_norm(Conv1d(ch, 1, 7, 1, padding=3))
        self.ups.apply(init_weights)
        self.conv_post.apply(init_weights)
//  [python]
def __init__(self, n_feats, conv):
        super(ESA, self).__init__()
        wn = lambda x: torch.nn.utils.weight_norm(x)
        f = n_feats // 4
        self.conv1 = wn(conv(n_feats, f, kernel_size=1))
        self.conv_f = wn(conv(f, f, kernel_size=1))
        self.conv_max = wn(conv(f, f, kernel_size=3, padding=1))
        self.conv2 = wn(conv(f, f, kernel_size=3, stride=2, padding=0))
        self.conv3 = wn(conv(f, f, kernel_size=3, padding=1))
        self.conv3_ = wn(conv(f, f, kernel_size=3, padding=1))
        self.conv4 = wn(conv(f, n_feats, kernel_size=1))
        self.sigmoid = nn.Sigmoid()
        self.relu = nn.SiLU(inplace=True)
//  [python]
def __init__(self, n_feats, conv):
        super(ESA, self).__init__()
        wn = lambda x: torch.nn.utils.parametrizations.weight_norm(x)
        f = n_feats // 4
        self.conv1 = wn(conv(n_feats, f, kernel_size=1))
        self.conv_f = wn(conv(f, f, kernel_size=1))
        self.conv_max = wn(conv(f, f, kernel_size=3, padding=1))
        self.conv2 = wn(conv(f, f, kernel_size=3, stride=2, padding=0))
        self.conv3 = wn(conv(f, f, kernel_size=3, padding=1))
        self.conv3_ = wn(conv(f, f, kernel_size=3, padding=1))
        self.conv4 = wn(conv(f, n_feats, kernel_size=1))
        self.sigmoid = nn.Sigmoid()
        self.relu = nn.SiLU(inplace=True)
//  [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()
        start = torch.nn.Conv1d(n_in_channels, n_channels, 1)
        start = torch.nn.utils.weight_norm(start, name='weight')
        self.start = start
        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 = 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)
            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()
        start = torch.nn.Conv1d(n_in_channels, n_channels, 1)
        start = torch.nn.utils.parametrizations.weight_norm(start, name='weight')
        self.start = start
        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.parametrizations.weight_norm(cond_layer, name='weight')
        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)
            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 __init__(
        self, channels: int, kernel_size: int = 3, dilation: typing.Tuple[int] = (1, 3)
    ):
        super(ResBlock2, self).__init__()
        self.LRELU_SLOPE = 0.1
        self.convs = nn.ModuleList(
            [
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[0],
                        padding=get_padding(kernel_size, dilation[0]),
                    )
                ),
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[1],
                        padding=get_padding(kernel_size, dilation[1]),
                    )
                ),
            ]
        )
        self.convs.apply(init_weights)
//  [python]
def __init__(
        self, channels: int, kernel_size: int = 3, dilation: typing.Tuple[int] = (1, 3)
    ):
        super(ResBlock2, self).__init__()
        self.LRELU_SLOPE = 0.1
        self.convs = nn.ModuleList(
            [
                parametrizations.weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[0],
                        padding=get_padding(kernel_size, dilation[0]),
                    )
                ),
                parametrizations.weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[1],
                        padding=get_padding(kernel_size, dilation[1]),
                    )
                ),
            ]
        )
        self.convs.apply(init_weights)
//  [python]
def __init__(self, args, alway_mask=False):
        super().__init__(args)
        self.args = args
        self.dropout = args.dropout
        self.embedding_dim = args.encoder_embed_dim
        self.feat_scale = math.sqrt(args.encoder_embed_dim)
        if args.no_scale_feature:
            self.feat_scale = 1.0
        subsample = ConvFeatureExtractionModel(conv_layers=eval(args.conv_feature_layers), dropout=0.0, mode=args.speech_extractor_mode, conv_bias=args.speech_conv_bias)
        self.feature_enc_layers = eval(args.conv_feature_layers)
        self.subsample = subsample
        self.feat_proj = nn.Linear(self.feature_enc_layers[-1][0], self.embedding_dim) if self.feature_enc_layers[-1][0] != self.embedding_dim else None
        self.feat_layer_norm = LayerNorm(self.feature_enc_layers[-1][0])
        self.embed_positions = nn.Conv1d(self.embedding_dim, self.embedding_dim, kernel_size=args.conv_pos, padding=args.conv_pos // 2, groups=args.conv_pos_groups)
        std = math.sqrt(4 / (args.conv_pos * self.embedding_dim))
        nn.init.normal_(self.embed_positions.weight, mean=0, std=std)
        nn.init.constant_(self.embed_positions.bias, 0)
        self.embed_positions = nn.utils.weight_norm(self.embed_positions, name='weight', dim=2)
        self.embed_positions = nn.Sequential(self.embed_positions, SamePad(args.conv_pos), nn.GELU())
        self.mask_prob = args.speech_mask_prob
        self.mask_selection = args.speech_mask_selection
        self.mask_other = args.speech_mask_other
        self.mask_length = args.speech_mask_length
        self.no_mask_overlap = args.speech_no_mask_overlap
        self.mask_min_space = args.speech_mask_min_space
        self.mask_channel_prob = args.speech_mask_channel_prob
        self.mask_channel_selection = args.speech_mask_channel_selection
        self.mask_channel_other = args.speech_mask_channel_other
        self.mask_channel_length = args.speech_mask_channel_length
        self.no_mask_channel_overlap = args.speech_no_mask_channel_overlap
        self.mask_channel_min_space = args.speech_mask_channel_min_space
        self.dropout_input = nn.Dropout(args.dropout_input)
        self.dropout_features = nn.Dropout(args.dropout_features)
        self.feature_grad_mult = args.feature_grad_mult
        self.mask_emb = nn.Parameter(torch.FloatTensor(args.encoder_embed_dim).uniform_())
        self.layers = nn.ModuleList([TransformerEncoderLayer(args) for _ in range(args.encoder_layers)])
        self.layer_norm = LayerNorm(args.encoder_embed_dim)
        self.normalize_before = args.encoder_normalize_before
        self.alway_mask = alway_mask
//  [python]
def __init__(self, args, alway_mask=False):
        super().__init__(args)
        self.args = args
        self.dropout = args.dropout
        self.embedding_dim = args.encoder_embed_dim
        self.feat_scale = math.sqrt(args.encoder_embed_dim)
        if args.no_scale_feature:
            self.feat_scale = 1.0
        subsample = ConvFeatureExtractionModel(conv_layers=eval(args.conv_feature_layers), dropout=0.0, mode=args.speech_extractor_mode, conv_bias=args.speech_conv_bias)
        self.feature_enc_layers = eval(args.conv_feature_layers)
        self.subsample = subsample
        self.feat_proj = nn.Linear(self.feature_enc_layers[-1][0], self.embedding_dim) if self.feature_enc_layers[-1][0] != self.embedding_dim else None
        self.feat_layer_norm = LayerNorm(self.feature_enc_layers[-1][0])
        self.embed_positions = nn.Conv1d(self.embedding_dim, self.embedding_dim, kernel_size=args.conv_pos, padding=args.conv_pos // 2, groups=args.conv_pos_groups)
        std = math.sqrt(4 / (args.conv_pos * self.embedding_dim))
        nn.init.normal_(self.embed_positions.weight, mean=0, std=std)
        nn.init.constant_(self.embed_positions.bias, 0)
        self.embed_positions = nn.utils.parametrizations.weight_norm(self.embed_positions, name='weight', dim=2)
        self.embed_positions = nn.Sequential(self.embed_positions, SamePad(args.conv_pos), nn.GELU())
        self.mask_prob = args.speech_mask_prob
        self.mask_selection = args.speech_mask_selection
        self.mask_other = args.speech_mask_other
        self.mask_length = args.speech_mask_length
        self.no_mask_overlap = args.speech_no_mask_overlap
        self.mask_min_space = args.speech_mask_min_space
        self.mask_channel_prob = args.speech_mask_channel_prob
        self.mask_channel_selection = args.speech_mask_channel_selection
        self.mask_channel_other = args.speech_mask_channel_other
        self.mask_channel_length = args.speech_mask_channel_length
        self.no_mask_channel_overlap = args.speech_no_mask_channel_overlap
        self.mask_channel_min_space = args.speech_mask_channel_min_space
        self.dropout_input = nn.Dropout(args.dropout_input)
        self.dropout_features = nn.Dropout(args.dropout_features)
        self.feature_grad_mult = args.feature_grad_mult
        self.mask_emb = nn.Parameter(torch.FloatTensor(args.encoder_embed_dim).uniform_())
        self.layers = nn.ModuleList([TransformerEncoderLayer(args) for _ in range(args.encoder_layers)])
        self.layer_norm = LayerNorm(args.encoder_embed_dim)
        self.normalize_before = args.encoder_normalize_before
        self.alway_mask = alway_mask
//  [python]
def __init__(
        self,
        hidden_channels: int,
        kernel_size: int,
        dilation_rate: int,
        n_layers: int,
        gin_channels: int = 0,
        p_dropout: float = 0,
    ):
        super().__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
            )
            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: int,
        kernel_size: int,
        dilation_rate: int,
        n_layers: int,
        gin_channels: int = 0,
        p_dropout: float = 0,
    ):
        super().__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
            )
            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,
        channels: int,
        kernel_size: int = 3,
        dilation: typing.Tuple[int] = (1, 3, 5),
    ):
        super(ResBlock1, self).__init__()
        self.LRELU_SLOPE = 0.1
        self.convs1 = nn.ModuleList(
            [
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[0],
                        padding=get_padding(kernel_size, dilation[0]),
                    )
                ),
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[1],
                        padding=get_padding(kernel_size, dilation[1]),
                    )
                ),
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[2],
                        padding=get_padding(kernel_size, dilation[2]),
                    )
                ),
            ]
        )
        self.convs1.apply(init_weights)

        self.convs2 = nn.ModuleList(
            [
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=1,
                        padding=get_padding(kernel_size, 1),
                    )
                ),
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=1,
                        padding=get_padding(kernel_size, 1),
                    )
                ),
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=1,
                        padding=get_padding(kernel_size, 1),
                    )
                ),
            ]
        )
        self.convs2.apply(init_weights)
//  [python]
def __init__(
        self,
        channels: int,
        kernel_size: int = 3,
        dilation: typing.Tuple[int] = (1, 3, 5),
    ):
        super(ResBlock1, self).__init__()
        self.LRELU_SLOPE = 0.1
        self.convs1 = nn.ModuleList(
            [
                parametrizations.weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[0],
                        padding=get_padding(kernel_size, dilation[0]),
                    )
                ),
                parametrizations.weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[1],
                        padding=get_padding(kernel_size, dilation[1]),
                    )
                ),
                parametrizations.weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[2],
                        padding=get_padding(kernel_size, dilation[2]),
                    )
                ),
            ]
        )
        self.convs1.apply(init_weights)

        self.convs2 = nn.ModuleList(
            [
                parametrizations.weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=1,
                        padding=get_padding(kernel_size, 1),
                    )
                ),
                parametrizations.weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=1,
                        padding=get_padding(kernel_size, 1),
                    )
                ),
                parametrizations.weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=1,
                        padding=get_padding(kernel_size, 1),
                    )
                ),
            ]
        )
        self.convs2.apply(init_weights)
//  [python]
def __init__(
        self, model_config: VocoderModelConfig, preprocess_config: PreprocessingConfig
    ):
        super(Generator, self).__init__()
        self.mel_channel = preprocess_config.stft.n_mel_channels
        self.noise_dim = model_config.gen.noise_dim
        self.hop_length = preprocess_config.stft.hop_length
        channel_size = model_config.gen.channel_size
        kpnet_conv_size = model_config.gen.kpnet_conv_size

        self.res_stack = nn.ModuleList()
        hop_length = 1
        for stride in model_config.gen.strides:
            hop_length = stride * hop_length
            self.res_stack.append(
                LVCBlock(
                    channel_size,
                    preprocess_config.stft.n_mel_channels,
                    stride=stride,
                    dilations=model_config.gen.dilations,
                    lReLU_slope=model_config.gen.lReLU_slope,
                    cond_hop_length=hop_length,
                    kpnet_conv_size=kpnet_conv_size,
                )
            )

        self.conv_pre = nn.utils.weight_norm(
            nn.Conv1d(
                model_config.gen.noise_dim,
                channel_size,
                7,
                padding=3,
                padding_mode="reflect",
            )
        )

        self.conv_post = nn.Sequential(
            nn.LeakyReLU(model_config.gen.lReLU_slope),
            nn.utils.weight_norm(
                nn.Conv1d(channel_size, 1, 7, padding=3, padding_mode="reflect")
            ),
            nn.Tanh(),
        )

        # Output of STFT(zeros)
        self.mel_mask_value = -11.5129
//  [python]
def __init__(
        self, model_config: VocoderModelConfig, preprocess_config: PreprocessingConfig
    ):
        super(Generator, self).__init__()
        self.mel_channel = preprocess_config.stft.n_mel_channels
        self.noise_dim = model_config.gen.noise_dim
        self.hop_length = preprocess_config.stft.hop_length
        channel_size = model_config.gen.channel_size
        kpnet_conv_size = model_config.gen.kpnet_conv_size

        self.res_stack = nn.ModuleList()
        hop_length = 1
        for stride in model_config.gen.strides:
            hop_length = stride * hop_length
            self.res_stack.append(
                LVCBlock(
                    channel_size,
                    preprocess_config.stft.n_mel_channels,
                    stride=stride,
                    dilations=model_config.gen.dilations,
                    lReLU_slope=model_config.gen.lReLU_slope,
                    cond_hop_length=hop_length,
                    kpnet_conv_size=kpnet_conv_size,
                )
            )

        self.conv_pre = nn.utils.parametrizations.weight_norm(
            nn.Conv1d(
                model_config.gen.noise_dim,
                channel_size,
                7,
                padding=3,
                padding_mode="reflect",
            )
        )

        self.conv_post = nn.Sequential(
            nn.LeakyReLU(model_config.gen.lReLU_slope),
            nn.utils.parametrizations.weight_norm(
                nn.Conv1d(channel_size, 1, 7, padding=3, padding_mode="reflect")
            ),
            nn.Tanh(),
        )

        # Output of STFT(zeros)
        self.mel_mask_value = -11.5129
//  [python]
def __init__(self, args):
        super(MODEL, self).__init__()
        # hyper-params
        self.args = args
        scale = args.scale[0]
        n_resblocks = args.n_resblocks
        n_feats = args.n_feats
        kernel_size = 3
        res_scale = args.res_scale
        n_awru = args.n_awru
        act = nn.ReLU(True)
        # wn = lambda x: x
        wn = lambda x: torch.nn.utils.weight_norm(x)

        self.rgb_mean = torch.autograd.Variable(torch.FloatTensor(
            [0.4488, 0.4371, 0.4040])).view([1, 3, 1, 1])

        # define head module
        # head = HEAD(args, n_feats, kernel_size, wn)
        head = []
        head.append(
            wn(nn.Conv2d(args.n_colors, n_feats, 3, padding=3//2)))

        # define body module
        body = []
        for i in range(n_resblocks):
            body.append(
                LFB(n_feats, kernel_size, args.block_feats, n_awru, wn=wn, res_scale=res_scale, act=act))

        # define tail module
        out_feats = scale*scale*args.n_colors
        tail = AWMS(args, scale, n_feats, kernel_size, wn)

        skip = []
        skip.append(
            wn(nn.Conv2d(args.n_colors, out_feats, 3, padding=3//2))
        )
        skip.append(nn.PixelShuffle(scale))

        # make object members
        self.head = nn.Sequential(*head)
        self.body = nn.Sequential(*body)
        self.tail = tail
        self.skip = nn.Sequential(*skip)
//  [python]
def __init__(self, args):
        super(MODEL, self).__init__()
        # hyper-params
        self.args = args
        scale = args.scale[0]
        n_resblocks = args.n_resblocks
        n_feats = args.n_feats
        kernel_size = 3
        res_scale = args.res_scale
        n_awru = args.n_awru
        act = nn.ReLU(True)
        # wn = lambda x: x
        wn = lambda x: torch.nn.utils.parametrizations.weight_norm(x)

        self.rgb_mean = torch.autograd.Variable(torch.FloatTensor(
            [0.4488, 0.4371, 0.4040])).view([1, 3, 1, 1])

        # define head module
        # head = HEAD(args, n_feats, kernel_size, wn)
        head = []
        head.append(
            wn(nn.Conv2d(args.n_colors, n_feats, 3, padding=3//2)))

        # define body module
        body = []
        for i in range(n_resblocks):
            body.append(
                LFB(n_feats, kernel_size, args.block_feats, n_awru, wn=wn, res_scale=res_scale, act=act))

        # define tail module
        out_feats = scale*scale*args.n_colors
        tail = AWMS(args, scale, n_feats, kernel_size, wn)

        skip = []
        skip.append(
            wn(nn.Conv2d(args.n_colors, out_feats, 3, padding=3//2))
        )
        skip.append(nn.PixelShuffle(scale))

        # make object members
        self.head = nn.Sequential(*head)
        self.body = nn.Sequential(*body)
        self.tail = tail
        self.skip = nn.Sequential(*skip)
//  [python]
def __init__(self, input_size, output_size, upsample_factor, hp=None, index=0):
        super(DownsampleNet, self).__init__()
        self.input_size = input_size
        self.output_size = output_size
        self.upsample_factor = upsample_factor
        self.skip_conv = nn.Conv1d(input_size, output_size, kernel_size=1)
        self.index = index
        layer = nn.Conv1d(
            input_size,
            output_size,
            kernel_size=upsample_factor * 2,
            stride=upsample_factor,
            padding=upsample_factor // 2 + upsample_factor % 2,
        )

        self.layer = nn.utils.weight_norm(layer)
//  [python]
def __init__(self, input_size, output_size, upsample_factor, hp=None, index=0):
        super(DownsampleNet, self).__init__()
        self.input_size = input_size
        self.output_size = output_size
        self.upsample_factor = upsample_factor
        self.skip_conv = nn.Conv1d(input_size, output_size, kernel_size=1)
        self.index = index
        layer = nn.Conv1d(
            input_size,
            output_size,
            kernel_size=upsample_factor * 2,
            stride=upsample_factor,
            padding=upsample_factor // 2 + upsample_factor % 2,
        )

        self.layer = nn.utils.parametrizations.weight_norm(layer)
//  [python]
def __init__(self, in_nc=3, nf=32, num_modules=6, out_nc=3, upscale=4):
        super(RLCSR, self).__init__()
        self.conv1_1 = nn.Conv2d(in_channels=in_nc, out_channels=nf, kernel_size=(1, 3), padding=(0, 1), groups=1,
                                 bias=False)
        wn = lambda x: torch.nn.utils.weight_norm(x)
        self.BAM= BAM(64,16)
        self.conv1_2 = nn.Conv2d(in_channels=in_nc, out_channels=nf, kernel_size=(3, 1), padding=(1, 0), groups=1,
                                 bias=False)
        self.conv1_3 = nn.Conv2d(in_channels=in_nc, out_channels=nf, kernel_size=(3, 3), padding=(1, 1), groups=1,
                                 bias=False)
        self.convl11 = nn.Conv2d(in_channels=nf, out_channels=nf, kernel_size=(1, 3), padding=(0, 1), groups=1,
                                 bias=False)
        self.convl22 = nn.Conv2d(in_channels=nf, out_channels=nf, kernel_size=(3, 1), padding=(1, 0), groups=1,
                                 bias=False)
        self.convl33 = nn.Conv2d(in_channels=nf, out_channels=nf, kernel_size=(3, 3), padding=(1, 1), groups=1,
                                 bias=False)
        # self.fea_conv = conv_layer1(in_nc, nf)
        self.convl1 = nn.Conv2d(in_channels=nf * 2, out_channels=nf, kernel_size=(1, 3), padding=(0, 1), groups=1,
                                bias=False)
        self.convl2 = nn.Conv2d(in_channels=nf * 2, out_channels=nf, kernel_size=(3, 1), padding=(1, 0), groups=1,
                                bias=False)
        self.convl3 = nn.Conv2d(in_channels=nf * 2, out_channels=nf, kernel_size=(3, 3), padding=(1, 1), groups=1,
                                bias=False)

        self.B1 = RFDB(in_channels=nf)
        self.B2 = RFDB(in_channels=nf)
        self.B3 = RFDB(in_channels=nf)
        self.B4 = RFDB(in_channels=nf)
        self.B5 = RFDB(in_channels=nf)
        self.B6 = RFDB(in_channels=nf)

        self.reduction1 = wn(nn.Conv2d(nf * 2, nf, 1))
        self.reduction2 = wn(nn.Conv2d(nf * 2, nf, 1))
        self.reduction3 = wn(nn.Conv2d(nf * 2, nf, 1))
        self.reduction4 = wn(nn.Conv2d(nf * 2, nf, 1))
        self.reduction5 = wn(nn.Conv2d(nf * 2, nf, 1))

        self.c = conv_block(nf * num_modules, nf, kernel_size=1, act_type='silu')
        self.la = LAM_Module(nf * num_modules)
        self.LR_conv = conv_layer(nf, nf, kernel_size=3)
        self.last_conv = nn.Conv2d(nf, nf, 3, 1, 1)
        upsample_block = pixelshuffle_block
        self.upsampler = upsample_block(nf, out_nc, upscale_factor=4)
        self.scale_idx = 0
        self.last = nn.Conv2d(nf * 2, nf, 3, 1, 1)
        self.res_scale = Scale(1)
        self.in_scale = Scale(1)
        self.outconv = nn.Conv2d(6,out_nc,3,1,1)
//  [python]
def __init__(self, in_nc=3, nf=32, num_modules=6, out_nc=3, upscale=4):
        super(RLCSR, self).__init__()
        self.conv1_1 = nn.Conv2d(in_channels=in_nc, out_channels=nf, kernel_size=(1, 3), padding=(0, 1), groups=1,
                                 bias=False)
        wn = lambda x: torch.nn.utils.parametrizations.weight_norm(x)
        self.BAM= BAM(64,16)
        self.conv1_2 = nn.Conv2d(in_channels=in_nc, out_channels=nf, kernel_size=(3, 1), padding=(1, 0), groups=1,
                                 bias=False)
        self.conv1_3 = nn.Conv2d(in_channels=in_nc, out_channels=nf, kernel_size=(3, 3), padding=(1, 1), groups=1,
                                 bias=False)
        self.convl11 = nn.Conv2d(in_channels=nf, out_channels=nf, kernel_size=(1, 3), padding=(0, 1), groups=1,
                                 bias=False)
        self.convl22 = nn.Conv2d(in_channels=nf, out_channels=nf, kernel_size=(3, 1), padding=(1, 0), groups=1,
                                 bias=False)
        self.convl33 = nn.Conv2d(in_channels=nf, out_channels=nf, kernel_size=(3, 3), padding=(1, 1), groups=1,
                                 bias=False)
        # self.fea_conv = conv_layer1(in_nc, nf)
        self.convl1 = nn.Conv2d(in_channels=nf * 2, out_channels=nf, kernel_size=(1, 3), padding=(0, 1), groups=1,
                                bias=False)
        self.convl2 = nn.Conv2d(in_channels=nf * 2, out_channels=nf, kernel_size=(3, 1), padding=(1, 0), groups=1,
                                bias=False)
        self.convl3 = nn.Conv2d(in_channels=nf * 2, out_channels=nf, kernel_size=(3, 3), padding=(1, 1), groups=1,
                                bias=False)

        self.B1 = RFDB(in_channels=nf)
        self.B2 = RFDB(in_channels=nf)
        self.B3 = RFDB(in_channels=nf)
        self.B4 = RFDB(in_channels=nf)
        self.B5 = RFDB(in_channels=nf)
        self.B6 = RFDB(in_channels=nf)

        self.reduction1 = wn(nn.Conv2d(nf * 2, nf, 1))
        self.reduction2 = wn(nn.Conv2d(nf * 2, nf, 1))
        self.reduction3 = wn(nn.Conv2d(nf * 2, nf, 1))
        self.reduction4 = wn(nn.Conv2d(nf * 2, nf, 1))
        self.reduction5 = wn(nn.Conv2d(nf * 2, nf, 1))

        self.c = conv_block(nf * num_modules, nf, kernel_size=1, act_type='silu')
        self.la = LAM_Module(nf * num_modules)
        self.LR_conv = conv_layer(nf, nf, kernel_size=3)
        self.last_conv = nn.Conv2d(nf, nf, 3, 1, 1)
        upsample_block = pixelshuffle_block
        self.upsampler = upsample_block(nf, out_nc, upscale_factor=4)
        self.scale_idx = 0
        self.last = nn.Conv2d(nf * 2, nf, 3, 1, 1)
        self.res_scale = Scale(1)
        self.in_scale = Scale(1)
        self.outconv = nn.Conv2d(6,out_nc,3,1,1)
//  [python]
def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        dilation=1,
        causal=False,
        mode="SAME",
    ):
        super(Conv, self).__init__()

        self.causal = causal
        self.mode = mode
        if self.causal and self.mode == "SAME":
            self.padding = dilation * (kernel_size - 1)
        elif self.mode == "SAME":
            self.padding = dilation * (kernel_size - 1) // 2
        else:
            self.padding = 0
        self.conv = nn.Conv1d(
            in_channels,
            out_channels,
            kernel_size,
            dilation=dilation,
            padding=self.padding,
        )
        self.conv = nn.utils.weight_norm(self.conv)
        nn.init.kaiming_normal_(self.conv.weight)