#!/usr/bin/env python3
from sys import argv, exit as sys_exit, stdin
-from os.path import dirname, basename, splitext, join as path_join, exists
+from os import listdir
+from os.path import dirname, basename, splitext, join as path_join
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
from random import randint
from datetime import datetime
def run():
- def save_path(count: int) -> str:
- n_iterations = len(args.gen_paramses) * args.quantity
- filename_count = f'_{count:08}' if n_iterations > 1 else ''
- # pylint: disable=possibly-used-before-assignment
- return path_join(dir_path,
- f'{filename_sans_ext}{filename_count}.{ext}')
-
# guard against overwriting
args = parse_args()
+ paths_to_write = []
if not args.list_schedulers:
- dir_path = dirname(args.output)
+ dir_path = dirname(args.output) if dirname(args.output) else '.'
filename = basename(args.output)
filename_sans_ext, ext = splitext(filename)
- ext = ext[1:] if ext else 'png'
+ ext = ext if ext else '.png'
+ filename_with_ext = f'{filename_sans_ext}{ext}'
+ start_at_idx = 0
+ for fn in list(listdir(dir_path)):
+ if filename_with_ext == fn:
+ start_at_idx = max(start_at_idx, 1)
+ fn_sans_ext, ext_ = splitext(fn)
+ if ext_ != ext:
+ continue
+ if fn_sans_ext[:-8] == f'{filename_sans_ext}_' and\
+ set(c in '0123456789' for c in fn_sans_ext[-8:]) == {1}:
+ start_at_idx = max(start_at_idx, int(fn_sans_ext[-8:]) + 1)
for n in range(args.quantity * len(args.gen_paramses)):
- path = save_path(n)
- if exists(path):
- raise Exception(f'Would overwrite file: {path}')
+ filename_count = ''
+ if start_at_idx + n > 0\
+ or args.quantity * len(args.gen_paramses) > 1:
+ filename_count = f'_{start_at_idx+n:08}'
+ filename = f'{filename_sans_ext}{filename_count}{ext}'
+ paths_to_write += [path_join(dir_path, filename)]
# pylint: disable=wrong-import-position, import-outside-toplevel
from stable.core import ImageMaker # noqa: E402
start_seed = start_seed if start_seed != 0 else randint(SEED_MIN,
SEED_MAX)
seed_corrector = 0
+ path_idx = 0
for n in range(args.quantity):
if 0 == start_seed + n + seed_corrector:
seed_corrector += 1
gen_params.seed = start_seed + n + seed_corrector
- path = save_path(i*args.quantity + n)
+ path = paths_to_write[path_idx]
maker.set_gen_params(gen_params)
index = i*args.quantity + n + 1
until_now = datetime.now() - start_time
f'estimated time left: {str(until_end)[:-7]}): '
f'{path}; {gen_params.to_str}')
maker.gen_image_to(path)
+ path_idx += 1
run()