Final rename for everything

hyperswap/src/dataset.py

@@ -0,0 +1,107 @@
+import glob
+import os
+import random
+from configparser import ConfigParser
+from typing import cast
+
+import albumentations
+from torch import Tensor
+from torch.utils.data import Dataset
+from torchvision import io, transforms
+
+from .helper import warp_tensor
+from .types import Batch, BatchMode, WarpTemplate
+
+
+class DynamicDataset(Dataset[Tensor]):
+	def __init__(self, config_parser : ConfigParser) -> None:
+		self.config_file_pattern = config_parser.get('training.dataset', 'file_pattern')
+		self.config_transform_size = config_parser.getint('training.dataset', 'transform_size')
+		self.config_batch_mode = cast(BatchMode, config_parser.get('training.dataset', 'batch_mode'))
+		self.config_batch_ratio = config_parser.getfloat('training.dataset', 'batch_ratio')
+		self.config_parser = config_parser
+		self.file_paths = glob.glob(self.config_file_pattern)
+		self.transforms = self.compose_transforms()
+
+	def __getitem__(self, index : int) -> Batch:
+		file_path = self.file_paths[index]
+
+		if random.random() < self.config_batch_ratio:
+			if self.config_batch_mode == 'equal':
+				return self.prepare_equal_batch(file_path)
+			if self.config_batch_mode == 'same':
+				return self.prepare_same_batch(file_path)
+
+		return self.prepare_different_batch(file_path)
+
+	def __len__(self) -> int:
+		return len(self.file_paths)
+
+	def compose_transforms(self) -> transforms:
+		return transforms.Compose(
+		[
+			AugmentTransform(),
+			transforms.ToPILImage(),
+			transforms.Resize((self.config_transform_size, self.config_transform_size), interpolation = transforms.InterpolationMode.BICUBIC),
+			transforms.ToTensor(),
+			WarpTransform(self.config_parser),
+			transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+		])
+
+	def prepare_different_batch(self, source_path : str) -> Batch:
+		target_path = random.choice(self.file_paths)
+		source_tensor = io.read_image(source_path)
+		source_tensor = self.transforms(source_tensor)
+		target_tensor = io.read_image(target_path)
+		target_tensor = self.transforms(target_tensor)
+		return source_tensor, target_tensor
+
+	def prepare_equal_batch(self, source_path : str) -> Batch:
+		source_tensor = io.read_image(source_path)
+		source_tensor = self.transforms(source_tensor)
+		return source_tensor, source_tensor
+
+	def prepare_same_batch(self, source_path : str) -> Batch:
+		target_directory_path = os.path.dirname(source_path)
+		target_file_name_and_extension = random.choice(os.listdir(target_directory_path))
+		target_path = os.path.join(target_directory_path, target_file_name_and_extension)
+		source_tensor = io.read_image(source_path)
+		source_tensor = self.transforms(source_tensor)
+		target_tensor = io.read_image(target_path)
+		target_tensor = self.transforms(target_tensor)
+		return source_tensor, target_tensor
+
+
+class AugmentTransform:
+	def __init__(self) -> None:
+		self.transforms = self.compose_transforms()
+
+	def __call__(self, input_tensor : Tensor) -> Tensor:
+		temp_tensor = input_tensor.numpy().transpose(1, 2, 0)
+		return self.transforms(image = temp_tensor).get('image')
+
+	@staticmethod
+	def compose_transforms() -> albumentations.Compose:
+		return albumentations.Compose(
+		[
+			albumentations.HorizontalFlip(),
+			albumentations.OneOf(
+			[
+				albumentations.MotionBlur(p = 0.1),
+				albumentations.ZoomBlur(max_factor = (1.0, 1.1), p = 0.1)
+			], p = 0.2),
+			albumentations.RandomBrightnessContrast(p = 0.7),
+			albumentations.ColorJitter(p = 0.2),
+			albumentations.RGBShift(p = 0.7),
+			albumentations.Illumination(p = 0.2),
+			albumentations.Affine(translate_percent = (-0.03, 0.03), scale = (0.98, 1.02), rotate = (-2, 2), border_mode = 1, p = 0.3)
+		])
+
+
+class WarpTransform:
+	def __init__(self, config_parser : ConfigParser) -> None:
+		self.config_warp_template = cast(WarpTemplate, config_parser.get('training.dataset', 'warp_template'))
+
+	def __call__(self, input_tensor : Tensor) -> Tensor:
+		temp_tensor = input_tensor.unsqueeze(0)
+		return warp_tensor(temp_tensor, self.config_warp_template).squeeze(0)

hyperswap/src/exporting.py

@@ -0,0 +1,47 @@
+import os
+from configparser import ConfigParser
+from typing import Tuple
+
+import torch
+from torch import Tensor, nn
+
+from .training import HyperSwapTrainer
+from .types import Embedding, Mask, Module
+
+CONFIG_PARSER = ConfigParser()
+CONFIG_PARSER.read('config.ini')
+
+
+class HalfPrecision(nn.Module):
+	def __init__(self, model : Module) -> None:
+		super().__init__()
+		self.model = model.half()
+
+	def forward(self, source_embedding : Embedding, target_tensor : Tensor) -> Tuple[Tensor, Mask]:
+		source_embedding = source_embedding.half()
+		target_tensor = target_tensor.half()
+		output_tensor, output_mask = self.model(source_embedding, target_tensor)
+		output_tensor = output_tensor.float()
+		output_mask = output_mask.float()
+		return output_tensor, output_mask
+
+
+def export() -> None:
+	config_directory_path = CONFIG_PARSER.get('exporting', 'directory_path')
+	config_source_path = CONFIG_PARSER.get('exporting', 'source_path')
+	config_target_path = CONFIG_PARSER.get('exporting', 'target_path')
+	config_target_size = CONFIG_PARSER.getint('exporting', 'target_size')
+	config_ir_version = CONFIG_PARSER.getint('exporting', 'ir_version')
+	config_opset_version = CONFIG_PARSER.getint('exporting', 'opset_version')
+	config_precision = CONFIG_PARSER.get('exporting', 'precision')
+
+	os.makedirs(config_directory_path, exist_ok = True)
+	model = HyperSwapTrainer.load_from_checkpoint(config_source_path, config_parser = CONFIG_PARSER, map_location ='cpu').eval()
+
+	if config_precision == 'half':
+		model = HalfPrecision(model).eval()
+
+	model.ir_version = torch.tensor(config_ir_version)
+	source_tensor = torch.randn(1, 512)
+	target_tensor = torch.randn(1, 3, config_target_size, config_target_size)
+	torch.onnx.export(model, (source_tensor, target_tensor), config_target_path, input_names = [ 'source', 'target' ], output_names = [ 'output', 'mask' ], opset_version = config_opset_version)

hyperswap/src/helper.py

@@ -0,0 +1,51 @@
+import torch
+from torch import Tensor, nn
+
+from .types import EmbedderModule, Embedding, Mask, Padding, WarpTemplate, WarpTemplateSet
+
+WARP_TEMPLATE_SET : WarpTemplateSet =\
+{
+	'arcface_128_v2_to_arcface_112_v2': torch.tensor(
+	[
+		[ 8.75000016e-01, -1.07193451e-08, 3.80446920e-10 ],
+		[ 1.07193451e-08, 8.75000016e-01, -1.25000007e-01 ]
+	]),
+	'ffhq_512_to_arcface_128_v2': torch.tensor(
+	[
+		[ 8.50048894e-01, -1.29486822e-04, 1.90956388e-03 ],
+		[ 1.29486822e-04, 8.50048894e-01, 9.56254653e-02 ]
+	]),
+	'vggfacehq_256_to_arcface_128_v2': torch.tensor(
+	[
+		[ 1.01305414, -0.00140513, -0.00585911 ],
+		[ 0.00140513, 1.01305414, 0.11169602 ]
+	])
+}
+
+
+def warp_tensor(input_tensor : Tensor, warp_template : WarpTemplate) -> Tensor:
+	normed_warp_template = WARP_TEMPLATE_SET.get(warp_template).repeat(input_tensor.shape[0], 1, 1)
+	affine_grid = nn.functional.affine_grid(normed_warp_template.to(input_tensor.device), list(input_tensor.shape))
+	output_tensor = nn.functional.grid_sample(input_tensor, affine_grid, padding_mode = 'reflection')
+	return output_tensor
+
+
+def calc_embedding(embedder : EmbedderModule, input_tensor : Tensor, padding : Padding) -> Embedding:
+	crop_tensor = warp_tensor(input_tensor, 'arcface_128_v2_to_arcface_112_v2')
+	crop_tensor = nn.functional.interpolate(crop_tensor, size = 112, mode = 'area')
+	crop_tensor[:, :, :padding[0], :] = 0
+	crop_tensor[:, :, 112 - padding[1]:, :] = 0
+	crop_tensor[:, :, :, :padding[2]] = 0
+	crop_tensor[:, :, :, 112 - padding[3]:] = 0
+
+	embedding = embedder(crop_tensor)
+	embedding = nn.functional.normalize(embedding, p = 2)
+	return embedding
+
+
+def overlay_mask(input_tensor : Tensor, input_mask : Mask) -> Tensor:
+	overlay_tensor = torch.zeros(*input_tensor.shape, dtype = input_tensor.dtype, device = input_tensor.device)
+	overlay_tensor[:, 2, :, :] = 1
+	input_mask = input_mask.repeat(1, 3, 1, 1).clamp(0, 0.8)
+	output_tensor = input_tensor * (1 - input_mask) + overlay_tensor * input_mask
+	return output_tensor

hyperswap/src/inferencing.py

@@ -0,0 +1,27 @@
+import configparser
+
+import torch
+from torchvision import io
+
+from .helper import calc_embedding
+from .training import HyperSwapTrainer
+
+CONFIG_PARSER = configparser.ConfigParser()
+CONFIG_PARSER.read('config.ini')
+
+
+def infer() -> None:
+	config_generator_path = CONFIG_PARSER.get('inferencing', 'generator_path')
+	config_embedder_path = CONFIG_PARSER.get('inferencing', 'embedder_path')
+	config_source_path = CONFIG_PARSER.get('inferencing', 'source_path')
+	config_target_path = CONFIG_PARSER.get('inferencing', 'target_path')
+	config_output_path = CONFIG_PARSER.get('inferencing', 'output_path')
+
+	generator = HyperSwapTrainer.load_from_checkpoint(config_generator_path, config_parser = CONFIG_PARSER, map_location ='cpu').eval()
+	embedder = torch.jit.load(config_embedder_path, map_location = 'cpu').eval()
+
+	source_tensor = io.read_image(config_source_path)
+	target_tensor = io.read_image(config_target_path)
+	source_embedding = calc_embedding(embedder, source_tensor, (0, 0, 0, 0))
+	output_tensor, _ = generator(source_embedding, target_tensor)
+	io.write_jpeg(output_tensor, config_output_path)

hyperswap/src/models/discriminator.py

@@ -0,0 +1,35 @@
+from configparser import ConfigParser
+from typing import List
+
+from torch import Tensor, nn
+
+from ..networks.nld import NLD
+
+
+class Discriminator(nn.Module):
+	def __init__(self, config_parser : ConfigParser) -> None:
+		super().__init__()
+		self.config_num_discriminators = config_parser.getint('training.model.discriminator', 'num_discriminators')
+		self.config_parser = config_parser
+		self.discriminators = self.create_discriminators()
+		self.avg_pool = nn.AvgPool2d(kernel_size = 3, stride = 2, padding = (1, 1), count_include_pad = False)
+
+	def create_discriminators(self) -> nn.ModuleList:
+		discriminators = nn.ModuleList()
+
+		for _ in range(self.config_num_discriminators):
+			discriminator = NLD(self.config_parser).sequences
+			discriminators.append(discriminator)
+
+		return discriminators
+
+	def forward(self, input_tensor : Tensor) -> List[Tensor]:
+		temp_tensor = input_tensor
+		output_tensors = []
+
+		for discriminator in self.discriminators:
+			output_tensor = discriminator(temp_tensor)
+			output_tensors.append(output_tensor)
+			temp_tensor = self.avg_pool(temp_tensor)
+
+		return output_tensors

hyperswap/src/models/generator.py

@@ -0,0 +1,42 @@
+from configparser import ConfigParser
+from typing import Tuple
+
+from torch import Tensor, nn
+
+from ..networks.aad import AAD
+from ..networks.masknet import MaskNet
+from ..networks.unet import UNet
+from ..types import Embedding, Feature, Mask
+
+
+class Generator(nn.Module):
+	def __init__(self, config_parser : ConfigParser) -> None:
+		super().__init__()
+		self.encoder = UNet(config_parser)
+		self.generator = AAD(config_parser)
+		self.masker = MaskNet(config_parser)
+		self.encoder.apply(init_weight)
+		self.generator.apply(init_weight)
+		self.masker.apply(init_weight)
+
+	def forward(self, source_embedding : Embedding, target_tensor : Tensor, target_features : Tuple[Feature, ...]) -> Tuple[Tensor, Mask]:
+		output_tensor = self.generator(source_embedding, target_features)
+		target_feature = target_features[-1]
+		output_mask = self.masker(target_tensor, target_feature)
+		output_tensor = output_tensor * output_mask + target_tensor * (1 - output_mask)
+		return output_tensor, output_mask
+
+	def encode_features(self, input_tensor : Tensor) -> Tuple[Feature, ...]:
+		return self.encoder(input_tensor)
+
+
+def init_weight(module : nn.Module) -> None:
+	if isinstance(module, nn.Linear):
+		module.weight.data.normal_(std = 0.001)
+		module.bias.data.zero_()
+
+	if isinstance(module, nn.Conv2d):
+		nn.init.xavier_normal_(module.weight.data)
+
+	if isinstance(module, nn.ConvTranspose2d):
+		nn.init.xavier_normal_(module.weight.data)

hyperswap/src/models/loss.py

@@ -0,0 +1,186 @@
+from configparser import ConfigParser
+from typing import List, Tuple
+
+import torch
+from pytorch_msssim import ssim
+from torch import Tensor, nn
+from torchvision import transforms
+
+from ..helper import calc_embedding
+from ..types import EmbedderModule, FaceMaskerModule, Feature, GazerModule, Loss, Mask
+
+
+class DiscriminatorLoss(nn.Module):
+	def __init__(self) -> None:
+		super().__init__()
+
+	def forward(self, discriminator_source_tensors : List[Tensor], discriminator_output_tensors : List[Tensor]) -> Loss:
+		positive_tensors = []
+		negative_tensors = []
+
+		for discriminator_source_tensor in discriminator_source_tensors:
+			positive_tensor = torch.relu(1 - discriminator_source_tensor).mean(dim = [ 1, 2, 3 ])
+			positive_tensors.append(positive_tensor)
+
+		for discriminator_output_tensor in discriminator_output_tensors:
+			negative_tensor = torch.relu(discriminator_output_tensor + 1).mean(dim = [ 1, 2, 3 ])
+			negative_tensors.append(negative_tensor)
+
+		positive_loss = torch.stack(positive_tensors).mean()
+		negative_loss = torch.stack(negative_tensors).mean()
+		discriminator_loss = (positive_loss + negative_loss) * 0.5
+		return discriminator_loss
+
+
+class AdversarialLoss(nn.Module):
+	def __init__(self, config_parser : ConfigParser) -> None:
+		super().__init__()
+		self.config_adversarial_weight = config_parser.getfloat('training.losses', 'adversarial_weight')
+
+	def forward(self, discriminator_output_tensors : List[Tensor]) -> Tuple[Loss, Loss]:
+		temp_tensors = []
+
+		for discriminator_output_tensor in discriminator_output_tensors:
+			temp_tensor = torch.relu(1 - discriminator_output_tensor).mean(dim = [ 1, 2, 3 ]).mean()
+			temp_tensors.append(temp_tensor)
+
+		adversarial_loss = torch.stack(temp_tensors).mean()
+		weighted_adversarial_loss = adversarial_loss * self.config_adversarial_weight
+		return adversarial_loss, weighted_adversarial_loss
+
+
+class CycleLoss(nn.Module):
+	def __init__(self, config_parser : ConfigParser) -> None:
+		super().__init__()
+		self.config_batch_size = config_parser.getint('training.loader', 'batch_size')
+		self.config_cycle_weight = config_parser.getfloat('training.losses', 'cycle_weight')
+		self.l1_loss = nn.L1Loss()
+
+	def forward(self, target_tensor : Tensor, cycle_tensor : Tensor, target_features : Tuple[Feature, ...], cycle_features : Tuple[Feature, ...]) -> Tuple[Loss, Loss]:
+		temp_tensors = []
+
+		for target_feature, output_feature in zip(target_features, cycle_features):
+			temp_tensor = torch.mean(torch.pow(output_feature - target_feature, 2).reshape(self.config_batch_size, -1), dim = 1).mean()
+			temp_tensors.append(temp_tensor)
+
+		feature_loss = torch.stack(temp_tensors).mean()
+		reconstruction_loss = self.l1_loss(target_tensor, cycle_tensor)
+		cycle_loss = (feature_loss + reconstruction_loss) * 0.5
+		weighted_feature_loss = cycle_loss * self.config_cycle_weight
+		return cycle_loss, weighted_feature_loss
+
+
+class FeatureLoss(nn.Module):
+	def __init__(self, config_parser : ConfigParser) -> None:
+		super().__init__()
+		self.config_batch_size = config_parser.getint('training.loader', 'batch_size')
+		self.config_feature_weight = config_parser.getfloat('training.losses', 'feature_weight')
+
+	def forward(self, target_features : Tuple[Feature, ...], output_features : Tuple[Feature, ...]) -> Tuple[Loss, Loss]:
+		temp_tensors = []
+
+		for target_feature, output_feature in zip(target_features, output_features):
+			temp_tensor = torch.mean(torch.pow(output_feature - target_feature, 2).reshape(self.config_batch_size, -1), dim = 1).mean()
+			temp_tensors.append(temp_tensor)
+
+		feature_loss = torch.stack(temp_tensors).mean() * 0.5
+		weighted_feature_loss = feature_loss * self.config_feature_weight
+		return feature_loss, weighted_feature_loss
+
+
+class ReconstructionLoss(nn.Module):
+	def __init__(self, config_parser : ConfigParser, embedder : EmbedderModule) -> None:
+		super().__init__()
+		self.config_reconstruction_weight = config_parser.getfloat('training.losses', 'reconstruction_weight')
+		self.embedder = embedder
+		self.mse_loss = nn.MSELoss()
+
+	def forward(self, source_tensor : Tensor, target_tensor : Tensor, output_tensor : Tensor) -> Tuple[Loss, Loss]:
+		with torch.no_grad():
+			source_embedding = calc_embedding(self.embedder, source_tensor, (0, 0, 0, 0))
+			target_embedding = calc_embedding(self.embedder, target_tensor, (0, 0, 0, 0))
+
+		has_similar_identity = torch.cosine_similarity(source_embedding, target_embedding) > 0.8
+
+		reconstruction_loss = torch.mean((source_tensor - target_tensor) ** 2, dim = (1, 2, 3))
+		reconstruction_loss = (reconstruction_loss * has_similar_identity).mean() * 0.5
+
+		data_range = float(torch.max(output_tensor) - torch.min(output_tensor))
+		visual_loss = 1 - ssim(output_tensor, target_tensor, data_range = data_range).mean()
+		reconstruction_loss = (reconstruction_loss + visual_loss) * 0.5
+		weighted_reconstruction_loss = reconstruction_loss * self.config_reconstruction_weight
+		return reconstruction_loss, weighted_reconstruction_loss
+
+
+class IdentityLoss(nn.Module):
+	def __init__(self, config_parser : ConfigParser, embedder : EmbedderModule) -> None:
+		super().__init__()
+		self.config_identity_weight = config_parser.getfloat('training.losses', 'identity_weight')
+		self.embedder = embedder
+
+	def forward(self, source_tensor : Tensor, output_tensor : Tensor) -> Tuple[Loss, Loss]:
+		output_embedding = calc_embedding(self.embedder, output_tensor, (30, 0, 10, 10))
+		source_embedding = calc_embedding(self.embedder, source_tensor, (30, 0, 10, 10))
+		identity_loss = (1 - torch.cosine_similarity(source_embedding, output_embedding)).mean()
+		weighted_identity_loss = identity_loss * self.config_identity_weight
+		return identity_loss, weighted_identity_loss
+
+
+class GazeLoss(nn.Module):
+	def __init__(self, config_parser : ConfigParser, gazer : GazerModule) -> None:
+		super().__init__()
+		self.config_gaze_weight = config_parser.getfloat('training.losses', 'gaze_weight')
+		self.config_output_size = config_parser.getint('training.model.generator', 'output_size')
+		self.gazer = gazer
+		self.l1_loss = nn.L1Loss()
+
+	def forward(self, target_tensor : Tensor, output_tensor : Tensor) -> Tuple[Loss, Loss]:
+		output_pitch, output_yaw = self.detect_gaze(output_tensor)
+		target_pitch, target_yaw = self.detect_gaze(target_tensor)
+
+		pitch_loss = self.l1_loss(output_pitch, target_pitch)
+		yaw_loss = self.l1_loss(output_yaw, target_yaw)
+
+		gaze_loss = (pitch_loss + yaw_loss) * 0.5
+		weighted_gaze_loss = gaze_loss * self.config_gaze_weight
+		return gaze_loss, weighted_gaze_loss
+
+	def detect_gaze(self, input_tensor : Tensor) -> Tuple[Tensor, Tensor]:
+		crop_sizes = (torch.tensor([ 0.235, 0.875, 0.0625, 0.8 ]) * self.config_output_size).int()
+		crop_tensor = input_tensor[:, :, crop_sizes[0]:crop_sizes[1], crop_sizes[2]:crop_sizes[3]]
+		crop_tensor = (crop_tensor + 1) * 0.5
+		crop_tensor = transforms.Normalize(mean = [ 0.485, 0.456, 0.406 ], std = [ 0.229, 0.224, 0.225 ])(crop_tensor)
+		crop_tensor = nn.functional.interpolate(crop_tensor, size = 448, mode = 'bicubic')
+
+		with torch.no_grad():
+			pitch, yaw = self.gazer(crop_tensor)
+
+		return pitch, yaw
+
+
+class MaskLoss(nn.Module):
+	def __init__(self, config_parser : ConfigParser, face_masker : FaceMaskerModule) -> None:
+		super().__init__()
+		self.config_mask_weight = config_parser.getfloat('training.losses', 'mask_weight')
+		self.config_output_size = config_parser.getint('training.model.generator', 'output_size')
+		self.face_masker = face_masker
+		self.mse_loss = nn.MSELoss()
+
+	def forward(self, target_tensor : Tensor, output_mask : Mask) -> Tuple[Loss, Loss]:
+		target_mask = self.calc_mask(target_tensor)
+		target_mask = target_mask.view(-1, self.config_output_size, self.config_output_size)
+		output_mask = output_mask.view(-1, self.config_output_size, self.config_output_size)
+		mask_loss = self.mse_loss(target_mask, output_mask)
+		weighted_mask_loss = mask_loss * self.config_mask_weight
+		return mask_loss, weighted_mask_loss
+
+	def calc_mask(self, target_tensor : Tensor) -> Tensor:
+		target_tensor = torch.nn.functional.interpolate(target_tensor, (256, 256), mode = 'bilinear')
+		target_tensor = (target_tensor.clip(-1, 1) + 1) * 0.5
+
+		with torch.no_grad():
+			output_tensor = self.face_masker(target_tensor)
+			output_tensor = output_tensor.clamp(0, 1)
+			output_tensor = torch.nn.functional.interpolate(output_tensor, (self.config_output_size, self.config_output_size), mode = 'bilinear')
+
+		return output_tensor

hyperswap/src/networks/aad.py

@@ -0,0 +1,191 @@
+from configparser import ConfigParser
+from typing import Tuple
+
+import torch
+from torch import Tensor, nn
+
+from ..types import Embedding, Feature
+
+
+class AAD(nn.Module):
+	def __init__(self, config_parser : ConfigParser) -> None:
+		super().__init__()
+		self.config_source_channels = config_parser.getint('training.model.generator', 'source_channels')
+		self.config_output_channels = config_parser.getint('training.model.generator', 'output_channels')
+		self.config_output_size = config_parser.getint('training.model.generator', 'output_size')
+		self.config_num_blocks = config_parser.getint('training.model.generator', 'num_blocks')
+		self.pixel_shuffle_up_sample = PixelShuffleUpSample(self.config_source_channels, self.config_output_channels)
+		self.layers = self.create_layers()
+
+	def create_layers(self) -> nn.ModuleList:
+		layers = nn.ModuleList()
+
+		if self.config_output_size == 128:
+			layers.extend(
+			[
+				AdaptiveFeatureModulation(512, 512, 512, self.config_source_channels, self.config_num_blocks),
+				AdaptiveFeatureModulation(512, 512, 1024, self.config_source_channels, self.config_num_blocks),
+				AdaptiveFeatureModulation(512, 512, 512, self.config_source_channels, self.config_num_blocks)
+			])
+
+		if self.config_output_size == 256:
+			layers.extend(
+			[
+				AdaptiveFeatureModulation(1024, 1024, 1024, self.config_source_channels, self.config_num_blocks),
+				AdaptiveFeatureModulation(1024, 1024, 2048, self.config_source_channels, self.config_num_blocks),
+				AdaptiveFeatureModulation(1024, 1024, 1024, self.config_source_channels, self.config_num_blocks),
+				AdaptiveFeatureModulation(1024, 512, 512, self.config_source_channels, self.config_num_blocks)
+			])
+
+		if self.config_output_size == 512:
+			layers.extend(
+			[
+				AdaptiveFeatureModulation(2048, 2048, 2048, self.config_source_channels, self.config_num_blocks),
+				AdaptiveFeatureModulation(2048, 2048, 4096, self.config_source_channels, self.config_num_blocks),
+				AdaptiveFeatureModulation(2048, 2048, 2048, self.config_source_channels, self.config_num_blocks),
+				AdaptiveFeatureModulation(2048, 1024, 1024, self.config_source_channels, self.config_num_blocks),
+				AdaptiveFeatureModulation(1024, 512, 512, self.config_source_channels, self.config_num_blocks)
+			])
+
+		if self.config_output_size == 1024:
+			layers.extend(
+			[
+				AdaptiveFeatureModulation(4096, 4096, 4096, self.config_source_channels, self.config_num_blocks),
+				AdaptiveFeatureModulation(4096, 4096, 8192, self.config_source_channels, self.config_num_blocks),
+				AdaptiveFeatureModulation(4096, 4096, 4096, self.config_source_channels, self.config_num_blocks),
+				AdaptiveFeatureModulation(4096, 2048, 2048, self.config_source_channels, self.config_num_blocks),
+				AdaptiveFeatureModulation(2048, 1024, 1024, self.config_source_channels, self.config_num_blocks),
+				AdaptiveFeatureModulation(1024, 512, 512, self.config_source_channels, self.config_num_blocks)
+			])
+
+		layers.extend(
+		[
+			AdaptiveFeatureModulation(512, 256, 256, self.config_source_channels, self.config_num_blocks),
+			AdaptiveFeatureModulation(256, 128, 128, self.config_source_channels, self.config_num_blocks),
+			AdaptiveFeatureModulation(128, 64, 64, self.config_source_channels, self.config_num_blocks),
+			AdaptiveFeatureModulation(64, 3, 64, self.config_source_channels, self.config_num_blocks)
+		])
+
+		return layers
+
+	def forward(self, source_embedding : Embedding, target_features : Tuple[Feature, ...]) -> Tensor:
+		temp_tensors = self.pixel_shuffle_up_sample(source_embedding)
+
+		for index, layer in enumerate(self.layers[:-1]):
+			target_feature = target_features[index]
+			temp_tensor = layer(temp_tensors, source_embedding, target_feature)
+			temp_tensors = nn.functional.interpolate(temp_tensor, scale_factor = 2, mode = 'bilinear', align_corners = False)
+
+		target_feature = target_features[-1]
+		temp_tensors = self.layers[-1](temp_tensors, source_embedding, target_feature)
+		output_tensor = torch.tanh(temp_tensors)
+		return output_tensor
+
+
+class AdaptiveFeatureModulation(nn.Module):
+	def __init__(self, input_channels : int, output_channels : int, target_channels : int, source_channels : int, num_blocks : int) -> None:
+		super().__init__()
+		self.context_input_channels = input_channels
+		self.context_output_channels = output_channels
+		self.context_target_channels = target_channels
+		self.context_source_channels = source_channels
+		self.context_num_blocks = num_blocks
+		self.primary_layers = self.create_primary_layers()
+		self.shortcut_layers = self.create_shortcut_layers()
+
+	def create_primary_layers(self) -> nn.ModuleList:
+		primary_layers = nn.ModuleList()
+
+		for index in range(self.context_num_blocks):
+			primary_layers.extend(
+			[
+				FeatureModulation(self.context_input_channels, self.context_target_channels, self.context_source_channels),
+				nn.ReLU(inplace = True)
+			])
+
+			if index < self.context_num_blocks - 1:
+				primary_layers.append(nn.Conv2d(self.context_input_channels, self.context_input_channels, kernel_size = 3, padding = 1, bias = False))
+			else:
+				primary_layers.append(nn.Conv2d(self.context_input_channels, self.context_output_channels, kernel_size = 3, padding = 1, bias = False))
+
+		return primary_layers
+
+	def create_shortcut_layers(self) -> nn.ModuleList:
+		shortcut_layers = nn.ModuleList()
+
+		if self.context_input_channels > self.context_output_channels:
+			shortcut_layers.extend(
+			[
+				FeatureModulation(self.context_input_channels, self.context_target_channels, self.context_source_channels),
+				nn.ReLU(inplace = True),
+				nn.Conv2d(self.context_input_channels, self.context_output_channels, kernel_size = 3, padding = 1, bias = False)
+			])
+
+		return shortcut_layers
+
+	def forward(self, input_tensor : Tensor, source_embedding : Embedding, target_feature : Feature) -> Tensor:
+		primary_tensor = input_tensor
+
+		for primary_layer in self.primary_layers:
+			if isinstance(primary_layer, FeatureModulation):
+				primary_tensor = primary_layer(primary_tensor, source_embedding, target_feature)
+			else:
+				primary_tensor = primary_layer(primary_tensor)
+
+		if self.context_input_channels > self.context_output_channels:
+			shortcut_tensor = input_tensor
+
+			for shortcut_layer in self.shortcut_layers:
+				if isinstance(shortcut_layer, FeatureModulation):
+					shortcut_tensor = shortcut_layer(shortcut_tensor, source_embedding, target_feature)
+				else:
+					shortcut_tensor = shortcut_layer(shortcut_tensor)
+
+			input_tensor = shortcut_tensor
+
+		return primary_tensor + input_tensor
+
+
+class FeatureModulation(nn.Module):
+	def __init__(self, input_channels : int, target_channels : int, source_channels : int) -> None:
+		super().__init__()
+		self.context_input_channels = input_channels
+		self.conv1 = nn.Conv2d(target_channels, input_channels, kernel_size = 1)
+		self.conv2 = nn.Conv2d(target_channels, input_channels, kernel_size = 1)
+		self.conv3 = nn.Conv2d(input_channels, 1, kernel_size = 1)
+		self.linear1 = nn.Linear(source_channels, input_channels)
+		self.linear2 = nn.Linear(source_channels, input_channels)
+		self.instance_norm = nn.InstanceNorm2d(input_channels)
+
+	def forward(self, input_tensor : Tensor, source_embedding : Embedding, target_feature : Feature) -> Tensor:
+		temp_tensor = self.instance_norm(input_tensor)
+
+		source_scale = self.linear2(source_embedding).reshape(temp_tensor.shape[0], self.context_input_channels, 1, 1).expand_as(temp_tensor)
+		source_shift = self.linear1(source_embedding).reshape(temp_tensor.shape[0], self.context_input_channels, 1, 1).expand_as(temp_tensor)
+		source_modulation = source_scale * temp_tensor + source_shift
+
+		target_scale = self.conv1(target_feature)
+		target_shift = self.conv2(target_feature)
+		target_modulation = target_scale * temp_tensor + target_shift
+
+		temp_mask = torch.sigmoid(self.conv3(temp_tensor))
+		output_tensor = (1 - temp_mask) * target_modulation + temp_mask * source_modulation
+		return output_tensor
+
+
+class PixelShuffleUpSample(nn.Module):
+	def __init__(self, input_channels : int, output_channels : int) -> None:
+		super().__init__()
+		self.sequences = self.create_sequences(input_channels, output_channels)
+
+	@staticmethod
+	def create_sequences(input_channels : int, output_channels : int) -> nn.Sequential:
+		return nn.Sequential(
+			nn.Conv2d(input_channels, output_channels, kernel_size = 3, padding = 1),
+			nn.PixelShuffle(upscale_factor = 2)
+		)
+
+	def forward(self, input_tensor : Tensor) -> Tensor:
+		temp_tensor = input_tensor.view(input_tensor.shape[0], -1, 1, 1)
+		output_tensor = self.sequences(temp_tensor)
+		return output_tensor

hyperswap/src/networks/masknet.py

@@ -0,0 +1,111 @@
+from configparser import ConfigParser
+
+import torch
+from torch import Tensor, nn
+
+from ..types import Feature, Mask
+
+
+class MaskNet(nn.Module):
+	def __init__(self, config_parser : ConfigParser) -> None:
+		super().__init__()
+		self.config_input_channels = config_parser.getint('training.model.masker', 'input_channels')
+		self.config_output_channels = config_parser.getint('training.model.masker', 'output_channels')
+		self.config_num_filters = config_parser.getint('training.model.masker', 'num_filters')
+		self.down_samples = self.create_down_samples(self.config_input_channels, self.config_num_filters)
+		self.up_samples = self.create_up_samples(self.config_num_filters)
+		self.bottleneck = BottleNeck(self.config_num_filters * 4)
+		self.conv = nn.Conv2d(self.config_num_filters, self.config_output_channels, kernel_size = 1)
+		self.sigmoid = nn.Sigmoid()
+
+	@staticmethod
+	def create_down_samples(input_channels : int, num_filters : int) -> nn.ModuleList:
+		return nn.ModuleList(
+		[
+			DownSample(input_channels, num_filters),
+			DownSample(num_filters, num_filters * 2),
+			DownSample(num_filters * 2, num_filters * 4)
+		])
+
+	@staticmethod
+	def create_up_samples(num_filters : int) -> nn.ModuleList:
+		return nn.ModuleList(
+		[
+			UpSample(num_filters * 4, num_filters * 2),
+			UpSample(num_filters * 2, num_filters),
+			UpSample(num_filters, num_filters)
+		])
+
+	def forward(self, input_tensor : Tensor, input_feature : Feature) -> Mask:
+		output_mask = torch.cat([ input_tensor, input_feature ], dim = 1)
+
+		for down_sample in self.down_samples:
+			output_mask = down_sample(output_mask)
+
+		output_mask = self.bottleneck(output_mask)
+
+		for up_sample in self.up_samples:
+			output_mask = up_sample(output_mask)
+
+		output_mask = self.conv(output_mask)
+		output_mask = self.sigmoid(output_mask)
+		return output_mask
+
+
+class BottleNeck(nn.Module):
+	def __init__(self, num_filters : int):
+		super().__init__()
+		self.sequences = self.create_sequences(num_filters)
+		self.relu = nn.ReLU(inplace = True)
+
+	@staticmethod
+	def create_sequences(num_filters : int) -> nn.Sequential:
+		return nn.Sequential(
+			nn.Conv2d(num_filters, num_filters, kernel_size = 3, padding = 1, bias = False),
+			nn.BatchNorm2d(num_filters),
+			nn.ReLU(inplace = True),
+			nn.Conv2d(num_filters, num_filters, kernel_size = 3, padding = 1, bias = False),
+			nn.BatchNorm2d(num_filters),
+			nn.ReLU(inplace = True)
+		)
+
+	def forward(self, input_tensor : Tensor) -> Tensor:
+		output_tensor = self.sequences(input_tensor) + input_tensor
+		output_tensor = self.relu(output_tensor)
+		return output_tensor
+
+
+class UpSample(nn.Module):
+	def __init__(self, input_channels : int, output_channels : int) -> None:
+		super().__init__()
+		self.sequences = self.create_sequences(input_channels, output_channels)
+
+	@staticmethod
+	def create_sequences(input_channels : int, output_channels : int) -> nn.Sequential:
+		return nn.Sequential(
+			nn.ConvTranspose2d(input_channels, output_channels, kernel_size = 2, stride = 2),
+			nn.ReLU(inplace = True)
+		)
+
+	def forward(self, input_tensor : Tensor) -> Tensor:
+		output_tensor = self.sequences(input_tensor)
+		return output_tensor
+
+
+class DownSample(nn.Module):
+	def __init__(self, input_channels : int, output_channels : int) -> None:
+		super().__init__()
+		self.sequences = self.create_sequences(input_channels, output_channels)
+
+	@staticmethod
+	def create_sequences(input_channels : int, output_channels : int) -> nn.Sequential:
+		return nn.Sequential(
+			nn.Conv2d(input_channels, output_channels, kernel_size = 3, padding = 1, bias = False),
+			nn.BatchNorm2d(output_channels),
+			nn.ReLU(inplace = True),
+			nn.MaxPool2d(2)
+		)
+
+	def forward(self, input_tensor : Tensor) -> Tensor:
+		output_tensor = self.sequences(input_tensor)
+		return output_tensor

hyperswap/src/networks/nld.py

@@ -0,0 +1,48 @@
+import math
+from configparser import ConfigParser
+
+from torch import Tensor, nn
+
+
+class NLD(nn.Module):
+	def __init__(self, config_parser : ConfigParser) -> None:
+		super().__init__()
+		self.config_input_channels = config_parser.getint('training.model.discriminator', 'input_channels')
+		self.config_num_filters = config_parser.getint('training.model.discriminator', 'num_filters')
+		self.config_kernel_size = config_parser.getint('training.model.discriminator', 'kernel_size')
+		self.config_num_layers = config_parser.getint('training.model.discriminator', 'num_layers')
+		self.layers = self.create_layers()
+		self.sequences = nn.Sequential(*self.layers)
+
+	def create_layers(self) -> nn.ModuleList:
+		padding = math.ceil((self.config_kernel_size - 1) / 2)
+		current_filters = self.config_num_filters
+		layers = nn.ModuleList(
+		[
+			nn.Conv2d(self.config_input_channels, current_filters, kernel_size = self.config_kernel_size, stride = 2, padding = padding),
+			nn.LeakyReLU(0.2, True)
+		])
+
+		for _ in range(1, self.config_num_layers):
+			previous_filters = current_filters
+			current_filters = min(current_filters * 2, 512)
+			layers +=\
+			[
+				nn.Conv2d(previous_filters, current_filters, kernel_size = self.config_kernel_size, stride = 2, padding = padding),
+				nn.InstanceNorm2d(current_filters),
+				nn.LeakyReLU(0.2, True)
+			]
+
+		previous_filters = current_filters
+		current_filters = min(current_filters * 2, 512)
+		layers +=\
+		[
+			nn.Conv2d(previous_filters, current_filters, kernel_size = self.config_kernel_size, padding = padding),
+			nn.InstanceNorm2d(current_filters),
+			nn.LeakyReLU(0.2, True),
+			nn.Conv2d(current_filters, 1, kernel_size = self.config_kernel_size, padding = padding)
+		]
+		return layers
+
+	def forward(self, input_tensor : Tensor) -> Tensor:
+		return self.sequences(input_tensor)

hyperswap/src/networks/unet.py

@@ -0,0 +1,157 @@
+from configparser import ConfigParser
+from typing import Tuple
+
+import torch
+from torch import Tensor, nn
+
+from ..types import Feature
+
+
+class UNet(nn.Module):
+	def __init__(self, config_parser : ConfigParser) -> None:
+		super().__init__()
+		self.config_output_size = config_parser.getint('training.model.generator', 'output_size')
+		self.down_samples = self.create_down_samples()
+		self.up_samples = self.create_up_samples()
+
+	def create_down_samples(self) -> nn.ModuleList:
+		down_samples = nn.ModuleList(
+		[
+			DownSample(3, 32),
+			DownSample(32, 64),
+			DownSample(64, 128),
+			DownSample(128, 256),
+			DownSample(256, 512)
+		])
+
+		if self.config_output_size == 128:
+			down_samples.extend(
+			[
+				DownSample(512, 512)
+			])
+
+		if self.config_output_size == 256:
+			down_samples.extend(
+			[
+				DownSample(512, 1024),
+				DownSample(1024, 1024)
+			])
+
+		if self.config_output_size == 512:
+			down_samples.extend(
+			[
+				DownSample(512, 1024),
+				DownSample(1024, 2048),
+				DownSample(2048, 2048)
+			])
+
+		if self.config_output_size == 1024:
+			down_samples.extend(
+			[
+				DownSample(512, 1024),
+				DownSample(1024, 2048),
+				DownSample(2048, 4096),
+				DownSample(4096, 4096)
+			])
+
+		return down_samples
+
+	def create_up_samples(self) -> nn.ModuleList:
+		up_samples = nn.ModuleList()
+
+		if self.config_output_size == 128:
+			up_samples.extend(
+			[
+				UpSample(512, 512)
+			])
+
+		if self.config_output_size == 256:
+			up_samples.extend(
+			[
+				UpSample(1024, 1024),
+				UpSample(2048, 512)
+			])
+
+		if self.config_output_size == 512:
+			up_samples.extend(
+			[
+				UpSample(2048, 2048),
+				UpSample(4096, 1024),
+				UpSample(2048, 512)
+			])
+
+		if self.config_output_size == 1024:
+			up_samples.extend(
+			[
+				UpSample(4096, 4096),
+				UpSample(8192, 2048),
+				UpSample(4096, 1024),
+				UpSample(2048, 512)
+			])
+
+		up_samples.extend(
+		[
+			UpSample(1024, 256),
+			UpSample(512, 128),
+			UpSample(256, 64),
+			UpSample(128, 32)
+		])
+
+		return up_samples
+
+	def forward(self, target_tensor : Tensor) -> Tuple[Feature, ...]:
+		down_features = []
+		up_features = []
+		temp_feature = target_tensor
+
+		for down_sample in self.down_samples:
+			temp_feature = down_sample(temp_feature)
+			down_features.append(temp_feature)
+
+		bottleneck_feature = down_features[-1]
+		temp_feature = bottleneck_feature
+
+		for index, up_sample in enumerate(self.up_samples):
+			skip_tensor = down_features[-(index + 2)]
+			temp_feature = up_sample(temp_feature, skip_tensor)
+			up_features.append(temp_feature)
+
+		final_feature = nn.functional.interpolate(temp_feature, scale_factor = 2, mode = 'bilinear', align_corners = False)
+		return bottleneck_feature, *up_features, final_feature
+
+
+class UpSample(nn.Module):
+	def __init__(self, input_channels : int, output_channels : int) -> None:
+		super().__init__()
+		self.sequences = self.create_sequences(input_channels, output_channels)
+
+	@staticmethod
+	def create_sequences(input_channels : int, output_channels : int) -> nn.Sequential:
+		return nn.Sequential(
+			nn.ConvTranspose2d(input_channels, output_channels, kernel_size = 4, stride = 2, padding = 1, bias = False),
+			nn.BatchNorm2d(output_channels),
+			nn.LeakyReLU(0.1, inplace = True)
+		)
+
+	def forward(self, input_tensor : Tensor, skip_tensor : Tensor) -> Tensor:
+		output_tensor = self.sequences(input_tensor)
+		output_tensor = torch.cat((output_tensor, skip_tensor), dim = 1)
+		return output_tensor
+
+
+class DownSample(nn.Module):
+	def __init__(self, input_channels : int, output_channels : int) -> None:
+		super().__init__()
+		self.sequences = self.create_sequences(input_channels, output_channels)
+
+	@staticmethod
+	def create_sequences(input_channels : int, output_channels : int) -> nn.Sequential:
+		return nn.Sequential(
+			nn.Conv2d(input_channels, output_channels, kernel_size = 4, stride = 2, padding = 1, bias = False),
+			nn.BatchNorm2d(output_channels),
+			nn.LeakyReLU(0.1, inplace = True)
+		)
+
+	def forward(self, input_tensor : Tensor) -> Tensor:
+		output_tensor = self.sequences(input_tensor)
+		return output_tensor

hyperswap/src/training.py

@@ -0,0 +1,248 @@
+import os
+import warnings
+from configparser import ConfigParser
+from typing import List, Tuple
+
+import torch
+import torchvision
+from lightning import LightningModule, Trainer
+from lightning.pytorch.callbacks import ModelCheckpoint
+from lightning.pytorch.loggers import TensorBoardLogger
+from torch import Tensor, nn
+from torch.utils.data import ConcatDataset, Dataset, random_split
+from torchdata.stateful_dataloader import StatefulDataLoader
+
+from .dataset import DynamicDataset
+from .helper import calc_embedding, overlay_mask
+from .models.discriminator import Discriminator
+from .models.generator import Generator
+from .models.loss import AdversarialLoss, CycleLoss, DiscriminatorLoss, FeatureLoss, GazeLoss, IdentityLoss, MaskLoss, ReconstructionLoss
+from .types import Batch, Embedding, Mask, OptimizerSet
+
+warnings.filterwarnings('ignore', category = UserWarning, module = 'torch')
+
+CONFIG_PARSER = ConfigParser()
+CONFIG_PARSER.read('config.ini')
+
+
+class HyperSwapTrainer(LightningModule):
+	def __init__(self, config_parser : ConfigParser) -> None:
+		super().__init__()
+		self.config_generator_embedder_path = config_parser.get('training.model', 'generator_embedder_path')
+		self.config_loss_embedder_path = config_parser.get('training.model', 'loss_embedder_path')
+		self.config_gazer_path = config_parser.get('training.model', 'gazer_path')
+		self.config_face_masker_path = config_parser.get('training.model', 'face_masker_path')
+		self.config_accumulate_size = config_parser.getfloat('training.trainer', 'accumulate_size')
+		self.config_learning_rate = config_parser.getfloat('training.trainer', 'learning_rate')
+		self.config_preview_frequency = config_parser.getint('training.trainer', 'preview_frequency')
+		self.generator_embedder = torch.jit.load(self.config_generator_embedder_path, map_location = 'cpu').eval()
+		self.loss_embedder = torch.jit.load(self.config_loss_embedder_path, map_location = 'cpu').eval()
+		self.gazer = torch.jit.load(self.config_gazer_path, map_location = 'cpu').eval()
+		self.face_masker = torch.jit.load(self.config_face_masker_path, map_location ='cpu').eval()
+		self.generator = Generator(config_parser)
+		self.discriminator = Discriminator(config_parser)
+		self.discriminator_loss = DiscriminatorLoss()
+		self.adversarial_loss = AdversarialLoss(config_parser)
+		self.cycle_loss = CycleLoss(config_parser)
+		self.feature_loss = FeatureLoss(config_parser)
+		self.reconstruction_loss = ReconstructionLoss(config_parser, self.loss_embedder)
+		self.identity_loss = IdentityLoss(config_parser, self.loss_embedder)
+		self.gaze_loss = GazeLoss(config_parser, self.gazer)
+		self.mask_loss = MaskLoss(config_parser, self.face_masker)
+		self.automatic_optimization = False
+
+	def forward(self, source_embedding : Embedding, target_tensor : Tensor) -> Tuple[Tensor, Mask]:
+		with torch.no_grad():
+			generator_target_features = self.generator.encode_features(target_tensor)
+			output_tensor, output_mask = self.generator(source_embedding, target_tensor, generator_target_features)
+
+		return output_tensor, output_mask
+
+	def configure_optimizers(self) -> Tuple[OptimizerSet, OptimizerSet]:
+		generator_optimizer = torch.optim.AdamW(self.generator.parameters(), lr = self.config_learning_rate, betas = (0.0, 0.999), weight_decay = 1e-4)
+		discriminator_optimizer = torch.optim.AdamW(self.discriminator.parameters(), lr = self.config_learning_rate, betas = (0.0, 0.999), weight_decay = 1e-4)
+		generator_scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(generator_optimizer, T_0 = 300, T_mult = 2)
+		discriminator_scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(discriminator_optimizer, T_0 = 300, T_mult = 2)
+
+		generator_config =\
+		{
+			'optimizer': generator_optimizer,
+			'lr_scheduler':
+			{
+				'scheduler': generator_scheduler,
+				'interval': 'step'
+			}
+		}
+		discriminator_config =\
+		{
+			'optimizer': discriminator_optimizer,
+			'lr_scheduler':
+			{
+				'scheduler': discriminator_scheduler,
+				'interval': 'step'
+			}
+		}
+		return generator_config, discriminator_config
+
+	def training_step(self, batch : Batch, batch_index : int) -> Tensor:
+		source_tensor, target_tensor = batch
+		do_update = (batch_index + 1) % self.config_accumulate_size == 0
+		generator_optimizer, discriminator_optimizer = self.optimizers() #type:ignore[attr-defined]
+
+		source_embedding = calc_embedding(self.generator_embedder, source_tensor, (0, 0, 0, 0))
+		target_embedding = calc_embedding(self.generator_embedder, target_tensor, (0, 0, 0, 0))
+		generator_target_features = self.generator.encode_features(target_tensor)
+		generator_output_tensor, generator_output_mask = self.generator(source_embedding, target_tensor, generator_target_features)
+		generator_output_features = self.generator.encode_features(generator_output_tensor)
+		cycle_output_tensor, cycle_output_mask = self.generator(target_embedding, generator_output_tensor, generator_output_features)
+		cycle_output_features = self.generator.encode_features(cycle_output_tensor)
+		discriminator_output_tensors = self.discriminator(generator_output_tensor)
+		adversarial_loss, weighted_adversarial_loss = self.adversarial_loss(discriminator_output_tensors)
+		cycle_loss, weighted_cycle_loss = self.cycle_loss(target_tensor, cycle_output_tensor, generator_target_features, cycle_output_features)
+		feature_loss, weighted_feature_loss = self.feature_loss(generator_target_features, generator_output_features)
+		reconstruction_loss, weighted_reconstruction_loss = self.reconstruction_loss(source_tensor, target_tensor, generator_output_tensor)
+		identity_loss, weighted_identity_loss = self.identity_loss(generator_output_tensor, source_tensor)
+		gaze_loss, weighted_gaze_loss = self.gaze_loss(target_tensor, generator_output_tensor)
+		mask_loss, weighted_mask_loss = self.mask_loss(target_tensor, generator_output_mask)
+		generator_loss = weighted_adversarial_loss + weighted_cycle_loss + weighted_feature_loss + weighted_reconstruction_loss + weighted_identity_loss + weighted_gaze_loss + weighted_mask_loss
+
+		discriminator_source_tensors = self.discriminator(source_tensor)
+		discriminator_output_tensors = self.discriminator(generator_output_tensor.detach())
+		discriminator_loss = self.discriminator_loss(discriminator_source_tensors, discriminator_output_tensors)
+
+		self.toggle_optimizer(generator_optimizer)
+		self.manual_backward(generator_loss)
+
+		if do_update:
+			generator_optimizer.step()
+			generator_optimizer.zero_grad()
+		self.untoggle_optimizer(generator_optimizer)
+
+		self.toggle_optimizer(discriminator_optimizer)
+		self.manual_backward(discriminator_loss)
+
+		if do_update:
+			discriminator_optimizer.step()
+			discriminator_optimizer.zero_grad()
+		self.untoggle_optimizer(discriminator_optimizer)
+
+		if self.global_step % self.config_preview_frequency == 0:
+			self.generate_preview(source_tensor, target_tensor, generator_output_tensor, generator_output_mask)
+
+		self.log('generator_loss', generator_loss, prog_bar = True)
+		self.log('discriminator_loss', discriminator_loss, prog_bar = True)
+		self.log('adversarial_loss', adversarial_loss)
+		self.log('cycle_loss', cycle_loss)
+		self.log('feature_loss', feature_loss)
+		self.log('reconstruction_loss', reconstruction_loss)
+		self.log('identity_loss', identity_loss)
+		self.log('gaze_loss', gaze_loss)
+		self.log('mask_loss', mask_loss)
+		return generator_loss
+
+	def validation_step(self, batch : Batch, batch_index : int) -> Tensor:
+		source_tensor, target_tensor = batch
+		source_embedding = calc_embedding(self.generator_embedder, source_tensor, (0, 0, 0, 0))
+		output_tensor, _ = self.forward(source_embedding, target_tensor)
+		output_embedding = calc_embedding(self.generator_embedder, output_tensor, (0, 0, 0, 0))
+		validation_score = (nn.functional.cosine_similarity(source_embedding, output_embedding).mean() + 1) * 0.5
+		self.log('validation_score', validation_score, sync_dist = True, prog_bar = True)
+		return validation_score
+
+	def generate_preview(self, source_tensor : Tensor, target_tensor : Tensor, output_tensor : Tensor, output_mask : Mask) -> None:
+		preview_limit = 8
+		preview_cells = []
+		overlay_tensor = overlay_mask(output_tensor, output_mask)
+
+		for source_tensor, target_tensor, output_tensor, overlay_tensor in zip(source_tensor[:preview_limit], target_tensor[:preview_limit], output_tensor[:preview_limit], overlay_tensor[:preview_limit]):
+			preview_cell = torch.cat([ source_tensor, target_tensor, output_tensor, overlay_tensor ], dim = 2)
+			preview_cells.append(preview_cell)
+
+		preview_cells = torch.cat(preview_cells, dim = 1).unsqueeze(0)
+		preview_grid = torchvision.utils.make_grid(preview_cells, normalize = True, scale_each = True)
+		self.logger.experiment.add_image('preview', preview_grid, self.global_step) # type:ignore[attr-defined]
+
+
+def create_loaders(dataset : Dataset[Tensor]) -> Tuple[StatefulDataLoader[Tensor], StatefulDataLoader[Tensor]]:
+	config_batch_size = CONFIG_PARSER.getint('training.loader', 'batch_size')
+	config_num_workers = CONFIG_PARSER.getint('training.loader', 'num_workers')
+
+	training_dataset, validate_dataset = split_dataset(dataset)
+	training_loader = StatefulDataLoader(training_dataset, batch_size = config_batch_size, shuffle = True, num_workers = config_num_workers, drop_last = True, pin_memory = True, persistent_workers = True)
+	validation_loader = StatefulDataLoader(validate_dataset, batch_size = config_batch_size, shuffle = False, num_workers = config_num_workers, pin_memory = True, persistent_workers = True)
+	return training_loader, validation_loader
+
+
+def split_dataset(dataset : Dataset[Tensor]) -> Tuple[Dataset[Tensor], Dataset[Tensor]]:
+	config_split_ratio = CONFIG_PARSER.getfloat('training.loader', 'split_ratio')
+
+	dataset_size = len(dataset) # type:ignore[arg-type]
+	training_size = int(dataset_size * config_split_ratio)
+	validation_size = int(dataset_size - training_size)
+	training_dataset, validate_dataset = random_split(dataset, [ training_size, validation_size ])
+	return training_dataset, validate_dataset
+
+
+def prepare_datasets(config_parser : ConfigParser) -> List[Dataset[Tensor]]:
+	datasets = []
+
+	for config_section in config_parser.sections():
+
+		if config_section.startswith('training.dataset'):
+			current_config_parser = ConfigParser()
+			current_config_parser.add_section('training.dataset')
+
+			for key, value in config_parser.items(config_section):
+				current_config_parser.set('training.dataset', key, value)
+
+			datasets.append(DynamicDataset(current_config_parser))
+
+	return datasets
+
+
+def create_trainer() -> Trainer:
+	config_max_epochs = CONFIG_PARSER.getint('training.trainer', 'max_epochs')
+	config_strategy = CONFIG_PARSER.get('training.trainer', 'strategy')
+	config_precision = CONFIG_PARSER.get('training.trainer', 'precision')
+	config_logger_path = CONFIG_PARSER.get('training.trainer', 'logger_path')
+	config_logger_name = CONFIG_PARSER.get('training.trainer', 'logger_name')
+	config_directory_path = CONFIG_PARSER.get('training.output', 'directory_path')
+	config_file_pattern = CONFIG_PARSER.get('training.output', 'file_pattern')
+	logger = TensorBoardLogger(config_logger_path, config_logger_name)
+
+	return Trainer(
+		logger = logger,
+		log_every_n_steps = 10,
+		max_epochs = config_max_epochs,
+		strategy = config_strategy,
+		precision = config_precision,
+		callbacks =
+		[
+			ModelCheckpoint(
+				monitor = 'generator_loss',
+				dirpath = config_directory_path,
+				filename = config_file_pattern,
+				every_n_train_steps = 1000,
+				save_top_k = 3,
+				save_last = True
+			)
+		],
+		val_check_interval = 1000
+	)
+
+
+def train() -> None:
+	config_resume_path = CONFIG_PARSER.get('training.output', 'resume_path')
+
+	if torch.cuda.is_available():
+		torch.set_float32_matmul_precision('high')
+
+	dataset = ConcatDataset(prepare_datasets(CONFIG_PARSER))
+	training_loader, validation_loader = create_loaders(dataset)
+	hyperswap_trainer = HyperSwapTrainer(CONFIG_PARSER)
+	trainer = create_trainer()
+
+	if os.path.isfile(config_resume_path):
+		trainer.fit(hyperswap_trainer, training_loader, validation_loader, ckpt_path = config_resume_path)
+	else:
+		trainer.fit(hyperswap_trainer, training_loader, validation_loader)

hyperswap/src/types.py

@@ -0,0 +1,24 @@
+from typing import Any, Dict, Literal, Tuple, TypeAlias
+
+from torch import Tensor
+from torch.nn import Module
+
+Batch : TypeAlias = Tuple[Tensor, Tensor]
+BatchMode = Literal['equal', 'same', 'different']
+
+Feature : TypeAlias = Tensor
+Embedding : TypeAlias = Tensor
+Mask : TypeAlias = Tensor
+Loss : TypeAlias = Tensor
+
+Padding : TypeAlias = Tuple[int, int, int, int]
+
+GeneratorModule : TypeAlias = Module
+EmbedderModule : TypeAlias = Module
+GazerModule : TypeAlias = Module
+FaceMaskerModule : TypeAlias = Module
+
+OptimizerSet : TypeAlias = Any
+
+WarpTemplate = Literal['arcface_128_v2_to_arcface_112_v2', 'ffhq_512_to_arcface_128_v2', 'vggfacehq_256_to_arcface_128_v2']
+WarpTemplateSet : TypeAlias = Dict[WarpTemplate, Tensor]