training/networks.py

# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.

import numpy as np
import torch
from torch_utils import misc
from torch_utils import persistence
from torch_utils.ops import conv2d_resample
from torch_utils.ops import upfirdn2d
from torch_utils.ops import bias_act
from torch_utils.ops import fma
from training.flow import DDSF
import torch.nn as nn
import torch.nn.functional as F
from collections import Counter
#----------------------------------------------------------------------------

@misc.profiled_function
def normalize_2nd_moment(x, dim=1, eps=1e-8):
    return x * (x.square().mean(dim=dim, keepdim=True) + eps).rsqrt()

#----------------------------------------------------------------------------

@misc.profiled_function
def modulated_conv2d(
    x,                          # Input tensor of shape [batch_size, in_channels, in_height, in_width].
    weight,                     # Weight tensor of shape [out_channels, in_channels, kernel_height, kernel_width].
    styles,                     # Modulation coefficients of shape [batch_size, in_channels].
    noise           = None,     # Optional noise tensor to add to the output activations.
    up              = 1,        # Integer upsampling factor.
    down            = 1,        # Integer downsampling factor.
    padding         = 0,        # Padding with respect to the upsampled image.
    resample_filter = None,     # Low-pass filter to apply when resampling activations. Must be prepared beforehand by calling upfirdn2d.setup_filter().
    demodulate      = True,     # Apply weight demodulation?
    flip_weight     = True,     # False = convolution, True = correlation (matches torch.nn.functional.conv2d).
    fused_modconv   = True,     # Perform modulation, convolution, and demodulation as a single fused operation?
):
    batch_size = x.shape[0]
    out_channels, in_channels, kh, kw = weight.shape
    misc.assert_shape(weight, [out_channels, in_channels, kh, kw]) # [OIkk]
    misc.assert_shape(x, [batch_size, in_channels, None, None]) # [NIHW]
    misc.assert_shape(styles, [batch_size, in_channels]) # [NI]

    # Pre-normalize inputs to avoid FP16 overflow.
    if x.dtype == torch.float16 and demodulate:
        weight = weight * (1 / np.sqrt(in_channels * kh * kw) / weight.norm(float('inf'), dim=[1,2,3], keepdim=True)) # max_Ikk
        styles = styles / styles.norm(float('inf'), dim=1, keepdim=True) # max_I

    # Calculate per-sample weights and demodulation coefficients.
    w = None
    dcoefs = None
    if demodulate or fused_modconv:
        w = weight.unsqueeze(0) # [NOIkk]
        w = w * styles.reshape(batch_size, 1, -1, 1, 1) # [NOIkk]
    if demodulate:
        dcoefs = (w.square().sum(dim=[2,3,4]) + 1e-8).rsqrt() # [NO]
    if demodulate and fused_modconv:
        w = w * dcoefs.reshape(batch_size, -1, 1, 1, 1) # [NOIkk]

    # Execute by scaling the activations before and after the convolution.
    if not fused_modconv:
        x = x * styles.to(x.dtype).reshape(batch_size, -1, 1, 1)
        x = conv2d_resample.conv2d_resample(x=x, w=weight.to(x.dtype), f=resample_filter, up=up, down=down, padding=padding, flip_weight=flip_weight)
        if demodulate and noise is not None:
            x = fma.fma(x, dcoefs.to(x.dtype).reshape(batch_size, -1, 1, 1), noise.to(x.dtype))
        elif demodulate:
            x = x * dcoefs.to(x.dtype).reshape(batch_size, -1, 1, 1)
        elif noise is not None:
            x = x.add_(noise.to(x.dtype))
        return x

    # Execute as one fused op using grouped convolution.
    with misc.suppress_tracer_warnings(): # this value will be treated as a constant
        batch_size = int(batch_size)
    misc.assert_shape(x, [batch_size, in_channels, None, None])
    x = x.reshape(1, -1, *x.shape[2:])
    w = w.reshape(-1, in_channels, kh, kw)
    x = conv2d_resample.conv2d_resample(x=x, w=w.to(x.dtype), f=resample_filter, up=up, down=down, padding=padding, groups=batch_size, flip_weight=flip_weight)
    x = x.reshape(batch_size, -1, *x.shape[2:])
    if noise is not None:
        x = x.add_(noise)
    return x

#----------------------------------------------------------------------------

@persistence.persistent_class
class FullyConnectedLayer(torch.nn.Module):
    def __init__(self,
        in_features,                # Number of input features.
        out_features,               # Number of output features.
        bias            = True,     # Apply additive bias before the activation function?
        activation      = 'linear', # Activation function: 'relu', 'lrelu', etc.
        lr_multiplier   = 1,        # Learning rate multiplier.
        bias_init       = 0,        # Initial value for the additive bias.
        init = 'randn',
    ):
        super().__init__()
        self.activation = activation
        self.in_features = in_features
        self.out_features = out_features
        self.lr_multiplier = lr_multiplier
        if init == 'randn':
            self.weight = torch.nn.Parameter(torch.randn([out_features, in_features]) / lr_multiplier)
        else:
            self.weight = torch.nn.Parameter(torch.full([out_features, in_features],0.) / lr_multiplier)
        self.bias = torch.nn.Parameter(torch.full([out_features], np.float32(bias_init))) if bias else None
        self.weight_gain = lr_multiplier / np.sqrt(in_features)
        self.bias_gain = lr_multiplier

    def forward(self, x):
        w = self.weight.to(x.dtype) * self.weight_gain
        b = self.bias
        if b is not None:
            b = b.to(x.dtype)
            if self.bias_gain != 1:
                b = b * self.bias_gain

        if self.activation == 'linear' and b is not None:
            x = torch.addmm(b.unsqueeze(0), x, w.t())
        else:
            x = x.matmul(w.t())
            x = bias_act.bias_act(x, b, act=self.activation)
        return x

    def __repr__(self):
        return self.__class__.__name__ + '(' + 'in:%s, out:%s, lr:%s, act:%s' % \
            (self.in_features, self.out_features, self.lr_multiplier, self.activation) + ')'


#----------------------------------------------------------------------------

@persistence.persistent_class
class Conv2dLayer(torch.nn.Module):
    def __init__(self,
        in_channels,                    # Number of input channels.
        out_channels,                   # Number of output channels.
        kernel_size,                    # Width and height of the convolution kernel.
        bias            = True,         # Apply additive bias before the activation function?
        activation      = 'linear',     # Activation function: 'relu', 'lrelu', etc.
        up              = 1,            # Integer upsampling factor.
        down            = 1,            # Integer downsampling factor.
        resample_filter = [1,3,3,1],    # Low-pass filter to apply when resampling activations.
        conv_clamp      = None,         # Clamp the output to +-X, None = disable clamping.
        channels_last   = False,        # Expect the input to have memory_format=channels_last?
        trainable       = True,         # Update the weights of this layer during training?
    ):
        super().__init__()
        self.activation = activation
        self.up = up
        self.down = down
        self.conv_clamp = conv_clamp
        self.register_buffer('resample_filter', upfirdn2d.setup_filter(resample_filter))
        self.padding = kernel_size // 2
        self.weight_gain = 1 / np.sqrt(in_channels * (kernel_size ** 2))
        self.act_gain = bias_act.activation_funcs[activation].def_gain

        memory_format = torch.channels_last if channels_last else torch.contiguous_format
        weight = torch.randn([out_channels, in_channels, kernel_size, kernel_size]).to(memory_format=memory_format)
        bias = torch.zeros([out_channels]) if bias else None
        if trainable:
            self.weight = torch.nn.Parameter(weight)
            self.bias = torch.nn.Parameter(bias) if bias is not None else None
        else:
            self.register_buffer('weight', weight)
            if bias is not None:
                self.register_buffer('bias', bias)
            else:
                self.bias = None

    def forward(self, x, gain=1):
        w = self.weight * self.weight_gain
        b = self.bias.to(x.dtype) if self.bias is not None else None
        flip_weight = (self.up == 1) # slightly faster
        x = conv2d_resample.conv2d_resample(x=x, w=w.to(x.dtype), f=self.resample_filter, up=self.up, down=self.down, padding=self.padding, flip_weight=flip_weight)

        act_gain = self.act_gain * gain
        act_clamp = self.conv_clamp * gain if self.conv_clamp is not None else None
        x = bias_act.bias_act(x, b, act=self.activation, gain=act_gain, clamp=act_clamp)
        return x

#----------------------------------------------------------------------------
def gumbel_sigmoid(logits, tau: float = 1, hard: bool = False, threshold: float = 0.5, eval=False):
    gumbels = (
        -torch.empty_like(logits, memory_format=torch.legacy_contiguous_format).exponential_().log()
    )  # ~Gumbel(0, 1)

    gumbels = (logits + gumbels) / tau  # ~Gumbel(logits, tau)
    y_soft = gumbels.sigmoid()

    indices = (y_soft > threshold).nonzero(as_tuple=True)
    y_hard = torch.zeros_like(logits, memory_format=torch.legacy_contiguous_format)
    y_hard[indices[0], indices[1]] = 1.0
    ret = y_hard - y_soft.detach() + y_soft
    return ret, y_soft



def topk_gumbel_sigmoid(logits, tau: float = 1, hard: bool = False, threshold: float = 0.5, eval=False, topk=2):
    soft_mask = torch.sigmoid(logits)
    gumbels = (
        -torch.empty_like(logits, memory_format=torch.legacy_contiguous_format).exponential_().log()
    )  # ~Gumbel(0, 1)
    gumbels = (logits + gumbels) / tau  # ~Gumbel(logits, tau)
    y_soft = gumbels.sigmoid()
    indices = (y_soft > threshold).nonzero(as_tuple=True)
    y_hard = torch.zeros_like(logits, memory_format=torch.legacy_contiguous_format)
    y_hard[indices[0], indices[1]] = 1.0

    topk_values, topk_indices = soft_mask.topk(topk, dim=-1)
    mask = torch.zeros_like(logits)
    mask.scatter_(1, topk_indices, 1)
    ret = y_hard * mask + y_soft - y_soft.detach()
    return ret, y_soft



def sample_gumbel(shape, eps=1e-20):
    U = torch.rand(shape)
    return -torch.autograd.Variable(torch.log(-torch.log(U + eps) + eps))

def gumbel_softmax_sample(logits, temperature, eval):
    gumbels = sample_gumbel(logits.size()).to(logits.device)
    if eval:
        y = logits
    else:
        y = logits + gumbels
    return F.softmax(y/temperature, dim=-1)

def hard_softmax(logits, temperature=1):
    y = F.softmax(logits / temperature, dim=1)
    shape = y.size()
    _, ind = y.max(dim=-1)
    y_hard = torch.zeros_like(y).view(-1, shape[-1])
    y_hard.scatter_(1, ind.view(-1, 1), 1)
    y_hard = y_hard.view(*shape)
    return (y_hard - y).detach() + y


def gumbel_softmax(logits, temperature, eval):
    """
    input: [*, n_class]
    return: [*, n_class] an one-hot vector
    """
    y = gumbel_softmax_sample(logits, temperature, eval)
    y_soft = y
    shape = y.size()
    _, ind = y.max(dim=-1)
    y_hard = torch.zeros_like(y).view(-1, shape[-1])
    y_hard.scatter_(1, ind.view(-1, 1), 1)
    y_hard = y_hard.view(*shape)
    return (y_hard - y).detach() + y, y_soft

def get_onehot(y):
    shape = y.size()
    _, ind = y.max(dim=-1)
    y_hard = torch.zeros_like(y).view(-1, shape[-1])
    y_hard.scatter_(1, ind.view(-1, 1), 1)
    y_hard = y_hard.view(*shape)
    return y_hard


def orthogonal_loss(x):
    hard = (x>0.5)
    hard_expanded = hard.unsqueeze(2).float()
    hard_expanded_t = hard.unsqueeze(1).float()
    filter = (hard_expanded * hard_expanded_t).float()

    x_expanded = x.unsqueeze(2)
    x_expanded_t = x.unsqueeze(1)
    l1_distance = torch.sum(torch.abs(x_expanded - x_expanded_t)*filter, dim=0)
    l1_distance = torch.triu(l1_distance, diagonal=1)

    loss = -torch.mean(l1_distance)
    return loss

def get_topk(logit, topk=2):
    y_sigmoid = torch.sigmoid(logit)
    topk_values, topk_indices = y_sigmoid.topk(topk, dim=-1)
    mask = torch.zeros_like(logit)
    mask.scatter_(1, topk_indices, 1)
    return mask*y_sigmoid


list = [[1,0,0,1,0,0,0,0,0,0,0,0,0],
        [1,0,0,0,1,0,0,0,0,0,0,0,0],
        [0,1,0,1,0,0,0,0,0,0,0,0,0],
        [0,1,0,0,1,0,0,0,0,0,0,0,0],
        [0,1,0,0,0,1,0,0,0,0,0,0,0],
        [0,1,0,0,0,0,1,0,0,0,0,0,0],
        [0,1,0,0,0,0,0,1,0,0,0,0,0],
        [0,1,0,0,0,0,0,0,1,0,0,0,0],
        [0,1,0,0,0,0,0,0,0,1,0,0,0],
        [0,1,0,0,0,0,0,0,0,0,1,0,0],
        [0,1,0,0,0,0,0,0,0,0,0,1,0],
        [0,1,0,0,0,0,0,0,0,0,0,0,1],
        [1,0,0,0,0,1,0,0,0,0,0,0,0],
        [0,0,1,1,0,0,0,0,0,0,0,0,0],
        [0,0,1,0,1,0,0,0,0,0,0,0,0],
        [0,0,1,0,0,1,0,0,0,0,0,0,0],
        [0,0,1,0,0,0,1,0,0,0,0,0,0],
        [0,0,1,0,0,0,0,1,0,0,0,0,0],
        [0,0,1,0,0,0,0,0,1,0,0,0,0],
        [0,0,1,0,0,0,0,0,0,1,0,0,0],
        [0,0,1,0,0,0,0,0,0,0,1,0,0],
        [0,0,1,0,0,0,0,0,0,0,0,1,0],
        [0,0,1,0,0,0,0,0,0,0,0,0,1],
        [1,0,0,0,0,0,1,0,0,0,0,0,0],
        [1,0,0,0,0,0,0,1,0,0,0,0,0],
        [1,0,0,0,0,0,0,0,1,0,0,0,0],
        [1,0,0,0,0,0,0,0,0,1,0,0,0],
        [1,0,0,0,0,0,0,0,0,0,1,0,0],
        [1,0,0,0,0,0,0,0,0,0,0,1,0],
        [1,0,0,0,0,0,0,0,0,0,0,0,1],
        ]
ground_truth = torch.tensor(list).float()
ORTHO = torch.tensor([[-0.1580, -0.0408, -0.1414,  0.1170,  0.1882, -0.1885, -0.0104,  0.1914,
          0.0509, -0.1441, -0.2948,  0.1251, -0.1669],
        [ 0.1422, -0.0457,  0.0277,  0.1680,  0.0185, -0.2122, -0.1040,  0.0041,
          0.2174, -0.0141,  0.3782,  0.0765, -0.0174],
        [ 0.2878,  0.0051, -0.2083, -0.2044, -0.1188,  0.0932,  0.2063,  0.0671,
          0.2095, -0.1755, -0.1170,  0.2188,  0.1370],
        [ 0.2111, -0.2580,  0.2177, -0.0026, -0.1872, -0.2515, -0.2492, -0.0404,
          0.0025, -0.0382,  0.0111, -0.0352,  0.0146],
        [-0.1945,  0.0164,  0.0754,  0.0189, -0.1675,  0.2053,  0.0155,  0.2645,
         -0.2101, -0.0701, -0.1351,  0.1118,  0.1543],
        [ 0.1182, -0.2252,  0.2003, -0.0680,  0.1287, -0.1387, -0.2203, -0.0670,
         -0.1644,  0.1111,  0.0084,  0.0150, -0.1737],
        [ 0.1151,  0.1054,  0.2348,  0.0781, -0.0773,  0.1233, -0.1349, -0.2370,
         -0.1207, -0.0505, -0.1951,  0.0498,  0.0416],
        [-0.3048, -0.0374, -0.2396,  0.0777, -0.3567, -0.2694, -0.2050, -0.0189,
          0.0231,  0.0996, -0.1684,  0.1197, -0.1330],
        [-0.1611,  0.1017,  0.2342, -0.0282,  0.0928, -0.0018,  0.1993, -0.1433,
         -0.0727, -0.1053, -0.2802, -0.2779, -0.0839],
        [ 0.1368, -0.0055, -0.0157,  0.1525,  0.2246,  0.1000, -0.2871, -0.1524,
          0.0486, -0.1028,  0.0836, -0.2623,  0.0859],
        [-0.0752, -0.2077, -0.2013,  0.0231,  0.1458, -0.1430,  0.1958, -0.3751,
         -0.2262,  0.0225, -0.0760,  0.0732,  0.0299],
        [-0.1740, -0.2048,  0.0612, -0.2909,  0.1386,  0.0709,  0.1275,  0.0899,
         -0.1081,  0.2981, -0.0162, -0.1816, -0.1441],
        [-0.0404, -0.1640,  0.1119,  0.2111,  0.0357,  0.0474,  0.1237, -0.0681,
          0.0655, -0.0334,  0.0118,  0.0699, -0.1380],
        [ 0.1107, -0.3759,  0.1586, -0.0606, -0.1565,  0.0428,  0.1039,  0.0608,
          0.1760,  0.0363, -0.0802,  0.0947, -0.0543],
        [-0.0961,  0.1985, -0.0256,  0.0054,  0.0669,  0.0624, -0.1788,  0.1346,
          0.1483, -0.0602,  0.0913, -0.1270,  0.2608],
        [ 0.0254,  0.2653,  0.0828,  0.1078,  0.2526, -0.1783, -0.0357,  0.1296,
          0.2522,  0.2490,  0.0941,  0.0162,  0.0863],
        [-0.0473, -0.0018, -0.0941,  0.0532, -0.1086,  0.0423,  0.0508, -0.0304,
         -0.2711, -0.1112,  0.0677, -0.1355,  0.3170],
        [ 0.0134, -0.2489, -0.1870,  0.0748, -0.0764, -0.2126, -0.0964, -0.0404,
          0.1241, -0.3033, -0.1322, -0.2577, -0.0973],
        [-0.2806, -0.1622,  0.0704, -0.1336,  0.0622,  0.1741,  0.0632,  0.1448,
         -0.0587, -0.1539,  0.1617, -0.0279,  0.0363],
        [ 0.1749,  0.2154,  0.0963, -0.1549, -0.3061, -0.1309, -0.1245, -0.1402,
         -0.1692, -0.1410,  0.0752, -0.2097, -0.2414],
        [ 0.1401,  0.1368, -0.0531, -0.0407,  0.0151,  0.0888,  0.1823,  0.0020,
          0.0277, -0.1578,  0.2602,  0.1338, -0.3879],
        [ 0.2555,  0.1114,  0.4299,  0.0905,  0.0623, -0.0975,  0.1664,  0.3538,
         -0.2299, -0.0775, -0.0997,  0.0390, -0.1250],
        [ 0.0677, -0.0907,  0.1597, -0.1814,  0.1798,  0.0861, -0.0755, -0.1682,
          0.1742, -0.2388, -0.2412, -0.0729,  0.2385],
        [ 0.0193,  0.0062,  0.0031,  0.2542, -0.0530,  0.0144,  0.0780, -0.1098,
         -0.3371, -0.3164,  0.1891,  0.3701,  0.1644],
        [ 0.0799,  0.0810, -0.2274,  0.0251,  0.2413, -0.2354,  0.2490,  0.0038,
         -0.0775,  0.1271,  0.0097, -0.0019,  0.0399],
        [ 0.1140, -0.0978, -0.0916, -0.0106, -0.1800,  0.2802,  0.1105, -0.2525,
         -0.1132,  0.1886,  0.2471, -0.1485,  0.0221],
        [-0.1721, -0.1800,  0.1773, -0.2792,  0.0896, -0.1571, -0.1935,  0.0797,
         -0.1068, -0.1013,  0.2210,  0.0663,  0.2321],
        [-0.1225, -0.2221,  0.0689, -0.1025,  0.0532,  0.0361, -0.0843, -0.0799,
          0.0211,  0.0660,  0.0502,  0.4224, -0.0926],
        [ 0.0038, -0.0429, -0.1468,  0.2407,  0.0102,  0.3251, -0.1180,  0.0744,
         -0.1236,  0.0411, -0.0278, -0.1991, -0.2766],
        [-0.0863, -0.0660,  0.0396, -0.1107, -0.2354, -0.2599,  0.2507,  0.0745,
          0.0313,  0.0755,  0.2473, -0.2315,  0.2333],
        [-0.0166,  0.1887, -0.1992, -0.3009, -0.1392, -0.1847, -0.0457,  0.1943,
         -0.1482, -0.2822, -0.0030, -0.0758, -0.1392],
        [ 0.0020, -0.0024,  0.0625,  0.1569,  0.1537, -0.1713, -0.0473,  0.0963,
         -0.2901,  0.0191,  0.0918,  0.0159, -0.0347],
        [ 0.3627, -0.0454, -0.0322, -0.0956, -0.1250, -0.0341,  0.0803,  0.1461,
         -0.0883,  0.2248, -0.2998,  0.0395,  0.2176],
        [-0.2538,  0.3111,  0.2308,  0.0210, -0.1303, -0.2091,  0.0714, -0.3307,
         -0.0379,  0.2251, -0.1021,  0.0730,  0.0378],
        [ 0.0450,  0.0436,  0.0105, -0.1059,  0.1118, -0.1238,  0.1615, -0.0552,
         -0.0455, -0.0910,  0.1500, -0.0011, -0.0275],
        [-0.1807, -0.0692,  0.0591,  0.0076,  0.0127,  0.0891, -0.0765,  0.2371,
         -0.0692, -0.0207,  0.1000, -0.1159, -0.0882],
        [ 0.0912, -0.2033, -0.1026,  0.1210,  0.2329, -0.1923,  0.2109, -0.0092,
         -0.1319, -0.1324, -0.0528, -0.2045,  0.0713],
        [ 0.0842,  0.0389, -0.2070,  0.0645,  0.1097,  0.0028, -0.3349,  0.0986,
         -0.2645,  0.1306, -0.1034,  0.0816,  0.1803],
        [-0.2133,  0.0506,  0.1856,  0.1944, -0.0223, -0.0026,  0.1085, -0.0809,
          0.2229, -0.3000, -0.0402,  0.0148,  0.0605],
        [-0.0125, -0.1673,  0.0749,  0.4695, -0.2778, -0.0443,  0.1531,  0.2020,
          0.0622,  0.1171,  0.0037, -0.1444,  0.1088]])



class ResidualLinearBlock(torch.nn.Module):
    def __init__(self, w_dim=512):
        super().__init__()
        self.fc1 = FullyConnectedLayer(w_dim, w_dim, activation='lrelu', lr_multiplier=1)
        self.norm1 = nn.LayerNorm(w_dim)
        self.fc2 = FullyConnectedLayer(w_dim, w_dim, activation='lrelu', lr_multiplier=1)
        self.norm2 = nn.LayerNorm(w_dim)

    def forward(self, x):
        h = self.norm1(x)
        h = self.fc1(h)
        h = self.norm2(h)
        h = self.fc2(h)
        return h + x


@persistence.persistent_class
class ConceptMaskNetwork(nn.Module):
    def __init__(self, c_dim, i_dim, w_dim=512, activation='lrelu'):
        super().__init__()
        self.mask_net = nn.Sequential(
            FullyConnectedLayer(c_dim, w_dim, activation=activation, lr_multiplier=1),
            nn.LayerNorm(w_dim),
            FullyConnectedLayer(w_dim, w_dim, activation=activation, lr_multiplier=1),
            nn.LayerNorm(w_dim),
            FullyConnectedLayer(w_dim, i_dim, activation='linear', lr_multiplier=1, init='zeros'),
        )
        #self.param_net = nn.Parameter(-1e8*torch.ones([c_dim, i_dim]))

        """
        self.register_buffer('use_param', torch.zeros([c_dim, i_dim]))
        self.register_buffer('target_value', torch.zeros([c_dim, i_dim]))
        print(self)
        """
        self.register_buffer('use_param', torch.ones([c_dim, i_dim]))
        target_value = torch.tensor([
            [1,0,0,0,1,1,0],
            [1,0,0,1,0,1,0],
            [0,1,0,0,1,0,1],
            [0,1,0,0,1,1,0],
            [0,1,0,1,0,0,1],
            [0,1,0,1,0,1,0],
            [0,0,1,0,1,0,1],
            [0,0,1,1,0,0,1],
        ])
        self.register_buffer('target_value', target_value.float())

    def forward(self, c=None):
        mlp_out = (torch.tanh(self.mask_net(c))+1)/2
        buffer_out = self.target_value[c.argmax(dim=1)]
        use_param = self.use_param[c.argmax(dim=1)]
        return mlp_out * (1-use_param) + (use_param)*buffer_out


import pickle


@persistence.persistent_class
class ConceptMappingNetwork(torch.nn.Module):
    def __init__(self,
        z_dim,                      # Input latent (Z) dimensionality, 0 = no latent.
        c_dim,                      # Conditioning label (C) dimensionality, 0 = no label.
        w_dim,                      # Intermediate latent (W) dimensionality.
        num_ws,                     # Number of intermediate latents to output, None = do not broadcast.
        num_layers      = 8,        # Number of mapping layers.
        embed_features  = None,     # Label embedding dimensionality, None = same as w_dim.
        layer_features  = None,     # Number of intermediate features in the mapping layers, None = same as w_dim.
        activation      = 'lrelu',  # Activation function: 'relu', 'lrelu', etc.
        lr_multiplier   = 0.01,     # Learning rate multiplier for the mapping layers.
        w_avg_beta      = 0.995,    # Decay for tracking the moving average of W during training, None = do not track.
        cond_mode       = 'concat', # mode of coonditioning, stylegan3 uses concatenation
        i_dim           = 4,
        p_dim           = 64,
        flow_blocks     = 2,
        flow_dim        = 10,
        flow_norm       = 1,
        use_label       = 0,
        temperature     = 0.07
    ):
        super().__init__()

        self.z_dim = z_dim
        self.c_dim = c_dim
        self.w_dim = w_dim
        self.i_dim = i_dim
        self.p_dim = p_dim
        self.temperature = temperature

        self.temperature = 1

        self.num_ws = num_ws
        self.num_layers = num_layers
        self.w_avg_beta = w_avg_beta
        self.cond_mode = cond_mode
        self.flow_norm = flow_norm
        self.use_label = use_label
        if embed_features is None:
            embed_features = w_dim
        if c_dim == 0:
            embed_features = 0
        if layer_features is None:
            layer_features = w_dim

        #embedding_path = 'rgbmnist_pretrained_embedding.pkl'
        #with open(embedding_path, 'rb') as f:
        #    self.pretrained_embedding = pickle.load(f)
        #print('pretrained embedding loaded >>>>>>>>>> ', self.pretrained_embedding.shape)

        self.p_dim = p_dim
        for i in range(i_dim):
            mlp_net = nn.Sequential(FullyConnectedLayer(p_dim, w_dim, activation=activation, lr_multiplier=lr_multiplier),
                                    FullyConnectedLayer(w_dim, p_dim, activation=activation, lr_multiplier=lr_multiplier),)
            setattr(self, f'map_net{i}', mlp_net)

        self.deactivate_map_net = nn.Sequential(FullyConnectedLayer(p_dim, w_dim, activation=activation, lr_multiplier=lr_multiplier),
                                                FullyConnectedLayer(w_dim, p_dim, activation=activation, lr_multiplier=lr_multiplier),)

        self.main_map_net = nn.Sequential(FullyConnectedLayer((z_dim-i_dim*p_dim), w_dim, activation=activation, lr_multiplier=lr_multiplier),
                                          FullyConnectedLayer(w_dim, z_dim-i_dim*p_dim, activation=activation, lr_multiplier=lr_multiplier),)


        print(self)
        if num_ws is not None and w_avg_beta is not None:
            self.register_buffer('w_avg', torch.zeros([w_dim]))

    def forward(self, z, soft_mask, mask_mode='gumbel_hard',truncation_psi=1, truncation_cutoff=None, skip_w_avg_update=False,
                sparse_loss=False, label=None, entropy_thr=0.5, temperature=1.0):
        # Embed, normalize, and concat inputs.
        self.temperature = temperature

        x = None
        progress = 0
        outs = []
        """
        print(c.size(), ' >>>>>>>>>>.c size ssss')
        # Get unique elements and their counts
        unique_elements, counts = torch.unique(c.argmax(dim=1), return_counts=True)

        # Sort the counts in descending order, and sort unique elements according to this order
        sorted_counts, sorted_indices = torch.sort(counts, descending=True)
        sorted_elements = unique_elements[sorted_indices]

        print()
        # Print the sorted elements and their counts as pairs
        for element, count in zip(sorted_elements, sorted_counts):
            print(f"Element: {element.item()}, Count: {count.item()}")
        print()
        """
        with (torch.autograd.profiler.record_function('input')):
            if self.z_dim > 0:
                assert soft_mask.size() == (len(z), self.i_dim)
                #soft_mask = (torch.tanh(mask_logit)+1)/2
                hard_version = (soft_mask > 0.5).float()
                hard_mask = hard_version - soft_mask.detach() + soft_mask
                for i in range(self.i_dim):
                    cur_z = normalize_2nd_moment(z[:, i*self.p_dim:(i+1)*self.p_dim])
                    cur_map_net = getattr(self, f'map_net{i}')
                    cur_act_out = cur_map_net(cur_z)
                    cur_deact_out = self.deactivate_map_net(cur_z)
                    cur_out = cur_act_out*hard_mask[:, i].view(-1,1) + cur_deact_out*(1-hard_mask[:, i].view(-1,1))
                    outs.append(cur_out)
                rest_z = normalize_2nd_moment(z[:, self.i_dim*self.p_dim:])
                x = self.main_map_net(rest_z)
                outs.append(x)
        x = torch.cat(outs, dim=1)
        old_ws = x
        # Update moving average of W.
        if self.w_avg_beta is not None and self.training and not skip_w_avg_update:
            with torch.autograd.profiler.record_function('update_w_avg'):
                self.w_avg.copy_(x.detach().mean(dim=0).lerp(self.w_avg, self.w_avg_beta))

        # Broadcast.
        if self.num_ws is not None:
            with torch.autograd.profiler.record_function('broadcast'):
                x = x.unsqueeze(1).repeat([1, self.num_ws, 1])

        # Apply truncation.
        if truncation_psi != 1:
            with torch.autograd.profiler.record_function('truncate'):
                assert self.w_avg_beta is not None
                if self.num_ws is None or truncation_cutoff is None:
                    x = self.w_avg.lerp(x, truncation_psi)
                else:
                    x[:, :truncation_cutoff] = self.w_avg.lerp(x[:, :truncation_cutoff], truncation_psi)

        if sparse_loss:
            if self.cond_mode in ['flow', 'mlp']:



                loss_sparse = 0
                loss_entropy = 0
                loss_ortho = 0
                loss_path = 0
                loss_epsilon = 0
                loss_var = 0
                loss_colvar = 0
                loss_rowvar = 0
                loss_equal = 0

                entropy = -(soft_mask*torch.log(soft_mask+1e-20) + (1-soft_mask)*torch.log(1-soft_mask+1e-20))

                ent_cri = ((soft_mask>=entropy_thr))
                loss_entropy = torch.mean(entropy*ent_cri)
                filter_soft_mask = soft_mask*(soft_mask>0.5).float()
                crit = torch.sum(filter_soft_mask, dim=1).detach().view(-1,1)
                loss_sparse = torch.mean(hard_mask*crit*(soft_mask>0.5).float())
                loss_sparse = torch.mean(crit*hard_mask*(soft_mask>0.5).float())
                loss_sparse = torch.mean(crit*hard_mask*(soft_mask>0.5).float())
                loss_sparse = torch.mean(hard_mask*(soft_mask>0.5).float())
                #loss_sparse = torch.mean(crit*hard_mask*(soft_mask>0.5).float())
                crit = torch.sum(hard_mask, dim=1).view(-1,1).detach()
                loss_sparse = torch.mean(soft_mask*(soft_mask>0.5).float()*(soft_mask<0.9))

                sum_vec = torch.sum(soft_mask*(soft_mask>0.5).float(), dim=1)
                act_sum = torch.var(sum_vec)
                loss_rowvar = act_sum

                filter_hard_mask = hard_mask.detach()*(soft_mask>0.9)+hard_mask*(soft_mask<=0.9)
                sum_vec = torch.sum(filter_hard_mask, dim=1)
                act_sum = torch.var(sum_vec)
                loss_colvar = act_sum


                """
                cin = torch.arange(self.c_dim)
                cin = F.one_hot(cin, num_classes=self.c_dim).float().to(z.device)
                whole_soft_mask = self.mask_net(cin)
                whole_soft_mask = torch.sigmoid(whole_soft_mask)
                whole_soft_mask = whole_soft_mask*(whole_soft_mask>0.5).float()
                ortho_mat = torch.matmul(whole_soft_mask.t(), whole_soft_mask)
                ortho_mat = ortho_mat * (1-torch.eye(self.i_dim).to(z.device))
                loss_ortho = torch.mean(ortho_mat)
                """
                loss_dict = {
                             }
                loss_dict['loss_sparse'] = loss_sparse
                loss_dict['loss_entropy'] = loss_entropy
                loss_dict['loss_ortho'] = loss_ortho
                loss_dict['loss_path'] = loss_path
                loss_dict['loss_epsilon'] = loss_epsilon
                loss_dict['loss_cls'] = 0
                loss_dict['loss_colvar'] = loss_colvar
                loss_dict['loss_rowvar'] = loss_rowvar
                loss_dict['loss_equal'] = loss_equal

                return x, loss_dict
            else:
                return x, torch.tensor(0.)
        else:
            return x


@persistence.persistent_class
class MappingNetwork(torch.nn.Module):
    def __init__(self,
        z_dim,                      # Input latent (Z) dimensionality, 0 = no latent.
        c_dim,                      # Conditioning label (C) dimensionality, 0 = no label.
        w_dim,                      # Intermediate latent (W) dimensionality.
        num_ws,                     # Number of intermediate latents to output, None = do not broadcast.
        num_layers      = 8,        # Number of mapping layers.
        embed_features  = None,     # Label embedding dimensionality, None = same as w_dim.
        layer_features  = None,     # Number of intermediate features in the mapping layers, None = same as w_dim.
        activation      = 'lrelu',  # Activation function: 'relu', 'lrelu', etc.
        lr_multiplier   = 0.01,     # Learning rate multiplier for the mapping layers.
        w_avg_beta      = 0.995,    # Decay for tracking the moving average of W during training, None = do not track.
    ):
        super().__init__()
        self.z_dim = z_dim
        self.c_dim = c_dim
        self.w_dim = w_dim
        self.num_ws = num_ws
        self.num_layers = num_layers
        self.w_avg_beta = w_avg_beta

        if embed_features is None:
            embed_features = w_dim
        if c_dim == 0:
            embed_features = 0
        if layer_features is None:
            layer_features = w_dim
        features_list = [z_dim + embed_features] + [layer_features] * (num_layers - 1) + [w_dim]

        if c_dim > 0:
            self.embed = FullyConnectedLayer(c_dim, embed_features)
        for idx in range(num_layers):
            in_features = features_list[idx]
            out_features = features_list[idx + 1]
            layer = FullyConnectedLayer(in_features, out_features, activation=activation, lr_multiplier=lr_multiplier)
            setattr(self, f'fc{idx}', layer)

        if num_ws is not None and w_avg_beta is not None:
            self.register_buffer('w_avg', torch.zeros([w_dim]))

    def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, skip_w_avg_update=False):
        # Embed, normalize, and concat inputs.
        x = None
        with torch.autograd.profiler.record_function('input'):
            if self.z_dim > 0:
                misc.assert_shape(z, [None, self.z_dim])
                x = normalize_2nd_moment(z.to(torch.float32))
            if self.c_dim > 0:
                misc.assert_shape(c, [None, self.c_dim])
                y = normalize_2nd_moment(self.embed(c.to(torch.float32)))
                x = torch.cat([x, y], dim=1) if x is not None else y

        # Main layers.
        for idx in range(self.num_layers):
            layer = getattr(self, f'fc{idx}')
            x = layer(x)

        # Update moving average of W.
        if self.w_avg_beta is not None and self.training and not skip_w_avg_update:
            with torch.autograd.profiler.record_function('update_w_avg'):
                self.w_avg.copy_(x.detach().mean(dim=0).lerp(self.w_avg, self.w_avg_beta))

        # Broadcast.
        if self.num_ws is not None:
            with torch.autograd.profiler.record_function('broadcast'):
                x = x.unsqueeze(1).repeat([1, self.num_ws, 1])

        # Apply truncation.
        if truncation_psi != 1:
            with torch.autograd.profiler.record_function('truncate'):
                assert self.w_avg_beta is not None
                if self.num_ws is None or truncation_cutoff is None:
                    x = self.w_avg.lerp(x, truncation_psi)
                else:
                    x[:, :truncation_cutoff] = self.w_avg.lerp(x[:, :truncation_cutoff], truncation_psi)
        return x

#----------------------------------------------------------------------------

@persistence.persistent_class
class SynthesisLayer(torch.nn.Module):
    def __init__(self,
        in_channels,                    # Number of input channels.
        out_channels,                   # Number of output channels.
        w_dim,                          # Intermediate latent (W) dimensionality.
        resolution,                     # Resolution of this layer.
        kernel_size     = 3,            # Convolution kernel size.
        up              = 1,            # Integer upsampling factor.
        use_noise       = True,         # Enable noise input?
        activation      = 'lrelu',      # Activation function: 'relu', 'lrelu', etc.
        resample_filter = [1,3,3,1],    # Low-pass filter to apply when resampling activations.
        conv_clamp      = None,         # Clamp the output of convolution layers to +-X, None = disable clamping.
        channels_last   = False,        # Use channels_last format for the weights?
    ):
        super().__init__()
        self.resolution = resolution
        self.up = up
        self.use_noise = use_noise
        self.activation = activation
        self.conv_clamp = conv_clamp
        self.register_buffer('resample_filter', upfirdn2d.setup_filter(resample_filter))
        self.padding = kernel_size // 2
        self.act_gain = bias_act.activation_funcs[activation].def_gain

        self.affine = FullyConnectedLayer(w_dim, in_channels, bias_init=1)
        memory_format = torch.channels_last if channels_last else torch.contiguous_format
        self.weight = torch.nn.Parameter(torch.randn([out_channels, in_channels, kernel_size, kernel_size]).to(memory_format=memory_format))
        if use_noise:
            self.register_buffer('noise_const', torch.randn([resolution, resolution]))
            self.noise_strength = torch.nn.Parameter(torch.zeros([]))
        self.bias = torch.nn.Parameter(torch.zeros([out_channels]))

    def forward(self, x, w, noise_mode='random', fused_modconv=True, gain=1):
        assert noise_mode in ['random', 'const', 'none']
        in_resolution = self.resolution // self.up
        misc.assert_shape(x, [None, self.weight.shape[1], in_resolution, in_resolution])
        styles = self.affine(w)

        noise = None
        if self.use_noise and noise_mode == 'random':
            noise = torch.randn([x.shape[0], 1, self.resolution, self.resolution], device=x.device) * self.noise_strength
        if self.use_noise and noise_mode == 'const':
            noise = self.noise_const * self.noise_strength

        flip_weight = (self.up == 1) # slightly faster
        x = modulated_conv2d(x=x, weight=self.weight, styles=styles, noise=noise, up=self.up,
            padding=self.padding, resample_filter=self.resample_filter, flip_weight=flip_weight, fused_modconv=fused_modconv)

        act_gain = self.act_gain * gain
        act_clamp = self.conv_clamp * gain if self.conv_clamp is not None else None
        x = bias_act.bias_act(x, self.bias.to(x.dtype), act=self.activation, gain=act_gain, clamp=act_clamp)
        return x

#----------------------------------------------------------------------------

@persistence.persistent_class
class ToRGBLayer(torch.nn.Module):
    def __init__(self, in_channels, out_channels, w_dim, kernel_size=1, conv_clamp=None, channels_last=False):
        super().__init__()
        self.conv_clamp = conv_clamp
        self.affine = FullyConnectedLayer(w_dim, in_channels, bias_init=1)
        memory_format = torch.channels_last if channels_last else torch.contiguous_format
        self.weight = torch.nn.Parameter(torch.randn([out_channels, in_channels, kernel_size, kernel_size]).to(memory_format=memory_format))
        self.bias = torch.nn.Parameter(torch.zeros([out_channels]))
        self.weight_gain = 1 / np.sqrt(in_channels * (kernel_size ** 2))

    def forward(self, x, w, fused_modconv=True):
        styles = self.affine(w) * self.weight_gain
        x = modulated_conv2d(x=x, weight=self.weight, styles=styles, demodulate=False, fused_modconv=fused_modconv)
        x = bias_act.bias_act(x, self.bias.to(x.dtype), clamp=self.conv_clamp)
        return x

#----------------------------------------------------------------------------

@persistence.persistent_class
class SynthesisBlock(torch.nn.Module):
    def __init__(self,
        in_channels,                        # Number of input channels, 0 = first block.
        out_channels,                       # Number of output channels.
        w_dim,                              # Intermediate latent (W) dimensionality.
        resolution,                         # Resolution of this block.
        img_channels,                       # Number of output color channels.
        is_last,                            # Is this the last block?
        architecture        = 'skip',       # Architecture: 'orig', 'skip', 'resnet'.
        resample_filter     = [1,3,3,1],    # Low-pass filter to apply when resampling activations.
        conv_clamp          = None,         # Clamp the output of convolution layers to +-X, None = disable clamping.
        use_fp16            = False,        # Use FP16 for this block?
        fp16_channels_last  = False,        # Use channels-last memory format with FP16?
        **layer_kwargs,                     # Arguments for SynthesisLayer.
    ):
        assert architecture in ['orig', 'skip', 'resnet']
        super().__init__()
        self.in_channels = in_channels
        self.w_dim = w_dim
        self.resolution = resolution
        self.img_channels = img_channels
        self.is_last = is_last
        self.architecture = architecture
        self.use_fp16 = use_fp16
        self.channels_last = (use_fp16 and fp16_channels_last)
        self.register_buffer('resample_filter', upfirdn2d.setup_filter(resample_filter))
        self.num_conv = 0
        self.num_torgb = 0

        if in_channels == 0:
            self.const = torch.nn.Parameter(torch.randn([out_channels, resolution, resolution]))

        if in_channels != 0:
            self.conv0 = SynthesisLayer(in_channels, out_channels, w_dim=w_dim, resolution=resolution, up=2,
                resample_filter=resample_filter, conv_clamp=conv_clamp, channels_last=self.channels_last, **layer_kwargs)
            self.num_conv += 1

        self.conv1 = SynthesisLayer(out_channels, out_channels, w_dim=w_dim, resolution=resolution,
            conv_clamp=conv_clamp, channels_last=self.channels_last, **layer_kwargs)
        self.num_conv += 1

        if is_last or architecture == 'skip':
            self.torgb = ToRGBLayer(out_channels, img_channels, w_dim=w_dim,
                conv_clamp=conv_clamp, channels_last=self.channels_last)
            self.num_torgb += 1

        if in_channels != 0 and architecture == 'resnet':
            self.skip = Conv2dLayer(in_channels, out_channels, kernel_size=1, bias=False, up=2,
                resample_filter=resample_filter, channels_last=self.channels_last)

    def forward(self, x, img, ws, force_fp32=False, fused_modconv=None, **layer_kwargs):
        misc.assert_shape(ws, [None, self.num_conv + self.num_torgb, self.w_dim])
        w_iter = iter(ws.unbind(dim=1))
        dtype = torch.float16 if self.use_fp16 and not force_fp32 else torch.float32
        memory_format = torch.channels_last if self.channels_last and not force_fp32 else torch.contiguous_format
        if fused_modconv is None:
            with misc.suppress_tracer_warnings(): # this value will be treated as a constant
                fused_modconv = (not self.training) and (dtype == torch.float32 or int(x.shape[0]) == 1)

        # Input.
        if self.in_channels == 0:
            x = self.const.to(dtype=dtype, memory_format=memory_format)
            x = x.unsqueeze(0).repeat([ws.shape[0], 1, 1, 1])
        else:
            misc.assert_shape(x, [None, self.in_channels, self.resolution // 2, self.resolution // 2])
            x = x.to(dtype=dtype, memory_format=memory_format)

        # Main layers.
        if self.in_channels == 0:
            x = self.conv1(x, next(w_iter), fused_modconv=fused_modconv, **layer_kwargs)
        elif self.architecture == 'resnet':
            y = self.skip(x, gain=np.sqrt(0.5))
            x = self.conv0(x, next(w_iter), fused_modconv=fused_modconv, **layer_kwargs)
            x = self.conv1(x, next(w_iter), fused_modconv=fused_modconv, gain=np.sqrt(0.5), **layer_kwargs)
            x = y.add_(x)
        else:
            x = self.conv0(x, next(w_iter), fused_modconv=fused_modconv, **layer_kwargs)
            x = self.conv1(x, next(w_iter), fused_modconv=fused_modconv, **layer_kwargs)

        # ToRGB.
        if img is not None:
            misc.assert_shape(img, [None, self.img_channels, self.resolution // 2, self.resolution // 2])
            img = upfirdn2d.upsample2d(img, self.resample_filter)
        if self.is_last or self.architecture == 'skip':
            y = self.torgb(x, next(w_iter), fused_modconv=fused_modconv)
            y = y.to(dtype=torch.float32, memory_format=torch.contiguous_format)
            img = img.add_(y) if img is not None else y

        assert x.dtype == dtype
        assert img is None or img.dtype == torch.float32
        return x, img

#----------------------------------------------------------------------------

@persistence.persistent_class
class SynthesisNetwork(torch.nn.Module):
    def __init__(self,
        w_dim,                      # Intermediate latent (W) dimensionality.
        img_resolution,             # Output image resolution.
        img_channels,               # Number of color channels.
        channel_base    = 32768,    # Overall multiplier for the number of channels.
        channel_max     = 512,      # Maximum number of channels in any layer.
        num_fp16_res    = 0,        # Use FP16 for the N highest resolutions.
        **block_kwargs,             # Arguments for SynthesisBlock.
    ):
        assert img_resolution >= 4 and img_resolution & (img_resolution - 1) == 0
        super().__init__()
        self.w_dim = w_dim
        self.img_resolution = img_resolution
        self.img_resolution_log2 = int(np.log2(img_resolution))
        self.img_channels = img_channels
        self.block_resolutions = [2 ** i for i in range(2, self.img_resolution_log2 + 1)]
        channels_dict = {res: min(channel_base // res, channel_max) for res in self.block_resolutions}
        fp16_resolution = max(2 ** (self.img_resolution_log2 + 1 - num_fp16_res), 8)

        self.num_ws = 0
        for res in self.block_resolutions:
            in_channels = channels_dict[res // 2] if res > 4 else 0
            out_channels = channels_dict[res]
            use_fp16 = (res >= fp16_resolution)
            is_last = (res == self.img_resolution)
            block = SynthesisBlock(in_channels, out_channels, w_dim=w_dim, resolution=res,
                img_channels=img_channels, is_last=is_last, use_fp16=use_fp16, **block_kwargs)
            self.num_ws += block.num_conv
            if is_last:
                self.num_ws += block.num_torgb
            setattr(self, f'b{res}', block)

    def forward(self, ws, **block_kwargs):
        block_ws = []
        with torch.autograd.profiler.record_function('split_ws'):
            misc.assert_shape(ws, [None, self.num_ws, self.w_dim])
            ws = ws.to(torch.float32)
            w_idx = 0
            for res in self.block_resolutions:
                block = getattr(self, f'b{res}')
                block_ws.append(ws.narrow(1, w_idx, block.num_conv + block.num_torgb))
                w_idx += block.num_conv

        x = img = None
        for res, cur_ws in zip(self.block_resolutions, block_ws):
            block = getattr(self, f'b{res}')
            x, img = block(x, img, cur_ws, **block_kwargs)
        return img

#----------------------------------------------------------------------------

@persistence.persistent_class
class Generator(torch.nn.Module):
    def __init__(self,
        z_dim,                      # Input latent (Z) dimensionality.
        c_dim,                      # Conditioning label (C) dimensionality.
        w_dim,                      # Intermediate latent (W) dimensionality.
        img_resolution,             # Output resolution.
        img_channels,               # Number of output color channels.
        mapping_kwargs      = {},   # Arguments for MappingNetwork.
        synthesis_kwargs    = {},   # Arguments for SynthesisNetwork.
    ):
        super().__init__()
        self.z_dim = z_dim
        self.c_dim = c_dim
        self.w_dim = w_dim
        self.img_resolution = img_resolution
        self.img_channels = img_channels
        self.synthesis = SynthesisNetwork(w_dim=w_dim, img_resolution=img_resolution, img_channels=img_channels, **synthesis_kwargs)
        self.num_ws = self.synthesis.num_ws
        self.mapping = ConceptMappingNetwork(z_dim=z_dim, c_dim=c_dim, w_dim=w_dim, num_ws=self.num_ws, **mapping_kwargs)

    def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, label=None, **synthesis_kwargs):
        ws = self.mapping(z, c, truncation_psi=truncation_psi, truncation_cutoff=truncation_cutoff, label=label)
        img = self.synthesis(ws, **synthesis_kwargs)
        return img

#----------------------------------------------------------------------------

@persistence.persistent_class
class DiscriminatorBlock(torch.nn.Module):
    def __init__(self,
        in_channels,                        # Number of input channels, 0 = first block.
        tmp_channels,                       # Number of intermediate channels.
        out_channels,                       # Number of output channels.
        resolution,                         # Resolution of this block.
        img_channels,                       # Number of input color channels.
        first_layer_idx,                    # Index of the first layer.
        architecture        = 'resnet',     # Architecture: 'orig', 'skip', 'resnet'.
        activation          = 'lrelu',      # Activation function: 'relu', 'lrelu', etc.
        resample_filter     = [1,3,3,1],    # Low-pass filter to apply when resampling activations.
        conv_clamp          = None,         # Clamp the output of convolution layers to +-X, None = disable clamping.
        use_fp16            = False,        # Use FP16 for this block?
        fp16_channels_last  = False,        # Use channels-last memory format with FP16?
        freeze_layers       = 0,            # Freeze-D: Number of layers to freeze.
    ):
        assert in_channels in [0, tmp_channels]
        assert architecture in ['orig', 'skip', 'resnet']
        super().__init__()
        self.in_channels = in_channels
        self.resolution = resolution
        self.img_channels = img_channels
        self.first_layer_idx = first_layer_idx
        self.architecture = architecture
        self.use_fp16 = use_fp16
        self.channels_last = (use_fp16 and fp16_channels_last)
        self.register_buffer('resample_filter', upfirdn2d.setup_filter(resample_filter))

        self.num_layers = 0
        def trainable_gen():
            while True:
                layer_idx = self.first_layer_idx + self.num_layers
                trainable = (layer_idx >= freeze_layers)
                self.num_layers += 1
                yield trainable
        trainable_iter = trainable_gen()

        if in_channels == 0 or architecture == 'skip':
            self.fromrgb = Conv2dLayer(img_channels, tmp_channels, kernel_size=1, activation=activation,
                trainable=next(trainable_iter), conv_clamp=conv_clamp, channels_last=self.channels_last)

        self.conv0 = Conv2dLayer(tmp_channels, tmp_channels, kernel_size=3, activation=activation,
            trainable=next(trainable_iter), conv_clamp=conv_clamp, channels_last=self.channels_last)

        self.conv1 = Conv2dLayer(tmp_channels, out_channels, kernel_size=3, activation=activation, down=2,
            trainable=next(trainable_iter), resample_filter=resample_filter, conv_clamp=conv_clamp, channels_last=self.channels_last)

        if architecture == 'resnet':
            self.skip = Conv2dLayer(tmp_channels, out_channels, kernel_size=1, bias=False, down=2,
                trainable=next(trainable_iter), resample_filter=resample_filter, channels_last=self.channels_last)

    def forward(self, x, img, force_fp32=False):
        dtype = torch.float16 if self.use_fp16 and not force_fp32 else torch.float32
        memory_format = torch.channels_last if self.channels_last and not force_fp32 else torch.contiguous_format

        # Input.
        if x is not None:
            misc.assert_shape(x, [None, self.in_channels, self.resolution, self.resolution])
            x = x.to(dtype=dtype, memory_format=memory_format)

        # FromRGB.
        if self.in_channels == 0 or self.architecture == 'skip':
            misc.assert_shape(img, [None, self.img_channels, self.resolution, self.resolution])
            img = img.to(dtype=dtype, memory_format=memory_format)
            y = self.fromrgb(img)
            x = x + y if x is not None else y
            img = upfirdn2d.downsample2d(img, self.resample_filter) if self.architecture == 'skip' else None

        # Main layers.
        if self.architecture == 'resnet':
            y = self.skip(x, gain=np.sqrt(0.5))
            x = self.conv0(x)
            x = self.conv1(x, gain=np.sqrt(0.5))
            x = y.add_(x)
        else:
            x = self.conv0(x)
            x = self.conv1(x)

        assert x.dtype == dtype
        return x, img

#----------------------------------------------------------------------------

@persistence.persistent_class
class MinibatchStdLayer(torch.nn.Module):
    def __init__(self, group_size, num_channels=1):
        super().__init__()
        self.group_size = group_size
        self.num_channels = num_channels

    def forward(self, x):
        N, C, H, W = x.shape
        with misc.suppress_tracer_warnings(): # as_tensor results are registered as constants
            G = torch.min(torch.as_tensor(self.group_size), torch.as_tensor(N)) if self.group_size is not None else N
        F = self.num_channels
        c = C // F

        y = x.reshape(G, -1, F, c, H, W)    # [GnFcHW] Split minibatch N into n groups of size G, and channels C into F groups of size c.
        y = y - y.mean(dim=0)               # [GnFcHW] Subtract mean over group.
        y = y.square().mean(dim=0)          # [nFcHW]  Calc variance over group.
        y = (y + 1e-8).sqrt()               # [nFcHW]  Calc stddev over group.
        y = y.mean(dim=[2,3,4])             # [nF]     Take average over channels and pixels.
        y = y.reshape(-1, F, 1, 1)          # [nF11]   Add missing dimensions.
        y = y.repeat(G, 1, H, W)            # [NFHW]   Replicate over group and pixels.
        x = torch.cat([x, y], dim=1)        # [NCHW]   Append to input as new channels.
        return x

#----------------------------------------------------------------------------

@persistence.persistent_class
class DiscriminatorEpilogue(torch.nn.Module):
    def __init__(self,
        in_channels,                    # Number of input channels.
        cmap_dim,                       # Dimensionality of mapped conditioning label, 0 = no label.
        resolution,                     # Resolution of this block.
        img_channels,                   # Number of input color channels.
        architecture        = 'resnet', # Architecture: 'orig', 'skip', 'resnet'.
        mbstd_group_size    = 4,        # Group size for the minibatch standard deviation layer, None = entire minibatch.
        mbstd_num_channels  = 1,        # Number of features for the minibatch standard deviation layer, 0 = disable.
        activation          = 'lrelu',  # Activation function: 'relu', 'lrelu', etc.
        conv_clamp          = None,     # Clamp the output of convolution layers to +-X, None = disable clamping.
    ):
        assert architecture in ['orig', 'skip', 'resnet']
        super().__init__()
        self.in_channels = in_channels
        self.cmap_dim = cmap_dim
        self.resolution = resolution
        self.img_channels = img_channels
        self.architecture = architecture

        if architecture == 'skip':
            self.fromrgb = Conv2dLayer(img_channels, in_channels, kernel_size=1, activation=activation)
        self.mbstd = MinibatchStdLayer(group_size=mbstd_group_size, num_channels=mbstd_num_channels) if mbstd_num_channels > 0 else None
        self.conv = Conv2dLayer(in_channels + mbstd_num_channels, in_channels, kernel_size=3, activation=activation, conv_clamp=conv_clamp)
        self.fc = FullyConnectedLayer(in_channels * (resolution ** 2), in_channels, activation=activation)
        self.out = FullyConnectedLayer(in_channels, 1 if cmap_dim == 0 else cmap_dim)
        #self.out = FullyConnectedLayer(in_channels, 1)

    def forward(self, x, img, cmap, force_fp32=False):
        misc.assert_shape(x, [None, self.in_channels, self.resolution, self.resolution]) # [NCHW]
        _ = force_fp32 # unused
        dtype = torch.float32
        memory_format = torch.contiguous_format

        # FromRGB.
        x = x.to(dtype=dtype, memory_format=memory_format)
        if self.architecture == 'skip':
            misc.assert_shape(img, [None, self.img_channels, self.resolution, self.resolution])
            img = img.to(dtype=dtype, memory_format=memory_format)
            x = x + self.fromrgb(img)

        # Main layers.
        if self.mbstd is not None:
            x = self.mbstd(x)
        x = self.conv(x)
        x = self.fc(x.flatten(1))
        x = self.out(x)

        # Conditioning.
        if self.cmap_dim > 0:
            misc.assert_shape(cmap, [None, self.cmap_dim])
            x = (x * cmap).sum(dim=1, keepdim=True) * (1 / np.sqrt(self.cmap_dim))
            #x = x + 0 * cmap.sum()
        assert x.dtype == dtype
        return x
#----------------------------------------------------------------------------

@persistence.persistent_class
class Discriminator(torch.nn.Module):
    def __init__(self,
        c_dim,                          # Conditioning label (C) dimensionality.
        img_resolution,                 # Input resolution.
        img_channels,                   # Number of input color channels.
        architecture        = 'resnet', # Architecture: 'orig', 'skip', 'resnet'.
        channel_base        = 32768,    # Overall multiplier for the number of channels.
        channel_max         = 512,      # Maximum number of channels in any layer.
        num_fp16_res        = 0,        # Use FP16 for the N highest resolutions.
        conv_clamp          = None,     # Clamp the output of convolution layers to +-X, None = disable clamping.
        cmap_dim            = None,     # Dimensionality of mapped conditioning label, None = default.
        block_kwargs        = {},       # Arguments for DiscriminatorBlock.
        mapping_kwargs      = {},       # Arguments for MappingNetwork.
        epilogue_kwargs     = {},       # Arguments for DiscriminatorEpilogue.
    ):
        super().__init__()
        self.c_dim = c_dim
        self.img_resolution = img_resolution
        self.img_resolution_log2 = int(np.log2(img_resolution))
        self.img_channels = img_channels
        self.block_resolutions = [2 ** i for i in range(self.img_resolution_log2, 2, -1)]
        channels_dict = {res: min(channel_base // res, channel_max) for res in self.block_resolutions + [4]}
        fp16_resolution = max(2 ** (self.img_resolution_log2 + 1 - num_fp16_res), 8)

        if cmap_dim is None:
            cmap_dim = channels_dict[4]
        if c_dim == 0:
            cmap_dim = 0

        common_kwargs = dict(img_channels=img_channels, architecture=architecture, conv_clamp=conv_clamp)
        cur_layer_idx = 0
        for res in self.block_resolutions:
            in_channels = channels_dict[res] if res < img_resolution else 0
            tmp_channels = channels_dict[res]
            out_channels = channels_dict[res // 2]
            use_fp16 = (res >= fp16_resolution)
            block = DiscriminatorBlock(in_channels, tmp_channels, out_channels, resolution=res,
                first_layer_idx=cur_layer_idx, use_fp16=use_fp16, **block_kwargs, **common_kwargs)
            setattr(self, f'b{res}', block)
            cur_layer_idx += block.num_layers
        if c_dim > 0:
            self.mapping = MappingNetwork(z_dim=0, c_dim=c_dim, w_dim=cmap_dim, num_ws=None, w_avg_beta=None, **mapping_kwargs)
        self.b4 = DiscriminatorEpilogue(channels_dict[4], cmap_dim=cmap_dim, resolution=4, **epilogue_kwargs, **common_kwargs)

    def forward(self, img, c, **block_kwargs):
        x = None
        for res in self.block_resolutions:
            block = getattr(self, f'b{res}')
            x, img = block(x, img, **block_kwargs)
            #print(img.size(), ' >>>>>>>> img sizeesssssssssss ', c.size(), ' ', res)

        cmap = None
        if self.c_dim > 0:
            cmap = self.mapping(None, c)
        x = self.b4(x, img, cmap)
        return x

#----------------------------------------------------------------------------