From 4ff9ab54f10fcc875f88cbe546af8571edf4e2dd Mon Sep 17 00:00:00 2001 From: Christian Heller 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