From 368da824aa971a29a22f9597ed8da6b6a1becfc6 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Thu, 6 Mar 2025 08:47:59 +0100 Subject: [PATCH] Fix CI --- embedding_converter/src/exporting.py | 2 +- face_swapper/src/training.py | 4 ++-- face_swapper/src/types.py | 3 ++- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/embedding_converter/src/exporting.py b/embedding_converter/src/exporting.py index b9f048d..0a39d69 100644 --- a/embedding_converter/src/exporting.py +++ b/embedding_converter/src/exporting.py @@ -19,7 +19,7 @@ def export() -> None: 'opset_version': CONFIG.getint('exporting', 'opset_version') } - makedirs(config.get('directory_path'), exist_ok = True) + makedirs(config.get('directory_path'), exist_ok = True) # type:ignore[arg-type] model = EmbeddingConverterTrainer.load_from_checkpoint(config.get('source_path'), map_location = 'cpu') model.eval() model.ir_version = torch.tensor(config.get('ir_version')) diff --git a/face_swapper/src/training.py b/face_swapper/src/training.py index 6472c90..4873913 100644 --- a/face_swapper/src/training.py +++ b/face_swapper/src/training.py @@ -18,7 +18,7 @@ from .helper import calc_embedding from .models.discriminator import Discriminator from .models.generator import Generator from .models.loss import AdversarialLoss, AttributeLoss, DiscriminatorLoss, GazeLoss, IdentityLoss, MotionLoss, ReconstructionLoss -from .types import Batch, BatchMode, Embedding, OptimizerConfig, WarpTemplate +from .types import Batch, BatchMode, Embedding, OptimizerSet, WarpTemplate warnings.filterwarnings('ignore', category = UserWarning, module = 'torch') @@ -52,7 +52,7 @@ class FaceSwapperTrainer(lightning.LightningModule): output_tensor = self.generator(source_embedding, target_tensor) return output_tensor - def configure_optimizers(self) -> Tuple[OptimizerConfig, OptimizerConfig]: + def configure_optimizers(self) -> Tuple[OptimizerSet, OptimizerSet]: learning_rate = CONFIG.getfloat('training.trainer', 'learning_rate') generator_optimizer = torch.optim.AdamW(self.generator.parameters(), lr = learning_rate, betas = (0.0, 0.999), weight_decay = 1e-4) discriminator_optimizer = torch.optim.AdamW(self.discriminator.parameters(), lr = learning_rate, betas = (0.0, 0.999), weight_decay = 1e-4) diff --git a/face_swapper/src/types.py b/face_swapper/src/types.py index 5aed9bd..34ebe12 100644 --- a/face_swapper/src/types.py +++ b/face_swapper/src/types.py @@ -17,7 +17,8 @@ EmbedderModule : TypeAlias = Module GazerModule : TypeAlias = Module MotionExtractorModule : TypeAlias = Module -OptimizerConfig : TypeAlias = Any +Config : TypeAlias = Dict[str, Any] +OptimizerSet : TypeAlias = Any WarpTemplate = Literal['vgg_face_hq_to_arcface_128_v2', 'arcface_128_v2_to_arcface_112_v2'] WarpTemplateSet : TypeAlias = Dict[WarpTemplate, Tensor]