From 80a88f5d322b0b73ac72bf4233cb85e84082fa44 Mon Sep 17 00:00:00 2001 From: Hendrik Schroeter Date: Mon, 12 Jun 2023 23:30:12 +0200 Subject: [PATCH] df: Add option to run only 1st stage of model via commandline --- DeepFilterNet/df/enhance.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/DeepFilterNet/df/enhance.py b/DeepFilterNet/df/enhance.py index 0eab4888f..9900f9e64 100644 --- a/DeepFilterNet/df/enhance.py +++ b/DeepFilterNet/df/enhance.py @@ -51,6 +51,7 @@ def main(args): log_level=args.log_level, config_allow_defaults=True, epoch=args.epoch, + mask_only=args.no_df_stage, ) suffix = suffix if args.suffix else None if args.output_dir is None: @@ -105,6 +106,7 @@ def init_df( config_allow_defaults: bool = False, epoch: Union[str, int, None] = "best", default_model: str = DEFAULT_MODEL, + mask_only: bool = False, ) -> Tuple[nn.Module, DF, str]: """Initializes and loads config, model and deep filtering state. @@ -161,7 +163,9 @@ def init_df( load_cp = epoch is not None and not (isinstance(epoch, str) and epoch.lower() == "none") if not load_cp: checkpoint_dir = None - mask_only = config("mask_only", cast=bool, section="train", default=False, save=False) + mask_only = mask_only or config( + "mask_only", cast=bool, section="train", default=False, save=False + ) model, epoch = load_model_cp(checkpoint_dir, df_state, epoch=epoch, mask_only=mask_only) if (epoch is None or epoch == 0) and load_cp: logger.error("Could not find a checkpoint") @@ -361,6 +365,7 @@ def run(): dest="suffix", help="Don't add the model suffix to the enhanced audio files", ) + parser.add_argument("--no-df-stage", action="store_true") args = parser.parse_args() main(args)