refactor: extract argument parsing into function

fboulnois · Dec 3, 2023 · 7d559de · 7d559de
1 parent e4227b0
commit 7d559de
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/docker-entrypoint.py b/docker-entrypoint.py
@@ -150,7 +150,7 @@ def stable_diffusion_inference(p):
     print("completed pipeline:", iso_date_time(), flush=True)
 
 
-def main():
+def parse_args():
     parser = argparse.ArgumentParser(description="Create images from a text prompt.")
     parser.add_argument(
         "--attention-slicing",
@@ -289,6 +289,10 @@ def main():
     if args.prompt0 is not None:
         args.prompt = args.prompt0
 
+    return args
+
+def main():
+    args = parse_args()
     pipeline = stable_diffusion_pipeline(args)
     stable_diffusion_inference(pipeline)