This commit is contained in:
harisreedhar
2025-03-11 22:57:08 +05:30
parent 4f67e045a0
commit 99931def84
2 changed files with 10 additions and 5 deletions
+4 -4
View File
@@ -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)
+6 -1
View File
@@ -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]: