mirror of
https://github.com/facefusion/facefusion-labs.git
synced 2026-06-25 07:59:55 +02:00
fix
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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]:
|
||||
|
||||
Reference in New Issue
Block a user