CalvinBackup/facefusion-labs
1
0
Fork 0
mirror of https://github.com/facefusion/facefusion-labs.git synced 2026-06-25 07:59:55 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #55 from facefusion/fix-export

Browse Source 
Fix export
This commit is contained in:
Henry Ruhs
2025-03-11 18:33:16 +01:00
committed by GitHub
parent 4f67e045a0 5567b49a6d
commit e85aa20602
2 changed files with 8 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
face_swapper/src/exporting.py
+2 -2
View File
@@ -19,8 +19,8 @@ def export() -> None:
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)
torch.onnx.export(model, (source_tensor, target_tensor), config_target_path, input_names = [ 'source', 'target' ], output_names = [ 'output', 'mask' ], opset_version = config_opset_version)
face_swapper/src/training.py
+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]: