diff --git a/.flake8 b/.flake8 index 544f9ff..3e399c6 100644 --- a/.flake8 +++ b/.flake8 @@ -1,5 +1,5 @@ [flake8] select = E22, E23, E24, E27, E3, E4, E7, F, I1, I2 plugins = flake8-import-order -application_import_names = embedding_converter, face_swapper +application_import_names = crossface, hyperswap import-order-style = pycharm diff --git a/.github/previews/embedding_converter.png b/.github/previews/crossface.png similarity index 100% rename from .github/previews/embedding_converter.png rename to .github/previews/crossface.png diff --git a/.github/previews/face_swapper.png b/.github/previews/hyperswap.png similarity index 100% rename from .github/previews/face_swapper.png rename to .github/previews/hyperswap.png diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c2e3cd2..c73f037 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -15,8 +15,8 @@ jobs: - run: pip install flake8 - run: pip install flake8-import-order - run: pip install mypy - - run: flake8 embedding_converter face_swapper - - run: mypy embedding_converter face_swapper + - run: flake8 crossface hyperswap + - run: mypy crossface hyperswap test: runs-on: ubuntu-latest steps: diff --git a/embedding_converter/LICENSE.md b/crossface/LICENSE.md similarity index 100% rename from embedding_converter/LICENSE.md rename to crossface/LICENSE.md diff --git a/embedding_converter/README.md b/crossface/README.md similarity index 69% rename from embedding_converter/README.md rename to crossface/README.md index 2d1bf18..faf9e87 100644 --- a/embedding_converter/README.md +++ b/crossface/README.md @@ -1,7 +1,7 @@ -Embedding Converter -=================== +CrossFace +========= -> Convert face embeddings between various models. +> Seamless transform face embeddings across embedder models. ![License](https://img.shields.io/badge/license-OpenRAIL--MS-green) @@ -9,7 +9,7 @@ Embedding Converter Preview ------- -![Preview](https://raw.githubusercontent.com/facefusion/facefusion-labs/master/.github/previews/embedding_converter.png?sanitize=true) +![Preview](https://raw.githubusercontent.com/facefusion/facefusion-labs/master/.github/previews/crossface.png?sanitize=true) Installation @@ -23,7 +23,7 @@ pip install -r requirements.txt Setup ----- -This `config.ini` utilizes the MegaFace dataset to train the Embedding Converter for SimSwap. +This `config.ini` utilizes the MegaFace dataset to train the CrossFace model for SimSwap. ``` [training.dataset] @@ -50,13 +50,13 @@ max_epochs = 4096 strategy = auto precision = 16-mixed logger_path = .logs -logger_name = arcface_converter_simswap +logger_name = crossface_simswap ``` ``` [training.output] directory_path = .outputs -file_pattern = arcface_converter_simswap_{epoch}_{step} +file_pattern = crossface_simswap_{epoch}_{step} resume_path = .outputs/last.ckpt ``` @@ -64,7 +64,7 @@ resume_path = .outputs/last.ckpt [exporting] directory_path = .exports source_path = .outputs/last.ckpt -target_path = .exports/arcface_converter_simswap.onnx +target_path = .exports/crossface_simswap.onnx ir_version = 10 opset_version = 15 ``` @@ -73,7 +73,7 @@ opset_version = 15 Training -------- -Train the Embedding Converter model. +Train the model. ``` python train.py diff --git a/embedding_converter/__init__.py b/crossface/__init__.py similarity index 100% rename from embedding_converter/__init__.py rename to crossface/__init__.py diff --git a/embedding_converter/config.ini b/crossface/config.ini similarity index 100% rename from embedding_converter/config.ini rename to crossface/config.ini diff --git a/embedding_converter/export.py b/crossface/export.py similarity index 100% rename from embedding_converter/export.py rename to crossface/export.py diff --git a/embedding_converter/src/__init__.py b/crossface/src/__init__.py similarity index 100% rename from embedding_converter/src/__init__.py rename to crossface/src/__init__.py diff --git a/embedding_converter/src/dataset.py b/crossface/src/dataset.py similarity index 100% rename from embedding_converter/src/dataset.py rename to crossface/src/dataset.py diff --git a/embedding_converter/src/exporting.py b/crossface/src/exporting.py similarity index 83% rename from embedding_converter/src/exporting.py rename to crossface/src/exporting.py index f7e7549..7b85fe3 100644 --- a/embedding_converter/src/exporting.py +++ b/crossface/src/exporting.py @@ -3,7 +3,7 @@ from configparser import ConfigParser import torch -from .training import EmbeddingConverterTrainer +from .training import CrossFaceTrainer CONFIG_PARSER = ConfigParser() CONFIG_PARSER.read('config.ini') @@ -17,7 +17,7 @@ def export() -> None: config_opset_version = CONFIG_PARSER.getint('exporting', 'opset_version') os.makedirs(config_directory_path, exist_ok = True) - model = EmbeddingConverterTrainer.load_from_checkpoint(config_source_path, map_location = 'cpu').eval() + model = CrossFaceTrainer.load_from_checkpoint(config_source_path, map_location ='cpu').eval() model.ir_version = torch.tensor(config_ir_version) input_tensor = torch.randn(1, 512) torch.onnx.export(model, input_tensor, config_target_path, input_names = [ 'input' ], output_names = [ 'output' ], opset_version = config_opset_version) diff --git a/embedding_converter/src/models/__init__.py b/crossface/src/models/__init__.py similarity index 100% rename from embedding_converter/src/models/__init__.py rename to crossface/src/models/__init__.py diff --git a/embedding_converter/src/models/embedding_converter.py b/crossface/src/models/crossface.py similarity index 94% rename from embedding_converter/src/models/embedding_converter.py rename to crossface/src/models/crossface.py index bc599c3..1edf564 100644 --- a/embedding_converter/src/models/embedding_converter.py +++ b/crossface/src/models/crossface.py @@ -2,7 +2,7 @@ import torch from torch import Tensor, nn -class EmbeddingConverter(nn.Module): +class CrossFace(nn.Module): def __init__(self) -> None: super().__init__() self.layers = self.create_layers() diff --git a/embedding_converter/src/training.py b/crossface/src/training.py similarity index 91% rename from embedding_converter/src/training.py rename to crossface/src/training.py index fbd5e39..f3b140c 100644 --- a/embedding_converter/src/training.py +++ b/crossface/src/training.py @@ -11,26 +11,26 @@ from torch.utils.data import Dataset, random_split from torchdata.stateful_dataloader import StatefulDataLoader from .dataset import StaticDataset -from .models.embedding_converter import EmbeddingConverter +from .models.crossface import CrossFace from .types import Batch, Embedding, OptimizerSet CONFIG_PARSER = ConfigParser() CONFIG_PARSER.read('config.ini') -class EmbeddingConverterTrainer(LightningModule): +class CrossFaceTrainer(LightningModule): def __init__(self, config_parser : ConfigParser) -> None: super().__init__() self.config_source_path = config_parser.get('training.model', 'source_path') self.config_target_path = config_parser.get('training.model', 'target_path') self.config_learning_rate = config_parser.getfloat('training.trainer', 'learning_rate') - self.embedding_converter = EmbeddingConverter() + self.crossface = CrossFace() self.source_embedder = torch.jit.load(self.config_source_path, map_location = 'cpu').eval() self.target_embedder = torch.jit.load(self.config_target_path, map_location = 'cpu').eval() self.mse_loss = nn.MSELoss() def forward(self, source_embedding : Embedding) -> Embedding: - return self.embedding_converter(source_embedding) + return self.crossface(source_embedding) def training_step(self, batch : Batch, batch_index : int) -> Tensor: with torch.no_grad(): @@ -125,10 +125,10 @@ def train() -> None: dataset = StaticDataset(CONFIG_PARSER) training_loader, validation_loader = create_loaders(dataset) - embedding_converter_trainer = EmbeddingConverterTrainer(CONFIG_PARSER) + crossface_trainer = CrossFaceTrainer(CONFIG_PARSER) trainer = create_trainer() if os.path.exists(config_resume_path): - trainer.fit(embedding_converter_trainer, training_loader, validation_loader, ckpt_path = config_resume_path) + trainer.fit(crossface_trainer, training_loader, validation_loader, ckpt_path = config_resume_path) else: - trainer.fit(embedding_converter_trainer, training_loader, validation_loader) + trainer.fit(crossface_trainer, training_loader, validation_loader) diff --git a/embedding_converter/src/types.py b/crossface/src/types.py similarity index 100% rename from embedding_converter/src/types.py rename to crossface/src/types.py diff --git a/embedding_converter/train.py b/crossface/train.py similarity index 100% rename from embedding_converter/train.py rename to crossface/train.py diff --git a/face_swapper/LICENSE.md b/hyperswap/LICENSE.md similarity index 100% rename from face_swapper/LICENSE.md rename to hyperswap/LICENSE.md diff --git a/face_swapper/README.md b/hyperswap/README.md similarity index 87% rename from face_swapper/README.md rename to hyperswap/README.md index 3923a27..d8a710a 100644 --- a/face_swapper/README.md +++ b/hyperswap/README.md @@ -1,7 +1,7 @@ -Face Swapper -============ +HyperSwap +========= -> Face shape and occlusion aware identity transfer. +> Hyper accurate face swapping for everyone. ![License](https://img.shields.io/badge/license-ResearchRAIL--MS-red) @@ -23,12 +23,12 @@ pip install -r requirements.txt Setup ----- -This `config.ini` utilizes the VGGFace2 dataset to train the Face Swapper model. +This `config.ini` utilizes the VGGFace2 dataset to train the HyperSwap model. ``` [training.dataset] file_pattern = .datasets/vggface2/**/*.jpg -warp_template = vgg_face_hq_to_arcface_128_v2 +warp_template = vggfacehq_256_to_arcface_128_v2 transform_size = 256 batch_mode = equal batch_ratio = 0.2 @@ -92,14 +92,14 @@ max_epochs = 50 strategy = auto precision = 16-mixed logger_path = .logs -logger_name = face_swapper +logger_name = hyperswap preview_frequency = 100 ``` ``` [training.output] directory_path = .outputs -file_pattern = face_swapper_{epoch}_{step} +file_pattern = hyperswap_{epoch}_{step} resume_path = .outputs/last.ckpt ``` @@ -127,7 +127,7 @@ output_path = .outputs/output.jpg Training -------- -Train the Face Swapper model. +Train the model. ``` python train.py diff --git a/face_swapper/__init__.py b/hyperswap/__init__.py similarity index 100% rename from face_swapper/__init__.py rename to hyperswap/__init__.py diff --git a/face_swapper/config.ini b/hyperswap/config.ini similarity index 100% rename from face_swapper/config.ini rename to hyperswap/config.ini diff --git a/face_swapper/export.py b/hyperswap/export.py similarity index 100% rename from face_swapper/export.py rename to hyperswap/export.py diff --git a/face_swapper/infer.py b/hyperswap/infer.py similarity index 100% rename from face_swapper/infer.py rename to hyperswap/infer.py diff --git a/face_swapper/src/__init__.py b/hyperswap/src/__init__.py similarity index 100% rename from face_swapper/src/__init__.py rename to hyperswap/src/__init__.py diff --git a/face_swapper/src/dataset.py b/hyperswap/src/dataset.py similarity index 100% rename from face_swapper/src/dataset.py rename to hyperswap/src/dataset.py diff --git a/face_swapper/src/exporting.py b/hyperswap/src/exporting.py similarity index 90% rename from face_swapper/src/exporting.py rename to hyperswap/src/exporting.py index 34916a0..20c5020 100644 --- a/face_swapper/src/exporting.py +++ b/hyperswap/src/exporting.py @@ -5,7 +5,7 @@ from typing import Tuple import torch from torch import Tensor, nn -from .training import FaceSwapperTrainer +from .training import HyperSwapTrainer from .types import Embedding, Mask, Module CONFIG_PARSER = ConfigParser() @@ -36,7 +36,7 @@ def export() -> None: config_precision = CONFIG_PARSER.get('exporting', 'precision') os.makedirs(config_directory_path, exist_ok = True) - model = FaceSwapperTrainer.load_from_checkpoint(config_source_path, config_parser = CONFIG_PARSER, map_location = 'cpu').eval() + model = HyperSwapTrainer.load_from_checkpoint(config_source_path, config_parser = CONFIG_PARSER, map_location ='cpu').eval() if config_precision == 'half': model = HalfPrecision(model).eval() diff --git a/face_swapper/src/helper.py b/hyperswap/src/helper.py similarity index 95% rename from face_swapper/src/helper.py rename to hyperswap/src/helper.py index 982d0d9..d4a2371 100644 --- a/face_swapper/src/helper.py +++ b/hyperswap/src/helper.py @@ -10,12 +10,12 @@ WARP_TEMPLATE_SET : WarpTemplateSet =\ [ 8.75000016e-01, -1.07193451e-08, 3.80446920e-10 ], [ 1.07193451e-08, 8.75000016e-01, -1.25000007e-01 ] ]), - 'ffhq_to_arcface_128_v2': torch.tensor( + '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 ] ]), - 'vgg_face_hq_to_arcface_128_v2': torch.tensor( + 'vggfacehq_256_to_arcface_128_v2': torch.tensor( [ [ 1.01305414, -0.00140513, -0.00585911 ], [ 0.00140513, 1.01305414, 0.11169602 ] diff --git a/face_swapper/src/inferencing.py b/hyperswap/src/inferencing.py similarity index 83% rename from face_swapper/src/inferencing.py rename to hyperswap/src/inferencing.py index 6f81b2d..b9f0c54 100644 --- a/face_swapper/src/inferencing.py +++ b/hyperswap/src/inferencing.py @@ -4,7 +4,7 @@ import torch from torchvision import io from .helper import calc_embedding -from .training import FaceSwapperTrainer +from .training import HyperSwapTrainer CONFIG_PARSER = configparser.ConfigParser() CONFIG_PARSER.read('config.ini') @@ -17,7 +17,7 @@ def infer() -> None: config_target_path = CONFIG_PARSER.get('inferencing', 'target_path') config_output_path = CONFIG_PARSER.get('inferencing', 'output_path') - generator = FaceSwapperTrainer.load_from_checkpoint(config_generator_path, config_parser = CONFIG_PARSER, map_location = 'cpu').eval() + 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) diff --git a/face_swapper/src/models/__init__.py b/hyperswap/src/models/__init__.py similarity index 100% rename from face_swapper/src/models/__init__.py rename to hyperswap/src/models/__init__.py diff --git a/face_swapper/src/models/discriminator.py b/hyperswap/src/models/discriminator.py similarity index 100% rename from face_swapper/src/models/discriminator.py rename to hyperswap/src/models/discriminator.py diff --git a/face_swapper/src/models/generator.py b/hyperswap/src/models/generator.py similarity index 100% rename from face_swapper/src/models/generator.py rename to hyperswap/src/models/generator.py diff --git a/face_swapper/src/models/loss.py b/hyperswap/src/models/loss.py similarity index 100% rename from face_swapper/src/models/loss.py rename to hyperswap/src/models/loss.py diff --git a/face_swapper/src/networks/__init__.py b/hyperswap/src/networks/__init__.py similarity index 100% rename from face_swapper/src/networks/__init__.py rename to hyperswap/src/networks/__init__.py diff --git a/face_swapper/src/networks/aad.py b/hyperswap/src/networks/aad.py similarity index 100% rename from face_swapper/src/networks/aad.py rename to hyperswap/src/networks/aad.py diff --git a/face_swapper/src/networks/masknet.py b/hyperswap/src/networks/masknet.py similarity index 100% rename from face_swapper/src/networks/masknet.py rename to hyperswap/src/networks/masknet.py diff --git a/face_swapper/src/networks/nld.py b/hyperswap/src/networks/nld.py similarity index 100% rename from face_swapper/src/networks/nld.py rename to hyperswap/src/networks/nld.py diff --git a/face_swapper/src/networks/unet.py b/hyperswap/src/networks/unet.py similarity index 100% rename from face_swapper/src/networks/unet.py rename to hyperswap/src/networks/unet.py diff --git a/face_swapper/src/training.py b/hyperswap/src/training.py similarity index 97% rename from face_swapper/src/training.py rename to hyperswap/src/training.py index 67662b4..acba51f 100644 --- a/face_swapper/src/training.py +++ b/hyperswap/src/training.py @@ -25,7 +25,7 @@ CONFIG_PARSER = ConfigParser() CONFIG_PARSER.read('config.ini') -class FaceSwapperTrainer(LightningModule): +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') @@ -239,10 +239,10 @@ def train() -> None: dataset = ConcatDataset(prepare_datasets(CONFIG_PARSER)) training_loader, validation_loader = create_loaders(dataset) - face_swapper_trainer = FaceSwapperTrainer(CONFIG_PARSER) + hyperswap_trainer = HyperSwapTrainer(CONFIG_PARSER) trainer = create_trainer() if os.path.isfile(config_resume_path): - trainer.fit(face_swapper_trainer, training_loader, validation_loader, ckpt_path = config_resume_path) + trainer.fit(hyperswap_trainer, training_loader, validation_loader, ckpt_path = config_resume_path) else: - trainer.fit(face_swapper_trainer, training_loader, validation_loader) + trainer.fit(hyperswap_trainer, training_loader, validation_loader) diff --git a/face_swapper/src/types.py b/hyperswap/src/types.py similarity index 91% rename from face_swapper/src/types.py rename to hyperswap/src/types.py index d840897..43e108f 100644 --- a/face_swapper/src/types.py +++ b/hyperswap/src/types.py @@ -20,5 +20,5 @@ FaceMaskerModule : TypeAlias = Module OptimizerSet : TypeAlias = Any -WarpTemplate = Literal['arcface_128_v2_to_arcface_112_v2', 'ffhq_to_arcface_128_v2', 'vgg_face_hq_to_arcface_128_v2'] +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] diff --git a/face_swapper/tests/test_networks.py b/hyperswap/tests/test_networks.py similarity index 89% rename from face_swapper/tests/test_networks.py rename to hyperswap/tests/test_networks.py index 654a172..9f52bbb 100644 --- a/face_swapper/tests/test_networks.py +++ b/hyperswap/tests/test_networks.py @@ -3,9 +3,9 @@ from configparser import ConfigParser import pytest import torch -from face_swapper.src.networks.aad import AAD -from face_swapper.src.networks.masknet import MaskNet -from face_swapper.src.networks.unet import UNet +from hyperswap.src.networks.aad import AAD +from hyperswap.src.networks.masknet import MaskNet +from hyperswap.src.networks.unet import UNet @pytest.mark.parametrize('output_size', [ 128, 256, 512 ]) diff --git a/face_swapper/train.py b/hyperswap/train.py similarity index 100% rename from face_swapper/train.py rename to hyperswap/train.py