diff --git a/face_swapper/README.md b/face_swapper/README.md index 2ef47cb..99cfe2b 100644 --- a/face_swapper/README.md +++ b/face_swapper/README.md @@ -112,7 +112,7 @@ target_path = .exports/face_swapper.onnx target_size = 256 ir_version = 10 opset_version = 15 -precision = fp16 +precision = full ``` ``` diff --git a/face_swapper/src/exporting.py b/face_swapper/src/exporting.py index cdcbd29..34916a0 100644 --- a/face_swapper/src/exporting.py +++ b/face_swapper/src/exporting.py @@ -12,7 +12,7 @@ CONFIG_PARSER = ConfigParser() CONFIG_PARSER.read('config.ini') -class HalfPrecisionModel(nn.Module): +class HalfPrecision(nn.Module): def __init__(self, model : Module) -> None: super().__init__() self.model = model.half() @@ -38,8 +38,8 @@ def export() -> None: os.makedirs(config_directory_path, exist_ok = True) model = FaceSwapperTrainer.load_from_checkpoint(config_source_path, config_parser = CONFIG_PARSER, map_location = 'cpu').eval() - if config_precision == 'fp16': - model = HalfPrecisionModel(model).eval() + if config_precision == 'half': + model = HalfPrecision(model).eval() model.ir_version = torch.tensor(config_ir_version) source_tensor = torch.randn(1, 512)