diff --git a/cli/modules/process.py b/cli/modules/process.py
index 10f72178d..67453ae4c 100755
--- a/cli/modules/process.py
+++ b/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
diff --git a/cli/modules/sdapi.py b/cli/modules/sdapi.py
index e677e8dea..a8930c20b 100755
--- a/cli/modules/sdapi.py
+++ b/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')
diff --git a/cli/train-cd.py b/cli/train-cd.py
deleted file mode 100644
index 5f75113ae..000000000
--- a/cli/train-cd.py
+++ /dev/null
@@ -1,11 +0,0 @@
-#!/bin/env python
-"""
-Custom Diffusion (CD) training script
-Based on:
--
--
--
-"""
-
-# TBD
-
diff --git a/cli/train-lora.py b/cli/train-lora.py
index 56a68443c..6f82063a6 100755
--- a/cli/train-lora.py
+++ b/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()
diff --git a/cli/train-ti.py b/cli/train-ti.py
index 43052e0de..37c52c883 100755
--- a/cli/train-ti.py
+++ b/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")
diff --git a/cli/train/console.py b/cli/train/console.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/cli/train/latents.py b/cli/train/latents.py
new file mode 100755
index 000000000..671c6bdcf
--- /dev/null
+++ b/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))
diff --git a/cli/train/options.py b/cli/train/options.py
new file mode 100644
index 000000000..7bd9ebb73
--- /dev/null
+++ b/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
+})
diff --git a/cli/train/process.py b/cli/train/process.py
new file mode 100644
index 000000000..8aedcd872
--- /dev/null
+++ b/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
+ 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
+ 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
diff --git a/cli/train/sdapi.py b/cli/train/sdapi.py
new file mode 100644
index 000000000..e85148e59
--- /dev/null
+++ b/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)
diff --git a/cli/train/train.py b/cli/train/train.py
new file mode 100755
index 000000000..0eddb2384
--- /dev/null
+++ b/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')
diff --git a/cli/train/util.py b/cli/train/util.py
new file mode 100644
index 000000000..8c9aeb6c9
--- /dev/null
+++ b/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
diff --git a/config.json b/config.json
index 0b18d24ad..be232f081 100644
--- a/config.json
+++ b/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",
diff --git a/extensions-builtin/sd-extension-system-info b/extensions-builtin/sd-extension-system-info
index d09a25b4b..29062d131 160000
--- a/extensions-builtin/sd-extension-system-info
+++ b/extensions-builtin/sd-extension-system-info
@@ -1 +1 @@
-Subproject commit d09a25b4bcef33933ff91f89b77834372c579c5a
+Subproject commit 29062d131324d0bd51a70c8e8502906ba0b5f461
diff --git a/extensions-builtin/sd-webui-controlnet b/extensions-builtin/sd-webui-controlnet
index 56194e110..de8fdeff3 160000
--- a/extensions-builtin/sd-webui-controlnet
+++ b/extensions-builtin/sd-webui-controlnet
@@ -1 +1 @@
-Subproject commit 56194e110e62b70f079fbaa0eb7718e0931b9917
+Subproject commit de8fdeff373695431440b9c29e9a6fe24155e795
diff --git a/extensions-builtin/seed_travel b/extensions-builtin/seed_travel
index 3d075a59e..d12646e13 160000
--- a/extensions-builtin/seed_travel
+++ b/extensions-builtin/seed_travel
@@ -1 +1 @@
-Subproject commit 3d075a59e5493af4d0cad34c98b22bc9302e70b2
+Subproject commit d12646e1351b6ee19171e034187da3ce880bdd0c
diff --git a/launch.py b/launch.py
index e2bef4ffe..db2145037 100644
--- a/launch.py
+++ b/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
diff --git a/modules/lora b/modules/lora
index 7c1cf7f4e..432353185 160000
--- a/modules/lora
+++ b/modules/lora
@@ -1 +1 @@
-Subproject commit 7c1cf7f4eaf011e3c90e163049f85bdbadb75ef2
+Subproject commit 432353185c89c761f2d710c136e3421e6d2a7761
diff --git a/modules/lycoris b/modules/lycoris
index 2389611cc..bd21392c1 160000
--- a/modules/lycoris
+++ b/modules/lycoris
@@ -1 +1 @@
-Subproject commit 2389611cc83e0a23499c97a56a9be666a4a7b604
+Subproject commit bd21392c12914d2d13b1f51080f3634f398b3266
diff --git a/modules/shared.py b/modules/shared.py
index bbf7624be..37083a97b 100644
--- a/modules/shared.py
+++ b/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
diff --git a/ui-config.json b/ui-config.json
index f8b078e5a..304bcb7ff 100644
--- a/ui-config.json
+++ b/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
}
\ No newline at end of file