Integrating a discriminator to guide the model toward generating more realistic facial details. This did introduce some texture to the faces.

Experimenting with Adversarial Loss/Discriminatorv3_3.py

@@ -0,0 +1,196 @@
+# a modified version of https://github.com/deepinsight/insightface/blob/master/recognition/arcface_torch/backbones/iresnet.py
+
+import torch
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+
+__all__ = ['iresnet18', 'iresnet34', 'iresnet50', 'iresnet100', 'iresnet200']
+using_ckpt = False
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes,
+                     out_planes,
+                     kernel_size=3,
+                     stride=stride,
+                     padding=dilation,
+                     groups=groups,
+                     bias=True,
+                     dilation=dilation)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes,
+                     out_planes,
+                     kernel_size=1,
+                     stride=stride,
+                     bias=True)
+
+
+class IBasicBlock(nn.Module):
+    expansion = 1
+    def __init__(self, inplanes, planes, stride=1, downsample=None,
+                 groups=1, base_width=64, dilation=1):
+        super(IBasicBlock, self).__init__()
+        if groups != 1 or base_width != 64:
+            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        self.bn1 = nn.BatchNorm2d(inplanes, eps=1e-05,)
+        self.conv1 = conv3x3(inplanes, planes)
+        self.bn2 = nn.BatchNorm2d(planes, eps=1e-05,)
+        self.prelu = nn.PReLU(planes)
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn3 = nn.BatchNorm2d(planes, eps=1e-05,)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward_impl(self, x):
+        identity = x
+        out = self.bn1(x)
+        out = self.conv1(out)
+        out = self.bn2(out)
+        out = self.prelu(out)
+        out = self.conv2(out)
+        out = self.bn3(out)
+        if self.downsample is not None:
+            identity = self.downsample(x)
+        out += identity
+        return out        
+
+    def forward(self, x):
+        if self.training and using_ckpt:
+            return checkpoint(self.forward_impl, x)
+        else:
+            return self.forward_impl(x)
+
+
+class IResNet(nn.Module):
+    fc_scale = 14 * 14
+    def __init__(self,
+                 block, layers, dropout=0, num_features=512, zero_init_residual=False,
+                 groups=1, width_per_group=64, replace_stride_with_dilation=None, fp16=False):
+        super(IResNet, self).__init__()
+        self.extra_gflops = 0.0
+        self.fp16 = fp16
+        self.inplanes = 64
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError("replace_stride_with_dilation should be None "
+                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
+        self.groups = groups
+        self.base_width = width_per_group
+        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=True)
+        self.bn1 = nn.BatchNorm2d(self.inplanes, eps=1e-05)
+        self.prelu = nn.PReLU(self.inplanes)
+        self.layer1 = self._make_layer(block, 64, layers[0], stride=2)
+        self.layer2 = self._make_layer(block,
+                                       128,
+                                       layers[1],
+                                       stride=2,
+                                       dilate=replace_stride_with_dilation[0])
+        self.layer3 = self._make_layer(block,
+                                       256,
+                                       layers[2],
+                                       stride=2,
+                                       dilate=replace_stride_with_dilation[1])
+        self.layer4 = self._make_layer(block,
+                                       512,
+                                       layers[3],
+                                       stride=2,
+                                       dilate=replace_stride_with_dilation[2])
+        self.bn2 = nn.BatchNorm2d(512 * block.expansion, eps=1e-05,)
+        self.dropout = nn.Dropout(p=dropout, inplace=True)
+        self.fc = nn.Linear(512 * block.expansion * self.fc_scale, num_features)
+        self.features = nn.BatchNorm1d(num_features, eps=1e-05)
+        nn.init.constant_(self.features.weight, 1.0)
+        self.features.weight.requires_grad = False
+
+        # for m in self.modules():
+        #     if isinstance(m, nn.Conv2d):
+        #         nn.init.normal_(m.weight, 0, 0.1)
+        #     elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+        #         nn.init.constant_(m.weight, 1)
+        #         nn.init.constant_(m.bias, 0)
+
+        # if zero_init_residual:
+        #     for m in self.modules():
+        #         if isinstance(m, IBasicBlock):
+        #             nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                nn.BatchNorm2d(planes * block.expansion, eps=1e-05, ),
+            )
+        layers = []
+        layers.append(
+            block(self.inplanes, planes, stride, downsample, self.groups,
+                  self.base_width, previous_dilation))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(
+                block(self.inplanes,
+                      planes,
+                      groups=self.groups,
+                      base_width=self.base_width,
+                      dilation=self.dilation))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        with torch.cuda.amp.autocast(self.fp16):
+            x = self.conv1(x)
+            x = self.bn1(x)
+            x = self.prelu(x)
+            x = self.layer1(x)
+            x = self.layer2(x)
+            x = self.layer3(x)
+            x = self.layer4(x)
+            x = self.bn2(x)
+            x = torch.flatten(x, 1)
+            x = self.dropout(x)
+        x = self.fc(x.float() if self.fp16 else x)
+        # x = self.features(x)
+        return x
+
+
+def _iresnet(arch, block, layers, pretrained, progress, **kwargs):
+    model = IResNet(block, layers, **kwargs)
+    if pretrained:
+        raise ValueError()
+    return model
+
+
+def iresnet18(pretrained=False, progress=True, **kwargs):
+    return _iresnet('iresnet18', IBasicBlock, [2, 2, 2, 2], pretrained,
+                    progress, **kwargs)
+
+
+def iresnet34(pretrained=False, progress=True, **kwargs):
+    return _iresnet('iresnet34', IBasicBlock, [3, 4, 6, 3], pretrained,
+                    progress, **kwargs)
+
+
+def iresnet50(pretrained=False, progress=True, **kwargs):
+    return _iresnet('iresnet50', IBasicBlock, [3, 4, 14, 3], pretrained,
+                    progress, **kwargs)
+
+
+def iresnet100(pretrained=False, progress=True, **kwargs):
+    return _iresnet('iresnet100', IBasicBlock, [3, 13, 30, 3], pretrained,
+                    progress, **kwargs)
+
+
+def iresnet200(pretrained=False, progress=True, **kwargs):
+    return _iresnet('iresnet200', IBasicBlock, [6, 26, 60, 6], pretrained,
+                    progress, **kwargs)

Experimenting with Adversarial Loss/discriminator-16796-16328-37280.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bd44167f0badfb6adfa6975b76c29ff360fb9e4eaa80b248768e15cb0145bec1
+size 328890738

Experimenting with Adversarial Loss/discriminator-580-596-640.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6f0cfcd96fcefa9b06d182d5a471ddf94db35c9143d0c3afb5b073502ec1cc07
+size 328887546

Experimenting with Adversarial Loss/reswapper-1679500.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0653401aad18b8c82a565ea4f954e044b2c2d72b5dde965b4c06e52abddac2cf
+size 553194302

Experimenting with Adversarial Loss/reswapper-1683150.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c13e960d555a7075fb38bfaa2b0fa59a6c84ca470fd85b3b0c54f526ecb32e8f
+size 553194302

Experimenting with Adversarial Loss/train_dis.3_3_1_Good_1.1.1.py

@@ -0,0 +1,435 @@
+from datetime import datetime
+import os
+import random
+import torch
+import torch.optim as optim
+import torch.nn.functional as F
+
+from Discriminatorv3_3 import iresnet50
+import Image
+import ModelFormat
+from StyleTransferLoss import StyleTransferLoss
+import onnxruntime as rt
+
+import cv2
+from insightface.data import get_image as ins_get_image
+from insightface.app import FaceAnalysis
+import face_align
+
+from StyleTransferModel_128 import StyleTransferModel
+from torch.utils.tensorboard import SummaryWriter
+
+inswapper_128_path = 'inswapper_128.onnx'
+img_size = 128
+
+providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+
+inswapperInferenceSession = rt.InferenceSession(inswapper_128_path, providers=providers)
+
+faceAnalysis = FaceAnalysis(name='buffalo_l')
+faceAnalysis.prepare(ctx_id=0, det_size=(512, 512))
+
+class FocalLoss(torch.nn.Module):
+    def __init__(self, gamma=0, eps=1e-7):
+        super(FocalLoss, self).__init__()
+        self.gamma = gamma
+        self.eps = eps
+        self.ce = torch.nn.CrossEntropyLoss()
+
+    def forward(self, input, target):
+        logp = self.ce(input, target)
+        p = torch.exp(-logp)
+        loss = (1 - p) ** self.gamma * logp
+        return loss.mean()
+
+def get_device():
+    return torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+style_loss_fn = StyleTransferLoss().to(get_device())
+
+def patchgan_prediction(pred, threshold=0.5):
+    """Process PatchGAN output to image-level decision"""
+    # pred shape: (batch_size, 1, 8, 8)
+    probabilities = torch.sigmoid(pred)
+    
+    # Two aggregation strategies
+    patch_confidence = probabilities.mean(dim=[1,2,3])  # Average all patches
+    any_patch_positive = (probabilities > threshold).any(dim=[1,2,3]).float()  # Any patch thinks it's real
+    
+    return patch_confidence, any_patch_positive
+
+def compute_gradient_penalty(D, real, fake):
+    alpha = torch.rand(real.size(0), 1, 1, 1).to(real.device)
+    interpolates = (alpha * real + (1 - alpha) * fake).requires_grad_(True)
+    d_interpolates = D(interpolates)
+    
+    gradients = torch.autograd.grad(
+        outputs=d_interpolates,
+        inputs=interpolates,
+        grad_outputs=torch.ones_like(d_interpolates),
+        create_graph=True,
+        retain_graph=True
+    )[0]
+    
+    gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()
+    return gradient_penalty
+
+def cosin_metric(x1,x2):
+    return torch.sum(x1 * x2, dim=1) / (torch.norm(x1, dim=1) * torch.norm(x2, dim=1))
+
+def createFakeImage(datasetDir, image, enableDataAugmentation, steps, resolution, device):
+    targetFaceIndex = random.randint(0, len(image)-1)
+    sourceFaceIndex = random.randint(0, len(image)-1)
+
+    target_img=cv2.imread(f"{datasetDir}/{image[targetFaceIndex]}")
+    if enableDataAugmentation and steps % 2 == 0:
+        target_img = cv2.cvtColor(target_img, cv2.COLOR_BGR2GRAY)
+        target_img = cv2.cvtColor(target_img, cv2.COLOR_GRAY2BGR)
+    faces = faceAnalysis.get(target_img)
+
+    if targetFaceIndex != sourceFaceIndex:
+        source_img = cv2.imread(f"{datasetDir}/{image[sourceFaceIndex]}")
+        faces2 = faceAnalysis.get(source_img)
+    else:
+        faces2 = faces
+
+    if len(faces) > 0 and len(faces2) > 0:
+        new_aligned_face, _ = face_align.norm_crop2(target_img, faces[0].kps, img_size)
+        blob = Image.getBlob(new_aligned_face)
+        latent = Image.getLatent(faces2[0])
+    else:
+        return createFakeImage(datasetDir, image, enableDataAugmentation, steps, resolution, device)
+
+    if targetFaceIndex != sourceFaceIndex:
+        input = {inswapperInferenceSession.get_inputs()[0].name: blob,
+                inswapperInferenceSession.get_inputs()[1].name: latent}
+
+        expected_output = inswapperInferenceSession.run([inswapperInferenceSession.get_outputs()[0].name], input)[0]
+    else:
+        expected_output = blob
+
+    expected_output_tensor = torch.from_numpy(expected_output).to(device)
+
+    if resolution != 128:
+        new_aligned_face, _ = face_align.norm_crop2(target_img, faces[0].kps, resolution)
+        blob = Image.getBlob(new_aligned_face, (resolution, resolution))
+
+    latent_tensor = torch.from_numpy(latent).to(device)
+    target_input_tensor = torch.from_numpy(blob).to(device)
+
+    return target_input_tensor, latent_tensor, expected_output_tensor
+
+def train(datasetDir, learning_rate=0.0001, model_path=None, outputModelFolder='', saveModelEachSteps = 1, stopAtSteps=None, logDir=None, previewDir=None, saveAs_onnx = False, resolutions = [128], enableDataAugmentation = False):
+    device = get_device()
+    print(f"Using device: {device}")
+    train_g = True #True
+    train_d = True #False
+
+    model = StyleTransferModel().to(device)
+    discriminator = iresnet50().to(device)  # Add discriminator
+    # discriminator.features.weight.requires_grad = True
+    optimizer_D = optim.Adam(discriminator.parameters(), lr=0.00005)  # S
+    fake_correct_count = 0
+    real_correct_count = 0
+    d_steps = 0
+
+    if model_path is not None:
+        model.load_state_dict(torch.load(model_path, map_location=device), strict=False)
+        lastSteps = 0
+        # lastSteps = 200
+        d_steps = 640
+        fake_correct_count=580
+        real_correct_count=596
+
+        discriminator.load_state_dict(torch.load(f"D:\\ReSwapper\\model\\discriminatorV4\\discriminator-{fake_correct_count}-{real_correct_count}-{d_steps}.pth", map_location=device), strict=False)
+        print(f"Loaded model from {model_path}")
+        if train_g:
+            lastSteps = int(model_path.split('-')[-1].split('.')[0])
+        print(f"Resuming training from step {lastSteps}")
+        d_steps *= 2
+    else:
+        lastSteps = 0
+
+    model.train()
+    model = model.to(device)
+    # criterion = FocalLoss(gamma=2).to(device)
+    # # criterion = torch.nn.CrossEntropyLoss().to(device)
+    # criterion = torch.nn.BCELoss().to(device)
+    criterion = torch.nn.BCELoss().to(device)
+
+    # Initialize optimizer
+    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
+    # torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50, eta_min=1e-6)
+
+    # Initialize TensorBoard writer
+    if logDir is not None:
+        train_writer = SummaryWriter(os.path.join(logDir, "training"))
+        val_writer = SummaryWriter(os.path.join(logDir, "validation"))
+
+    steps = 0
+
+    image = os.listdir(datasetDir)
+
+    resolutionIndex = 0
+
+
+    batch_size = 5
+    # Training loop
+    while True:
+        start_time = datetime.now()
+        
+        resolution = resolutions[resolutionIndex%len(resolutions)]
+        optimizer.zero_grad()
+
+        # if steps % 100 == 0 or True:
+        real_images_list = []
+
+        fake_images_list = []
+        while len(real_images_list)!=batch_size:
+            realFaceIndex = random.randint(0, len(image)-1)
+            real_img = cv2.imread(f"{datasetDir}/{image[realFaceIndex]}")
+            faces3 = faceAnalysis.get(real_img)
+            if len(faces3) == 0 : continue
+
+            aligned_real_face, _ = face_align.norm_crop2(real_img, faces3[0].kps, resolution)
+            real_images = torch.from_numpy(Image.getBlob(aligned_real_face, (resolution, resolution))).to(device)
+
+            real_images = F.interpolate(real_images, size=(224, 224), mode='bilinear', align_corners=False)
+            real_images_list.append(real_images)
+
+        while len(fake_images_list)!=batch_size:
+            target_input_tensor, latent_tensor, expected_output_tensor = createFakeImage(datasetDir, image, enableDataAugmentation, steps, resolution, device)
+
+            with torch.no_grad():
+                output = model(target_input_tensor, latent_tensor)
+
+            fake_images = output.detach()  # Detach to avoid backprop through generator
+            fake_images = F.interpolate(fake_images, size=(224, 224), mode='bilinear', align_corners=False)
+
+            fake_images_list.append(fake_images)
+
+        if train_d and resolution == 256:
+            # ---------------------
+            # Train Discriminator
+            # ---------------------
+            optimizer_D.zero_grad()
+
+            # Use ground truth as real samples
+            fake_images_list = torch.stack(fake_images_list, 1).to(device)
+            real_images_list = torch.stack(real_images_list, 1).to(device)
+
+            real_pred = discriminator(real_images_list[0])
+            # real_label = torch.from_numpy([1]) * 1
+            # real_label = real_label.float().to(device)
+            
+            d_loss_real = F.binary_cross_entropy_with_logits(real_pred, torch.ones_like(real_pred))
+            # d_loss_real= 1- F.cosine_similarity(real_pred, torch.ones_like(real_pred))
+
+            # old
+            # if real_pred.mean() > 0.5:
+                # real_correct_count += 1
+            #new
+            mean_per_real_sample = real_pred.mean(dim=1)
+
+            # Create a boolean mask where mean > 0.5
+            real_mask = mean_per_real_sample > 0.5
+
+            # Sum the True values to get the count
+            real_correct_count += real_mask.sum().item()
+            #new end
+
+            # Use generator output as fake samples
+            # fake_images = output.detach()  # Detach to avoid backprop through generator
+            # fake_images = F.interpolate(fake_images, size=(224, 224), mode='bilinear', align_corners=False)
+            fake_pred = discriminator(fake_images_list[0])
+
+            # fake_label = [0] * 1
+            # fake_label = fake_label.float().to(device)
+
+            # if fake_pred.mean() < 0.5:
+            #     fake_correct_count += 1
+            mean_per_fake_sample = fake_pred.mean(dim=1)
+
+            # Create a boolean mask where mean > 0.5
+            fake_mask = mean_per_fake_sample < 0.5
+
+            # Sum the True values to get the count
+            fake_correct_count += fake_mask.sum().item()
+
+            d_loss_fake = F.binary_cross_entropy_with_logits(fake_pred, torch.zeros_like(real_pred) * -1)
+            # d_loss_fake= 1- F.cosine_similarity(fake_pred, torch.zeros_like(real_pred))
+            # d_loss_fake_v2 = 1 - cosin_metric(fake_pred[0], torch.zeros_like(real_pred)[0])
+            d_loss = d_loss_real + d_loss_fake
+            d_loss.backward()
+            optimizer_D.step()
+            d_steps += batch_size * 2
+
+            # real, p = patchgan_prediction(real_pred)
+            # fake, p2 = patchgan_prediction(fake_pred)
+
+        #Train Gen
+        if train_g:
+
+            target_input_tensor, latent_tensor, expected_output_tensor = createFakeImage(datasetDir, image, enableDataAugmentation, steps, resolution, device)
+
+            output = model(target_input_tensor, latent_tensor)
+
+            if (resolution != 128):
+                output_128 = F.interpolate(output, size=(128, 128), mode='bilinear', align_corners=False)
+
+            content_loss, identity_loss = style_loss_fn(output_128, expected_output_tensor)
+            # Adversarial loss
+            output_224 = F.interpolate(output, size=(224, 224), mode='bilinear', align_corners=False)
+
+            fake_pred = discriminator(output_224)
+            adversarial_loss = F.binary_cross_entropy_with_logits(fake_pred, torch.ones_like(fake_pred))
+
+            loss = content_loss + adversarial_loss
+
+            if identity_loss is not None:
+                loss +=identity_loss
+                
+            
+            loss.backward()
+
+            optimizer.step()
+
+        steps += 1
+        totalSteps = steps + lastSteps
+
+        acc = (fake_correct_count+real_correct_count)/ d_steps
+
+        if logDir is not None:
+            if train_g:
+                train_writer.add_scalar("Loss/total", loss.item(), totalSteps)
+                train_writer.add_scalar("Loss/content_loss", content_loss.item(), totalSteps)
+                train_writer.add_scalar("Loss/adversarial_loss", adversarial_loss.item(), totalSteps)
+
+                if identity_loss is not None:
+                    train_writer.add_scalar("Loss/identity_loss", identity_loss.item(), totalSteps)
+
+            if train_d:
+                train_writer.add_scalar("Loss/d_acc", acc, totalSteps)
+
+                train_writer.add_scalar("Loss/d_loss", d_loss.item(), totalSteps)
+                train_writer.add_scalar("Loss/d_loss_fake", d_loss_fake.item(), totalSteps)
+                train_writer.add_scalar("Loss/d_loss_real", d_loss_real.item(), totalSteps)
+
+        elapsed_time = datetime.now() - start_time
+
+        if train_d:
+            print(f"Total Steps: {totalSteps}, Step: {steps}, D_Loss: {d_loss.item():.4f}, d_loss_real: {d_loss_real.item():.4f}, d_loss_fake: {d_loss_fake.item():.4f}, acc: {(acc):.4f}, Elapsed time: {elapsed_time}")
+        if train_g:
+            print(f"Total Steps: {totalSteps}, Step: {steps}, G_Loss: {loss.item():.4f}, Elapsed time: {elapsed_time}")
+
+        if steps % saveModelEachSteps == 0:
+            if train_g:
+                outputModelPath = f"reswapper-{totalSteps}.pth"
+                if outputModelFolder != '':
+                    outputModelPath = f"{outputModelFolder}/{outputModelPath}"
+                saveModel(model, outputModelPath)
+            if train_d:
+                discriminatorModelPath = f"discriminator-{fake_correct_count}-{real_correct_count}-{d_steps}.pth"
+                if outputModelFolder != '':
+                    discriminatorModelPath = f"{outputModelFolder}/{discriminatorModelPath}"
+                saveModel(discriminator, discriminatorModelPath)
+
+            if train_g:
+                validation_total_loss, validation_content_loss, validation_identity_loss, swapped_face, swapped_face_256 = validate(outputModelPath)
+                if previewDir is not None:
+                    cv2.imwrite(f"{previewDir}/{totalSteps}.jpg", swapped_face)
+                    cv2.imwrite(f"{previewDir}/{totalSteps}_256.jpg", swapped_face_256)
+
+                if logDir is not None:
+                    val_writer.add_scalar("Loss/total", validation_total_loss.item(), totalSteps)
+                    val_writer.add_scalar("Loss/content_loss", validation_content_loss.item(), totalSteps)
+                    if validation_identity_loss is not None:
+                        val_writer.add_scalar("Loss/identity_loss", validation_identity_loss.item(), totalSteps)
+
+            if saveAs_onnx :
+                ModelFormat.save_as_onnx_model(outputModelPath)
+
+        if stopAtSteps is not None and steps == stopAtSteps:
+            exit()
+
+        resolutionIndex += 1
+
+def saveModel(model, outputModelPath):
+    torch.save(model.state_dict(), outputModelPath)
+
+def load_model(model_path):
+    device = get_device()
+    model = StyleTransferModel().to(device)
+    model.load_state_dict(torch.load(model_path, map_location=device), strict=False)
+
+    model.eval()
+    return model
+
+def swap_face(model, target_face, source_face_latent):
+    device = get_device()
+
+    target_tensor = torch.from_numpy(target_face).to(device)
+    source_tensor = torch.from_numpy(source_face_latent).to(device)
+
+    with torch.no_grad():
+        swapped_tensor = model(target_tensor, source_tensor)
+
+    swapped_face = Image.postprocess_face(swapped_tensor)
+    
+    return swapped_face, swapped_tensor
+
+# test image
+test_img = ins_get_image('t1')
+
+test_faces = faceAnalysis.get(test_img)
+test_faces = sorted(test_faces, key = lambda x : x.bbox[0])
+test_target_face, _ = face_align.norm_crop2(test_img, test_faces[0].kps, img_size)
+test_target_face = Image.getBlob(test_target_face)
+test_l = Image.getLatent(test_faces[2])
+
+test_target_face_256, _ = face_align.norm_crop2(test_img, test_faces[0].kps, 256)
+test_target_face_256 = Image.getBlob(test_target_face_256, (256, 256))
+
+test_input = {inswapperInferenceSession.get_inputs()[0].name: test_target_face,
+        inswapperInferenceSession.get_inputs()[1].name: test_l}
+
+test_inswapperOutput = inswapperInferenceSession.run([inswapperInferenceSession.get_outputs()[0].name], test_input)[0]
+
+def validate(modelPath):
+    model = load_model(modelPath)
+    swapped_face, swapped_tensor= swap_face(model, test_target_face, test_l)
+    swapped_face_256, _= swap_face(model, test_target_face_256, test_l)
+
+    validation_content_loss, validation_identity_loss = style_loss_fn(swapped_tensor, torch.from_numpy(test_inswapperOutput).to(get_device()))
+
+    validation_total_loss = validation_content_loss
+    if validation_identity_loss is not None:
+        validation_total_loss += validation_identity_loss
+
+    return validation_total_loss, validation_content_loss, validation_identity_loss, swapped_face, swapped_face_256
+
+def main():
+    outputModelFolder = "model/discriminatorV4"
+    modelPath = None
+    modelPath = f"model/discriminatorV4/reswapper-1679500.pth"
+
+    logDir = "training/log/moreRes"
+    previewDir = "training/preview/moreRes"
+    datasetDir = "FFHQ"
+
+    os.makedirs(outputModelFolder, exist_ok=True)
+    os.makedirs(previewDir, exist_ok=True)
+
+    train(
+        datasetDir=datasetDir,
+        model_path=modelPath,
+        learning_rate=0.000001,
+        resolutions = [256],
+        enableDataAugmentation=True,
+        outputModelFolder=outputModelFolder,
+        saveModelEachSteps = 100,
+        stopAtSteps = 70000,
+        logDir=f"{logDir}/{datetime.now().strftime('%Y%m%d %H%M%S')}",
+        previewDir=previewDir)
+                    
+if __name__ == "__main__":
+    main()