implement new training master script

cli/modules/process.py

@@ -432,7 +432,7 @@ def process_file(f: str, dst: str = None, preview: bool = False, offline: bool =
     if params.keep_original:
         resized = save_original(image)
         fn = save(resized, f, 'original')
-        log.info({ 'keep original': fn })
+        log.info({ 'original': fn })
 
     image.close()
     return i, metadata

cli/modules/sdapi.py

@@ -129,6 +129,12 @@ async def progress():
     return res
 
 
+def progresssync():
+    res = getsync('/sdapi/v1/progress?skip_current_image=true')
+    log.debug({ 'progress': res })
+    return res
+
+
 def options():
     options = getsync('/sdapi/v1/options')
     flags = getsync('/sdapi/v1/cmd-flags')

cli/train-cd.py

@@ -1,11 +0,0 @@
-#!/bin/env python
-"""
-Custom Diffusion (CD) training script
-Based on:
-- <https://www.cs.cmu.edu/~custom-diffusion/>
-- <https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/6751>
-- <https://github.com/guaneec/custom-diffusion-webui>
-"""
-
-# TBD
-

cli/train-lora.py

@@ -190,7 +190,6 @@ if __name__ == '__main__':
     options.unet_lr = args.unetlr
     options.text_encoder_lr = args.textlr
     options.train_batch_size = args.batch
-    options.network_alpha = args.alpha
     log.info({ 'train lora args': vars(options) })
     transformers.logging.set_verbosity_error()
     mem_stats()

cli/train-ti.py

@@ -93,7 +93,7 @@ args = Map({
       "tag_drop_out": 0,
       "clip_grad_mode": "disabled",
       "clip_grad_value": "0.1",
-      "latent_sampling_method": "deterministic",
+      "latent_sampling_method": "once",
       "create_image_every": -1,
       "save_embedding_every": -1,
       "save_image_with_stored_embedding": False,
@@ -107,6 +107,7 @@ args = Map({
       "preview_width": 512,
       "preview_height": 512,
       "varsize": False,
+      "use_weight": False,
   },
 })
 
@@ -498,7 +499,7 @@ async def monitor(params):
 
 
 async def main():
-    parser = argparse.ArgumentParser(description="sd train pipeline")
+    parser = argparse.ArgumentParser(description="sd train ti pipeline")
     parser.add_argument("--name", type = str, required = True, help = "embedding name, set to auto to use src folder name")
     parser.add_argument("--src", type = str, required = True, help = "source image folder or movie file")
     parser.add_argument("--init", type = str, default = "person", required = False, help = "initialization class, default: %(default)s")

cli/train/latents.py

@@ -0,0 +1,161 @@
+#!/bin/env python
+
+import os
+import sys
+import json
+import pathlib
+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 Map
+
+from rich.pretty import install as pretty_install
+from rich.traceback import install as traceback_install
+from rich.console import Console
+console = Console(log_time=True, log_time_format='%H:%M:%S-%f')
+pretty_install(console=console)
+traceback_install(console=console, extra_lines=1, width=console.width, word_wrap=False, indent_guides=False)
+
+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})
+    console.log(f'create vae latents args: {args}')
+    image_paths = train_util.glob_images(args.input)
+    if os.path.exists(args.json):
+        with open(args.json, 'rt', encoding='utf-8') as f:
+            metadata = json.load(f)
+    else:
+        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()
+    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 = 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)
+        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:
+            console.log(f'vae latents bucket: {i+1}/{len(bucket_manager.resos)} resolution: {reso} images: {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/train/options.py

@@ -0,0 +1,141 @@
+from util import Map
+
+embedding = Map({
+    "id_task": 0,
+    "embedding_name": "",
+    "learn_rate": -1,
+    "batch_size": 1,
+    "steps": 500,
+    "data_root": "",
+    "log_directory": "train/log",
+    "template_filename": "subject_filewords.txt",
+    "gradient_step": 20,
+    "training_width": 512,
+    "training_height": 512,
+    "shuffle_tags": False,
+    "tag_drop_out": 0,
+    "clip_grad_mode": "disabled",
+    "clip_grad_value": "0.1",
+    "latent_sampling_method": "deterministic",
+    "create_image_every": 0,
+    "save_embedding_every": 0,
+    "save_image_with_stored_embedding": False,
+    "preview_from_txt2img": False,
+    "preview_prompt": "",
+    "preview_negative_prompt": "blurry, duplicate, ugly, deformed, low res, watermark, text",
+    "preview_steps": 20,
+    "preview_sampler_index": 0,
+    "preview_cfg_scale": 6,
+    "preview_seed": -1,
+    "preview_width": 512,
+    "preview_height": 512,
+    "varsize": False,
+    "use_weight": False,
+})
+
+lora = Map({
+    "bucket_no_upscale": False,
+    "bucket_reso_steps": 64,
+    "cache_latents": True,
+    "caption_dropout_every_n_epochs": None,
+    "caption_dropout_rate": 0.0,
+    "caption_extension": ".txt",
+    "caption_extention": ".txt",
+    "caption_tag_dropout_rate": 0.0,
+    "clip_skip": None,
+    "color_aug": False,
+    "dataset_repeats": 1,
+    "debug_dataset": False,
+    "enable_bucket": False,
+    "face_crop_aug_range": None,
+    "flip_aug": False,
+    "full_fp16": False,
+    "gradient_accumulation_steps": 1,
+    "gradient_checkpointing": False,
+    "in_json": "",
+    "keep_tokens": None,
+    "learning_rate": 5e-05,
+    "log_prefix": None,
+    "logging_dir": None,
+    "lr_scheduler_num_cycles": 1,
+    "lr_scheduler_power": 1,
+    "lr_scheduler": "cosine",
+    "lr_warmup_steps": 0,
+    "max_bucket_reso": 1024,
+    "max_data_loader_n_workers": 8,
+    "max_grad_norm": 0.0,
+    "max_token_length": None,
+    "max_train_epochs": None,
+    "max_train_steps": 2500,
+    "mem_eff_attn": False,
+    "min_bucket_reso": 256,
+    "mixed_precision": "fp16",
+    "network_alpha": 1.0,
+    "network_args": None,
+    "network_dim": 16,
+    "network_module": "networks.lora",
+    "network_train_text_encoder_only": False,
+    "network_train_unet_only": False,
+    "network_weights": None,
+    "no_metadata": False,
+    "output_dir": "",
+    "output_name": "",
+    "persistent_data_loader_workers": False,
+    "pretrained_model_name_or_path": "",
+    "prior_loss_weight": 1.0,
+    "random_crop": False,
+    "reg_data_dir": None,
+    "resolution": "512,512",
+    "resume": None,
+    "save_every_n_epochs": None,
+    "save_last_n_epochs_state": None,
+    "save_last_n_epochs": None,
+    "save_model_as": "ckpt",
+    "save_n_epoch_ratio": None,
+    "save_precision": "fp16",
+    "save_state": False,
+    "seed": 42,
+    "shuffle_caption": False,
+    "text_encoder_lr": 5e-05,
+    "train_batch_size": 1,
+    "train_data_dir": "",
+    "training_comment": "mood-magic",
+    "unet_lr": 1e-04,
+    "use_8bit_adam": False,
+    "v_parameterization": False,
+    "v2": False,
+    "vae": None,
+    "xformers": False,
+})
+
+process = Map({
+    # general settings, do not modify
+    'format': '.jpg', # image format
+    'target_size': 512, # target resolution
+    'segmentation_model': 0, # segmentation model 0/general 1/landscape
+    'segmentation_background': (192, 192, 192), # segmentation background color
+    'blur_score': 1.8, # max score for face blur detection
+    'blur_samplesize': 60, # sample size to use for blur detection
+    'similarity_score': 0.8, # maximum similarity score before image is discarded
+    'similarity_size': 64, # base similarity detection on reduced images
+    'range_score': 0.15, # min score for face color dynamicrange detection
+    # face processing settings
+    'face_score': 0.7, # min face detection score
+    'face_pad': 0.1, # pad face image percentage
+    'face_model': 1, # which face model to use 0/close-up 1/standard
+    # body processing settings
+    '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
+    'body_pad': 0.2,  # pad body image percentage
+    'body_model': 2, # body model to use 0/low 1/medium 2/high
+    # similarity detection settings
+    # interrogate settings
+    'interrogate': False, # interrogate images
+    'interrogate_model': ['clip', 'deepdanbooru'], # interrogate models
+    'tag_limit': 5, # number of tags to extract
+    # validations
+    # tbd
+    'face_segmentation': False, # segmentation enabled
+    'body_segmentation': False, # segmentation enabled
+})

cli/train/process.py

@@ -0,0 +1,326 @@
+import os
+import sys
+import io
+import math
+import base64
+import pathlib
+import numpy as np
+import mediapipe as mp
+from PIL import Image, ImageOps
+from skimage.metrics import structural_similarity as ssim
+from scipy.stats import beta
+sys.path.append(os.path.join(os.path.dirname(__file__)))
+
+import util
+import sdapi
+import options
+
+face_model = None
+body_model = None
+segmentation_model = None
+all_images = []
+all_images_by_type = {}
+
+
+class Result(object):
+    def __init__(self, type: str, input: str, tag: str = None, requested: list = []):
+        self.type = type
+        self.input = input
+        self.output = ''
+        self.basename = ''
+        self.message = ''
+        self.image = None
+        self.caption = ''
+        self.tag = tag
+        self.tags = []
+        self.ops = []
+        self.steps = requested
+
+
+def detect_blur(image: Image):
+    # based on <https://github.com/karthik9319/Blur-Detection/>
+    bw = ImageOps.grayscale(image)
+    cx, cy = image.size[0] // 2, image.size[1] // 2
+    fft = np.fft.fft2(bw)
+    fftShift = np.fft.fftshift(fft)
+    fftShift[cy - options.process.blur_samplesize: cy + options.process.blur_samplesize, cx - options.process.blur_samplesize: cx + options.process.blur_samplesize] = 0
+    fftShift = np.fft.ifftshift(fftShift)
+    recon = np.fft.ifft2(fftShift)
+    magnitude = np.log(np.abs(recon))
+    mean = round(np.mean(magnitude), 2)
+    return mean
+
+
+def detect_dynamicrange(image: Image):
+    # based on <https://towardsdatascience.com/measuring-enhancing-image-quality-attributes-234b0f250e10>
+    data = np.asarray(image)
+    image = np.float32(data)
+    RGB = [0.299, 0.587, 0.114]
+    height, width = image.shape[:2]
+    brightness_image = np.sqrt(image[..., 0] ** 2 * RGB[0] + image[..., 1] ** 2 * RGB[1] + image[..., 2] ** 2 * RGB[2])
+    hist, _ = np.histogram(brightness_image, bins=256, range=(0, 255))
+    img_brightness_pmf = hist / (height * width)
+    dist = beta(2, 2)
+    ys = dist.pdf(np.linspace(0, 1, 256))
+    ref_pmf = ys / np.sum(ys)
+    dot_product = np.dot(ref_pmf, img_brightness_pmf)
+    squared_dist_a = np.sum(ref_pmf ** 2)
+    squared_dist_b = np.sum(img_brightness_pmf ** 2)
+    res = dot_product / math.sqrt(squared_dist_a * squared_dist_b)
+    return round(res, 2)
+
+
+def detect_simmilar(image: Image):
+    img = image.resize((options.process.similarity_size, options.process.similarity_size))
+    img = ImageOps.grayscale(img)
+    data = np.array(img)
+    similarity = 0
+    for i in all_images:
+        val = ssim(data, i, data_range=255, channel_axis=None, gradient=False, full=False)
+        if val > similarity:
+            similarity = val
+    all_images.append(data)
+    return similarity
+
+
+def segmentation(res: Result):
+    global segmentation_model
+    if segmentation_model is None:
+        segmentation_model = mp.solutions.selfie_segmentation.SelfieSegmentation(model_selection=options.process.segmentation_model)
+    data = np.array(res.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[:] = options.process.segmentation_background
+    data = np.where(condition, data, background) # consider using a joint bilateral filter instead of pure combine
+    segmented = Image.fromarray(data)
+    res.image = segmented
+    res.ops.append('segmentation')
+    return res
+
+
+def unload():
+    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
+
+
+def encode(img):
+    with io.BytesIO() as stream:
+        img.save(stream, 'JPEG')
+        values = stream.getvalue()
+        encoded = base64.b64encode(values).decode()
+        return encoded
+
+
+def reset():
+    unload()
+    global all_images_by_type
+    all_images_by_type = {}
+    global all_images
+    all_images = []
+
+
+def upscale_restore_image(res: Result, upscale: bool = False, restore: bool = False):
+    kwargs = util.Map({
+        'image': encode(res.image),
+        'codeformer_visibility': 0.0,
+        'codeformer_weight': 0.0,
+    })
+    if res.image.width >= options.process.target_size and res.image.height >= options.process.target_size:
+        upscale = False
+    if upscale:
+        kwargs.upscaler_1 = 'SwinIR_4x'
+        kwargs.upscaling_resize = 2
+        res.ops.append('upscale')
+    if restore:
+        kwargs.codeformer_visibility = 1.0
+        kwargs.codeformer_weight: 0.2
+        res.ops.append('restore')
+    if upscale or restore:
+        result = sdapi.postsync('/sdapi/v1/extra-single-image', kwargs)
+        if 'image' not in result:
+            res.message = 'failed to upscale/restore image'
+        else:
+            res.image = Image.open(io.BytesIO(base64.b64decode(result['image'])))
+    return res
+
+
+def interrogate_image(res: Result, tag: str = None):
+    caption = ''
+    tags = []
+    for model in options.process.interrogate_model:
+        json = util.Map({ 'image': encode(res.image), 'model': model })
+        result = sdapi.postsync('/sdapi/v1/interrogate', json)
+        if model == 'clip':
+            caption = result.caption if 'caption' in result else ''
+            caption = caption.split(',')[0].replace('a ', '')
+            if tag is not None:
+                caption = res.tag + ', ' + caption
+        if model == 'deepdanbooru':
+            tag = result.caption if 'caption' in result else ''
+            tags = tag.split(',')
+            tags = [t.replace('(', '').replace(')', '').replace('\\', '').split(':')[0].strip() for t in tags]
+            if tag is not None:
+                for t in res.tag.split(',')[::-1]:
+                    tags.insert(0, t.strip())
+    pos = 0 if len(tags) == 0 else 1
+    tags.insert(pos, caption.split(' ')[1])
+    if len(tags) > options.process.tag_limit:
+        tags = tags[:options.process.tag_limit]
+    res.caption = caption
+    res.tags = tags
+    res.ops.append('interrogate')
+    return res
+
+
+def resize_image(res: Result):
+    resized = res.image
+    resized.thumbnail((options.process.target_size, options.process.target_size), Image.HAMMING)
+    res.image = resized
+    res.ops.append('resize')
+    return res
+
+
+def square_image(res: Result):
+    size = max(res.image.width, res.image.height)
+    squared = Image.new('RGB', (size, size))
+    squared.paste(res.image, ((size - res.image.width) // 2, (size - res.image.height) // 2))
+    res.image = squared
+    res.ops.append('square')
+    return res
+
+
+def process_face(res: Result):
+    res.ops.append('face')
+    global face_model
+    if face_model is None:
+        face_model = mp.solutions.face_detection.FaceDetection(min_detection_confidence=options.process.face_score, model_selection=options.process.face_model)
+    results = face_model.process(np.array(res.image))
+    if results.detections is None:
+        res.message = 'no face detected'
+        res.image = None
+        return res
+    box = results.detections[0].location_data.relative_bounding_box
+    if box.xmin < 0 or box.ymin < 0 or (box.width - box.xmin) > 1 or (box.height - box.ymin) > 1:
+        res.message = 'face out of frame'
+        res.image = None
+        return res
+    x = max(0, (box.xmin - options.process.face_pad / 2) * res.image.width)
+    y = max(0, (box.ymin - options.process.face_pad / 2)* res.image.height)
+    w = min(res.image.width, (box.width + options.process.face_pad) * res.image.width)
+    h = min(res.image.height, (box.height + options.process.face_pad) * res.image.height)
+    x = max(0, x)
+    res.image = res.image.crop((x, y, x + w, y + h))
+    return res
+
+
+def process_body(res: Result):
+    res.ops.append('body')
+    global body_model
+    if body_model is None:
+        body_model = mp.solutions.pose.Pose(static_image_mode=True, min_detection_confidence=options.process.body_score, model_complexity=options.process.body_model)
+    results = body_model.process(np.array(res.image))
+    if results.pose_landmarks is None:
+        res.message = 'no body detected'
+        res.image = None
+        return res
+    x0 = [res.image.width * (i.x - options.process.body_pad / 2) for i in results.pose_landmarks.landmark if i.visibility > options.process.body_visibility]
+    y0 = [res.image.height * (i.y - options.process.body_pad / 2) for i in results.pose_landmarks.landmark if i.visibility > options.process.body_visibility]
+    x1 = [res.image.width * (i.x + options.process.body_pad / 2) for i in results.pose_landmarks.landmark if i.visibility > options.process.body_visibility]
+    y1 = [res.image.height * (i.y + options.process.body_pad / 2) for i in results.pose_landmarks.landmark if i.visibility > options.process.body_visibility]
+    if len(x0) < options.process.body_parts:
+        res.message = f'insufficient body parts detected: {len(x0)}'
+        res.image = None
+        return res
+    res.image = res.image.crop((max(0, min(x0)), max(0, min(y0)), min(res.image.width, max(x1)), min(res.image.height, max(y1))))
+    return res
+
+
+def process_original(res: Result):
+    res.ops.append('original')
+    return res
+
+
+def save_image(res: Result, folder: str):
+    if res.image is None or folder is None:
+        return res
+    all_images_by_type[res.type] = all_images_by_type.get(res.type, 0) + 1
+    res.basename = os.path.basename(res.input).split('.')[0]
+    res.basename = str(all_images_by_type[res.type]).rjust(3, '0') + '-' + res.type + '-' + res.basename
+    res.basename = os.path.join(folder, res.basename)
+    res.output = res.basename + options.process.format
+    res.image.save(res.output)
+    res.image.close()
+    res.ops.append('save')
+    return res
+
+
+def file(filename: str, folder: str, tag = None, requested = []):
+    # initialize result dict
+    res = Result(input = filename, type='unknown', tag=tag, requested = requested)
+    # open image
+    try:
+        res.image = Image.open(filename)
+        if res.image.mode == 'RGBA':
+            res.image = res.image.convert('RGB')
+        res.image = ImageOps.exif_transpose(res.image) # rotate image according to EXIF orientation
+    except Exception as e:
+        res.message = f'error opening: {e}'
+        return res
+    # primary steps
+    if 'face' in requested:
+        res.type = 'face'
+        res = process_face(res)
+    elif 'body' in requested:
+        res.type = 'body'
+        res = process_body(res)
+    elif 'original' in requested:
+        res.type = 'original'
+        res = process_original(res)
+    # validation steps
+    if res.image is None:
+        return res
+    if 'blur' in requested:
+        res.ops.append('blur')
+        val = detect_blur(res.image)
+        if val > options.process.blur_score:
+            res.message = f'blur check failed: {val}'
+            res.image = None
+    if 'range' in requested:
+        res.ops.append('range')
+        val = detect_dynamicrange(res.image)
+        if val < options.process.range_score:
+            res.message = f'dynamic range check failed: {val}'
+            res.image = None
+    if 'similarity' in requested:
+        res.ops.append('similarity')
+        val = detect_simmilar(res.image)
+        if val > options.process.similarity_score:
+            res.message = f'dynamic range check failed: {val}'
+            res.image = None
+    if res.image is None:
+        return res
+    # post processing steps
+    res = upscale_restore_image(res, 'upscale' in requested, 'restore' in requested)
+    if res.image.width < options.process.target_size or res.image.height < options.process.target_size:
+        res.message = f'low resolution: [{res.image.width}, {res.image.height}]'
+        res.image = None
+        return res
+    if 'interrogate' in requested:
+        res = interrogate_image(res, tag)
+    if 'resize' in requested:
+        res = resize_image(res)
+    if 'square' in requested:
+        res = square_image(res)
+    if 'segment' in requested:
+        res = segmentation(res)
+    # finally save image
+    res = save_image(res, folder)
+    return res

cli/train/sdapi.py

@@ -0,0 +1,111 @@
+import sys
+import json
+import aiohttp
+import asyncio
+import requests
+from util import Map
+
+
+sd_url = "http://127.0.0.1:7860" # automatic1111 api url root
+use_session = True
+timeout = aiohttp.ClientTimeout(total = None, sock_connect = 10, sock_read = None) # default value is 5 minutes, we need longer for training
+sess = None
+quiet = False
+
+
+async def result(req):
+    if req.status != 200:
+        if not use_session and sess is not None:
+            await sess.close()
+        return Map({ 'error': req.status, 'reason': req.reason, 'url': req.url })
+    else:
+        json = await req.json()
+        if type(json) == list:
+            res = json
+        elif json is None:
+            res = {}
+        else:
+            res = Map(json)
+        return res
+
+
+def resultsync(req: requests.Response):
+    if req.status_code != 200:
+        return Map({ 'error': req.status_code, 'reason': req.reason, 'url': req.url })
+    else:
+        json = req.json()
+        if type(json) == list:
+            res = json
+        elif json is None:
+            res = {}
+        else:
+            res = Map(json)
+        return res
+
+
+async def get(endpoint: str, json: dict = None):
+    global sess # pylint: disable=global-statement
+    sess = sess if sess is not None else await session()
+    async with sess.get(url = endpoint, json = json) as req:
+        res = await result(req)
+        return res
+
+
+def getsync(endpoint: str, json: dict = None):
+    req = requests.get(f'{sd_url}{endpoint}', json = json) # pylint: disable=missing-timeout
+    res = resultsync(req)
+    return res
+
+
+async def post(endpoint: str, json: dict = None):
+    global sess # pylint: disable=global-statement
+    # sess = sess if sess is not None else await session()
+    if sess and not sess.closed:
+        await sess.close()
+    sess = await session()
+    async with sess.post(url = endpoint, json = json) as req:
+        res = await result(req)
+        return res
+
+
+def postsync(endpoint: str, json: dict = None):
+    req = requests.post(f'{sd_url}{endpoint}', json = json) # pylint: disable=missing-timeout
+    res = resultsync(req)
+    return res
+
+
+def interrupt():
+    res = getsync('/sdapi/v1/progress?skip_current_image=true')
+    if 'state' in res and res.state.job_count > 0:
+        res = postsync('/sdapi/v1/interrupt')
+        return res
+    else:
+        return { 'interrupt': 'idle' }
+
+
+def progress():
+    res = getsync('/sdapi/v1/progress?skip_current_image=true')
+    return res
+
+
+def options():
+    options = getsync('/sdapi/v1/options')
+    flags = getsync('/sdapi/v1/cmd-flags')
+    return { 'options': options, 'flags': flags }
+
+
+def shutdown():
+    postsync('/sdapi/v1/shutdown')
+
+
+async def session():
+    global sess # pylint: disable=global-statement
+    time = aiohttp.ClientTimeout(total = None, sock_connect = 10, sock_read = None) # default value is 5 minutes, we need longer for training
+    sess = aiohttp.ClientSession(timeout = time, base_url = sd_url)
+    return sess
+
+
+async def close():
+    if sess is not None:
+        await asyncio.sleep(0)
+        await sess.__aexit__(None, None, None)

cli/train/train.py

@@ -0,0 +1,323 @@
+#!/bin/env python
+
+# system imports
+import os
+import re
+import gc
+import sys
+import json
+import shutil
+import pathlib
+import asyncio
+import tempfile
+import argparse
+
+# 3rd party imports
+import filetype
+from tqdm.rich import tqdm
+
+# local imports
+import util
+import sdapi
+import process
+import latents
+import options
+
+# console handler
+from rich import print
+from rich.pretty import install as pretty_install
+from rich.traceback import install as traceback_install
+from rich.console import Console
+console = Console(log_time=True, log_time_format='%H:%M:%S-%f')
+pretty_install(console=console)
+import torch, accelerate, diffusers, requests, urllib3, http
+traceback_install(console=console, extra_lines=1, width=console.width, word_wrap=False, indent_guides=False, suppress=[torch,accelerate,diffusers,asyncio,http,urllib3,requests])
+
+# lora imports
+lora_path = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, 'modules', 'lora'))
+sys.path.append(lora_path)
+lycoris_path = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, 'modules', 'lycoris'))
+sys.path.append(lycoris_path)
+import train_network
+
+
+# globals
+args = None
+valid_steps = ['original', 'face', 'body', 'blur', 'range', 'upscale', 'restore', 'interrogate', 'resize', 'square', 'segment']
+
+# methods
+
+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 = util.get_memory()
+    peak = { 'active': mem['gpu-active']['peak'], 'allocated': mem['gpu-allocated']['peak'], 'reserved': mem['gpu-reserved']['peak'] }
+    console.log(f"memory cpu: {mem.ram} gpu current: {mem.gpu} gpu peak: {peak}")
+
+
+def parse_args():
+    global args
+    parser = argparse.ArgumentParser(description = 'train lora')
+    # basic section
+    parser.add_argument('--output', '--name', type=str, default=None, required=True, help='output filename')
+    parser.add_argument('--type', type=str, choices=['embedding', 'lora', 'lycoris', 'dreambooth'], default=None, required=True, help='training type')
+    parser.add_argument('--tag', type=str, default='person', required=False, help='primary tag, default: %(default)s')
+    parser.add_argument('--process', type=str, default='original,interrogate,resize,square', required=False, help=f'list of possible processing steps: {valid_steps}, default: %(default)s')
+    parser.add_argument('--dir', type=str, default='', required=False, help='where to store processed images, default is system temp/train')
+    parser.add_argument('--input', '--dataset', type=str, default=None, required=True, help='input folder with training images')
+    parser.add_argument('--overwrite', default = False, action='store_true', help = "overwrite existing training, default: %(default)s")
+
+    # global params
+    parser.add_argument('--gradient', type=int, default=1, required=False, help='gradient accumulation steps, default: %(default)s')
+    parser.add_argument('--steps', type=int, default=2500, required=False, help='training steps, 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('--dim', '--vectors', type=int, default=40, required=False, help='network dimension, default: %(default)s')
+
+    # lora params
+    parser.add_argument('--repeats', type=int, default=10, required=False, help='number of repeats per image, default: %(default)s')
+    parser.add_argument('--alpha', type=float, default=0, required=False, help='alpha for weights scaling, default: half of dim')
+
+    args = parser.parse_args()
+
+
+def prepare_server():
+    try:
+        server_status = util.Map(sdapi.progress())
+        server_state = server_status['state']
+    except:
+        console.log('server error:', server_status)
+        exit(1)
+    if server_state['job_count'] > 0:
+        console.log('server not idle:', server_state)
+        exit(1)
+
+    server_options = util.Map(sdapi.options())
+    server_options.options.save_training_settings_to_txt = False
+    server_options.options.training_enable_tensorboard = False
+    server_options.options.training_tensorboard_save_images = False
+    server_options.options.pin_memory = True
+    server_options.options.save_optimizer_state = False
+    server_options.options.training_image_repeats_per_epoch = args.repeats
+    server_options.options.training_write_csv_every = 0
+    server_options.options.training_xattention_optimizations = False
+    sdapi.postsync('/sdapi/v1/options', server_options.options)
+    console.log(f'updated server options')
+
+
+def verify_args():
+    global args
+    server_options = util.Map(sdapi.options())
+    args.model = server_options.options['sd_model_checkpoint'].split(' [')[0]
+    args.lora_dir = server_options.flags['lora_dir']
+    if not os.path.isabs(args.model) and not os.path.exists(args.model):
+        args.model = os.path.abspath(os.path.join(args.lora_dir, os.pardir, 'Stable-diffusion', args.model))
+
+    if not os.path.exists(args.model) or not os.path.isfile(args.model):
+        console.log('cannot find model:', args.model)
+        exit(1)
+    if not os.path.exists(args.input) or not os.path.isdir(args.input):
+        console.log('cannot find training folder:', args.input)
+        exit(1)
+    if not os.path.exists(args.lora_dir) or not os.path.isdir(args.lora_dir):
+        console.log('cannot find lora folder:', args.dir)
+        exit(1)
+    if args.dir != '':
+        args.process_dir = args.dir
+    else:
+        args.process_dir = os.path.join(tempfile.gettempdir(), 'train', args.output)
+    console.log(f'args: {vars(args)}')
+
+
+async def training_loop():
+    async def async_train():
+        res = await sdapi.post('/sdapi/v1/train/embedding', options.embedding)
+        console.log(f'train embedding result: {res}')
+
+    async def async_monitor():
+        await asyncio.sleep(3)
+        res = util.Map(sdapi.progress())
+        with tqdm(desc='train embedding', total=res.state.job_count) as pbar:
+            while res.state.job_no < res.state.job_count and not res.state.interrupted and not res.state.skipped:
+                await asyncio.sleep(2)
+                prev_job = res.state.job_no
+                res = util.Map(sdapi.progress())
+                loss = re.search(r"Loss: (.*?)(?=\<)", res.textinfo)
+                if loss:
+                    pbar.set_postfix({ 'loss': loss.group(0) })
+                    pbar.update(res.state.job_no - prev_job)
+
+    a = asyncio.create_task(async_train())
+    b = asyncio.create_task(async_monitor())
+    await asyncio.gather(a, b) # wait for both pipeline and monitor to finish
+
+
+def train_embedding():
+    console.log(f'{args.type} options: {options.embedding}')
+    create_options = util.Map({
+        "name": args.output,
+        "num_vectors_per_token": args.dim,
+        "overwrite_old": False,
+        "init_text": args.tag,
+    })
+    server_options = util.Map(sdapi.options())
+    fn = os.path.join(server_options.flags.embeddings_dir, args.output) + '.pt'
+    if os.path.exists(fn) and args.overwrite:
+        console.log(f'delete existing embedding {fn}')
+        os.remove(fn)
+    else:
+        console.log(f'embedding exists {fn}')
+        return
+    console.log(f'create embedding {create_options}')
+    res = sdapi.postsync('/sdapi/v1/create/embedding', create_options)
+    if 'info' in res and 'error' in res['info']: # formatted error
+        console.log(res.info)
+    elif 'info' in res: # no error
+        asyncio.run(training_loop())
+    else: # unknown error
+        console.log(f'create embedding error {res}')
+
+
+def train_lora():
+    fn = os.path.join(args.lora_dir, args.output)
+    for ext in ['.ckpt', '.pt', '.safetensors']:
+        if os.path.exists(fn + ext):
+            if args.overwrite:
+                console.log(f'delete existing lora: {fn + ext}')
+                os.remove(fn + ext)
+            else:
+                console.log(f'lora exists: {fn + ext}')
+                return
+    console.log(f'{args.type} options: {options.lora}')
+    train_network.train(options.lora)
+
+
+def prepare_options():
+    # lora specific
+    options.lora.pretrained_model_name_or_path = args.model
+    options.lora.output_dir = args.lora_dir
+    options.lora.output_name = args.output
+    options.lora.max_train_steps = args.steps
+    options.lora.network_dim = args.dim
+    options.lora.network_alpha = args.dim // 2 if args.alpha == 0 else args.alpha
+    options.lora.gradient_accumulation_steps = args.gradient
+    options.lora.learning_rate = args.lr
+    options.lora.train_batch_size = args.batch
+    options.lora.network_alpha = args.dim // 2 if args.alpha == 0 else args.alpha
+    options.lora.train_data_dir = args.process_dir
+    if args.type == 'lycoris':
+        console.log('train using lycoris network')
+        options.lora.network_module = 'lycoris.kohya'
+        options.lora.in_json = os.path.join(args.process_dir, args.output + '.json')
+    if args.type == 'dreambooth':
+        console.log('train using dreambooth style training')
+        options.lora.in_json = None
+    if args.type == 'lora':
+        console.log('train using lora style training')
+        options.lora.in_json = os.path.join(args.process_dir, args.output + '.json')
+    if args.type == 'embedding':
+        console.log('train embedding')
+        options.lora.in_json = None
+        pass
+    # embedding specific
+    options.embedding.embedding_name = args.output
+    options.embedding.learn_rate = str(args.lr)
+    options.embedding.batch_size = args.batch
+    options.embedding.steps = args.steps
+    options.embedding.data_root = args.process_dir
+    options.embedding.log_directory = os.path.join(args.process_dir, 'log')
+    options.embedding.gradient_step = args.gradient
+
+
+def process_inputs():
+    pathlib.Path(args.process_dir).mkdir(parents=True, exist_ok=True)
+    processing_options = args.process.split(',') if isinstance(args.process, str) else args.process
+    processing_options = [opt.strip() for opt in re.split(',| ', args.process)]
+    console.log(f'processing steps: {processing_options}')
+    for step in processing_options:
+        if step not in valid_steps:
+            console.log(f'invalid processing step: {[step]}')
+            exit(1)
+    for root, _sub_dirs, folder in os.walk(args.input):
+        files = [os.path.join(root, f) for f in folder if filetype.is_image(os.path.join(root, f))]
+    console.log(f'processing input images: {len(files)}')
+    if os.path.exists(args.process_dir):
+        console.log('removing existing processed folder:', args.process_dir)
+        shutil.rmtree(args.process_dir, ignore_errors=True)
+    steps = [step for step in processing_options if step in ['face', 'body', 'original']]
+    process.reset()
+    metadata = {}
+    for step in steps:
+        if step == 'face':
+            opts = [step for step in processing_options if step not in ['body', 'original']]
+        if step == 'body':
+            opts = [step for step in processing_options if step not in ['face', 'original', 'upscale', 'restore']] # body does not perform upscale or restore
+        if step == 'original':
+            opts = [step for step in processing_options if step not in ['face', 'body', 'upscale', 'restore', 'blur', 'range', 'segment']] # original does not perform most steps
+        console.log(f'processing current step: {opts}')
+        tag = step
+        if tag == 'original' and args.tag is not None:
+            concept = args.tag.split(',')[0].strip()
+        else:
+            concept = step
+        if args.type in ['lora', 'lycoris', 'dreambooth']:
+            dir = os.path.join(args.process_dir, str(args.repeats) + '_' + concept) # separate concepts per folder
+        if args.type in ['embedding']:
+            dir = os.path.join(args.process_dir) # everything into same folder
+        console.log('processing concept:', concept)
+        console.log('processing output folder:', dir)
+        pathlib.Path(dir).mkdir(parents=True, exist_ok=True)
+        results = {}
+        for f in files:
+            res = process.file(filename = f, folder = dir, tag = args.tag, requested = opts)
+            if res.image: # valid result
+                results[res.type] = results.get(res.type, 0) + 1
+                results['total'] = results.get('total', 0) + 1
+                rel_path = res.basename.replace(os.path.commonpath([res.basename, args.process_dir]), '')
+                if rel_path.startswith(os.path.sep): rel_path = rel_path[1:]
+                metadata[rel_path] = { 'caption': res.caption, 'tags': ','.join(res.tags) }
+                if options.lora.in_json is None:
+                    with open(res.output.replace(options.process.format, '.txt'), "w") as outfile:
+                        outfile.write(res.caption)
+            console.log(f"processing {'saved' if res.image is not None else 'skipped'}: {f} => {res.output} {res.ops} {res.message}")
+    dirs = [os.path.join(args.process_dir, dir) for dir in os.listdir(args.process_dir) if os.path.isdir(os.path.join(args.process_dir, dir))]
+    console.log(f'input datasets {dirs}')
+    if options.lora.in_json is not None:
+        with open(options.lora.in_json, "w") as outfile: # write json at the end only
+            outfile.write(json.dumps(metadata, indent=2))
+        for dir in dirs: # create latents
+            latents.create_vae_latents(util.Map({ 'input': dir, 'json': options.lora.in_json }))
+            latents.unload_vae()
+    r = { 'inputs': len(files), 'outputs': results, 'metadata': options.lora.in_json }
+    console.log(f'processing steps result: {r}')
+    if args.gradient < 0:
+        console.log(f"setting gradient accumulation to number of images: {results['total']}")
+        options.lora.gradient_accumulation_steps = results['total']
+        options.embedding.gradient_step = results['total']
+    process.unload()
+
+
+if __name__ == '__main__':
+    console.log('train script for stable diffusion')
+    parse_args()
+    prepare_server()
+    verify_args()
+    prepare_options()
+    mem_stats()
+    process_inputs()
+    mem_stats()
+    try:
+        if args.type == 'embedding':
+            train_embedding()
+        if args.type == 'lora' or args.type == 'lycoris' or args.type == 'dreambooth':
+            train_lora()
+    except KeyboardInterrupt as e:
+        console.log('interrupt requested')
+        sdapi.interrupt()
+    mem_stats()
+    console.log('done')

cli/train/util.py

@@ -0,0 +1,85 @@
+#!/bin/env python
+import os
+
+import transformers
+transformers.logging.set_verbosity_error()
+
+
+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):
+    __slots__ = ('__dict__')
+    def __init__(self, *args, **kwargs):
+        super(Map, self).__init__(*args, **kwargs)
+        for arg in args:
+            if isinstance(arg, dict):
+                for k, v in arg.items():
+                    if isinstance(v, dict):
+                        v = Map(v)
+                    if isinstance(v, list):
+                        self.__convert(v)
+                    self[k] = v
+        if kwargs:
+            for k, v in kwargs.items():
+                if isinstance(v, dict):
+                    v = Map(v)
+                elif isinstance(v, list):
+                    self.__convert(v)
+                self[k] = v
+    def __convert(self, v):
+        for elem in range(0, len(v)): # pylint: disable=consider-using-enumerate
+            if isinstance(v[elem], dict):
+                v[elem] = Map(v[elem])
+            elif isinstance(v[elem], list):
+                self.__convert(v[elem])
+    def __getattr__(self, attr):
+        return self.get(attr)
+    def __setattr__(self, key, value):
+        self.__setitem__(key, value)
+    def __setitem__(self, key, value):
+        super(Map, self).__setitem__(key, value)
+        self.__dict__.update({key: value})
+    def __delattr__(self, item):
+        self.__delitem__(item)
+    def __delitem__(self, key):
+        super(Map, self).__delitem__(key)
+        del self.__dict__[key]
+
+
+if __name__ == "__main__":
+    pass

config.json

@@ -172,7 +172,7 @@
     "save_optimizer_state": false,
     "save_selected_only": true,
     "save_to_dirs": false,
-    "save_training_settings_to_txt": true,
+    "save_training_settings_to_txt": false,
     "save_txt": false,
     "sd_checkpoint_cache": 0,
     "sd_checkpoint_hash": "cc6cb27103417325ff94f52b7a5d2dde45a7515b25c255d8e396c90014281516",
@@ -196,10 +196,10 @@
     "target_side_length": 4000.0,
     "temp_dir": "",
     "training_enable_tensorboard": false,
-    "training_image_repeats_per_epoch": 1,
+    "training_image_repeats_per_epoch": 10,
     "training_tensorboard_flush_every": 120,
     "training_tensorboard_save_images": false,
-    "training_write_csv_every": 1.0,
+    "training_write_csv_every": 0.0,
     "training_xattention_optimizations": false,
     "ui_extra_networks_tab_reorder": "",
     "ui_reorder": "sampler, dimensions, cfg, seed, checkboxes, hires_fix, batch, scripts",

launch.py

@@ -8,6 +8,7 @@ import platform
 import argparse
 import json
 import warnings
+from rich import print
 
 parser = argparse.ArgumentParser(add_help=False)
 parser.add_argument("--ui-settings-file", type=str, default='config.json')
@@ -282,10 +283,12 @@ def start():
 
     rich_installed = False
     try:
-        from rich.traceback import install
+        from rich.pretty import install as pretty_install
+        from rich.traceback import install as traceback_install
         from rich.console import Console
-        console = Console()
-        install(show_locals=True, max_frames=2, extra_lines=1, word_wrap=False, width=min([console.width, 200]))
+        console = Console(log_time=True, log_time_format='%H:%M:%S-%f')
+        pretty_install(console=console)
+        traceback_install(console=console, extra_lines=1, width=console.width, word_wrap=False, indent_guides=False, show_locals=True, max_frames=2)
         rich_installed = True
     except:
         import traceback

modules/shared.py

@@ -14,6 +14,7 @@ import modules.styles
 import modules.devices as devices
 from modules import localization, extensions, script_loading, errors, ui_components, shared_items
 from modules.paths import models_path, script_path, data_path
+from rich import print
 
 
 demo = None
@@ -744,8 +745,12 @@ def html(filename):
     return ""
 
 try:
+    from rich.pretty import install as pretty_install
+    from rich.traceback import install as traceback_install
     from rich.console import Console
-    console = Console()
+    console = Console(log_time=True, log_time_format='%H:%M:%S-%f')
+    pretty_install(console=console)
+    traceback_install(console=console, extra_lines=1, width=console.width, word_wrap=False, indent_guides=False, show_locals=True, max_frames=2)
 except:
     console = None
     import traceback

ui-config.json

@@ -1441,5 +1441,21 @@
     "customscript/seed_travel.py/img2img/Desired min SSIM threshold (% of threshold)/value": 75,
     "customscript/seed_travel.py/img2img/Desired min SSIM threshold (% of threshold)/minimum": 0,
     "customscript/seed_travel.py/img2img/Desired min SSIM threshold (% of threshold)/maximum": 100,
-    "customscript/seed_travel.py/img2img/Desired min SSIM threshold (% of threshold)/step": 1
+    "customscript/seed_travel.py/img2img/Desired min SSIM threshold (% of threshold)/step": 1,
+    "customscript/seed_travel.py/txt2img/Frames per second (0 to disable video)/visible": true,
+    "customscript/seed_travel.py/txt2img/Frames per second (0 to disable video)/value": 30.0,
+    "customscript/seed_travel.py/txt2img/RIFE passes/visible": true,
+    "customscript/seed_travel.py/txt2img/RIFE passes/value": 0.0,
+    "customscript/seed_travel.py/txt2img/Drop original frames/visible": true,
+    "customscript/seed_travel.py/txt2img/Drop original frames/value": false,
+    "customscript/seed_travel.py/img2img/Frames per second (0 to disable video)/visible": true,
+    "customscript/seed_travel.py/img2img/Frames per second (0 to disable video)/value": 30.0,
+    "customscript/seed_travel.py/img2img/RIFE passes/visible": true,
+    "customscript/seed_travel.py/img2img/RIFE passes/value": 0.0,
+    "customscript/seed_travel.py/img2img/Drop original frames/visible": true,
+    "customscript/seed_travel.py/img2img/Drop original frames/value": false,
+    "customscript/movie2movie.py/txt2img/Save preprocessed/visible": true,
+    "customscript/movie2movie.py/txt2img/Save preprocessed/value": false,
+    "customscript/movie2movie.py/img2img/Save preprocessed/visible": true,
+    "customscript/movie2movie.py/img2img/Save preprocessed/value": false
 }