From b0e017a16c51bede2c744aa05b233a6cee002fb5 Mon Sep 17 00:00:00 2001 From: madroid Date: Sat, 12 Oct 2024 18:44:47 +0800 Subject: [PATCH] FLUX: split small functions --- flux/dreambooth.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/flux/dreambooth.py b/flux/dreambooth.py index 4a4dbb08..59fc38a5 100644 --- a/flux/dreambooth.py +++ b/flux/dreambooth.py @@ -157,7 +157,8 @@ def save_adapters(iteration, flux, args): ) -if __name__ == "__main__": +def setup_arg_parser(): + """Set up and return the argument parser.""" parser = argparse.ArgumentParser( description="Finetune Flux to generate images with a specific subject" ) @@ -247,7 +248,11 @@ if __name__ == "__main__": ) parser.add_argument("dataset") + return parser + +if __name__ == "__main__": + parser = setup_arg_parser() args = parser.parse_args() # Load the model and set it up for LoRA training. We use the same random