methods/alphaface/Models/Swapper_AlphaFace.py

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

import pytorch_lightning as pl
from omegaconf import OmegaConf
from torch import nn
from Models.Swapper_Units_AdIN import Encoder, Discriminator, Decoder, Encoder_noBNIN, VGGPerceptualLoss
from Models.arcface_resnet import resnet50
# from Model.MultiScaleDiscriminator import MultiscaleDiscriminator
from Models.iresnet import iresnet100
from backbones import get_model
import numpy as np

# from Model.loss import GANLoss, AEI_Loss
batch_size = 1


class Normalize(nn.Module):
    def __init__(self, mean, std):
        super(Normalize, self).__init__()
        self.mean = mean
        self.std = std

    def forward(self, x):
        x = x - self.mean
        x = x / self.std
        return x

class ResNet_ID_encoder(pl.LightningModule):
    def __init__(self, ema_path):
        super(ResNet_ID_encoder, self).__init__()
        self.resnet = resnet50()
        #pdb.set_trace()

        self.ema = nn.Linear(512, 512)
        self.ema.weight.data = torch.nn.Parameter(torch.from_numpy(np.load(ema_path)).permute(1,0)).cuda()
        self.ema.bias.data.fill_(0.0) 
        #self.ema = torch.from_numpy(np.load(ema_path)).cuda()

    def processing_ema_only(self,latent):
        x = self.ema(latent)
        # x = torch.matmul(x,self.ema)
        out = torch.div(x, torch.linalg.norm(x, dim=1, keepdim=True))
        return out

    def forward(self, inputs):
        x = self.resnet(inputs)
        x = self.ema(x)
        #x = torch.matmul(x,self.ema)
        out = torch.div(x, torch.linalg.norm(x, dim=1, keepdim=True))
        return out

class Swapper(nn.Module): #class Swapper_Enlarge(pl.LightningModule):
    def __init__(self, source_dim,exist_BN=True):
        super(Swapper, self).__init__()
        self.source_dim = source_dim
        if exist_BN:
            self.E = Encoder(self.source_dim)
        else:
            self.E = Encoder_noBNIN(self.source_dim)
        self.G = Decoder(1024, 3)
        self.dis = None
        self.train_adv = None


    def forward(self, target, source,get_latent=False):
        # source = np.dot(source, self.ema)
        # source /= np.linalg.norm(source)
        output = self.E(target, source)
        if get_latent==True:
            output,latent = self.G(output,get_latent=get_latent)
            return output,latent 
        else:
            output = self.G(output)
        return output

    def forward(self, target, source,get_latent=False):
        # source = np.dot(source, self.ema)
        # source /= np.linalg.norm(source)
        output = self.E(target, source)
        if get_latent==True:
            output,latent = self.G(output,get_latent=get_latent)
            return output,latent 
        else:
            output = self.G(output)
        return output
    


class AlphaFace(pl.LightningModule):
    def __init__(self, Swapper, id_encoder,fine_tune=False):
        super(AlphaFace, self).__init__()
        self.Swapper = Swapper.cuda()
        #self.Id_encoder = None
        self.fine_tune = fine_tune
        self.Id_encoder = id_encoder.cuda()

        self.dis = None
        self.train_adv = False
        
        self.feats_extractor = None

    def Prepareing_adversrial_learning(self,img_size,max_conv_size):
        self.dis = Discriminator(img_size=img_size, max_conv_dim=max_conv_size).cuda()
        self.train_adv = True
        
    def Preparing_VGG_percet_loss(self):
        self.feats_extractor = VGGPerceptualLoss(layer_ids=[3, 8, 15, 22]).cuda()
        

    def set_grads(self):
        swapper_param_list = list(self.Swapper.parameters())

        for param in swapper_param_list:
            param.requires_grad = True

        
        idc_param_list = list(self.Id_encoder.parameters())
        for param in idc_param_list:
            param.requires_grad = False

        # if self.fine_tune==False:
        #     idc_param_list = list(self.Id_encoder.parameters())
        #     for param in idc_param_list:
        #         param.requires_grad = False

        if self.train_adv==True:
            adv_param_list = list(self.dis.parameters())
            for param in adv_param_list:
                param.requires_grad = True

    def save(self, fname, step):
        print('Saving checkpoint into %s...' % fname)
        PATH_Swapper = fname + '_swapper_%d' % (step) + '.pt'
        PATH_IDC = fname + '_idc_%d' % (step) + '.pt'
        PATH_DIS = fname + '_dis_%d' % (step) + '.pt'

        PATH_Swapper_last = fname + '_swapper_last' + '.pt'
        PATH_IDC_last = fname + '_idc_last' + '.pt'
        PATH_DIS_last = fname + '_dis_last' + '.pt'

        # Save the mask module
        torch.save(self.Swapper.state_dict(), PATH_Swapper)
        torch.save(self.Swapper.state_dict(), PATH_Swapper_last)
        torch.save(self.ID_encoder.state_dict(), PATH_IDC)
        torch.save(self.ID_encoder.state_dict(), PATH_IDC_last)

        # Save discriminator
        if self.train_adv==True:
            torch.save(self.discriminator.state_dict(), PATH_DIS)
            torch.save(self.discriminator.state_dict(), PATH_DIS_last)

        # Save the mask module


    def load(self, fname, step=None):
        print('Loading checkpoint from %s...' % fname)
        if step == None:
            PATH_Swapper = fname + '_swapper_last' + '.pt'
            PATH_IDC = fname + '_idc_last' + '.pt'
            PATH_DIS = fname + '_dis_last' + '.pt'
        else:
            PATH_Swapper = fname + '_swapper_%d' % (step) + '.pt'
            PATH_IDC = fname + '_idc_%d' % (step) + '.pt'
            PATH_DIS = fname + '_dis_%d' % (step) + '.pt'

        # Save the mask module
        self.Swapper.load_state_dict(torch.load(PATH_Swapper))
        # Save discriminator
        self.dis.load_state_dict(torch.load(PATH_DIS))
        if self.fine_tune==False:
            self.Id_encoder.load_state_dict(torch.load(PATH_IDC))

    def get_id_code(self,source):
        return self.Id_encoder(source).detach()


    def forward(self, target, source,get_latent=False):
        #if it's the fine-tunning process, only identity code (512D) will be provided, so, no need to use Id_encoder
        if self.fine_tune==False:
            source = self.Id_encoder(source)
        #Else
        output = self.Swapper(target, source,get_latent=get_latent)
        return output


    def add_batch_instant_norm2swapper(self):
        bn_layer = nn.BatchNorm2d(1024)  # BatchNorm after fusion
        in_layer = nn.InstanceNorm2d(1024, affine=True)  # InstanceNorm after fusion
        #name, module = self.Swapper.E.Encoder['layer_3'].items()
        self.Swapper.E.Encoder['layer_3']  = nn.Sequential(
                self.Swapper.E.Encoder['layer_3'] ,  # Original Feature_Fusion_Block
                bn_layer,  # BatchNorm
                in_layer,  # InstanceNorm
            )
        #pdb.set_trace()
        for i in range(5):  # Since we need norm layers between 6 fusion layers, we only need 5 sets of norms
            out_channels = 1024  # Assuming fusion layer output size is 2048
            # Directly add BatchNorm and InstanceNorm as attributes to the model
            setattr(self.Swapper.E, f'batch_norm_{i}', nn.BatchNorm2d(out_channels))
            setattr(self.Swapper.E, f'instance_norm_{i}', nn.InstanceNorm2d(out_channels, affine=True))



def build_AlphaFace(config=None,fine_tune=False,adv_train=True,new_id_model=False):
    #Swapper
    
    swapper = Swapper(512)
    print('Loading pre-trained model for the ID encoder')
    if new_id_model==False:
        ID_encoder = ResNet_ID_encoder(ema_path='./Models/emp.npy')
        ID_encoder.resnet.load_state_dict(torch.load('./Models/arcface_w600k_r50_pytorch.pt', map_location=torch.device('cuda')))
    else:
        weight = torch.load(config.id_network_path)
        ID_encoder = get_model(config.id_network, dropout=0, fp16=False).cuda()
        ID_encoder.load_state_dict(weight)
        
    #Framework
    DPG = AlphaFace(swapper, ID_encoder, fine_tune=fine_tune)
    if adv_train==True:
        DPG.Prepareing_adversrial_learning(img_size=256,max_conv_size=512)
        DPG.Preparing_VGG_percet_loss()
    return DPG

if __name__ == '__main__':
    build_models()