# Combined loss
g_loss = g_A_loss_adv + g_B_loss_adv + loss_rec + loss_idt
if self.checkpoint or self.direct:
    g_loss = g_A_loss_adv + g_B_loss_adv + loss_rec + loss_idt + g_A_loss_his + g_B_loss_his

# GAN loss D_A(G_A(A))
fake_B = self.G_A(org_A)
pred_fake = self.D_A(fake_B)
g_A_loss_adv =  self.criterionGAN(pred_fake, True)
#g_loss_adv = self.get_G_loss(out)

# GAN loss D_B(G_B(B))
fake_A = self.G_B(ref_B)
pred_fake = self.D_B(fake_A)
g_B_loss_adv = self.criterionGAN(pred_fake, True)

# identity loss
if self.lambda_idt > 0:
    # G should be identity if ref_B or org_A is fed
    idt_A1, idt_A2 = self.G(org_A, org_A)
    idt_B1, idt_B2 = self.G(ref_B, ref_B)
    loss_idt_A1 = self.criterionL1(idt_A1, org_A) * self.lambda_A * self.lambda_idt
    loss_idt_A2 = self.criterionL1(idt_A2, org_A) * self.lambda_A * self.lambda_idt
    loss_idt_B1 = self.criterionL1(idt_B1, ref_B) * self.lambda_B * self.lambda_idt
    loss_idt_B2 = self.criterionL1(idt_B2, ref_B) * self.lambda_B * self.lambda_idt
    # loss_idt
    loss_idt = (loss_idt_A1 + loss_idt_A2 + loss_idt_B1 + loss_idt_B2) * 0.5
else:
    loss_idt = 0

# vgg loss
vgg_org = self.vgg(org_A, self.content_layer)[0]
vgg_org = Variable(vgg_org.data).detach()
vgg_fake_A = self.vgg(fake_A, self.content_layer)[0]
g_loss_A_vgg = self.criterionL2(vgg_fake_A, vgg_org) * self.lambda_A * self.lambda_vgg

vgg_ref = self.vgg(ref_B, self.content_layer)[0]
vgg_ref = Variable(vgg_ref.data).detach()
vgg_fake_B = self.vgg(fake_B, self.content_layer)[0]
g_loss_B_vgg = self.criterionL2(vgg_fake_B, vgg_ref) * self.lambda_B * self.lambda_vgg

loss_rec = (g_loss_rec_A + g_loss_rec_B + g_loss_A_vgg + g_loss_B_vgg) * 0.5

# Forward cycle loss
rec_A = self.G_B(fake_B)
g_loss_rec_A = self.criterionL1(rec_A, org_A) * self.lambda_A

# Backward cycle loss
rec_B = self.G_A(fake_A)
g_loss_rec_B = self.criterionL1(rec_B, ref_B) * self.lambda_B

for self.i, (img_A, img_B, mask_A, mask_B) in enumerate(self.data_loader_train):
    # Convert tensor to variable
    # mask attribute: 0:background 1:face 2:left-eyebrown 3:right-eyebrown 4:left-eye 5: right-eye 6: nose 
    # 7: upper-lip 8: teeth 9: under-lip 10:hair 11: left-ear 12: right-ear 13: neck
    if self.checkpoint or self.direct:
        if self.lips==True:
            mask_A_lip = (mask_A==7).float() + (mask_A==9).float()
            mask_B_lip = (mask_B==7).float() + (mask_B==9).float()
            mask_A_lip, mask_B_lip, index_A_lip, index_B_lip = self.mask_preprocess(mask_A_lip, mask_B_lip)
        if self.skin==True:
            mask_A_skin = (mask_A==1).float() + (mask_A==6).float() + (mask_A==13).float()
            mask_B_skin = (mask_B==1).float() + (mask_B==6).float() + (mask_B==13).float()
            mask_A_skin, mask_B_skin, index_A_skin, index_B_skin = self.mask_preprocess(mask_A_skin, mask_B_skin)
        if self.eye==True:
            mask_A_eye_left = (mask_A==4).float()
            mask_A_eye_right = (mask_A==5).float()
            mask_B_eye_left = (mask_B==4).float()
            mask_B_eye_right = (mask_B==5).float()
            mask_A_face = (mask_A==1).float() + (mask_A==6).float()
            mask_B_face = (mask_B==1).float() + (mask_B==6).float()
            # avoid the situation that images with eye closed
            if not ((mask_A_eye_left>0).any() and (mask_B_eye_left>0).any() and \
                (mask_A_eye_right > 0).any() and (mask_B_eye_right > 0).any()):
                continue
            mask_A_eye_left, mask_A_eye_right = self.rebound_box(mask_A_eye_left, mask_A_eye_right, mask_A_face)
            mask_B_eye_left, mask_B_eye_right = self.rebound_box(mask_B_eye_left, mask_B_eye_right, mask_B_face)
            mask_A_eye_left, mask_B_eye_left, index_A_eye_left, index_B_eye_left = \
                self.mask_preprocess(mask_A_eye_left, mask_B_eye_left)
            mask_A_eye_right, mask_B_eye_right, index_A_eye_right, index_B_eye_right = \
                self.mask_preprocess(mask_A_eye_right, mask_B_eye_right)

# color_histogram loss
g_A_loss_his = 0
g_B_loss_his = 0
if self.checkpoint or self.direct:
    if self.lips==True:
        g_A_lip_loss_his = self.criterionHis(fake_A, ref_B, mask_A_lip, mask_B_lip, index_A_lip) * self.lambda_his_lip
        g_B_lip_loss_his = self.criterionHis(fake_B, org_A, mask_B_lip, mask_A_lip, index_B_lip) * self.lambda_his_lip
        g_A_loss_his += g_A_lip_loss_his
        g_B_loss_his += g_B_lip_loss_his
    if self.skin==True:
        g_A_skin_loss_his = self.criterionHis(fake_A, ref_B, mask_A_skin, mask_B_skin, index_A_skin) * self.lambda_his_skin_1
        g_B_skin_loss_his = self.criterionHis(fake_B, org_A, mask_B_skin, mask_A_skin, index_B_skin) * self.lambda_his_skin_2
        g_A_loss_his += g_A_skin_loss_his
        g_B_loss_his += g_B_skin_loss_his
    if self.eye==True:
        g_A_eye_left_loss_his = self.criterionHis(fake_A, ref_B, mask_A_eye_left, mask_B_eye_left, index_A_eye_left) * self.lambda_his_eye
        g_B_eye_left_loss_his = self.criterionHis(fake_B, org_A, mask_B_eye_left, mask_A_eye_left, index_B_eye_left) * self.lambda_his_eye
        g_A_eye_right_loss_his = self.criterionHis(fake_A, ref_B, mask_A_eye_right, mask_B_eye_right, index_A_eye_right) * self.lambda_his_eye
        g_B_eye_right_loss_his = self.criterionHis(fake_B, org_A, mask_B_eye_right, mask_A_eye_right, index_B_eye_right) * self.lambda_his_eye
        g_A_loss_his += g_A_eye_left_loss_his + g_A_eye_right_loss_his
        g_B_loss_his += g_B_eye_left_loss_his + g_B_eye_right_loss_his

class Generator(nn.Module):
    """Generator. Encoder-Decoder Architecture."""
    def __init__(self, conv_dim=64, repeat_num=6):
        super(Generator, self).__init__()

        layers = []
        layers.append(nn.Conv2d(3, conv_dim, kernel_size=7, stride=1, padding=3, bias=False))
        layers.append(nn.InstanceNorm2d(conv_dim, affine=True))
        layers.append(nn.ReLU(inplace=True))

        # Down-Sampling
        curr_dim = conv_dim
        for i in range(2):
            layers.append(nn.Conv2d(curr_dim, curr_dim*2, kernel_size=4, stride=2, padding=1, bias=False))
            layers.append(nn.InstanceNorm2d(curr_dim*2, affine=True))
            layers.append(nn.ReLU(inplace=True))
            curr_dim = curr_dim * 2

        # Bottleneck
        for i in range(repeat_num):
            layers.append(ResidualBlock(dim_in=curr_dim, dim_out=curr_dim))

        # Up-Sampling
        for i in range(2):
            layers.append(nn.ConvTranspose2d(curr_dim, curr_dim//2, kernel_size=4, stride=2, padding=1, bias=False))
            layers.append(nn.InstanceNorm2d(curr_dim//2, affine=True))
            layers.append(nn.ReLU(inplace=True))
            curr_dim = curr_dim // 2

        layers.append(nn.Conv2d(curr_dim, 3, kernel_size=7, stride=1, padding=3, bias=False))
        layers.append(nn.Tanh())
        self.main = nn.Sequential(*layers)

    def forward(self, x):
        out = self.main(x)
        return out

class Generator_makeup(nn.Module):
    """Generator. Encoder-Decoder Architecture."""
    # input 2 images and output 2 images as well
    def __init__(self, conv_dim=64, repeat_num=6, input_nc=6):
        super(Generator_makeup, self).__init__()

        layers = []
        layers.append(nn.Conv2d(input_nc, conv_dim, kernel_size=7, stride=1, padding=3, bias=False))
        layers.append(nn.InstanceNorm2d(conv_dim, affine=True))
        layers.append(nn.ReLU(inplace=True))

        # Down-Sampling
        curr_dim = conv_dim
        for i in range(2):
            layers.append(nn.Conv2d(curr_dim, curr_dim*2, kernel_size=4, stride=2, padding=1, bias=False))
            layers.append(nn.InstanceNorm2d(curr_dim*2, affine=True))
            layers.append(nn.ReLU(inplace=True))
            curr_dim = curr_dim * 2

        # Bottleneck
        for i in range(repeat_num):
            layers.append(ResidualBlock(dim_in=curr_dim, dim_out=curr_dim))

        # Up-Sampling
        for i in range(2):
            layers.append(nn.ConvTranspose2d(curr_dim, curr_dim//2, kernel_size=4, stride=2, padding=1, bias=False))
            layers.append(nn.InstanceNorm2d(curr_dim//2, affine=True))
            layers.append(nn.ReLU(inplace=True))
            curr_dim = curr_dim // 2

        self.main = nn.Sequential(*layers)

        layers_1 = []
        layers_1.append(nn.Conv2d(curr_dim, 3, kernel_size=7, stride=1, padding=3, bias=False))
        layers_1.append(nn.Tanh())
        self.branch_1 = nn.Sequential(*layers_1)
        layers_2 = []
        layers_2.append(nn.Conv2d(curr_dim, 3, kernel_size=7, stride=1, padding=3, bias=False))
        layers_2.append(nn.Tanh())
        self.branch_2 = nn.Sequential(*layers_2)

    def forward(self, x, y):
        input_x = torch.cat((x, y), dim=1)
        out = self.main(input_x)
        out_A = self.branch_1(out)
        out_B = self.branch_2(out)
        return out_A, out_B


class Generator_branch(nn.Module):
    """Generator. Encoder-Decoder Architecture."""
    # input 2 images and output 2 images as well
    def __init__(self, conv_dim=64, repeat_num=6, input_nc=3):
        super(Generator_branch, self).__init__()

        # Branch input
        layers_branch = []
        layers_branch.append(nn.Conv2d(input_nc, conv_dim, kernel_size=7, stride=1, padding=3, bias=False))
        layers_branch.append(nn.InstanceNorm2d(conv_dim, affine=True))
        layers_branch.append(nn.ReLU(inplace=True))
        layers_branch.append(nn.Conv2d(conv_dim, conv_dim*2, kernel_size=4, stride=2, padding=1, bias=False))
        layers_branch.append(nn.InstanceNorm2d(conv_dim*2, affine=True))
        layers_branch.append(nn.ReLU(inplace=True))
        self.Branch_0 = nn.Sequential(*layers_branch)

        # Branch input
        layers_branch = []
        layers_branch.append(nn.Conv2d(input_nc, conv_dim, kernel_size=7, stride=1, padding=3, bias=False))
        layers_branch.append(nn.InstanceNorm2d(conv_dim, affine=True))
        layers_branch.append(nn.ReLU(inplace=True))
        layers_branch.append(nn.Conv2d(conv_dim, conv_dim*2, kernel_size=4, stride=2, padding=1, bias=False))
        layers_branch.append(nn.InstanceNorm2d(conv_dim*2, affine=True))
        layers_branch.append(nn.ReLU(inplace=True))
        self.Branch_1 = nn.Sequential(*layers_branch)

        # Down-Sampling, branch merge
        layers = []
        curr_dim = conv_dim*2
        layers.append(nn.Conv2d(curr_dim*2, curr_dim*2, kernel_size=4, stride=2, padding=1, bias=False))
        layers.append(nn.InstanceNorm2d(curr_dim*2, affine=True))
        layers.append(nn.ReLU(inplace=True))
        curr_dim = curr_dim * 2
     
        # Bottleneck
        for i in range(repeat_num):
            layers.append(ResidualBlock(dim_in=curr_dim, dim_out=curr_dim))

        # Up-Sampling
        for i in range(2):
            layers.append(nn.ConvTranspose2d(curr_dim, curr_dim//2, kernel_size=4, stride=2, padding=1, bias=False))
            layers.append(nn.InstanceNorm2d(curr_dim//2, affine=True))
            layers.append(nn.ReLU(inplace=True))
            curr_dim = curr_dim // 2

        self.main = nn.Sequential(*layers)

        layers_1 = []
        layers_1.append(nn.Conv2d(curr_dim, curr_dim, kernel_size=3, stride=1, padding=1, bias=False))
        layers_1.append(nn.InstanceNorm2d(curr_dim, affine=True))
        layers_1.append(nn.ReLU(inplace=True))
        layers_1.append(nn.Conv2d(curr_dim, curr_dim, kernel_size=3, stride=1, padding=1, bias=False))
        layers_1.append(nn.InstanceNorm2d(curr_dim, affine=True))
        layers_1.append(nn.ReLU(inplace=True))
        layers_1.append(nn.Conv2d(curr_dim, 3, kernel_size=7, stride=1, padding=3, bias=False))
        layers_1.append(nn.Tanh())
        self.branch_1 = nn.Sequential(*layers_1)
        layers_2 = []
        layers_2.append(nn.Conv2d(curr_dim, curr_dim, kernel_size=3, stride=1, padding=1, bias=False))
        layers_2.append(nn.InstanceNorm2d(curr_dim, affine=True))
        layers_2.append(nn.ReLU(inplace=True))
        layers_2.append(nn.Conv2d(curr_dim, curr_dim, kernel_size=3, stride=1, padding=1, bias=False))
        layers_2.append(nn.InstanceNorm2d(curr_dim, affine=True))
        layers_2.append(nn.ReLU(inplace=True))
        layers_2.append(nn.Conv2d(curr_dim, 3, kernel_size=7, stride=1, padding=3, bias=False))
        layers_2.append(nn.Tanh())
        self.branch_2 = nn.Sequential(*layers_2)

    def forward(self, x, y):
        input_x = self.Branch_0(x)
        input_y = self.Branch_1(y)
        input_fuse = torch.cat((input_x, input_y), dim=1)
        out = self.main(input_fuse)
        out_A = self.branch_1(out)
        out_B = self.branch_2(out)
        return out_A, out_B

class Discriminator(nn.Module):
    """Discriminator. PatchGAN."""
    def __init__(self, image_size=128, conv_dim=64, repeat_num=3, norm='SN'):
        super(Discriminator, self).__init__()

        layers = []
        if norm=='SN':
            layers.append(SpectralNorm(nn.Conv2d(3, conv_dim, kernel_size=4, stride=2, padding=1)))
        else:
            layers.append(nn.Conv2d(3, conv_dim, kernel_size=4, stride=2, padding=1))
        layers.append(nn.LeakyReLU(0.01, inplace=True))

        curr_dim = conv_dim
        for i in range(1, repeat_num):
            if norm=='SN':
                layers.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim*2, kernel_size=4, stride=2, padding=1)))
            else:
                layers.append(nn.Conv2d(curr_dim, curr_dim*2, kernel_size=4, stride=2, padding=1))
            layers.append(nn.LeakyReLU(0.01, inplace=True))
            curr_dim = curr_dim * 2

        #k_size = int(image_size / np.power(2, repeat_num))
        if norm=='SN':
            layers.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim*2, kernel_size=4, stride=1, padding=1)))
        else:
            layers.append(nn.Conv2d(curr_dim, curr_dim*2, kernel_size=4, stride=1, padding=1))
        layers.append(nn.LeakyReLU(0.01, inplace=True))
        curr_dim = curr_dim *2

        self.main = nn.Sequential(*layers)
        if norm=='SN':
            self.conv1 = SpectralNorm(nn.Conv2d(curr_dim, 1, kernel_size=4, stride=1, padding=1, bias=False))
        else:
            self.conv1 = nn.Conv2d(curr_dim, 1, kernel_size=4, stride=1, padding=1, bias=False)

        # conv1 remain the last square size, 256*256-->30*30
        #self.conv2 = SpectralNorm(nn.Conv2d(curr_dim, 1, kernel_size=k_size, bias=False))
        #conv2 output a single number

    def forward(self, x):
        h = self.main(x)
        #out_real = self.conv1(h)
        out_makeup = self.conv1(h)
        #return out_real.squeeze(), out_makeup.squeeze()
        return out_makeup.squeeze()

def l2normalize(v, eps=1e-12):
    return v / (v.norm() + eps)

class SpectralNorm(object):
    def __init__(self):
        self.name = "weight"
        #print(self.name)
        self.power_iterations = 1

    def compute_weight(self, module):
        u = getattr(module, self.name + "_u")
        v = getattr(module, self.name + "_v")
        w = getattr(module, self.name + "_bar")

        height = w.data.shape[0]
        for _ in range(self.power_iterations):
            v.data = l2normalize(torch.mv(torch.t(w.view(height,-1).data), u.data))
            u.data = l2normalize(torch.mv(w.view(height,-1).data, v.data))
        # sigma = torch.dot(u.data, torch.mv(w.view(height,-1).data, v.data))
        sigma = u.dot(w.view(height, -1).mv(v))
        return w / sigma.expand_as(w)

    @staticmethod
    def apply(module):
        name = "weight"
        fn = SpectralNorm()

        try:
            u = getattr(module, name + "_u")
            v = getattr(module, name + "_v")
            w = getattr(module, name + "_bar")
        except AttributeError:
            w = getattr(module, name)
            height = w.data.shape[0]
            width = w.view(height, -1).data.shape[1]
            u = Parameter(w.data.new(height).normal_(0, 1), requires_grad=False)
            v = Parameter(w.data.new(width).normal_(0, 1), requires_grad=False)
            w_bar = Parameter(w.data)

            #del module._parameters[name]

            module.register_parameter(name + "_u", u)
            module.register_parameter(name + "_v", v)
            module.register_parameter(name + "_bar", w_bar)

        # remove w from parameter list
        del module._parameters[name]

        setattr(module, name, fn.compute_weight(module))

        # recompute weight before every forward()
        module.register_forward_pre_hook(fn)

        return fn

    def remove(self, module):
        weight = self.compute_weight(module)
        delattr(module, self.name)
        del module._parameters[self.name + '_u']
        del module._parameters[self.name + '_v']
        del module._parameters[self.name + '_bar']
        module.register_parameter(self.name, Parameter(weight.data))

    def __call__(self, module, inputs):
        setattr(module, self.name, self.compute_weight(module))

def spectral_norm(module):
    SpectralNorm.apply(module)
    return module

def remove_spectral_norm(module):
    name = 'weight'
    for k, hook in module._forward_pre_hooks.items():
        if isinstance(hook, SpectralNorm) and hook.name == name:
            hook.remove(module)
            del module._forward_pre_hooks[k]
            return module

    raise ValueError("spectral_norm of '{}' not found in {}"
                     .format(name, module))

parser.add_argument('--lambda_cls', default='1', type=float, help='the lambda_cls weight') 
parser.add_argument('--lambda_rec', default='10', type=int, help='lambda_A and lambda_B')
parser.add_argument('--lambda_vgg', default='5e-3', type=float, help='the param of vgg loss')
parser.add_argument('--lambda_his', default='1', type=float, help='histogram loss on lips') 
parser.add_argument('--lambda_eye', default='1', type=float, help='histogram loss on eyes equals to lambda_his*lambda_eye') 
parser.add_argument('--lambda_skin_1', default='0.1', type=float, help='histogram loss on skin equals to lambda_his* lambda_skin') 
parser.add_argument('--lambda_skin_2', default='0.1', type=float, help='histogram loss on skin equals to lambda_his* lambda_skin') 

parser.add_argument('--batch_size', default='1', type=int, help='batch_size')
parser.add_argument('--LR', default="2e-4", type=float, help='Learning rate')

def train_net():
    # enable cudnn
    cudnn.benchmark = True

    data_loaders = get_loader(dataset_config, config, mode="train")    # return train&test
    #get the solver
    if args.model == 'cycleGAN':
        solver = Solver_cycleGAN(data_loaders, config, dataset_config)
    elif args.model =='makeupGAN':
        solver = Solver_makeupGAN(data_loaders, config, dataset_config)
    else:
        print("model that not support")
        exit()
    solver.train()

import torch
import torch.nn as nn
from torch.autograd import Variable

class GANLoss(nn.Module):
    def __init__(self, use_lsgan=True, target_real_label=1.0, target_fake_label=0.0,
                 tensor=torch.FloatTensor):
        super(GANLoss, self).__init__()
        self.real_label = target_real_label
        self.fake_label = target_fake_label
        self.real_label_var = None
        self.fake_label_var = None
        self.Tensor = tensor
        if use_lsgan:
            self.loss = nn.MSELoss()
        else:
            self.loss = nn.BCELoss()

    def get_target_tensor(self, input, target_is_real):
        target_tensor = None
        if target_is_real:
            create_label = ((self.real_label_var is None) or
                            (self.real_label_var.numel() != input.numel()))
            if create_label:
                real_tensor = self.Tensor(input.size()).fill_(self.real_label)
                self.real_label_var = Variable(real_tensor, requires_grad=False)
            target_tensor = self.real_label_var
        else:
            create_label = ((self.fake_label_var is None) or
                            (self.fake_label_var.numel() != input.numel()))
            if create_label:
                fake_tensor = self.Tensor(input.size()).fill_(self.fake_label)
                self.fake_label_var = Variable(fake_tensor, requires_grad=False)
            target_tensor = self.fake_label_var
        return target_tensor

    def __call__(self, input, target_is_real):
        target_tensor = self.get_target_tensor(input, target_is_real)
        return self.loss(input, target_tensor)

def build_model(self):
    # Define generators and discriminators
    self.G_A = net.Generator(self.g_conv_dim, self.g_repeat_num) 
    self.G_B = net.Generator(self.g_conv_dim, self.g_repeat_num)
    self.D_A = net.Discriminator(self.img_size, self.d_conv_dim, self.d_repeat_num)
    self.D_B = net.Discriminator(self.img_size, self.d_conv_dim, self.d_repeat_num)
    self.criterionL1 = torch.nn.L1Loss()
    self.criterionGAN = GANLoss(use_lsgan=True, tensor =torch.cuda.FloatTensor)

    # Optimizers
    self.g_optimizer = torch.optim.Adam(itertools.chain(self.G_A.parameters(), self.G_B.parameters()),
                                            self.g_lr, [self.beta1, self.beta2])
    self.d_A_optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, self.D_A.parameters()), self.d_lr, [self.beta1, self.beta2])
    self.d_B_optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, self.D_B.parameters()), self.d_lr, [self.beta1, self.beta2])

    self.G_A.apply(self.weights_init_xavier)
    self.D_A.apply(self.weights_init_xavier)
    self.G_B.apply(self.weights_init_xavier)
    self.D_B.apply(self.weights_init_xavier)

    # Print networks
    #  self.print_network(self.E, 'E')
    self.print_network(self.G_A, 'G_A')
    self.print_network(self.D_A, 'D_A')
    self.print_network(self.G_B, 'G_B')
    self.print_network(self.D_B, 'D_B')

    if torch.cuda.is_available():
        self.G_A.cuda()
        self.G_B.cuda()
        self.D_A.cuda()
        self.D_B.cuda()


def train(self):
    """Train StarGAN within a single dataset."""
    # The number of iterations per epoch
    self.iters_per_epoch = len(self.data_loader_train)
    # Start with trained model if exists
    g_lr = self.g_lr
    d_lr = self.d_lr
    if self.checkpoint:
        start = int(self.checkpoint.split('_')[0])
    else:
        start = 0
    # Start training
    self.start_time = time.time()
    for self.e in range(start, self.num_epochs):
        for self.i, (img_A, img_B, _, _) in enumerate(self.data_loader_train):
            # Convert tensor to variable
            org_A = self.to_var(img_A, requires_grad=False)
            ref_B = self.to_var(img_B, requires_grad=False)

            # ================== Train D ================== #
            # training D_A
            # Real
            out = self.D_A(ref_B)
            d_loss_real = self.criterionGAN(out, True)
            # Fake
            fake = self.G_A(org_A)
            fake = Variable(fake.data)
            fake = fake.detach()
            out = self.D_A(fake)
            #d_loss_fake = self.get_D_loss(out, "fake")
            d_loss_fake =  self.criterionGAN(out, False)
            
            # Backward + Optimize
            d_loss = (d_loss_real + d_loss_fake) * 0.5
            self.d_A_optimizer.zero_grad()
            d_loss.backward(retain_graph=True)
            self.d_A_optimizer.step()

            # Logging
            self.loss = {}
            self.loss['D-A-loss_real'] = d_loss_real.item()

            # training D_B
            # Real
            out = self.D_B(org_A)
            d_loss_real = self.criterionGAN(out, True)
            # Fake
            fake = self.G_B(ref_B)
            fake = Variable(fake.data)
            fake = fake.detach()
            out = self.D_B(fake)
            #d_loss_fake = self.get_D_loss(out, "fake")
            d_loss_fake =  self.criterionGAN(out, False)
            
            # Backward + Optimize
            d_loss = (d_loss_real + d_loss_fake) * 0.5
            self.d_B_optimizer.zero_grad()
            d_loss.backward(retain_graph=True)
            self.d_B_optimizer.step()

            # Logging
            self.loss['D-B-loss_real'] = d_loss_real.item()

            # ================== Train G ================== #
            if (self.i + 1) % self.ndis == 0:
                # adversarial loss, i.e. L_trans,v in the paper 

                # identity loss
                if self.lambda_idt > 0:
                    # G_A should be identity if ref_B is fed
                    idt_A = self.G_A(ref_B)
                    loss_idt_A = self.criterionL1(idt_A, ref_B) * self.lambda_B * self.lambda_idt
                    # G_B should be identity if org_A is fed
                    idt_B = self.G_B(org_A)
                    loss_idt_B = self.criterionL1(idt_B, org_A) * self.lambda_A * self.lambda_idt
                    g_loss_idt = loss_idt_A + loss_idt_B
                else:
                    g_loss_idt = 0
                    
                # GAN loss D_A(G_A(A))
                fake_B = self.G_A(org_A)
                pred_fake = self.D_A(fake_B)
                g_A_loss_adv =  self.criterionGAN(pred_fake, True)
                #g_loss_adv = self.get_G_loss(out)

                # GAN loss D_B(G_B(B))
                fake_A = self.G_B(ref_B)
                pred_fake = self.D_B(fake_A)
                g_B_loss_adv = self.criterionGAN(pred_fake, True)

                # Forward cycle loss
                rec_A = self.G_B(fake_B)
                g_loss_rec_A = self.criterionL1(rec_A, org_A) * self.lambda_A

                # Backward cycle loss
                rec_B = self.G_A(fake_A)
                g_loss_rec_B = self.criterionL1(rec_B, ref_B) * self.lambda_B

                # Combined loss
                g_loss = g_A_loss_adv + g_B_loss_adv + g_loss_rec_A + g_loss_rec_B + g_loss_idt
                
                self.g_optimizer.zero_grad()
                g_loss.backward(retain_graph=True)
                self.g_optimizer.step()

                # Logging
                self.loss['G-A-loss_adv'] = g_A_loss_adv.item()
                self.loss['G-B-loss_adv'] = g_A_loss_adv.item()
                self.loss['G-loss_org'] = g_loss_rec_A.item()
                self.loss['G-loss_ref'] = g_loss_rec_B.item()
                self.loss['G-loss_idt'] = g_loss_idt.item()

            # Print out log info
            if (self.i + 1) % self.log_step == 0:
                self.log_terminal()

            #plot the figures
            for key_now in self.loss.keys():
                plot_fig.plot(key_now, self.loss[key_now])

            #save the images
            if (self.i + 1) % self.vis_step == 0:
                print("Saving middle output...")
                self.vis_train([org_A, ref_B, fake_A, fake_B, rec_A, rec_B])
                self.vis_test()

            # Save model checkpoints
            if (self.i + 1) % self.snapshot_step == 0:
                self.save_models()

            if (self.i % 100 == 99):
                plot_fig.flush(self.task_name)

            plot_fig.tick()
        
        # Decay learning rate
        if (self.e+1) > (self.num_epochs - self.num_epochs_decay):
            g_lr -= (self.g_lr / float(self.num_epochs_decay))
            d_lr -= (self.d_lr / float(self.num_epochs_decay))
            self.update_lr(g_lr, d_lr)
            print('Decay learning rate to g_lr: {}, d_lr:{}.'.format(g_lr, d_lr))

def build_model(self):
    # Define generators and discriminators
    if self.whichG=='normal':
        self.G = net.Generator_makeup(self.g_conv_dim, self.g_repeat_num)
    if self.whichG=='branch':
        self.G = net.Generator_branch(self.g_conv_dim, self.g_repeat_num)
    for i in self.cls:
        setattr(self, "D_" + i, net.Discriminator(self.img_size, self.d_conv_dim, self.d_repeat_num, self.norm))

    self.criterionL1 = torch.nn.L1Loss()
    self.criterionL2 = torch.nn.MSELoss()
    self.criterionGAN = GANLoss(use_lsgan=True, tensor =torch.cuda.FloatTensor)
    self.vgg = net.VGG()
    self.vgg.load_state_dict(torch.load('addings/vgg_conv.pth'))
    # Optimizers
    self.g_optimizer = torch.optim.Adam(self.G.parameters(), self.g_lr, [self.beta1, self.beta2])
    for i in self.cls:
        setattr(self, "d_" + i + "_optimizer", \
                torch.optim.Adam(filter(lambda p: p.requires_grad, getattr(self, "D_" + i).parameters()), \
                                    self.d_lr, [self.beta1, self.beta2]))

    # Weights initialization
    self.G.apply(self.weights_init_xavier)
    for i in self.cls:
        getattr(self, "D_" + i).apply(self.weights_init_xavier)

    # Print networks
    self.print_network(self.G, 'G')
    for i in self.cls:
        self.print_network(getattr(self, "D_" + i), "D_" + i)

    if torch.cuda.is_available():
        self.G.cuda()
        self.vgg.cuda()
        for i in self.cls:
            getattr(self, "D_" + i).cuda()


def train(self):
    """Train StarGAN within a single dataset."""
    # The number of iterations per epoch
    self.iters_per_epoch = len(self.data_loader_train)
    # Start with trained model if exists
    cls_A = self.cls[0]
    cls_B = self.cls[1]
    g_lr = self.g_lr
    d_lr = self.d_lr
    if self.checkpoint:
        start = int(self.checkpoint.split('_')[0])
        self.vis_test()
    else:
        start = 0
    # Start training
    self.start_time = time.time()
    for self.e in range(start, self.num_epochs):
        for self.i, (img_A, img_B, mask_A, mask_B) in enumerate(self.data_loader_train):
            # Convert tensor to variable
            # mask attribute: 0:background 1:face 2:left-eyebrown 3:right-eyebrown 4:left-eye 5: right-eye 6: nose 
            # 7: upper-lip 8: teeth 9: under-lip 10:hair 11: left-ear 12: right-ear 13: neck
            if self.checkpoint or self.direct:
                if self.lips==True:
                    mask_A_lip = (mask_A==7).float() + (mask_A==9).float()
                    mask_B_lip = (mask_B==7).float() + (mask_B==9).float()
                    mask_A_lip, mask_B_lip, index_A_lip, index_B_lip = self.mask_preprocess(mask_A_lip, mask_B_lip)
                if self.skin==True:
                    mask_A_skin = (mask_A==1).float() + (mask_A==6).float() + (mask_A==13).float()
                    mask_B_skin = (mask_B==1).float() + (mask_B==6).float() + (mask_B==13).float()
                    mask_A_skin, mask_B_skin, index_A_skin, index_B_skin = self.mask_preprocess(mask_A_skin, mask_B_skin)
                if self.eye==True:
                    mask_A_eye_left = (mask_A==4).float()
                    mask_A_eye_right = (mask_A==5).float()
                    mask_B_eye_left = (mask_B==4).float()
                    mask_B_eye_right = (mask_B==5).float()
                    mask_A_face = (mask_A==1).float() + (mask_A==6).float()
                    mask_B_face = (mask_B==1).float() + (mask_B==6).float()
                    # avoid the situation that images with eye closed
                    if not ((mask_A_eye_left>0).any() and (mask_B_eye_left>0).any() and \
                        (mask_A_eye_right > 0).any() and (mask_B_eye_right > 0).any()):
                        continue
                    mask_A_eye_left, mask_A_eye_right = self.rebound_box(mask_A_eye_left, mask_A_eye_right, mask_A_face)
                    mask_B_eye_left, mask_B_eye_right = self.rebound_box(mask_B_eye_left, mask_B_eye_right, mask_B_face)
                    mask_A_eye_left, mask_B_eye_left, index_A_eye_left, index_B_eye_left = \
                        self.mask_preprocess(mask_A_eye_left, mask_B_eye_left)
                    mask_A_eye_right, mask_B_eye_right, index_A_eye_right, index_B_eye_right = \
                        self.mask_preprocess(mask_A_eye_right, mask_B_eye_right)

            org_A = self.to_var(img_A, requires_grad=False)
            ref_B = self.to_var(img_B, requires_grad=False)
            # ================== Train D ================== #
            # training D_A, D_A aims to distinguish class B
            # Real
            out = getattr(self, "D_" + cls_A)(ref_B)
            d_loss_real = self.criterionGAN(out, True)
            # Fake
            fake_A, fake_B = self.G(org_A, ref_B)
            fake_A = Variable(fake_A.data).detach()
            fake_B = Variable(fake_B.data).detach()
            out = getattr(self, "D_" + cls_A)(fake_A)
            #d_loss_fake = self.get_D_loss(out, "fake")
            d_loss_fake =  self.criterionGAN(out, False)
            
            # Backward + Optimize
            d_loss = (d_loss_real + d_loss_fake) * 0.5
            getattr(self, "d_" + cls_A + "_optimizer").zero_grad()
            d_loss.backward(retain_graph=True)
            getattr(self, "d_" + cls_A + "_optimizer").step()

            # Logging
            self.loss = {}
            self.loss['D-A-loss_real'] = d_loss_real.item()

            # training D_B, D_B aims to distinguish class A
            # Real
            out = getattr(self, "D_" + cls_B)(org_A)
            d_loss_real = self.criterionGAN(out, True)
            # Fake
            out = getattr(self, "D_" + cls_B)(fake_B)
            #d_loss_fake = self.get_D_loss(out, "fake")
            d_loss_fake =  self.criterionGAN(out, False)
            
            # Backward + Optimize
            d_loss = (d_loss_real + d_loss_fake) * 0.5
            getattr(self, "d_" + cls_B + "_optimizer").zero_grad()
            d_loss.backward(retain_graph=True)
            getattr(self, "d_" + cls_B + "_optimizer").step()

            # Logging
            self.loss['D-B-loss_real'] = d_loss_real.item()

            # ================== Train G ================== #
            if (self.i + 1) % self.ndis == 0:
                # adversarial loss, i.e. L_trans,v in the paper 

                # identity loss
                if self.lambda_idt > 0:
                    # G should be identity if ref_B or org_A is fed
                    idt_A1, idt_A2 = self.G(org_A, org_A)
                    idt_B1, idt_B2 = self.G(ref_B, ref_B)
                    loss_idt_A1 = self.criterionL1(idt_A1, org_A) * self.lambda_A * self.lambda_idt
                    loss_idt_A2 = self.criterionL1(idt_A2, org_A) * self.lambda_A * self.lambda_idt
                    loss_idt_B1 = self.criterionL1(idt_B1, ref_B) * self.lambda_B * self.lambda_idt
                    loss_idt_B2 = self.criterionL1(idt_B2, ref_B) * self.lambda_B * self.lambda_idt
                    # loss_idt
                    loss_idt = (loss_idt_A1 + loss_idt_A2 + loss_idt_B1 + loss_idt_B2) * 0.5
                else:
                    loss_idt = 0
                    
                # GAN loss D_A(G_A(A))
                # fake_A in class B, 
                fake_A, fake_B = self.G(org_A, ref_B)
                pred_fake = getattr(self, "D_" + cls_A)(fake_A)
                g_A_loss_adv = self.criterionGAN(pred_fake, True)
                #g_loss_adv = self.get_G_loss(out)
                # GAN loss D_B(G_B(B))
                pred_fake = getattr(self, "D_" + cls_B)(fake_B)
                g_B_loss_adv = self.criterionGAN(pred_fake, True)
                rec_B, rec_A = self.G(fake_B, fake_A)

                # color_histogram loss
                g_A_loss_his = 0
                g_B_loss_his = 0
                if self.checkpoint or self.direct:
                    if self.lips==True:
                        g_A_lip_loss_his = self.criterionHis(fake_A, ref_B, mask_A_lip, mask_B_lip, index_A_lip) * self.lambda_his_lip
                        g_B_lip_loss_his = self.criterionHis(fake_B, org_A, mask_B_lip, mask_A_lip, index_B_lip) * self.lambda_his_lip
                        g_A_loss_his += g_A_lip_loss_his
                        g_B_loss_his += g_B_lip_loss_his
                    if self.skin==True:
                        g_A_skin_loss_his = self.criterionHis(fake_A, ref_B, mask_A_skin, mask_B_skin, index_A_skin) * self.lambda_his_skin_1
                        g_B_skin_loss_his = self.criterionHis(fake_B, org_A, mask_B_skin, mask_A_skin, index_B_skin) * self.lambda_his_skin_2
                        g_A_loss_his += g_A_skin_loss_his
                        g_B_loss_his += g_B_skin_loss_his
                    if self.eye==True:
                        g_A_eye_left_loss_his = self.criterionHis(fake_A, ref_B, mask_A_eye_left, mask_B_eye_left, index_A_eye_left) * self.lambda_his_eye
                        g_B_eye_left_loss_his = self.criterionHis(fake_B, org_A, mask_B_eye_left, mask_A_eye_left, index_B_eye_left) * self.lambda_his_eye
                        g_A_eye_right_loss_his = self.criterionHis(fake_A, ref_B, mask_A_eye_right, mask_B_eye_right, index_A_eye_right) * self.lambda_his_eye
                        g_B_eye_right_loss_his = self.criterionHis(fake_B, org_A, mask_B_eye_right, mask_A_eye_right, index_B_eye_right) * self.lambda_his_eye
                        g_A_loss_his += g_A_eye_left_loss_his + g_A_eye_right_loss_his
                        g_B_loss_his += g_B_eye_left_loss_his + g_B_eye_right_loss_his

                # cycle loss
                g_loss_rec_A = self.criterionL1(rec_A, org_A) * self.lambda_A
                g_loss_rec_B = self.criterionL1(rec_B, ref_B) * self.lambda_B

                # vgg loss
                vgg_org = self.vgg(org_A, self.content_layer)[0]
                vgg_org = Variable(vgg_org.data).detach()
                vgg_fake_A = self.vgg(fake_A, self.content_layer)[0]
                g_loss_A_vgg = self.criterionL2(vgg_fake_A, vgg_org) * self.lambda_A * self.lambda_vgg
                
                vgg_ref = self.vgg(ref_B, self.content_layer)[0]
                vgg_ref = Variable(vgg_ref.data).detach()
                vgg_fake_B = self.vgg(fake_B, self.content_layer)[0]
                g_loss_B_vgg = self.criterionL2(vgg_fake_B, vgg_ref) * self.lambda_B * self.lambda_vgg
                
                loss_rec = (g_loss_rec_A + g_loss_rec_B + g_loss_A_vgg + g_loss_B_vgg) * 0.5
                
                # Combined loss
                g_loss = g_A_loss_adv + g_B_loss_adv + loss_rec + loss_idt
                if self.checkpoint or self.direct:
                    g_loss = g_A_loss_adv + g_B_loss_adv + loss_rec + loss_idt + g_A_loss_his + g_B_loss_his
                
                self.g_optimizer.zero_grad()
                g_loss.backward(retain_graph=True)
                self.g_optimizer.step()

                # Logging
                self.loss['G-A-loss-adv'] = g_A_loss_adv.item()
                self.loss['G-B-loss-adv'] = g_A_loss_adv.item()
                self.loss['G-loss-org'] = g_loss_rec_A.item()
                self.loss['G-loss-ref'] = g_loss_rec_B.item()
                self.loss['G-loss-idt'] = loss_idt.item()
                self.loss['G-loss-img-rec'] = (g_loss_rec_A + g_loss_rec_B).item()
                self.loss['G-loss-vgg-rec'] = (g_loss_A_vgg + g_loss_B_vgg).item()
                if self.direct:
                    self.loss['G-A-loss-his'] = g_A_loss_his.item()
                    self.loss['G-B-loss-his'] = g_B_loss_his.item()

            # Print out log info
            if (self.i + 1) % self.log_step == 0:
                self.log_terminal()

            #plot the figures
            for key_now in self.loss.keys():
                plot_fig.plot(key_now, self.loss[key_now])

            #save the images
            if (self.i + 1) % self.vis_step == 0:
                print("Saving middle output...")
                self.vis_train([org_A, ref_B, fake_A, fake_B, rec_A, rec_B])


            # Save model checkpoints
            if (self.i + 1) % self.snapshot_step == 0:
                self.save_models()

            if (self.i % 100 == 99):
                plot_fig.flush(self.task_name)

            plot_fig.tick()
        
        # Decay learning rate
        if (self.e+1) > (self.num_epochs - self.num_epochs_decay):
            g_lr -= (self.g_lr / float(self.num_epochs_decay))
            d_lr -= (self.d_lr / float(self.num_epochs_decay))
            self.update_lr(g_lr, d_lr)
            print('Decay learning rate to g_lr: {}, d_lr:{}.'.format(g_lr, d_lr))

import torch
import torch.nn as nn
import torch.nn.functional as F

from ops.spectral_norm import spectral_norm as SpectralNorm

# Defines the GAN loss which uses either LSGAN or the regular GAN.
# When LSGAN is used, it is basically same as MSELoss,
# but it abstracts away the need to create the target label tensor
# that has the same size as the input

class ResidualBlock(nn.Module):
    """Residual Block."""
    def __init__(self, dim_in, dim_out):
        super(ResidualBlock, self).__init__()
        self.main = nn.Sequential(
            nn.Conv2d(dim_in, dim_out, kernel_size=3, stride=1, padding=1, bias=False),
            nn.InstanceNorm2d(dim_out, affine=True),
            nn.ReLU(inplace=True),
            nn.Conv2d(dim_out, dim_out, kernel_size=3, stride=1, padding=1, bias=False),
            nn.InstanceNorm2d(dim_out, affine=True))

    def forward(self, x):
        return x + self.main(x)


class Generator(nn.Module):
    """Generator. Encoder-Decoder Architecture."""
    def __init__(self, conv_dim=64, repeat_num=6):
        super(Generator, self).__init__()

        layers = []
        layers.append(nn.Conv2d(3, conv_dim, kernel_size=7, stride=1, padding=3, bias=False))
        layers.append(nn.InstanceNorm2d(conv_dim, affine=True))
        layers.append(nn.ReLU(inplace=True))

        # Down-Sampling
        curr_dim = conv_dim
        for i in range(2):
            layers.append(nn.Conv2d(curr_dim, curr_dim*2, kernel_size=4, stride=2, padding=1, bias=False))
            layers.append(nn.InstanceNorm2d(curr_dim*2, affine=True))
            layers.append(nn.ReLU(inplace=True))
            curr_dim = curr_dim * 2

        # Bottleneck
        for i in range(repeat_num):
            layers.append(ResidualBlock(dim_in=curr_dim, dim_out=curr_dim))

        # Up-Sampling
        for i in range(2):
            layers.append(nn.ConvTranspose2d(curr_dim, curr_dim//2, kernel_size=4, stride=2, padding=1, bias=False))
            layers.append(nn.InstanceNorm2d(curr_dim//2, affine=True))
            layers.append(nn.ReLU(inplace=True))
            curr_dim = curr_dim // 2

        layers.append(nn.Conv2d(curr_dim, 3, kernel_size=7, stride=1, padding=3, bias=False))
        layers.append(nn.Tanh())
        self.main = nn.Sequential(*layers)

    def forward(self, x):
        out = self.main(x)
        return out

class Generator_makeup(nn.Module):
    """Generator. Encoder-Decoder Architecture."""
    # input 2 images and output 2 images as well
    def __init__(self, conv_dim=64, repeat_num=6, input_nc=6):
        super(Generator_makeup, self).__init__()

        layers = []
        layers.append(nn.Conv2d(input_nc, conv_dim, kernel_size=7, stride=1, padding=3, bias=False))
        layers.append(nn.InstanceNorm2d(conv_dim, affine=True))
        layers.append(nn.ReLU(inplace=True))

        # Down-Sampling
        curr_dim = conv_dim
        for i in range(2):
            layers.append(nn.Conv2d(curr_dim, curr_dim*2, kernel_size=4, stride=2, padding=1, bias=False))
            layers.append(nn.InstanceNorm2d(curr_dim*2, affine=True))
            layers.append(nn.ReLU(inplace=True))
            curr_dim = curr_dim * 2

        # Bottleneck
        for i in range(repeat_num):
            layers.append(ResidualBlock(dim_in=curr_dim, dim_out=curr_dim))

        # Up-Sampling
        for i in range(2):
            layers.append(nn.ConvTranspose2d(curr_dim, curr_dim//2, kernel_size=4, stride=2, padding=1, bias=False))
            layers.append(nn.InstanceNorm2d(curr_dim//2, affine=True))
            layers.append(nn.ReLU(inplace=True))
            curr_dim = curr_dim // 2

        self.main = nn.Sequential(*layers)

        layers_1 = []
        layers_1.append(nn.Conv2d(curr_dim, 3, kernel_size=7, stride=1, padding=3, bias=False))
        layers_1.append(nn.Tanh())
        self.branch_1 = nn.Sequential(*layers_1)
        layers_2 = []
        layers_2.append(nn.Conv2d(curr_dim, 3, kernel_size=7, stride=1, padding=3, bias=False))
        layers_2.append(nn.Tanh())
        self.branch_2 = nn.Sequential(*layers_2)

    def forward(self, x, y):
        input_x = torch.cat((x, y), dim=1)
        out = self.main(input_x)
        out_A = self.branch_1(out)
        out_B = self.branch_2(out)
        return out_A, out_B


class Generator_branch(nn.Module):
    """Generator. Encoder-Decoder Architecture."""
    # input 2 images and output 2 images as well
    def __init__(self, conv_dim=64, repeat_num=6, input_nc=3):
        super(Generator_branch, self).__init__()

        # Branch input
        layers_branch = []
        layers_branch.append(nn.Conv2d(input_nc, conv_dim, kernel_size=7, stride=1, padding=3, bias=False))
        layers_branch.append(nn.InstanceNorm2d(conv_dim, affine=True))
        layers_branch.append(nn.ReLU(inplace=True))
        layers_branch.append(nn.Conv2d(conv_dim, conv_dim*2, kernel_size=4, stride=2, padding=1, bias=False))
        layers_branch.append(nn.InstanceNorm2d(conv_dim*2, affine=True))
        layers_branch.append(nn.ReLU(inplace=True))
        self.Branch_0 = nn.Sequential(*layers_branch)

        # Branch input
        layers_branch = []
        layers_branch.append(nn.Conv2d(input_nc, conv_dim, kernel_size=7, stride=1, padding=3, bias=False))
        layers_branch.append(nn.InstanceNorm2d(conv_dim, affine=True))
        layers_branch.append(nn.ReLU(inplace=True))
        layers_branch.append(nn.Conv2d(conv_dim, conv_dim*2, kernel_size=4, stride=2, padding=1, bias=False))
        layers_branch.append(nn.InstanceNorm2d(conv_dim*2, affine=True))
        layers_branch.append(nn.ReLU(inplace=True))
        self.Branch_1 = nn.Sequential(*layers_branch)

        # Down-Sampling, branch merge
        layers = []
        curr_dim = conv_dim*2
        layers.append(nn.Conv2d(curr_dim*2, curr_dim*2, kernel_size=4, stride=2, padding=1, bias=False))
        layers.append(nn.InstanceNorm2d(curr_dim*2, affine=True))
        layers.append(nn.ReLU(inplace=True))
        curr_dim = curr_dim * 2
     
        # Bottleneck
        for i in range(repeat_num):
            layers.append(ResidualBlock(dim_in=curr_dim, dim_out=curr_dim))

        # Up-Sampling
        for i in range(2):
            layers.append(nn.ConvTranspose2d(curr_dim, curr_dim//2, kernel_size=4, stride=2, padding=1, bias=False))
            layers.append(nn.InstanceNorm2d(curr_dim//2, affine=True))
            layers.append(nn.ReLU(inplace=True))
            curr_dim = curr_dim // 2

        self.main = nn.Sequential(*layers)

        layers_1 = []
        layers_1.append(nn.Conv2d(curr_dim, curr_dim, kernel_size=3, stride=1, padding=1, bias=False))
        layers_1.append(nn.InstanceNorm2d(curr_dim, affine=True))
        layers_1.append(nn.ReLU(inplace=True))
        layers_1.append(nn.Conv2d(curr_dim, curr_dim, kernel_size=3, stride=1, padding=1, bias=False))
        layers_1.append(nn.InstanceNorm2d(curr_dim, affine=True))
        layers_1.append(nn.ReLU(inplace=True))
        layers_1.append(nn.Conv2d(curr_dim, 3, kernel_size=7, stride=1, padding=3, bias=False))
        layers_1.append(nn.Tanh())
        self.branch_1 = nn.Sequential(*layers_1)
        layers_2 = []
        layers_2.append(nn.Conv2d(curr_dim, curr_dim, kernel_size=3, stride=1, padding=1, bias=False))
        layers_2.append(nn.InstanceNorm2d(curr_dim, affine=True))
        layers_2.append(nn.ReLU(inplace=True))
        layers_2.append(nn.Conv2d(curr_dim, curr_dim, kernel_size=3, stride=1, padding=1, bias=False))
        layers_2.append(nn.InstanceNorm2d(curr_dim, affine=True))
        layers_2.append(nn.ReLU(inplace=True))
        layers_2.append(nn.Conv2d(curr_dim, 3, kernel_size=7, stride=1, padding=3, bias=False))
        layers_2.append(nn.Tanh())
        self.branch_2 = nn.Sequential(*layers_2)

    def forward(self, x, y):
        input_x = self.Branch_0(x)
        input_y = self.Branch_1(y)
        input_fuse = torch.cat((input_x, input_y), dim=1)
        out = self.main(input_fuse)
        out_A = self.branch_1(out)
        out_B = self.branch_2(out)
        return out_A, out_B

class Discriminator(nn.Module):
    """Discriminator. PatchGAN."""
    def __init__(self, image_size=128, conv_dim=64, repeat_num=3, norm='SN'):
        super(Discriminator, self).__init__()

        layers = []
        if norm=='SN':
            layers.append(SpectralNorm(nn.Conv2d(3, conv_dim, kernel_size=4, stride=2, padding=1)))
        else:
            layers.append(nn.Conv2d(3, conv_dim, kernel_size=4, stride=2, padding=1))
        layers.append(nn.LeakyReLU(0.01, inplace=True))

        curr_dim = conv_dim
        for i in range(1, repeat_num):
            if norm=='SN':
                layers.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim*2, kernel_size=4, stride=2, padding=1)))
            else:
                layers.append(nn.Conv2d(curr_dim, curr_dim*2, kernel_size=4, stride=2, padding=1))
            layers.append(nn.LeakyReLU(0.01, inplace=True))
            curr_dim = curr_dim * 2

        #k_size = int(image_size / np.power(2, repeat_num))
        if norm=='SN':
            layers.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim*2, kernel_size=4, stride=1, padding=1)))
        else:
            layers.append(nn.Conv2d(curr_dim, curr_dim*2, kernel_size=4, stride=1, padding=1))
        layers.append(nn.LeakyReLU(0.01, inplace=True))
        curr_dim = curr_dim *2

        self.main = nn.Sequential(*layers)
        if norm=='SN':
            self.conv1 = SpectralNorm(nn.Conv2d(curr_dim, 1, kernel_size=4, stride=1, padding=1, bias=False))
        else:
            self.conv1 = nn.Conv2d(curr_dim, 1, kernel_size=4, stride=1, padding=1, bias=False)

        # conv1 remain the last square size, 256*256-->30*30
        #self.conv2 = SpectralNorm(nn.Conv2d(curr_dim, 1, kernel_size=k_size, bias=False))
        #conv2 output a single number

    def forward(self, x):
        h = self.main(x)
        #out_real = self.conv1(h)
        out_makeup = self.conv1(h)
        #return out_real.squeeze(), out_makeup.squeeze()
        return out_makeup.squeeze()

class VGG(nn.Module):
    def __init__(self, pool='max'):
        super(VGG, self).__init__()
        # vgg modules
        self.conv1_1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
        self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, padding=1)
        self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        self.conv3_3 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        self.conv3_4 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, padding=1)
        self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.conv4_3 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.conv4_4 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.conv5_1 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.conv5_2 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.conv5_3 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.conv5_4 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        if pool == 'max':
            self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
            self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
            self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
            self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)
            self.pool5 = nn.MaxPool2d(kernel_size=2, stride=2)
        elif pool == 'avg':
            self.pool1 = nn.AvgPool2d(kernel_size=2, stride=2)
            self.pool2 = nn.AvgPool2d(kernel_size=2, stride=2)
            self.pool3 = nn.AvgPool2d(kernel_size=2, stride=2)
            self.pool4 = nn.AvgPool2d(kernel_size=2, stride=2)
            self.pool5 = nn.AvgPool2d(kernel_size=2, stride=2)

    def forward(self, x, out_keys):
        out = {}
        out['r11'] = F.relu(self.conv1_1(x))
        out['r12'] = F.relu(self.conv1_2(out['r11']))
        out['p1'] = self.pool1(out['r12'])
        out['r21'] = F.relu(self.conv2_1(out['p1']))
        out['r22'] = F.relu(self.conv2_2(out['r21']))
        out['p2'] = self.pool2(out['r22'])
        out['r31'] = F.relu(self.conv3_1(out['p2']))
        out['r32'] = F.relu(self.conv3_2(out['r31']))
        out['r33'] = F.relu(self.conv3_3(out['r32']))
        out['r34'] = F.relu(self.conv3_4(out['r33']))
        out['p3'] = self.pool3(out['r34'])
        out['r41'] = F.relu(self.conv4_1(out['p3']))
        
        out['r42'] = F.relu(self.conv4_2(out['r41']))
        out['r43'] = F.relu(self.conv4_3(out['r42']))
        out['r44'] = F.relu(self.conv4_4(out['r43']))
        out['p4'] = self.pool4(out['r44'])
        out['r51'] = F.relu(self.conv5_1(out['p4']))
        out['r52'] = F.relu(self.conv5_2(out['r51']))
        out['r53'] = F.relu(self.conv5_3(out['r52']))
        out['r54'] = F.relu(self.conv5_4(out['r53']))
        out['p5'] = self.pool5(out['r54'])
        
        return [out[key] for key in out_keys]

class WeightNet(nn.Module):

    def __init__(self, in_channel, out_channel, hidden_unit = [8, 8]):
        super(WeightNet, self).__init__()

        self.mlp_convs = nn.ModuleList()
        self.mlp_bns = nn.ModuleList()
        if hidden_unit is None or len(hidden_unit) == 0:
            self.mlp_convs.append(nn.Conv2d(in_channel, out_channel, 1))
            self.mlp_bns.append(nn.BatchNorm2d(out_channel))
        else:
            self.mlp_convs.append(nn.Conv2d(in_channel, hidden_unit[0], 1))
            self.mlp_bns.append(nn.BatchNorm2d(hidden_unit[0]))
            for i in range(1, len(hidden_unit)):
                self.mlp_convs.append(nn.Conv2d(hidden_unit[i - 1], hidden_unit[i], 1))
                self.mlp_bns.append(nn.BatchNorm2d(hidden_unit[i]))
            self.mlp_convs.append(nn.Conv2d(hidden_unit[-1], out_channel, 1))
            self.mlp_bns.append(nn.BatchNorm2d(out_channel))
        
    def forward(self, localized_xyz):
        #xyz : BxCxKxN

        weights = localized_xyz
        for i, conv in enumerate(self.mlp_convs):
            bn = self.mlp_bns[i]
            weights =  F.relu(bn(conv(weights)))

        return weights

def compute_density(xyz, bandwidth):
    '''
    xyz: input points position data, [B, N, C]
    '''
    #import ipdb; ipdb.set_trace()
    B, N, C = xyz.shape
    sqrdists = square_distance(xyz, xyz)
    gaussion_density = torch.exp(- sqrdists / (2.0 * bandwidth * bandwidth)) / (2.5 * bandwidth)
    xyz_density = gaussion_density.mean(dim = -1)

    return xyz_density

class DensityNet(nn.Module):
    def __init__(self, hidden_unit = [16, 8]):
        super(DensityNet, self).__init__()
        self.mlp_convs = nn.ModuleList()
        self.mlp_bns = nn.ModuleList() 

        self.mlp_convs.append(nn.Conv2d(1, hidden_unit[0], 1))
        self.mlp_bns.append(nn.BatchNorm2d(hidden_unit[0]))
        for i in range(1, len(hidden_unit)):
            self.mlp_convs.append(nn.Conv2d(hidden_unit[i - 1], hidden_unit[i], 1))
            self.mlp_bns.append(nn.BatchNorm2d(hidden_unit[i]))
        self.mlp_convs.append(nn.Conv2d(hidden_unit[-1], 1, 1))
        self.mlp_bns.append(nn.BatchNorm2d(1))

    def forward(self, density_scale):
        for i, conv in enumerate(self.mlp_convs):
            bn = self.mlp_bns[i]
            density_scale =  bn(conv(density_scale))
            if i == len(self.mlp_convs):
                density_scale = F.sigmoid(density_scale)
            else:
                density_scale = F.relu(density_scale)
        
        return density_scale

class PointConvDensitySetAbstraction(nn.Module):
    def __init__(self, npoint, nsample, in_channel, mlp, bandwidth, group_all):
        super(PointConvDensitySetAbstraction, self).__init__()
        self.npoint = npoint
        self.nsample = nsample
        self.mlp_convs = nn.ModuleList()
        self.mlp_bns = nn.ModuleList()
        last_channel = in_channel
        for out_channel in mlp:
            self.mlp_convs.append(nn.Conv2d(last_channel, out_channel, 1))
            self.mlp_bns.append(nn.BatchNorm2d(out_channel))
            last_channel = out_channel

        self.weightnet = WeightNet(3, 16)
        self.linear = nn.Linear(16 * mlp[-1], mlp[-1])
        self.bn_linear = nn.BatchNorm1d(mlp[-1])
        self.densitynet = DensityNet()
        self.group_all = group_all
        self.bandwidth = bandwidth

    def forward(self, xyz, points):
        """
        Input:
            xyz: input points position data, [B, C, N]
            points: input points data, [B, D, N]
        Return:
            new_xyz: sampled points position data, [B, C, S]
            new_points_concat: sample points feature data, [B, D', S]
        """
        B = xyz.shape[0]
        N = xyz.shape[2]
        xyz = xyz.permute(0, 2, 1)
        if points is not None:
            points = points.permute(0, 2, 1)

        xyz_density = compute_density(xyz, self.bandwidth)
        inverse_density = 1.0 / xyz_density 

        if self.group_all:
            new_xyz, new_points, grouped_xyz_norm, grouped_density = sample_and_group_all(xyz, points, inverse_density.view(B, N, 1))
        else:
            new_xyz, new_points, grouped_xyz_norm, _, grouped_density = sample_and_group(self.npoint, self.nsample, xyz, points, inverse_density.view(B, N, 1))
        # new_xyz: sampled points position data, [B, npoint, C]
        # new_points: sampled points data, [B, npoint, nsample, C+D]
        new_points = new_points.permute(0, 3, 2, 1) # [B, C+D, nsample,npoint]
        for i, conv in enumerate(self.mlp_convs):
            bn = self.mlp_bns[i]
            new_points =  F.relu(bn(conv(new_points)))

        inverse_max_density = grouped_density.max(dim = 2, keepdim=True)[0]
        density_scale = grouped_density / inverse_max_density
        density_scale = self.densitynet(density_scale.permute(0, 3, 2, 1))
        new_points = new_points * density_scale

        grouped_xyz = grouped_xyz_norm.permute(0, 3, 2, 1)
        weights = self.weightnet(grouped_xyz)     
        new_points = torch.matmul(input=new_points.permute(0, 3, 1, 2), other = weights.permute(0, 3, 2, 1)).view(B, self.npoint, -1)
        new_points = self.linear(new_points)
        new_points = self.bn_linear(new_points.permute(0, 2, 1))
        new_points = F.relu(new_points)
        new_xyz = new_xyz.permute(0, 2, 1)

        return new_xyz, new_points

class PointConvDensityClsSsg(nn.Module):
    def __init__(self, num_classes = 40):
        super(PointConvDensityClsSsg, self).__init__()
        feature_dim = 3
        self.sa1 = PointConvDensitySetAbstraction(npoint=512, nsample=32, in_channel=feature_dim + 3, mlp=[64, 64, 128], bandwidth = 0.1, group_all=False)
        self.sa2 = PointConvDensitySetAbstraction(npoint=128, nsample=64, in_channel=128 + 3, mlp=[128, 128, 256], bandwidth = 0.2, group_all=False)
        self.sa3 = PointConvDensitySetAbstraction(npoint=1, nsample=None, in_channel=256 + 3, mlp=[256, 512, 1024], bandwidth = 0.4, group_all=True)
        self.fc1 = nn.Linear(1024, 512)
        self.bn1 = nn.BatchNorm1d(512)
        self.drop1 = nn.Dropout(0.7)
        self.fc2 = nn.Linear(512, 256)
        self.bn2 = nn.BatchNorm1d(256)
        self.drop2 = nn.Dropout(0.7)
        self.fc3 = nn.Linear(256, num_classes)

    def forward(self, xyz, feat):
        B, _, _ = xyz.shape
        l1_xyz, l1_points = self.sa1(xyz, feat)
        l2_xyz, l2_points = self.sa2(l1_xyz, l1_points)
        l3_xyz, l3_points = self.sa3(l2_xyz, l2_points)
        x = l3_points.view(B, 1024)
        x = self.drop1(F.relu(self.bn1(self.fc1(x))))
        x = self.drop2(F.relu(self.bn2(self.fc2(x))))
        x = self.fc3(x)
        x = F.log_softmax(x, -1)
        return x