From d3d22ce5a894becb951eec03e663951b28d45135 Mon Sep 17 00:00:00 2001 From: Brian Whicheloe Date: Wed, 26 Oct 2022 09:46:56 -0700 Subject: [PATCH] Small modification to enable usage by external scripts (#956) * Make training code usable by external scripts Add parameter inputs to training and argument parsing function to allow this script to be used by an external call. * Apply suggestions from code review Co-authored-by: Patrick von Platen --- examples/dreambooth/train_dreambooth.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/examples/dreambooth/train_dreambooth.py b/examples/dreambooth/train_dreambooth.py index 6dbba9af56..6d3880b1a7 100644 --- a/examples/dreambooth/train_dreambooth.py +++ b/examples/dreambooth/train_dreambooth.py @@ -26,7 +26,7 @@ from transformers import CLIPTextModel, CLIPTokenizer logger = get_logger(__name__) -def parse_args(): +def parse_args(input_args): parser = argparse.ArgumentParser(description="Simple example of a training script.") parser.add_argument( "--pretrained_model_name_or_path", @@ -196,7 +196,11 @@ def parse_args(): ) parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") - args = parser.parse_args() + if input_args is not None: + args = parser.parse_args(input_args) + else: + args = parser.parse_args() + env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) if env_local_rank != -1 and env_local_rank != args.local_rank: args.local_rank = env_local_rank @@ -319,8 +323,7 @@ def get_full_repo_name(model_id: str, organization: Optional[str] = None, token: return f"{organization}/{model_id}" -def main(): - args = parse_args() +def main(args): logging_dir = Path(args.output_dir, args.logging_dir) accelerator = Accelerator( @@ -653,4 +656,5 @@ def main(): if __name__ == "__main__": - main() + args = parse_args() + main(args)