add lora training

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)

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))

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))

cli/modules/train-lora.py

@@ -3,16 +3,36 @@
 """
 Extract approximating LoRA by SVD from two SD models
 Based on: <https://github.com/kohya-ss/sd-scripts/blob/main/networks/train_network.py>
+
+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/
-"""
+    mem_stats()

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)

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
 }