jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(self, in_channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0, sigmoid_scale=False):
    super().__init__()
    self.in_channels = in_channels
    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.sigmoid_scale = sigmoid_scale

    start = torch.nn.Conv1d(in_channels//2, hidden_channels, 1)
    start = torch.nn.utils.parametrizations.weight_norm(start)
    self.start = start
    # Initializing last layer to 0 makes the affine coupling layers
    # do nothing at first.  It helps to stabilze training.
    end = torch.nn.Conv1d(hidden_channels, in_channels, 1)
    end.weight.data.zero_()
    end.bias.data.zero_()
    self.end = end

    self.wn = modules.WN(in_channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels, p_dropout)
//  [python]
def __init__(self, channel, rate):
        super(PatchAttn, self).__init__()
        wn = lambda x: torch.nn.utils.weight_norm(x)
        self.softmax = nn.Softmax(dim=-1)
        self.channel = channel
        self.depth_conv = wn(
            nn.Conv2d(self.channel * 4, self.channel * 4, 3, padding=rate, dilation=rate, groups=self.channel * 4))
        self.point_conv = wn(nn.Conv2d(self.channel * 4, self.channel, 1, groups=1))
//  [python]
def __init__(self, channel, rate):
        super(PatchAttn, self).__init__()
        wn = lambda x: torch.nn.utils.parametrizations.weight_norm(x)
        self.softmax = nn.Softmax(dim=-1)
        self.channel = channel
        self.depth_conv = wn(
            nn.Conv2d(self.channel * 4, self.channel * 4, 3, padding=rate, dilation=rate, groups=self.channel * 4))
        self.point_conv = wn(nn.Conv2d(self.channel * 4, self.channel, 1, groups=1))
//  [python]
def split_cond_layers(model):
    for WN in model.WN:
        if hasattr(WN, "cond_layer"):
            n_layers = len(WN.res_skip_layers)
            conv_weights = WN.cond_layer.weight
            conv_bias = WN.cond_layer.bias
            conv_stride = WN.cond_layer.stride
            conv_dilation = WN.cond_layer.dilation
            conv_padding = WN.cond_layer.padding
            num_in_channels = conv_weights.size(1)
            num_out_channels = conv_weights.size(0)//n_layers
            kernel_size = conv_weights.size(2)
            WN.cond_layers = []
            for i in range(n_layers):
                layer = torch.nn.Conv1d(
                    in_channels=num_in_channels,
                    out_channels=num_out_channels,
                    kernel_size=kernel_size,
                    stride=conv_stride,
                    padding=conv_padding,
                    dilation=conv_dilation)
                layer.weight.data[:, :, :] = conv_weights.data[
                        i*num_out_channels:(i+1)*num_out_channels, :, :]
                layer.bias.data[:] = conv_bias.data[
                        i*num_out_channels:(i+1)*num_out_channels]
                layer = torch.nn.utils.weight_norm(layer, name='weight')
                WN.cond_layers.append(layer)
    return model
//  [python]
def split_cond_layers(model):
    for WN in model.WN:
        if hasattr(WN, "cond_layer"):
            n_layers = len(WN.res_skip_layers)
            conv_weights = WN.cond_layer.weight
            conv_bias = WN.cond_layer.bias
            conv_stride = WN.cond_layer.stride
            conv_dilation = WN.cond_layer.dilation
            conv_padding = WN.cond_layer.padding
            num_in_channels = conv_weights.size(1)
            num_out_channels = conv_weights.size(0)//n_layers
            kernel_size = conv_weights.size(2)
            WN.cond_layers = []
            for i in range(n_layers):
                layer = torch.nn.Conv1d(
                    in_channels=num_in_channels,
                    out_channels=num_out_channels,
                    kernel_size=kernel_size,
                    stride=conv_stride,
                    padding=conv_padding,
                    dilation=conv_dilation)
                layer.weight.data[:, :, :] = conv_weights.data[
                        i*num_out_channels:(i+1)*num_out_channels, :, :]
                layer.bias.data[:] = conv_bias.data[
                        i*num_out_channels:(i+1)*num_out_channels]
                layer = torch.nn.utils.parametrizations.weight_norm(layer, name='weight')
                WN.cond_layers.append(layer)
    return model
//  [python]
def __init__(self, in_channels, out_channels, groups=3):
        super(ConvBlock, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.groups = groups
        wn = lambda x: torch.nn.utils.weight_norm(x)
        self.group_conv = wn(nn.Conv2d(self.in_channels, self.in_channels, 1, groups=self.groups))
        self.depth_conv = wn(nn.Conv2d(self.in_channels, self.in_channels, 3, padding=1, groups=in_channels))
        self.point_conv = wn(nn.Conv2d(self.in_channels, self.out_channels, 1, groups=1))
//  [python]
def __init__(self, in_channels, out_channels, groups=3):
        super(ConvBlock, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.groups = groups
        wn = lambda x: torch.nn.utils.parametrizations.weight_norm(x)
        self.group_conv = wn(nn.Conv2d(self.in_channels, self.in_channels, 1, groups=self.groups))
        self.depth_conv = wn(nn.Conv2d(self.in_channels, self.in_channels, 3, padding=1, groups=in_channels))
        self.point_conv = wn(nn.Conv2d(self.in_channels, self.out_channels, 1, groups=1))
//  [python]
def __init__(self, in_channels, out_channels, groups=3, ker_size=2):
        super(ConvBlockD, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.groups = groups
        wn = lambda x: torch.nn.utils.weight_norm(x)
        self.group_conv = wn(nn.Conv2d(self.in_channels, self.in_channels, 1, groups=self.groups))
        self.depth_conv = wn(
            nn.Conv2d(self.in_channels, self.in_channels, 3, padding=ker_size, dilation=ker_size, groups=in_channels))
        self.point_conv = wn(nn.Conv2d(self.in_channels, self.out_channels, 1, groups=1))
//  [python]
def __init__(self, in_channels, out_channels, groups=3, ker_size=2):
        super(ConvBlockD, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.groups = groups
        wn = lambda x: torch.nn.utils.parametrizations.weight_norm(x)
        self.group_conv = wn(nn.Conv2d(self.in_channels, self.in_channels, 1, groups=self.groups))
        self.depth_conv = wn(
            nn.Conv2d(self.in_channels, self.in_channels, 3, padding=ker_size, dilation=ker_size, groups=in_channels))
        self.point_conv = wn(nn.Conv2d(self.in_channels, self.out_channels, 1, groups=1))
//  [python]
def __init__(self, n_feats):
        super(MIRB1, self).__init__()
        self.c_out =  n_feats // 2
        wn = lambda x: torch.nn.utils.weight_norm(x)
        self.conv3_1 = ConvBlock( n_feats, self.c_out)
        self.convd_1 = ConvBlock( n_feats, self.c_out)

        self.conv3_2 = ConvBlock( n_feats, self.c_out)
        self.convd_2 = ConvBlock( n_feats, self.c_out)

        self.conv3_3 = ConvBlock( n_feats, self.c_out)
        self.convd_3 = ConvBlock( n_feats, self.c_out)
        self.conv_last = wn(nn.Conv2d( n_feats,  n_feats, 1))
        # self.calayer = CALayer( n_feats)
        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
//  [python]
def __init__(self, n_feats):
        super(MIRB1, self).__init__()
        self.c_out =  n_feats // 2
        wn = lambda x: torch.nn.utils.parametrizations.weight_norm(x)
        self.conv3_1 = ConvBlock( n_feats, self.c_out)
        self.convd_1 = ConvBlock( n_feats, self.c_out)

        self.conv3_2 = ConvBlock( n_feats, self.c_out)
        self.convd_2 = ConvBlock( n_feats, self.c_out)

        self.conv3_3 = ConvBlock( n_feats, self.c_out)
        self.convd_3 = ConvBlock( n_feats, self.c_out)
        self.conv_last = wn(nn.Conv2d( n_feats,  n_feats, 1))
        # self.calayer = CALayer( n_feats)
        self.lrelu = nn.LeakyReLU(negative_slope=0.2, 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

        # 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 = 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)

            # 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()

        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

        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)

            # 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 __init__(self, n_feats):
        super(MIRB2, self).__init__()
        self.c_out =  n_feats // 2
        wn = lambda x: torch.nn.utils.weight_norm(x)

        self.conv3_1 = ConvBlock( n_feats, self.c_out)
        self.convd_1 = ConvBlockD( n_feats, self.c_out, ker_size=2)

        self.conv3_2 = ConvBlock( n_feats, self.c_out)
        self.convd_2 = ConvBlockD( n_feats, self.c_out, ker_size=2)

        self.conv3_3 = ConvBlock( n_feats, self.c_out)
        self.convd_3 = ConvBlockD( n_feats, self.c_out, ker_size=2)
        self.conv_last = wn(nn.Conv2d( n_feats,  n_feats, 1))
        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
//  [python]
def __init__(self, n_feats):
        super(MIRB2, self).__init__()
        self.c_out =  n_feats // 2
        wn = lambda x: torch.nn.utils.parametrizations.weight_norm(x)

        self.conv3_1 = ConvBlock( n_feats, self.c_out)
        self.convd_1 = ConvBlockD( n_feats, self.c_out, ker_size=2)

        self.conv3_2 = ConvBlock( n_feats, self.c_out)
        self.convd_2 = ConvBlockD( n_feats, self.c_out, ker_size=2)

        self.conv3_3 = ConvBlock( n_feats, self.c_out)
        self.convd_3 = ConvBlockD( n_feats, self.c_out, ker_size=2)
        self.conv_last = wn(nn.Conv2d( n_feats,  n_feats, 1))
        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
//  [python]
def __init__(self, args):
        super().__init__()

        self.dropout = args.dropout
        self.embedding_dim = args.encoder_embed_dim

        PosConv = Fp32Conv1d if args.fp32_pos_conv else nn.Conv1d

        self.pos_conv = PosConv(
            self.embedding_dim,
            self.embedding_dim,
            kernel_size=args.conv_pos,
            padding=args.conv_pos // 2,
            groups=args.conv_pos_groups,
        )
        dropout = 0
        std = math.sqrt((4 * (1.0 - dropout)) / (args.conv_pos * self.embedding_dim))
        nn.init.normal_(self.pos_conv.weight, mean=0, std=std)
        nn.init.constant_(self.pos_conv.bias, 0)

        self.pos_conv = nn.utils.weight_norm(self.pos_conv, name="weight", dim=2)
        self.pos_conv = nn.Sequential(self.pos_conv, SamePad(args.conv_pos), nn.GELU())

        def create_decoder_layers(n_layers):
            return nn.ModuleList([
                TransformerSentenceEncoderLayer(
                    embedding_dim=self.embedding_dim,
                    ffn_embedding_dim=args.encoder_ffn_embed_dim,
                    num_attention_heads=args.encoder_attention_heads,
                    dropout=self.dropout,
                    attention_dropout=args.attention_dropout,
                    activation_dropout=args.activation_dropout,
                    activation_fn=args.activation_fn,
                    layer_norm_first=args.layer_norm_first,
                    rotary_embeddings=args.rotary_embeddings,
                    mha=args.mha,
                    fp32_transformer_layernorm=args.fp32_transformer_layernorm,
                    fp32_mha_softmax=args.fp32_mha_softmax,
                )
                for _ in range(n_layers)
            ])

        if args.hourglass_transformer:
            n_pre, (n_hourglass, self.shorten_factor), n_post = eval(
                args.hourglass_transformer)

            self.layers = create_decoder_layers(n_pre)
            self.hourglass_layers = create_decoder_layers(n_hourglass)
            self.post_layers = create_decoder_layers(n_post)

            assert args.hourglass_resample in ['linear', 'naive']
            # otherwise i want to resample before merging resutls
            assert not args.layer_norm_first

            kw = {'emb_dim': self.embedding_dim, 'factor': self.shorten_factor,
                  'mode': args.hourglass_resample}
            self.upsample_layer = Upsampler(**kw)
            self.downsample_layer = Downsampler(**kw)
        else:
            self.layers = create_decoder_layers(args.encoder_layers)
            self.hourglass_layers = None
            self.post_layers = None

        self.layer_norm_first = args.layer_norm_first
        self.layer_norm = LayerNorm(self.embedding_dim)
        self.layerdrop = args.encoder_layerdrop

        self.apply(init_bert_params)
//  [python]
def __init__(self, args):
        super().__init__()

        self.dropout = args.dropout
        self.embedding_dim = args.encoder_embed_dim

        PosConv = Fp32Conv1d if args.fp32_pos_conv else nn.Conv1d

        self.pos_conv = PosConv(
            self.embedding_dim,
            self.embedding_dim,
            kernel_size=args.conv_pos,
            padding=args.conv_pos // 2,
            groups=args.conv_pos_groups,
        )
        dropout = 0
        std = math.sqrt((4 * (1.0 - dropout)) / (args.conv_pos * self.embedding_dim))
        nn.init.normal_(self.pos_conv.weight, mean=0, std=std)
        nn.init.constant_(self.pos_conv.bias, 0)

        self.pos_conv = nn.utils.parametrizations.weight_norm(self.pos_conv, name="weight", dim=2)
        self.pos_conv = nn.Sequential(self.pos_conv, SamePad(args.conv_pos), nn.GELU())

        def create_decoder_layers(n_layers):
            return nn.ModuleList([
                TransformerSentenceEncoderLayer(
                    embedding_dim=self.embedding_dim,
                    ffn_embedding_dim=args.encoder_ffn_embed_dim,
                    num_attention_heads=args.encoder_attention_heads,
                    dropout=self.dropout,
                    attention_dropout=args.attention_dropout,
                    activation_dropout=args.activation_dropout,
                    activation_fn=args.activation_fn,
                    layer_norm_first=args.layer_norm_first,
                    rotary_embeddings=args.rotary_embeddings,
                    mha=args.mha,
                    fp32_transformer_layernorm=args.fp32_transformer_layernorm,
                    fp32_mha_softmax=args.fp32_mha_softmax,
                )
                for _ in range(n_layers)
            ])

        if args.hourglass_transformer:
            n_pre, (n_hourglass, self.shorten_factor), n_post = eval(
                args.hourglass_transformer)

            self.layers = create_decoder_layers(n_pre)
            self.hourglass_layers = create_decoder_layers(n_hourglass)
            self.post_layers = create_decoder_layers(n_post)

            assert args.hourglass_resample in ['linear', 'naive']
            # otherwise i want to resample before merging resutls
            assert not args.layer_norm_first

            kw = {'emb_dim': self.embedding_dim, 'factor': self.shorten_factor,
                  'mode': args.hourglass_resample}
            self.upsample_layer = Upsampler(**kw)
            self.downsample_layer = Downsampler(**kw)
        else:
            self.layers = create_decoder_layers(args.encoder_layers)
            self.hourglass_layers = None
            self.post_layers = None

        self.layer_norm_first = args.layer_norm_first
        self.layer_norm = LayerNorm(self.embedding_dim)
        self.layerdrop = args.encoder_layerdrop

        self.apply(init_bert_params)
//  [python]
def __init__(self, n_feats):
        super(MMFB, self).__init__()
        wn = lambda x: torch.nn.utils.weight_norm(x)
        # self.n_feats =  n_feats
        self.bs1 = MIRB1(n_feats)
        self.bs11 = MIRB1(n_feats)
        self.bs2 = MIRB2(n_feats)
        self.bs22 = MIRB2(n_feats)
        self.bs3 = MIRB3(n_feats)
        self.bs33 = MIRB3(n_feats)
//  [python]
def __init__(self, n_feats):
        super(MMFB, self).__init__()
        wn = lambda x: torch.nn.utils.parametrizations.weight_norm(x)
        # self.n_feats =  n_feats
        self.bs1 = MIRB1(n_feats)
        self.bs11 = MIRB1(n_feats)
        self.bs2 = MIRB2(n_feats)
        self.bs22 = MIRB2(n_feats)
        self.bs3 = MIRB3(n_feats)
        self.bs33 = MIRB3(n_feats)
//  [python]
def __init__(self, n_feats):
        super(MIRB3, self).__init__()
        self.c_out =  n_feats // 2
        wn = lambda x: torch.nn.utils.weight_norm(x)

        self.conv3_1 = ConvBlock( n_feats, self.c_out)
        self.convd_1 = ConvBlockD( n_feats, self.c_out, ker_size=3)

        self.conv3_2 = ConvBlock( n_feats, self.c_out)
        self.convd_2 = ConvBlockD( n_feats, self.c_out, ker_size=3)

        self.conv3_3 = ConvBlock( n_feats, self.c_out)
        self.convd_3 = ConvBlockD( n_feats, self.c_out, ker_size=3)
        self.conv_last = wn(nn.Conv2d( n_feats,  n_feats, 1))
        # self.calayer = CALayer( n_feats)
        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
//  [python]
def __init__(self, n_feats):
        super(MIRB3, self).__init__()
        self.c_out =  n_feats // 2
        wn = lambda x: torch.nn.utils.parametrizations.weight_norm(x)

        self.conv3_1 = ConvBlock( n_feats, self.c_out)
        self.convd_1 = ConvBlockD( n_feats, self.c_out, ker_size=3)

        self.conv3_2 = ConvBlock( n_feats, self.c_out)
        self.convd_2 = ConvBlockD( n_feats, self.c_out, ker_size=3)

        self.conv3_3 = ConvBlock( n_feats, self.c_out)
        self.convd_3 = ConvBlockD( n_feats, self.c_out, ker_size=3)
        self.conv_last = wn(nn.Conv2d( n_feats,  n_feats, 1))
        # self.calayer = CALayer( n_feats)
        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
//  [python]
def __init__(self, n_feats, scale, n_colors):
        super(MDAB, self).__init__()
        self.n_feats = n_feats
        self.scale = scale
        self.out_feats = self.scale * self.scale * n_colors
        self.im_feats = self.out_feats * 2  # 2:24,3:54,4:96

        wn = lambda x: torch.nn.utils.weight_norm(x)
        self.tail1 = wn(nn.Conv2d(n_feats, n_feats // 2, 1))
        self.tail2 = ConvBlock(n_feats, n_feats // 2)

        self.conv = wn(nn.Conv2d(n_feats, n_feats, 1))
        self.conv3 = ConvBlock(n_feats, n_feats)
        self.soft_c = nn.Softmax(dim=1)
        self.soft_hw = nn.Softmax(dim=2)
        self.conv_end = wn(nn.Conv2d(n_feats, n_feats, 1))
//  [python]
def __init__(self, n_feats, scale, n_colors):
        super(MDAB, self).__init__()
        self.n_feats = n_feats
        self.scale = scale
        self.out_feats = self.scale * self.scale * n_colors
        self.im_feats = self.out_feats * 2  # 2:24,3:54,4:96

        wn = lambda x: torch.nn.utils.parametrizations.weight_norm(x)
        self.tail1 = wn(nn.Conv2d(n_feats, n_feats // 2, 1))
        self.tail2 = ConvBlock(n_feats, n_feats // 2)

        self.conv = wn(nn.Conv2d(n_feats, n_feats, 1))
        self.conv3 = ConvBlock(n_feats, n_feats)
        self.soft_c = nn.Softmax(dim=1)
        self.soft_hw = nn.Softmax(dim=2)
        self.conv_end = wn(nn.Conv2d(n_feats, n_feats, 1))
//  [python]
def apply_weight_norm(self):
        self.res_block = weight_norm(self.res_block)
        for idx, layer in enumerate(self.main_block):
            if len(layer.state_dict()) != 0:
               self.main_block[idx] = weight_norm(layer)
//  [python]
def apply_weight_norm(self):
        self.res_block = torch.nn.utils.parametrizations.weight_norm(self.res_block)
        for idx, layer in enumerate(self.main_block):
            if len(layer.state_dict()) != 0:
               self.main_block[idx] = torch.nn.utils.parametrizations.weight_norm(layer)
//  [python]
def __init__(self, args):
        super(MAFFSRN, self).__init__()

        # hyper-params
        self.scale = args.scale
        n_FFGs = args.n_FFGs
        n_feats = args.n_feats 
        kernel_size = 3
        act = nn.LeakyReLU(True)

        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.append(
            wn(nn.Conv2d(3, n_feats, 3, padding=3//2)))

        # define body module
        body = []
        for i in range(n_FFGs):
            body.append(
                FFG(n_feats, wn=wn, act=act))

        # define tail module
        tail = Tail(args, self.scale[0], n_feats, kernel_size, wn)

        # make object members
        self.head = nn.Sequential(*head)
        self.body = nn.Sequential(*body)
        self.tail = tail
//  [python]
def __init__(self, args):
        super(MAFFSRN, self).__init__()

        # hyper-params
        self.scale = args.scale
        n_FFGs = args.n_FFGs
        n_feats = args.n_feats 
        kernel_size = 3
        act = nn.LeakyReLU(True)

        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.append(
            wn(nn.Conv2d(3, n_feats, 3, padding=3//2)))

        # define body module
        body = []
        for i in range(n_FFGs):
            body.append(
                FFG(n_feats, wn=wn, act=act))

        # define tail module
        tail = Tail(args, self.scale[0], n_feats, kernel_size, wn)

        # make object members
        self.head = nn.Sequential(*head)
        self.body = nn.Sequential(*body)
        self.tail = tail
//  [python]
def __init__(self, n_feats=48, scale=4, n_colors=3):
        super(MDAN, self).__init__()
        # n_resgroups = args.n_resgroups
        self.n_feats = n_feats
        self.scale = scale
        self.out_feats = self.scale * self.scale * n_colors
        self.im_feats = self.out_feats * 2  # 2:24,3:54,4:96

        wn = lambda x: torch.nn.utils.weight_norm(x)
        self.sub_mean = MeanShift(3)
        self.up_bic = nn.Upsample(scale_factor=self.scale, mode='bicubic', align_corners=False)
        self.upsample = nn.Upsample(scale_factor=2, mode='bicubic', align_corners=True)
        self.conv_first = wn(nn.Conv2d(3, self.n_feats, 3, padding=1))

        self.BS1 = MMFB(n_feats)
        self.BS2 = MMFB(n_feats)
        self.BS3 = MMFB(n_feats)
        self.upb1 = MDAB(n_feats, scale, n_colors)
        self.upb2 = MDAB(n_feats, scale, n_colors)
        self.upb3 = MDAB(n_feats, scale, n_colors)

        self.scale1 = Scale(0.3)
        self.scale2 = Scale(0.3)
        self.scale3 = Scale(0.4)
        self.conv_add = wn(nn.Conv2d(n_feats * 3, n_feats, 1))
        #self.out = wn(nn.Conv2d(n_feats, self.scale * self.scale * 3, 3, padding=1)) #2x
        self.out1 = wn(nn.Conv2d(n_feats, self.scale * self.scale * 3, 3, padding=1)) # 3x pre_train
        self.pixelshuffle = nn.PixelShuffle(self.scale)

        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
        self.add_mean = MeanShift(3, sign=1)
//  [python]
def __init__(self, n_feats=48, scale=4, n_colors=3):
        super(MDAN, self).__init__()
        # n_resgroups = args.n_resgroups
        self.n_feats = n_feats
        self.scale = scale
        self.out_feats = self.scale * self.scale * n_colors
        self.im_feats = self.out_feats * 2  # 2:24,3:54,4:96

        wn = lambda x: torch.nn.utils.parametrizations.weight_norm(x)
        self.sub_mean = MeanShift(3)
        self.up_bic = nn.Upsample(scale_factor=self.scale, mode='bicubic', align_corners=False)
        self.upsample = nn.Upsample(scale_factor=2, mode='bicubic', align_corners=True)
        self.conv_first = wn(nn.Conv2d(3, self.n_feats, 3, padding=1))

        self.BS1 = MMFB(n_feats)
        self.BS2 = MMFB(n_feats)
        self.BS3 = MMFB(n_feats)
        self.upb1 = MDAB(n_feats, scale, n_colors)
        self.upb2 = MDAB(n_feats, scale, n_colors)
        self.upb3 = MDAB(n_feats, scale, n_colors)

        self.scale1 = Scale(0.3)
        self.scale2 = Scale(0.3)
        self.scale3 = Scale(0.4)
        self.conv_add = wn(nn.Conv2d(n_feats * 3, n_feats, 1))
        #self.out = wn(nn.Conv2d(n_feats, self.scale * self.scale * 3, 3, padding=1)) #2x
        self.out1 = wn(nn.Conv2d(n_feats, self.scale * self.scale * 3, 3, padding=1)) # 3x pre_train
        self.pixelshuffle = nn.PixelShuffle(self.scale)

        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
        self.add_mean = MeanShift(3, sign=1)
//  [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(BlendingRenderingNetwork, self).__init__()

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

        self.embedder = None
        if multires_view > 0:
            self.embedder = Embedding(3, multires_view)
            dims[0] += (self.embedder.out_channels - 3)

        self.num_layers = len(dims)

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

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

            setattr(self, "lin" + str(l), lin)

        self.relu = nn.ReLU()

        self.color_volume = None

        self.softmax = nn.Softmax(dim=1)

        self.type = 'blending'
//  [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(BlendingRenderingNetwork, self).__init__()

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

        self.embedder = None
        if multires_view > 0:
            self.embedder = Embedding(3, multires_view)
            dims[0] += (self.embedder.out_channels - 3)

        self.num_layers = len(dims)

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

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

            setattr(self, "lin" + str(l), lin)

        self.relu = nn.ReLU()

        self.color_volume = None

        self.softmax = nn.Softmax(dim=1)

        self.type = 'blending'
//  [python]
def __init__(self,
                 in_channels,
                 hidden_channels,
                 kernel_size,
                 dilation_rate,
                 num_layers,
                 c_in_channels=0,
                 dropout_p=0,
                 weight_norm=True):
        super().__init__()
        assert kernel_size % 2 == 1
        assert hidden_channels % 2 == 0
        self.in_channels = in_channels
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size
        self.dilation_rate = dilation_rate
        self.num_layers = num_layers
        self.c_in_channels = c_in_channels
        self.dropout_p = dropout_p

        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.dropout = nn.Dropout(dropout_p)

        # init conditioning layer
        if c_in_channels > 0:
            cond_layer = torch.nn.Conv1d(c_in_channels,
                                         2 * hidden_channels * num_layers, 1)
            self.cond_layer = torch.nn.utils.weight_norm(cond_layer,
                                                         name='weight')
        # intermediate layers
        for i in range(num_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)

            if i < num_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)
        # setup weight norm
        if not weight_norm:
            self.remove_weight_norm()
//  [python]
def __init__(self,
                 in_channels,
                 hidden_channels,
                 kernel_size,
                 dilation_rate,
                 num_layers,
                 c_in_channels=0,
                 dropout_p=0,
                 weight_norm=True):
        super().__init__()
        assert kernel_size % 2 == 1
        assert hidden_channels % 2 == 0
        self.in_channels = in_channels
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size
        self.dilation_rate = dilation_rate
        self.num_layers = num_layers
        self.c_in_channels = c_in_channels
        self.dropout_p = dropout_p

        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.dropout = nn.Dropout(dropout_p)

        # init conditioning layer
        if c_in_channels > 0:
            cond_layer = torch.nn.Conv1d(c_in_channels,
                                         2 * hidden_channels * num_layers, 1)
            self.cond_layer = torch.nn.utils.parametrizations.weight_norm(cond_layer,
                                                         name='weight')
        # intermediate layers
        for i in range(num_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)

            if i < num_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)
        # setup weight norm
        if not weight_norm:
            self.remove_weight_norm()
//  [python]
def __init__(self, in_dim, mid_dim, num_layers, mask, init_zeros=False, dropout=False):
        
        super(CouplingLayerTabular, self).__init__()
        self.mask = mask
        self.st_net = nn.Sequential(nn.Linear(in_dim, mid_dim),
                                    nn.ReLU(),
                                    nn.Dropout(.5) if dropout else nn.Sequential(),
                                    *self._inner_seq(num_layers, mid_dim),
                                    nn.Linear(mid_dim, in_dim*2))

        if init_zeros:
                # init w zeros to init a flow w identity mapping
                torch.nn.init.zeros_(self.st_net[-1].weight)
                torch.nn.init.zeros_(self.st_net[-1].bias)

        self.rescale = nn.utils.weight_norm(RescaleTabular(in_dim))
//  [python]
def __init__(self, in_dim, mid_dim, num_layers, mask, init_zeros=False, dropout=False):
        
        super(CouplingLayerTabular, self).__init__()
        self.mask = mask
        self.st_net = nn.Sequential(nn.Linear(in_dim, mid_dim),
                                    nn.ReLU(),
                                    nn.Dropout(.5) if dropout else nn.Sequential(),
                                    *self._inner_seq(num_layers, mid_dim),
                                    nn.Linear(mid_dim, in_dim*2))

        if init_zeros:
                # init w zeros to init a flow w identity mapping
                torch.nn.init.zeros_(self.st_net[-1].weight)
                torch.nn.init.zeros_(self.st_net[-1].bias)

        self.rescale = torch.nn.utils.parametrizations.weight_norm(RescaleTabular(in_dim))
//  [python]
def __init__(self, in_channel, cin_channel, res_channel, n_height, n_flow, n_layer, layers_per_dilation_h_cycle,
                 size_flow_embed, coupling_type, num_bin=32, tail_bound=5.0, use_weightnorm_embed=False):
        super().__init__()
        self.in_channel = in_channel
        self.cin_channel = cin_channel
        self.res_channel = res_channel
        self.n_height = n_height
        self.n_flow = n_flow
        self.n_layer = n_layer

        self.layers_per_dilation_h_cycle = layers_per_dilation_h_cycle

        self.coupling_type = coupling_type
        if self.coupling_type == 'affine':
            pass
        elif self.coupling_type == 'nsf':
            # nsf params
            self.num_bin = num_bin
            self.tail_bound = tail_bound
        else:
            raise ValueError("unknown coupling type")

        # flow indicator conditioner
        self.size_flow_embed = size_flow_embed
        self.flow_embedding = nn.Embedding(self.n_flow, self.size_flow_embed)

        self.use_weightnorm_embed = use_weightnorm_embed
        if self.use_weightnorm_embed:
            print("INFO: using weightnorm at embedding layer")
            self.flow_embedding = nn.utils.weight_norm(self.flow_embedding)
        nn.init.orthogonal_(self.flow_embedding.weight)

        self.estimator = WaveNet2DHyperDensityEstimator(self.in_channel, self.cin_channel, self.size_flow_embed,
                                                        self.res_channel, self.n_layer, self.n_height,
                                                        self.layers_per_dilation_h_cycle)

        self.flows = nn.ModuleList()
        for i in range(self.n_flow):
            self.flows.append(
                FlowWithHyperSharedEstimator(self.in_channel, self.cin_channel, filter_size=self.res_channel,
                                             num_layer=self.n_layer, num_height=self.n_height,
                                             layers_per_dilation_h_cycle=self.layers_per_dilation_h_cycle,
                                             coupling_type=self.coupling_type,
                                             num_bin=num_bin, tail_bound=tail_bound))

        self.upsample_conv = nn.ModuleList()
        for s in [16, 16]:
            convt = nn.ConvTranspose2d(1, 1, (3, 2 * s), padding=(1, s // 2), stride=(1, s))
            convt = nn.utils.weight_norm(convt)
            nn.init.kaiming_normal_(convt.weight)
            self.upsample_conv.append(convt)
            self.upsample_conv.append(nn.LeakyReLU(0.4))

        self.upsample_conv_kernel_size = (2*s)**2
        self.upsample_conv_stride = s**2
//  [python]
def __init__(self, in_channel, cin_channel, res_channel, n_height, n_flow, n_layer, layers_per_dilation_h_cycle,
                 size_flow_embed, coupling_type, num_bin=32, tail_bound=5.0, use_weightnorm_embed=False):
        super().__init__()
        self.in_channel = in_channel
        self.cin_channel = cin_channel
        self.res_channel = res_channel
        self.n_height = n_height
        self.n_flow = n_flow
        self.n_layer = n_layer

        self.layers_per_dilation_h_cycle = layers_per_dilation_h_cycle

        self.coupling_type = coupling_type
        if self.coupling_type == 'affine':
            pass
        elif self.coupling_type == 'nsf':
            # nsf params
            self.num_bin = num_bin
            self.tail_bound = tail_bound
        else:
            raise ValueError("unknown coupling type")

        # flow indicator conditioner
        self.size_flow_embed = size_flow_embed
        self.flow_embedding = nn.Embedding(self.n_flow, self.size_flow_embed)

        self.use_weightnorm_embed = use_weightnorm_embed
        if self.use_weightnorm_embed:
            print("INFO: using weightnorm at embedding layer")
            self.flow_embedding = nn.utils.parametrizations.weight_norm(self.flow_embedding)
        nn.init.orthogonal_(self.flow_embedding.weight)

        self.estimator = WaveNet2DHyperDensityEstimator(self.in_channel, self.cin_channel, self.size_flow_embed,
                                                        self.res_channel, self.n_layer, self.n_height,
                                                        self.layers_per_dilation_h_cycle)

        self.flows = nn.ModuleList()
        for i in range(self.n_flow):
            self.flows.append(
                FlowWithHyperSharedEstimator(self.in_channel, self.cin_channel, filter_size=self.res_channel,
                                             num_layer=self.n_layer, num_height=self.n_height,
                                             layers_per_dilation_h_cycle=self.layers_per_dilation_h_cycle,
                                             coupling_type=self.coupling_type,
                                             num_bin=num_bin, tail_bound=tail_bound))

        self.upsample_conv = nn.ModuleList()
        for s in [16, 16]:
            convt = nn.ConvTranspose2d(1, 1, (3, 2 * s), padding=(1, s // 2), stride=(1, s))
            convt = nn.utils.parametrizations.weight_norm(convt)
            nn.init.kaiming_normal_(convt.weight)
            self.upsample_conv.append(convt)
            self.upsample_conv.append(nn.LeakyReLU(0.4))

        self.upsample_conv_kernel_size = (2*s)**2
        self.upsample_conv_stride = s**2
//  [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.time_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, 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)

            time_layer = nn.Linear(1, 2*n_channels)
            time_layer = torch.nn.utils.weight_norm(time_layer, name='weight')
            self.time_layers.append(time_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, 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.time_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

        # 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, 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 = 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.parametrizations.weight_norm(in_layer, name='weight')
            self.in_layers.append(in_layer)

            time_layer = nn.Linear(1, 2*n_channels)
            time_layer = torch.nn.utils.parametrizations.weight_norm(time_layer, name='weight')
            self.time_layers.append(time_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 __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)
            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)