From 99931def84e02798b47f148518e9362e5606b0e8 Mon Sep 17 00:00:00 2001 From: harisreedhar Date: Tue, 11 Mar 2025 22:57:08 +0530 Subject: [PATCH] fix --- face_swapper/src/exporting.py | 8 ++++---- face_swapper/src/training.py | 7 ++++++- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/face_swapper/src/exporting.py b/face_swapper/src/exporting.py index ae0cc4d..3241ad3 100644 --- a/face_swapper/src/exporting.py +++ b/face_swapper/src/exporting.py @@ -14,13 +14,13 @@ 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_output_size = CONFIG_PARSER.getint('training.model.generator', 'output_size') config_ir_version = CONFIG_PARSER.getint('exporting', 'ir_version') config_opset_version = CONFIG_PARSER.getint('exporting', 'opset_version') os.makedirs(config_directory_path, exist_ok = True) - model = FaceSwapperTrainer.load_from_checkpoint(config_source_path, map_location = 'cpu').eval() + model = FaceSwapperTrainer.load_from_checkpoint(config_source_path, config_parser = CONFIG_PARSER, map_location = 'cpu').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' ], opset_version = config_opset_version) + target_tensor = torch.randn(1, 3, config_output_size, config_output_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) diff --git a/face_swapper/src/training.py b/face_swapper/src/training.py index 927510b..1c8f1ee 100644 --- a/face_swapper/src/training.py +++ b/face_swapper/src/training.py @@ -54,7 +54,12 @@ class FaceSwapperTrainer(LightningModule): self.automatic_optimization = False def forward(self, source_embedding : Embedding, target_tensor : Tensor) -> Tuple[Tensor, Tensor]: - output_tensor, mask_tensor = self.generator(source_embedding, target_tensor) + + with torch.no_grad(): + output_tensor = self.generator(source_embedding, target_tensor) + target_attributes = self.generator.get_attributes(target_tensor) + mask_tensor = self.masker(target_tensor, target_attributes[-1]) + return output_tensor, mask_tensor def configure_optimizers(self) -> Tuple[OptimizerSet, OptimizerSet, OptimizerSet]: