From 4c54da5749a45f37f9d8b1641f2381534735ee9c Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 14 Feb 2023 08:57:57 -0500 Subject: [PATCH] update lora processing --- cli/modules/lora-latents.py | 4 +++- cli/modules/process.py | 21 +++++++++++++++------ cli/train-lora.py | 31 +++++++++++++++++++------------ extensions-builtin/seed_travel | 2 +- 4 files changed, 38 insertions(+), 20 deletions(-) diff --git a/cli/modules/lora-latents.py b/cli/modules/lora-latents.py index 95b933ac6..d556d596b 100755 --- a/cli/modules/lora-latents.py +++ b/cli/modules/lora-latents.py @@ -3,6 +3,7 @@ import os import sys import json +import pathlib import argparse import warnings @@ -100,7 +101,8 @@ def create_vae_latents(params): image = transforms.functional.to_pil_image(img_tensor) else: image = Image.open(image_path) - image_key = image_path + image_key = os.path.basename(image_path) + image_key = os.path.join(os.path.basename(pathlib.Path(image_path).parent), pathlib.Path(image_path).stem) if image_key not in metadata: metadata[image_key] = {} reso, resized_size, ar_error = bucket_manager.select_bucket(image.width, image.height) diff --git a/cli/modules/process.py b/cli/modules/process.py index 817f97a98..4c501475f 100755 --- a/cli/modules/process.py +++ b/cli/modules/process.py @@ -52,8 +52,10 @@ params = Map({ 'segmentation_background': (192, 192, 192), # segmentation background color 'blur_samplesize': 60, # sample size to use for blur detection 'similarity_size': 64, # base similarity detection on reduced images + # original image processing settings + 'keep_original': True, # keep original image # face processing settings - 'extract_face': True, # extract face from image + 'extract_face': False, # extract face from image 'face_score': 0.7, # min face detection score 'face_pad': 0.2, # pad face image percentage 'face_model': 1, # which face model to use 0/close-up 1/standard @@ -63,7 +65,7 @@ params = Map({ 'face_upscale': True, # attempt to scale small faces 'face_segmentation': False, # segmentation enabled # body processing settings - 'extract_body': True, # extract face from image + 'extract_body': False, # extract face from image 'body_score': 0.9, # min body detection score 'body_visibility': 0.5, # min visibility score for each detected body part 'body_parts': 15, # min number of detected body parts with sufficient visibility @@ -342,14 +344,17 @@ def process_file(f: str, dst: str = None, preview: bool = False, offline: bool = else: dir = dst base = os.path.basename(f).split('.')[0] - fn = os.path.join(dir, str(i[what]).rjust(3, '0') + '-' + what + '-' + base + params.format) + parent = os.path.basename(pathlib.Path(dir)) + basename = str(i[what]).rjust(3, '0') + '-' + what + '-' + base + fn = basename + params.format # log.debug({ 'save': fn }) caption = '' + tags = '' if not preview: - img.save(fn) + img.save(os.path.join(dir, fn)) if not offline: - caption, tags = interrogate(img, fn) - metadata[fn] = { 'caption': caption, 'tags': tags } + caption, tags = interrogate(img, os.path.join(dir, fn)) + metadata[os.path.join(parent, basename)] = { 'caption': caption, 'tags': ','.join(tags) } return fn log.info({ 'processing': f }) @@ -385,6 +390,10 @@ def process_file(f: str, dst: str = None, preview: bool = False, offline: bool = else: log.debug({ 'no body': f }) + if params.keep_original: + fn = save(image, f, 'original') + log.info({ 'original': fn }) + image.close() return i, metadata diff --git a/cli/train-lora.py b/cli/train-lora.py index 7246a988a..ad8f0056f 100755 --- a/cli/train-lora.py +++ b/cli/train-lora.py @@ -135,9 +135,11 @@ if __name__ == '__main__': parser.add_argument('--nocaptions', default = False, action='store_true', help = 'skip creating captions and tags') parser.add_argument('--nolatents', default = False, action='store_true', help = 'skip generating vae latents') parser.add_argument('--offline', default = False, action='store_true', help = 'do not use webui server for processing') + parser.add_argument('--shutdown', default = False, action='store_true', help = 'shutdown webui server') parser.add_argument('--gradient', type=int, default=1, required=False, help='gradient accumulation steps, default: %(default)s') parser.add_argument('--steps', type=int, default=5000, required=False, help='training steps, default: %(default)s') parser.add_argument('--dim', type=int, default=128, required=False, help='network dimension, default: %(default)s') + parser.add_argument('--repeats', type=int, default=10, required=False, help='number of repeats per image, default: %(default)s') parser.add_argument('--batch', type=int, default=1, required=False, help='batch size, default: %(default)s') parser.add_argument('--lr', type=float, default=1e-04, required=False, help='model learning rate, default: %(default)s') parser.add_argument('--unetlr', type=float, default=1e-04, required=False, help='unet learning rate, default: %(default)s') @@ -171,37 +173,32 @@ if __name__ == '__main__': transformers.logging.set_verbosity_error() mem_stats() - dir = os.path.join(tempfile.gettempdir(), args.output, '10_processed') + dir = os.path.join(tempfile.gettempdir(), args.output, str(args.repeats) + '_processed') Path(dir).mkdir(parents=True, exist_ok=True) - json_file = os.path.join(dir, args.output + '.json') + json_file = os.path.join(tempfile.gettempdir(), args.output, args.output + '.json') options.train_data_dir = os.path.join(tempfile.gettempdir(), args.output) options.in_json = json_file for root, _sub_dirs, folder in os.walk(args.input): files = [os.path.join(root, f) for f in folder] + if not args.noprocess: # preprocess for f in files: try: res, metadata = modules.process.process_file(f = f, dst = dir, preview = False, offline = args.offline, txt = False) + with open(json_file, "w") as outfile: + outfile.write(json.dumps(metadata, indent=2)) except ValueError as e: exit(1) modules.process.unload_models() mem_stats() if args.tag is not None: for name, item in metadata.items(): - item['tags'].insert(0, args.tag) - item['tags'] = ', '.join(item['tags']) + item['caption'] = args.tag + ',' + item['caption'] + item['tags'] = args.tag + ',' + item['tags'] with open(json_file, "w") as outfile: outfile.write(json.dumps(metadata, indent=2)) - log.info({ 'processed': res, 'inputs': len(files), 'metadata': json_file, 'path': dir }) - else: - log.info({ 'skip processing': len(files), 'metadata': json_file, 'path': dir }) - - if not args.notrain: - log.info({ 'server shutdown required': True }) - modules.sdapi.shutdown() - time.sleep(1) if not args.nolatents and json_file is not None: # create latents @@ -209,5 +206,15 @@ if __name__ == '__main__': latents.unload_vae() mem_stats() + log.info({ 'processed': res, 'inputs': len(files), 'metadata': json_file, 'path': dir }) + else: + log.info({ 'skip processing': len(files), 'metadata': json_file, 'path': dir }) + + if args.shutdown: + log.info({ 'server shutdown required': True }) + modules.sdapi.shutdown() + time.sleep(1) + + if not args.notrain: train(options) mem_stats() diff --git a/extensions-builtin/seed_travel b/extensions-builtin/seed_travel index 185ae54e9..791b43b5a 160000 --- a/extensions-builtin/seed_travel +++ b/extensions-builtin/seed_travel @@ -1 +1 @@ -Subproject commit 185ae54e93c560111a7053bfd31be187b9ed1246 +Subproject commit 791b43b5a2737ab2550e6eb18a8873ebd5b77b27