From 94ac152340280503917d30ca0fef85e7ce7ee204 Mon Sep 17 00:00:00 2001 From: Victor Hall Date: Thu, 12 Jan 2023 13:17:38 -0500 Subject: [PATCH] bugfix --- train.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/train.py b/train.py index b6fe28a..04194f1 100644 --- a/train.py +++ b/train.py @@ -267,7 +267,7 @@ def setup_args(args): args.lr = args.lr * (total_batch_size**0.55) logging.info(f"{Fore.CYAN} * Scaling learning rate {tmp_lr} by {total_batch_size**0.5}, new value: {args.lr}{Style.RESET_ALL}") - if not os.path.exists(args.save_ckpt_dir): + if args.save_ckpt_dir is not None and not os.path.exists(args.save_ckpt_dir): os.makedirs(args.save_ckpt_dir) return args @@ -410,8 +410,8 @@ def main(args): clean_prompt = clean_filename(prompt) - result.save(f"{log_folder}/samples/gs{gs:05}-{clean_prompt[:100]}.jpg", format="JPEG", quality=95, optimize=True, progressive=False) - with open(f"{log_folder}/samples/gs{gs:05}-{clean_prompt[:100]}.txt", "w", encoding='utf-8') as f: + result.save(f"{log_folder}/samples/gs{gs:05}-{i}-{clean_prompt[:100]}.jpg", format="JPEG", quality=95, optimize=True, progressive=False) + with open(f"{log_folder}/samples/gs{gs:05}-{i}-{clean_prompt[:100]}.txt", "w", encoding='utf-8') as f: f.write(prompt) f.write(f"\n seed: {seed}")