From 0ec6ded0cb44224a42b6b1ee32b1a56f5e8ed9fd Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Thu, 9 Feb 2023 15:25:36 -0500 Subject: [PATCH] add lora training --- cli/modules/interrogate_git.py | 102 +++++++++++ cli/modules/lora_latents.py | 159 +++++++++++++++++ cli/modules/process.py | 61 +++++-- cli/modules/train-lora.py | 167 ++++++++++-------- cli/modules/util.py | 38 ++++ config.json | 5 +- extensions-builtin/sd-dynamic-thresholding | 2 +- extensions-builtin/sd-extension-system-info | 2 +- .../stable-diffusion-webui-images-browser | 2 +- modules/lora | 2 +- 10 files changed, 445 insertions(+), 95 deletions(-) create mode 100755 cli/modules/interrogate_git.py create mode 100755 cli/modules/lora_latents.py diff --git a/cli/modules/interrogate_git.py b/cli/modules/interrogate_git.py new file mode 100755 index 000000000..e81e50baf --- /dev/null +++ b/cli/modules/interrogate_git.py @@ -0,0 +1,102 @@ +#!/bin/env python + +import os +import json +import argparse +import torch +import filetype +from PIL import Image +from transformers import AutoProcessor, AutoModelForCausalLM +from util import log, Map + + +git_processor = None +git_model = None +device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + + +options = Map({ + 'input': '', + 'model': 'microsoft/git-large-textcaps', + 'length': 256, + 'json': '', + 'txt': False, + 'tag': '', +}) + +def cleanup(s: str): + s = s.split('"')[0].split('.')[0].split(' that')[0] + s = s.split(' with a letter')[0].split(' with the number')[0].split(' with the word')[0] + return s.replace('a ', '') + + +def load_model(args): + global git_processor + global git_model + if git_processor is None: + git_processor = AutoProcessor.from_pretrained(args.model) + if git_model is None: + git_model = AutoModelForCausalLM.from_pretrained(args.model) + git_model.to(device) + log.info( { 'interrogate loaded model': args.model }) + + +def interrogate_files(params, files): + args = Map({**options, **params}) + data = [f for f in files if filetype.is_image(f)] + log.info({ 'interrogate files': len(files), 'images': len(data), 'args': args }) + load_model(args) + metadata = {} + for image_path in data: + image = Image.open(image_path) + inputs = git_processor(images=[image], return_tensors="pt").to(device) + generated_ids = git_model.generate(pixel_values=inputs.pixel_values, max_length=args.length) + caption = git_processor.batch_decode(generated_ids, skip_special_tokens=True)[0] + caption = cleanup(caption) + tags = '' + if args.tag != '': + tags += args.tag + ',' + tags += caption.split(' ')[0] + if args.txt: + with open(os.path.splitext(image_path)[0] + '.txt', "wt", encoding='utf-8') as f: + f.write(caption + "\n") + metadata[image_path] = { 'caption': caption, 'tags': tags } + log.info({ 'interrogate image': image_path, 'caption': caption, 'tags': tags }) + + git_model.to('cpu') + if args.json != '': + with open(args.json, "wt", encoding='utf-8') as f: + f.write(json.dumps(metadata, indent=2) + "\n") + return metadata + + +def unload_git(): + global git_processor + global git_model + del git_processor + del git_model + git_processor = None + git_model = None + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description = 'image interrogate') + parser.add_argument('input', type=str, nargs='*', help='input file or directory') + parser.add_argument("--model", type=str, default="microsoft/git-large-textcaps", help="model id for GIT in HuggingFace") + parser.add_argument("--length", type=int, default=256, help="max length of caption") + parser.add_argument("--json", type=str, default='', help="output json file") + parser.add_argument("--tag", type=str, default='', help="append tag") + parser.add_argument('--txt', default = False, action='store_true', help = "write captions to text files") + params = parser.parse_args() + log.info({ 'interrogate args': vars(params) }) + if len(params.input) == 0: + parser.print_help() + exit(1) + files = [] + for loc in params.input: + if os.path.isfile(loc): + files.append(loc) + elif os.path.isdir(loc): + for root, _sub_dirs, dir in os.walk(loc): + files = [os.path.join(root, f) for f in dir] + metadata = interrogate_files(vars(params), files) diff --git a/cli/modules/lora_latents.py b/cli/modules/lora_latents.py new file mode 100755 index 000000000..a48b83dbf --- /dev/null +++ b/cli/modules/lora_latents.py @@ -0,0 +1,159 @@ +#!/bin/env python + +import os +import sys +import json +import argparse +import warnings + +import cv2 +import numpy as np +import torch +from PIL import Image +from torchvision import transforms +from tqdm import tqdm +from util import log, Map + +sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..', 'modules', 'lora')) +import library.model_util as model_util +import library.train_util as train_util + + +warnings.filterwarnings('ignore') +device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +options = Map({ + 'batch': 1, + 'input': '', + 'json': '', + 'max': 1024, + 'min': 256, + 'noupscale': False, + 'precision': 'fp32', + 'resolution': '512,512', + 'steps': 64, + 'vae': 'stabilityai/sd-vae-ft-mse' +}) +vae = None + + +def get_latents(vae, images, weight_dtype): + image_transforms = transforms.Compose([ transforms.ToTensor(), transforms.Normalize([0.5], [0.5]) ]) + img_tensors = [image_transforms(image) for image in images] + img_tensors = torch.stack(img_tensors) + img_tensors = img_tensors.to(device, weight_dtype) + with torch.no_grad(): + latents = vae.encode(img_tensors).latent_dist.sample().float().to('cpu').numpy() + return latents + + +def get_npz_filename_wo_ext(data_dir, image_key): + return os.path.join(data_dir, os.path.splitext(os.path.basename(image_key))[0]) + + +def create_vae_latents(params): + args = Map({**options, **params}) + log.info({ 'latents args': args }) + if args.steps % 8 > 0: + log.warning({ 'latents': 'resolution is not multiple of 8' }) + image_paths = train_util.glob_images(args.input) + if os.path.exists(args.json): + log.info({ 'latents metadata': args.json, 'images': len(image_paths) }) + with open(args.json, 'rt', encoding='utf-8') as f: + metadata = json.load(f) + else: + log.error({ 'latents metadata missing': args.json }) + return + if args.precision == 'fp16': + weight_dtype = torch.float16 + elif args.precision == 'bf16': + weight_dtype = torch.bfloat16 + else: + weight_dtype = torch.float32 + global vae + if vae is None: + vae = model_util.load_vae(args.vae, weight_dtype) + vae.eval() + vae.to(device, dtype=weight_dtype) + max_reso = tuple([int(t) for t in args.resolution.split(',')]) + assert len(max_reso) == 2, f'illegal resolution: {args.resolution}' + bucket_manager = train_util.BucketManager(args.noupscale, max_reso, args.min, args.max, args.steps) + if not args.noupscale: + bucket_manager.make_buckets() + else: + log.warning({ 'latents': 'min and max are ignored if noupscale is set' }) + img_ar_errors = [] + def process_batch(is_last): + for bucket in bucket_manager.buckets: + if (is_last and len(bucket) > 0) or len(bucket) >= args.batch: + latents = get_latents(vae, [img for _, img in bucket], weight_dtype) + assert latents.shape[2] == bucket[0][1].shape[0] // 8 and latents.shape[3] == bucket[0][1].shape[1] // 8, f'latent shape {latents.shape}, {bucket[0][1].shape}' + for (image_key, _), latent in zip(bucket, latents): + npz_file_name = get_npz_filename_wo_ext(args.input, image_key) + np.savez(npz_file_name, latent) + bucket.clear() + data = [[(None, ip)] for ip in image_paths] + bucket_counts = {} + for data_entry in tqdm(data, smoothing=0.0): + if data_entry[0] is None: + continue + img_tensor, image_path = data_entry[0] + if img_tensor is not None: + image = transforms.functional.to_pil_image(img_tensor) + else: + image = Image.open(image_path) + image_key = image_path + if image_key not in metadata: + metadata[image_key] = {} + reso, resized_size, ar_error = bucket_manager.select_bucket(image.width, image.height) + img_ar_errors.append(abs(ar_error)) + bucket_counts[reso] = bucket_counts.get(reso, 0) + 1 + metadata[image_key]['train_resolution'] = (reso[0] - reso[0] % 8, reso[1] - reso[1] % 8) + if not args.noupscale: + assert resized_size[0] == reso[0] or resized_size[1] == reso[1], f'internal error, resized size not match: {reso}, {resized_size}, {image.width}, {image.height}' + assert resized_size[0] >= reso[0] and resized_size[1] >= reso[1], f'internal error, resized size too small: {reso}, {resized_size}, {image.width}, {image.height}' + assert resized_size[0] >= reso[0] and resized_size[1] >= reso[1], f'internal error resized size is small: {resized_size}, {reso}' + image = np.array(image) + if resized_size[0] != image.shape[1] or resized_size[1] != image.shape[0]: + image = cv2.resize(image, resized_size, interpolation=cv2.INTER_AREA) + if resized_size[0] > reso[0]: + trim_size = resized_size[0] - reso[0] + image = image[:, trim_size//2:trim_size//2 + reso[0]] + if resized_size[1] > reso[1]: + trim_size = resized_size[1] - reso[1] + image = image[trim_size//2:trim_size//2 + reso[1]] + assert image.shape[0] == reso[1] and image.shape[1] == reso[0], f'internal error, illegal trimmed size: {image.shape}, {reso}' + bucket_manager.add_image(reso, (image_key, image)) + process_batch(False) + + process_batch(True) + vae.to('cpu') + + bucket_manager.sort() + img_ar_errors = np.array(img_ar_errors) + for i, reso in enumerate(bucket_manager.resos): + count = bucket_counts.get(reso, 0) + if count > 0: + log.info({ 'latents bucket': i, 'resolution': reso, 'count': count, 'mean ar error': np.mean(img_ar_errors) }) + with open(args.json, 'wt', encoding='utf-8') as f: + json.dump(metadata, f, indent=2) + + +def unload_vae(): + global vae + vae = None + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('input', type=str, help='directory for train images') + parser.add_argument('--json', type=str, required=True, help='metadata file to input') + parser.add_argument('--vae', type=str, required=True, help='model name or path to encode latents') + parser.add_argument('--batch', type=int, default=1, help='batch size in inference') + parser.add_argument('--resolution', type=str, default='512,512', help='max resolution in fine tuning (width,height)') + parser.add_argument('--min', type=int, default=256, help='minimum resolution for buckets') + parser.add_argument('--max', type=int, default=1024, help='maximum resolution for buckets') + parser.add_argument('--steps', type=int, default=64, help='steps of resolution for buckets, divisible by 8') + parser.add_argument('--noupscale', action='store_true', help='make bucket for each image without upscaling') + parser.add_argument('--precision', type=str, default='fp32', choices=['fp32', 'fp16', 'bf16'], help='use precision') + params = parser.parse_args() + create_vae_latents(vars(params)) diff --git a/cli/modules/process.py b/cli/modules/process.py index 75737331c..68a8b8306 100755 --- a/cli/modules/process.py +++ b/cli/modules/process.py @@ -18,8 +18,10 @@ process people images but does not have to pass all other checks as body performs its own checks - runs clip interrogation on extracted images to generate filewords """ + import os import io +import json import math import base64 import pathlib @@ -65,6 +67,9 @@ params = Map({ 'similarity_size': 64, # base similarity detection on reduced images 'interrogate_model': 'clip' # interrogate model }) +face_model = None +body_model = None +segmentation_model = None def detect_blur(image): @@ -115,15 +120,17 @@ def detect_simmilar(image): def segmentation(image): - with mp.solutions.selfie_segmentation.SelfieSegmentation(model_selection=params.segmentation_model) as selfie_segmentation: - data = np.array(image) - results = selfie_segmentation.process(data) - condition = np.stack((results.segmentation_mask,) * 3, axis=-1) > 0.1 - background = np.zeros(data.shape, dtype=np.uint8) - background[:] = params.segmentation_background - data = np.where(condition, data, background) # consider using a joint bilateral filter instead of pure combine - segmented = Image.fromarray(data) - return segmented + global segmentation_model + if segmentation_model is None: + segmentation_model = mp.solutions.selfie_segmentation.SelfieSegmentation(model_selection=params.segmentation_model) + data = np.array(image) + results = segmentation_model.process(data) + condition = np.stack((results.segmentation_mask,) * 3, axis=-1) > 0.1 + background = np.zeros(data.shape, dtype=np.uint8) + background[:] = params.segmentation_background + data = np.where(condition, data, background) # consider using a joint bilateral filter instead of pure combine + segmented = Image.fromarray(data) + return segmented def extract_face(img): @@ -135,8 +142,10 @@ def extract_face(img): resized = img.copy() resized.thumbnail((params.target_size, params.target_size), Image.HAMMING) - with mp.solutions.face_detection.FaceDetection(min_detection_confidence=params.face_score, model_selection=params.face_model) as face: - results = face.process(np.array(resized)) + global face_model + if face_model is None: + face_model = mp.solutions.face_detection.FaceDetection(min_detection_confidence=params.face_score, model_selection=params.face_model) + results = face_model.process(np.array(resized)) if results.detections is None: return None, False box = results.detections[0].location_data.relative_bounding_box @@ -196,8 +205,11 @@ def extract_body(img): scale = max(img.size[0], img.size[1]) / params.target_size resized = img.copy() resized.thumbnail((params.target_size, params.target_size), Image.HAMMING) - with mp.solutions.pose.Pose(static_image_mode=True, min_detection_confidence=params.body_score, model_complexity=params.body_model) as pose: - results = pose.process(np.array(resized)) + + global body_model + if body_model is None: + body_model = mp.solutions.pose.Pose(static_image_mode=True, min_detection_confidence=params.body_score, model_complexity=params.body_model) + results = body_model.process(np.array(resized)) if results.pose_landmarks is None: return None, False x = [resized.width * (i.x - params.body_pad / 2) for i in results.pose_landmarks.landmark if i.visibility > params.body_visibility] @@ -266,9 +278,12 @@ def interrogate(img, fn): f = open(file, 'w') f.write(caption) f.close() + return caption i = {} +metadata = Map({}) + def process_file(f: str, dst: str = None, preview: bool = False, offline: bool = False): def save(img, f, what): i[what] = i.get(what, 0) + 1 @@ -279,10 +294,12 @@ def process_file(f: str, dst: str = None, preview: bool = False, offline: bool = base = os.path.basename(f).split('.')[0] fn = os.path.join(dir, str(i[what]).rjust(3, '0') + '-' + what + '-' + base + params.format) # log.debug({ 'save': fn }) + caption = '' if not preview: img.save(fn) if not offline: - interrogate(img, fn) + caption = interrogate(img, fn) + metadata[fn] = { 'caption': caption, 'tags': [] } return fn log.info({ 'processing': f }) @@ -317,7 +334,7 @@ def process_file(f: str, dst: str = None, preview: bool = False, offline: bool = log.debug({ 'no body': f }) image.close() - return i + return i, metadata def process_images(src: str, dst: str, args = None): params.src = src @@ -339,6 +356,19 @@ def process_images(src: str, dst: str, args = None): process_file(os.path.join(root, f), dst) return i + +def unload_models(): + global face_model + if face_model is not None: + face_model = None + global body_model + if body_model is not None: + body_model = None + global segmentation_model + if segmentation_model is not None: + segmentation_model = None + + if __name__ == '__main__': # log.setLevel(logging.DEBUG) parser = argparse.ArgumentParser(description = 'image watermarking') @@ -366,3 +396,4 @@ if __name__ == '__main__': for f in files: process_file(f, params.dst, args.preview, args.offline) log.info({ 'processed': i, 'inputs': len(files) }) + # print(json.dumps(metadata, indent=2)) diff --git a/cli/modules/train-lora.py b/cli/modules/train-lora.py index 835c39a28..7ad7ef660 100755 --- a/cli/modules/train-lora.py +++ b/cli/modules/train-lora.py @@ -3,16 +3,36 @@ """ Extract approximating LoRA by SVD from two SD models Based on: + +Train LoRA with custom preprocessing, tagging and bucketing + +Disabled/broken: +- `accelerate` with *dynamo* enabled +- `xformers` due to *faketensors* requirement +- `mem_eff_attn` due to *forwardfunc* mismatch +- 'use_8bit_adam` due to *bitsandbyttes* CUDA errors + +Example: +train-lora.py --name=ryan-faf-v0 --model=/mnt/d/Models/stable-diffusion/sd-v15-runwayml.ckpt --dir=/mnt/d/Models/lora --input=~/generative/Input/ryanreid/fuckafan --dim 4 --steps 4000 +train-lora.py --name=ryan-palmsprings-v0 --model=/mnt/d/Models/stable-diffusion/sd-v15-runwayml.ckpt --dir=/mnt/d/Models/lora --input=~/generative/Input/ryanreid/palmsprings --dim 16 --steps 6000 +train-lora.py --name=ryan-random-v0 --model=/mnt/d/Models/stable-diffusion/sd-v15-runwayml.ckpt --dir=/mnt/d/Models/lora --input=~/generative/Input/ryanreid/random --dim 16 --steps 6000 +train-lora.py --name=ryan-miami-v0 --model=/mnt/d/Models/stable-diffusion/sd-v15-runwayml.ckpt --dir=/mnt/d/Models/lora --input=~/generative/Input/ryanreid/miami --dim 64 --steps 8000 +train-lora.py --name=ryan-all-v0 --model=/mnt/d/Models/stable-diffusion/sd-v15-runwayml.ckpt --dir=/mnt/d/Models/lora --input=~/generative/Input/ryanreid/all --dim 128 --steps 10000 """ import os +import gc import sys +import json import argparse import tempfile +import torch import transformers from pathlib import Path -from util import log, Map -from process import process_file +from util import log, Map, get_memory +from process import process_file, unload_models +from interrogate_git import interrogate_files, unload_git +from lora_latents import create_vae_latents, unload_vae sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..', 'modules', 'lora')) from train_network import train @@ -21,8 +41,8 @@ from train_network import train options = Map({ "v2": False, "v_parameterization": False, - "pretrained_model_name_or_path": "/mnt/d/Models/stable-diffusion/sd-v15-runwayml.ckpt", - "train_data_dir": "/tmp/rreid/img", + "pretrained_model_name_or_path": "", + "train_data_dir": "", "shuffle_caption": False, "caption_extension": ".txt", "caption_extention": None, @@ -40,12 +60,12 @@ options = Map({ "bucket_reso_steps": 64, "bucket_no_upscale": False, "reg_data_dir": None, - "in_json": "/tmp/rreid/rreid.json", + "in_json": "", "dataset_repeats": 1, - "output_dir": "/mnt/d/Models/lora/", - "output_name": "lora-rreid-random-v1", + "output_dir": "", + "output_name": "", "save_precision": "fp16", - "save_every_n_epochs": 1, + "save_every_n_epochs": None, "save_n_epoch_ratio": None, "save_last_n_epochs": None, "save_last_n_epochs_state": None, @@ -57,8 +77,8 @@ options = Map({ "mem_eff_attn": False, "xformers": False, "vae": None, - "learning_rate": 1e-05, - "max_train_steps": 5000, + "learning_rate": 1e-04, + "max_train_steps": 8000, "max_train_epochs": None, "max_data_loader_n_workers": 8, "persistent_data_loader_workers": False, @@ -86,18 +106,41 @@ options = Map({ "network_args": None, "network_train_unet_only": False, "network_train_text_encoder_only": False, - "training_comment": "mood-magic" + "training_comment": "mood-magic", + "caption_dropout_rate": 0.0, + "caption_dropout_every_n_epochs": None, + "caption_tag_dropout_rate": 0.0, }) + +def mem_stats(): + gc.collect() + if torch.cuda.is_available(): + with torch.no_grad(): + torch.cuda.empty_cache() + with torch.cuda.device('cuda'): + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + mem = get_memory() + log.info({ 'memory': { 'ram': mem.ram, 'gpu': mem.gpu } }) + + if __name__ == '__main__': parser = argparse.ArgumentParser(description = 'train lora') parser.add_argument('--model', type=str, default=None, required=True, help='original model to use a base for training') parser.add_argument('--input', type=str, default=None, required=True, help='input folder with training images') parser.add_argument('--dir', type=str, default=None, required=True, help='folder containing lora checkpoints') parser.add_argument('--name', type=str, default=None, required=True, help='lora name') - parser.add_argument('--steps', type=int, default=5000, required=False, help='training steps') - parser.add_argument('--dim', type=int, default=16, required=False, help='network dimension') - parser.add_argument("--noprocess", default = False, action='store_true', help = "skip processing and use existing input data") + parser.add_argument('--interim', type=int, default=0, help = 'save interim checkpoints after n epoch') + parser.add_argument('--noprocess', default = False, action='store_true', help = 'skip processing and use existing input data') + 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('--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('--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') + parser.add_argument('--textlr', type=float, default=5e-05, required=False, help='text encoder learning rate, default: %(default)s') args = parser.parse_args() if not os.path.exists(args.model) or not os.path.isfile(args.model): log.error({ 'lora cannot find model': args.model }) @@ -113,74 +156,50 @@ if __name__ == '__main__': options.output_name = args.name options.max_train_steps = args.steps options.network_dim = args.dim + options.gradient_accumulation_steps = args.gradient + options.save_every_n_epochs = args.interim if args.interim > 0 else None + options.learning_rate = args.lr + options.unet_lr = args.unetlr + options.text_encoder_lr = args.textlr log.info({ 'train lora args': vars(options) }) transformers.logging.set_verbosity_error() + mem_stats() if args.noprocess: - options.train_data_dir = args.input + dir = args.input + options.train_data_dir = dir + options.in_json = None else: dir = os.path.join(tempfile.gettempdir(), args.name, '10_processed') Path(dir).mkdir(parents=True, exist_ok=True) - files = [] - json_data = {} + + # preprocess for root, _sub_dirs, folder in os.walk(args.input): - for f in folder: - files.append(os.path.join(root, f)) + files = [os.path.join(root, f) for f in folder] for f in files: - res = process_file(f = f, dst = dir, preview = False, offline = True) - - log.info({ 'processed': res, 'inputs': len(files) }) - options.train_data_dir = args.input - dir = os.path.join(tempfile.gettempdir(), args.name) + res, metadata = process_file(f = f, dst = dir, preview = False, offline = True) + unload_models() + options.train_data_dir = os.path.join(tempfile.gettempdir(), args.name) + mem_stats() + + if not args.nocaptions: + # interrogate + for root, _sub_dirs, folder in os.walk(dir): + files = [os.path.join(root, f) for f in folder] + metadata = interrogate_files(Map({ 'input': dir, 'json': '', 'tag': args.name }), files) + json_file = os.path.join(dir, args.name + '.json') + with open(json_file, "w") as outfile: + outfile.write(json.dumps(metadata, indent=2)) + unload_git() + mem_stats() + options.in_json = json_file + + log.info({ 'processed': res, 'inputs': len(files), 'metadata': json_file }) + + if not args.nolatents: + # create latents + create_vae_latents(Map({ 'input': dir, 'json': json_file })) + mem_stats() train(options) - - -""" -- cannot use `accelerate` with *dynamo* enabled -- cannot use `xformers` due to *faketensors* requirement -- cannot use `mem_eff_attn` due to *forwardfunc* mismatch - -TODO - ---gradient_checkpointing ---gradient_accumulation_steps=10 ---caption_extension=txt ---in_json - -WORKING - -process.py --output "/tmp/rreid/img/10_processed" /home/vlado/generative/Input/ryanreid/random --offline - -accelerate launch --no_python --quiet --num_cpu_threads_per_process=16 python /home/vlado/dev/automatic/modules/lora/train_network.py \ ---pretrained_model_name_or_path="/mnt/d/Models/stable-diffusion/sd-v15-runwayml.ckpt" \ ---train_data_dir="/tmp/rreid/img" \ ---logging_dir="/tmp/rreid/logging" \ ---output_dir="/mnt/d/Models/lora/" \ ---output_name="lora-rreid-random-v1" \ ---resolution=512,512 \ ---learning_rate=1e-5 \ ---unet_lr=1e-3 \ ---text_encoder_lr=5e-5 \ ---lr_scheduler_num_cycles=1 \ ---lr_scheduler=cosine \ ---max_train_steps=5000 \ ---network_alpha=1 \ ---network_dim=16 \ ---network_module=networks.lora \ ---save_every_n_epochs=1 \ ---save_model_as=ckpt \ ---save_precision=fp16 \ ---mixed_precision=fp16 \ ---seed=42 \ ---train_batch_size=1 \ ---cache_latents \ - -metadata { image_key: img_md: { caption: str, tags: [] } } - -abs_path = glob_images(train_data_dir, image_key) - -}} - -./train-lora.py --model /mnt/d/Models/stable-diffusion/sd-v15-runwayml.ckpt --name rreid --dir /mnt/d/Models/lora --input ~/generative/Input/ryanreid/random/ -""" \ No newline at end of file + mem_stats() diff --git a/cli/modules/util.py b/cli/modules/util.py index 6c8c4fa1e..7241b0f83 100755 --- a/cli/modules/util.py +++ b/cli/modules/util.py @@ -3,6 +3,7 @@ generic helper methods """ +import os import string import logging @@ -28,6 +29,43 @@ def safestring(text: str): return res[:1000] +def get_memory(): + def gb(val: float): + return round(val / 1024 / 1024 / 1024, 2) + mem = {} + try: + import psutil + process = psutil.Process(os.getpid()) + res = process.memory_info() + ram_total = 100 * res.rss / process.memory_percent() + ram = { 'free': gb(ram_total - res.rss), 'used': gb(res.rss), 'total': gb(ram_total) } + mem.update({ 'ram': ram }) + except Exception as e: + mem.update({ 'ram': e }) + try: + import torch + if torch.cuda.is_available(): + s = torch.cuda.mem_get_info() + gpu = { 'free': gb(s[0]), 'used': gb(s[1] - s[0]), 'total': gb(s[1]) } + s = dict(torch.cuda.memory_stats('cuda')) + allocated = { 'current': gb(s['allocated_bytes.all.current']), 'peak': gb(s['allocated_bytes.all.peak']) } + reserved = { 'current': gb(s['reserved_bytes.all.current']), 'peak': gb(s['reserved_bytes.all.peak']) } + active = { 'current': gb(s['active_bytes.all.current']), 'peak': gb(s['active_bytes.all.peak']) } + inactive = { 'current': gb(s['inactive_split_bytes.all.current']), 'peak': gb(s['inactive_split_bytes.all.peak']) } + warnings = { 'retries': s['num_alloc_retries'], 'oom': s['num_ooms'] } + mem.update({ + 'gpu': gpu, + 'gpu-active': active, + 'gpu-allocated': allocated, + 'gpu-reserved': reserved, + 'gpu-inactive': inactive, + 'events': warnings, + }) + except: + pass + return Map(mem) + + class Map(dict): def __init__(self, *args, **kwargs): super(Map, self).__init__(*args, **kwargs) diff --git a/config.json b/config.json index 0ea27004f..2241a15eb 100644 --- a/config.json +++ b/config.json @@ -128,7 +128,7 @@ "save_training_settings_to_txt": true, "save_txt": false, "sd_checkpoint_cache": 0, - "sd_checkpoint_hash": "cc6cb27103417325ff94f52b7a5d2dde45a7515b25c255d8e396c90014281516", + "sd_checkpoint_hash": "9dc4131ec86df35efc444e0d0f128800f787a9547eedb75afb1fc8fa20df1cbb", "sd_hypernetwork_strength": 1.0, "sd_hypernetwork": "None", "sd_lora": "", @@ -187,5 +187,6 @@ "additional_networks_sort_models_by": "name", "additional_networks_model_name_filter": "", "additional_networks_xy_grid_model_metadata": "", - "additional_networks_hash_thread_count": 1.0 + "additional_networks_hash_thread_count": 1.0, + "images_txt_files": false } \ No newline at end of file diff --git a/extensions-builtin/sd-dynamic-thresholding b/extensions-builtin/sd-dynamic-thresholding index ff14beafd..0e324249e 160000 --- a/extensions-builtin/sd-dynamic-thresholding +++ b/extensions-builtin/sd-dynamic-thresholding @@ -1 +1 @@ -Subproject commit ff14beafd30865e8ecedc8b50b9015482083dbfb +Subproject commit 0e324249e0c867b42bcb9c0897e26814ad1c7512 diff --git a/extensions-builtin/sd-extension-system-info b/extensions-builtin/sd-extension-system-info index 50e4a2d86..b5d8e6a2e 160000 --- a/extensions-builtin/sd-extension-system-info +++ b/extensions-builtin/sd-extension-system-info @@ -1 +1 @@ -Subproject commit 50e4a2d86f8eb95a48799b4c3c3c31d930698c67 +Subproject commit b5d8e6a2e2dd1eafef8cf101786c243feefd3ced diff --git a/extensions-builtin/stable-diffusion-webui-images-browser b/extensions-builtin/stable-diffusion-webui-images-browser index 3390e353f..aafc81fef 160000 --- a/extensions-builtin/stable-diffusion-webui-images-browser +++ b/extensions-builtin/stable-diffusion-webui-images-browser @@ -1 +1 @@ -Subproject commit 3390e353fd06d63a0906219c01577d53390b5d6f +Subproject commit aafc81fef837e53f7b89db4281415cc63f567ca7 diff --git a/modules/lora b/modules/lora index d59189104..6b790bace 160000 --- a/modules/lora +++ b/modules/lora @@ -1 +1 @@ -Subproject commit d591891048d925e7070375fe8ac1f190d3cc9130 +Subproject commit 6b790bace69c1deb7ed110c2ea119765de30090d