From 4ff9ab54f10fcc875f88cbe546af8571edf4e2dd Mon Sep 17 00:00:00 2001
From: Christian Heller <c.heller@plomlompom.de>
Date: Mon, 19 Aug 2024 06:43:12 +0200
Subject: [PATCH] Add basic exif Comment parsing.

---
 stable.py      | 50 +++++++++++++++++++++++++++++++++-----------------
 stable/core.py |  4 ++--
 2 files changed, 35 insertions(+), 19 deletions(-)

diff --git a/stable.py b/stable.py
index 26bb8e4..75eddd2 100755
--- a/stable.py
+++ b/stable.py
@@ -16,11 +16,14 @@ def save_path(count: int) -> str:
 
 def parse_args():
     parser = ArgumentParser(add_help=False)
-    parser.add_argument('-m', '--model', required=True)
+    parser.add_argument('-m', '--model')
     parser.add_argument('-o', '--output')
     parser.add_argument('-p', '--prompt')
     parser.add_argument('-H', '--help', action='help')
     parser.add_argument('-S', '--list-schedulers', action='store_true')
+    parser.add_argument('-C', '--copy-params')
+    parser.add_argument('-P', '--model_path_prefix', default='',
+                        help='useful if -m is set from -C as mere filename')
     parser.add_argument('-h', '--height', default=512, type=int,
                         help='default: 512')
     parser.add_argument('-w', '--width', default=512, type=int,
@@ -36,26 +39,39 @@ def parse_args():
     parser.add_argument('-s', '--scheduler', default=DEFAULT_SCHEDULER,
                         help=f'default: {DEFAULT_SCHEDULER}')
     parsed_args = parser.parse_args()
-    if not parsed_args.list_schedulers:
-        prefix = f'{argv[0]}: error: unless calling with -H/--help or '\
-                '-S/--list-schedulers, requiring '
-        suffix = ''
-        if (not parsed_args.output) and (not parsed_args.prompt):
-            suffix = '-o/--output and -p/--prompt'
-        elif not parsed_args.output:
-            suffix = '-o/--output'
-        elif not parsed_args.prompt:
-            suffix = '-p/--prompt'
-        if (not parsed_args.output) or (not parsed_args.prompt):
-            parser.print_usage()
-            print(f'{prefix}{suffix}')
-            sys_exit(1)
+    if parsed_args.copy_params:
+        for section in parsed_args.copy_params.split('; '):
+            key, val = section.split(': ', maxsplit=1)
+            key = key.lower()
+            key = 'randomness_seed' if key == 'seed' else key
+            if key in {'randomness_seed', 'height', 'width', 'n_steps'}:
+                val = int(val)
+            elif key in {'guidance'}:
+                val = float(val)
+            setattr(parsed_args, key, val)
+    prefix = f'{argv[0]}: error: '
+    if parsed_args.list_schedulers:
+        required = {'model': 'm'}
+        prefix += 'for -S/--list-schedulers '
+    else:
+        required = {'output': 'o', 'prompt': 'p', 'model': 'm'}
+        prefix += 'unless calling with -H/--help or -S/--list-schedulers, '
+    prefix += 'requiring '
+    suffixes = []
+    for k, v in required.items():
+        if not getattr(parsed_args, k):
+            suffixes += [f'-{v}/--{k}']
+    if suffixes:
+        parser.print_usage()
+        print(f'{prefix}{", ".join(suffixes)}')
+        sys_exit(1)
     return parsed_args
 
 
 args = parse_args()
+model_path = f'{args.model_path_prefix}{args.model}'
 if args.list_schedulers:
-    maker = ImageMaker(args.model)
+    maker = ImageMaker(model_path)
     print(f'AVAILABLE SCHEDULERS FOR MODEL {args.model}:\n')
     for name in maker.compatible_schedulers:
         print(name)
@@ -70,7 +86,7 @@ for n in range(args.quantity):
     if exists(path):
         raise Exception(f'Would overwrite file: {path}')
 
-maker = ImageMaker(args.model)
+maker = ImageMaker(model_path)
 start_seed = args.randomness_seed
 start_seed = start_seed if start_seed != 0 else randint(-(2**31-1), 2**31)
 for n in range(args.quantity):
diff --git a/stable/core.py b/stable/core.py
index 32a9372..d89d4a9 100644
--- a/stable/core.py
+++ b/stable/core.py
@@ -71,11 +71,11 @@ class ImageMaker:
     def gen_params_to_exif_comment(self):
         return f'SEED: {self.seed}; ' +\
                 f'GUIDANCE: {self.guidance}; ' +\
-                f'NUMBER OF STEPS: {self.n_steps}; ' +\
+                f'N_STEPS: {self.n_steps}; ' +\
                 f'HEIGHT: {self.height}; ' +\
                 f'WIDTH: {self.width}; ' +\
-                f'MODEL_FILE: {self.model_filename}; ' +\
                 f'SCHEDULER: {self.pipe.scheduler.__class__.__name__}; ' +\
+                f'MODEL: {self.model_filename}; ' +\
                 f'PROMPT: {self.prompt}'
 
     def gen_image_to(self, path):
-- 
2.30.2