add train preprocess options

This commit is contained in:
Vladimir Mandic
2023-05-05 09:05:59 -04:00
parent 19007ca695
commit fe496f4ebc
9 changed files with 56 additions and 47 deletions
-1
View File
@@ -24,7 +24,6 @@ import latents
import options
# console handler
from rich import print # pylint: disable=redefined-builtin
from rich.pretty import install as pretty_install
from rich.traceback import install as traceback_install
from rich.console import Console
+5 -6
View File
@@ -20,7 +20,7 @@ class Dot(dict): # dot notation access to dictionary attributes
log = logging.getLogger("sd")
args = Dot({ 'debug': False, 'upgrade': False, 'skip_update': False, 'skip_extensions': False, 'skip_requirements': False, 'skip_git': False, 'reset': False, 'use_directml': False, 'use_ipex': False, 'experimental': False, 'test': False })
args = Dot({ 'debug': False, 'upgrade': False, 'skip_update': False, 'skip_extensions': False, 'skip_requirements': False, 'skip_git': False, 'reset': False, 'use_directml': False, 'use_ipex': False, 'experimental': False, 'test': False, 'tls_selfsign': False })
quick_allowed = True
errors = 0
opts = {}
@@ -146,9 +146,9 @@ def update(folder):
log.debug(f'Setting branch: {folder} / {branch}')
git(f'checkout {branch}', folder)
if branch is None:
git('pull --autostash --rebase', folder)
git('pull --autostash --rebase --force', folder)
else:
git(f'pull origin {branch} --autostash --rebase', folder)
git(f'pull origin {branch} --autostash --rebase --force', folder)
# branch = git('branch', folder)
@@ -239,7 +239,6 @@ def check_torch():
log.info(f'Torch backend: DirectML ({version})')
for i in range(0, torch_directml.device_count()):
log.info(f'Torch detected GPU: {torch_directml.device_name(i)}')
log.info(f'DirectML default device: {torch_directml.device_name(torch_directml.default_device())}')
except:
log.warning("Torch repoorts CUDA not available")
except Exception as e:
@@ -435,7 +434,7 @@ def check_extensions():
# check version of the main repo and optionally upgrade it
def check_version(offline=False):
def check_version(offline=False): # pylint: disable=unused-argument
if not os.path.exists('.git'):
log.error('Not a git repository')
if not args.ignore:
@@ -464,7 +463,7 @@ def check_version(offline=False):
try:
git('add .')
git('stash')
update('.')
update('.') # TODO: can fail
# git('git stash pop')
ver = git('log -1 --pretty=format:"%h %ad"')
log.info(f'Upgraded to version: {ver}')
+19 -16
View File
@@ -2,10 +2,9 @@ import os
import numpy as np
from PIL import Image, ImageOps, ImageFilter, ImageEnhance, ImageChops, UnidentifiedImageError
import modules.scripts
from modules import sd_samplers
from modules import sd_samplers, shared
from modules.generation_parameters_copypaste import create_override_settings_dict
from modules.processing import Processed, StableDiffusionProcessingImg2Img, process_images
from modules.shared import opts, debug, state, listfiles, sd_model, log
from modules.ui import plaintext_to_html
import modules.processing as processing
from modules.memstats import memory_stats
@@ -13,23 +12,23 @@ from modules.memstats import memory_stats
def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args):
processing.fix_seed(p)
images = listfiles(input_dir)
images = shared.listfiles(input_dir)
is_inpaint_batch = False
if inpaint_mask_dir:
inpaint_masks = listfiles(inpaint_mask_dir)
inpaint_masks = shared.listfiles(inpaint_mask_dir)
is_inpaint_batch = len(inpaint_masks) > 0
if is_inpaint_batch:
log.info(f"\nInpaint batch is enabled. {len(inpaint_masks)} masks found.")
log.info(f"Will process {len(images)} images, creating {p.n_iter * p.batch_size} new images for each.")
shared.log.info(f"\nInpaint batch is enabled. {len(inpaint_masks)} masks found.")
shared.log.info(f"Will process {len(images)} images, creating {p.n_iter * p.batch_size} new images for each.")
save_normally = output_dir == ''
p.do_not_save_grid = True
p.do_not_save_samples = not save_normally
state.job_count = len(images) * p.n_iter
shared.state.job_count = len(images) * p.n_iter
for i, image in enumerate(images):
state.job = f"{i+1} out of {len(images)}"
if state.skipped:
state.skipped = False
if state.interrupted:
shared.state.job = f"{i+1} out of {len(images)}"
if shared.state.skipped:
shared.state.skipped = False
if shared.state.interrupted:
break
try:
img = Image.open(image)
@@ -61,11 +60,15 @@ def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args):
if processed_image.mode == 'RGBA':
processed_image = processed_image.convert("RGB")
processed_image.save(os.path.join(output_dir, filename))
debug(f'Processed: {len(images)} Memory: {memory_stats()} batch')
shared.debug(f'Processed: {len(images)} Memory: {memory_stats()} batch')
def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_styles, init_img, sketch, init_img_with_mask, inpaint_color_sketch, inpaint_color_sketch_orig, init_img_inpaint, init_mask_inpaint, steps: int, sampler_index: int, mask_blur: int, mask_alpha: float, inpainting_fill: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, image_cfg_scale: float, denoising_strength: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, selected_scale_tab: int, height: int, width: int, scale_by: float, resize_mode: int, inpaint_full_res: bool, inpaint_full_res_padding: int, inpainting_mask_invert: int, img2img_batch_input_dir: str, img2img_batch_output_dir: str, img2img_batch_inpaint_mask_dir: str, override_settings_texts, *args): # pylint: disable=unused-argument
if shared.sd_model is None:
shared.log.warning('Model not loaded')
return
override_settings = create_override_settings_dict(override_settings_texts)
is_batch = mode == 5
@@ -105,9 +108,9 @@ def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_s
assert 0. <= denoising_strength <= 1., 'can only work with strength in [0.0, 1.0]'
p = StableDiffusionProcessingImg2Img(
sd_model=sd_model,
outpath_samples=opts.outdir_samples or opts.outdir_img2img_samples,
outpath_grids=opts.outdir_grids or opts.outdir_img2img_grids,
sd_model=shared.sd_model,
outpath_samples=shared.opts.outdir_samples or shared.opts.outdir_img2img_samples,
outpath_grids=shared.opts.outdir_grids or shared.opts.outdir_img2img_grids,
prompt=prompt,
negative_prompt=negative_prompt,
styles=prompt_styles,
@@ -151,5 +154,5 @@ def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_s
processed = process_images(p)
p.close()
generation_info_js = processed.js()
debug(f'Processed: {len(processed.images)} Memory: {memory_stats()} img')
shared.debug(f'Processed: {len(processed.images)} Memory: {memory_stats()} img')
return processed.images, generation_info_js, plaintext_to_html(processed.info), plaintext_to_html(processed.comments)
+1 -3
View File
@@ -23,6 +23,7 @@ model_path = os.path.abspath(os.path.join(paths.models_path, model_dir))
checkpoints_list = {}
checkpoint_aliases = {}
checkpoints_loaded = collections.OrderedDict()
skip_next_load = False
class CheckpointInfo:
@@ -380,7 +381,6 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, timer=None)
sd_model = instantiate_from_config(sd_config.model)
except Exception:
sd_model = instantiate_from_config(sd_config.model)
# sd_model = instantiate_from_config(sd_config.model)
sd_model.used_config = checkpoint_config
timer.record("create")
load_model_weights(sd_model, checkpoint_info, state_dict, timer)
@@ -406,9 +406,7 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, timer=None)
shared.log.info(f"Model loaded in {timer.summary()}")
gc.collect()
shared.debug(f'Model load finished: {memory_stats()}')
return sd_model
skip_next_load = False
def reload_model_weights(sd_model=None, info=None):
global skip_next_load # pylint: disable=global-statement
+1 -1
View File
@@ -360,7 +360,7 @@ options_templates.update(options_section(('face-restoration', "Face restoration"
options_templates.update(options_section(('training', "Training"), {
"unload_models_when_training": OptionInfo(False, "Move VAE and CLIP to RAM when training if possible"),
"pin_memory": OptionInfo(True, "Turn on pin_memory for DataLoader"),
"pin_memory": OptionInfo(True, "Pin training dataset to memory"),
"save_optimizer_state": OptionInfo(False, "Saves resumable optimizer state when training embedding or hypernetwork"),
"save_training_settings_to_txt": OptionInfo(True, "Save textual inversion and hypernet settings to a text file whenever training starts"),
"dataset_filename_word_regex": OptionInfo("", "Filename word regex"),
+13 -10
View File
@@ -1,12 +1,12 @@
import os
import math
import tqdm
from tqdm.rich import tqdm
from PIL import Image, ImageOps
from modules import paths, shared, images, deepbooru
from modules.textual_inversion import autocrop
def preprocess(id_task, process_src, process_dst, process_width, process_height, preprocess_txt_action, process_keep_original_size, process_flip, process_split, process_caption, process_caption_deepbooru=False, split_threshold=0.5, overlap_ratio=0.2, process_focal_crop=False, process_focal_crop_face_weight=0.9, process_focal_crop_entropy_weight=0.3, process_focal_crop_edges_weight=0.5, process_focal_crop_debug=False, process_multicrop=None, process_multicrop_mindim=None, process_multicrop_maxdim=None, process_multicrop_minarea=None, process_multicrop_maxarea=None, process_multicrop_objective=None, process_multicrop_threshold=None): # pylint: disable=unused-argument
def preprocess(id_task, process_src, process_dst, process_width, process_height, preprocess_txt_action, process_keep_original_size=False, process_keep_channels=False, process_flip=False, process_split=False, process_caption_only=False, process_caption=False, process_caption_deepbooru=False, split_threshold=0.5, overlap_ratio=0.2, process_focal_crop=False, process_focal_crop_face_weight=0.9, process_focal_crop_entropy_weight=0.3, process_focal_crop_edges_weight=0.5, process_focal_crop_debug=False, process_multicrop=None, process_multicrop_mindim=None, process_multicrop_maxdim=None, process_multicrop_minarea=None, process_multicrop_maxarea=None, process_multicrop_objective=None, process_multicrop_threshold=None): # pylint: disable=unused-argument
try:
if process_caption:
shared.interrogator.load()
@@ -14,7 +14,7 @@ def preprocess(id_task, process_src, process_dst, process_width, process_height,
if process_caption_deepbooru:
deepbooru.model.start()
preprocess_work(process_src, process_dst, process_width, process_height, preprocess_txt_action, process_keep_original_size, process_flip, process_split, process_caption, process_caption_deepbooru, split_threshold, overlap_ratio, process_focal_crop, process_focal_crop_face_weight, process_focal_crop_entropy_weight, process_focal_crop_edges_weight, process_focal_crop_debug, process_multicrop, process_multicrop_mindim, process_multicrop_maxdim, process_multicrop_minarea, process_multicrop_maxarea, process_multicrop_objective, process_multicrop_threshold)
preprocess_work(process_src, process_dst, process_width, process_height, preprocess_txt_action, process_keep_original_size, process_keep_channels, process_flip, process_split, process_caption, process_caption_deepbooru, process_caption_only, split_threshold, overlap_ratio, process_focal_crop, process_focal_crop_face_weight, process_focal_crop_entropy_weight, process_focal_crop_edges_weight, process_focal_crop_debug, process_multicrop, process_multicrop_mindim, process_multicrop_maxdim, process_multicrop_minarea, process_multicrop_maxarea, process_multicrop_objective, process_multicrop_threshold)
finally:
@@ -34,6 +34,7 @@ class PreprocessParams:
dstdir = None
subindex = 0
flip = False
process_caption_only = False
process_caption = False
process_caption_deepbooru = False
preprocess_txt_action = None
@@ -55,7 +56,8 @@ def save_pic_with_caption(image, index, params: PreprocessParams, existing_capti
filename_part = os.path.basename(filename_part)
basename = f"{index:05}-{params.subindex}-{filename_part}"
image.save(os.path.join(params.dstdir, f"{basename}.png"))
if not params.process_caption_only:
image.save(os.path.join(params.dstdir, f"{basename}.png"))
if params.preprocess_txt_action == 'prepend' and existing_caption:
caption = existing_caption + ' ' + caption
@@ -75,7 +77,6 @@ def save_pic_with_caption(image, index, params: PreprocessParams, existing_capti
def save_pic(image, index, params, existing_caption=None):
save_pic_with_caption(image, index, params, existing_caption=existing_caption)
if params.flip:
save_pic_with_caption(ImageOps.mirror(image), index, params, existing_caption=existing_caption)
@@ -117,7 +118,7 @@ def center_crop(image: Image, w: int, h: int):
def multicrop_pic(image: Image, mindim, maxdim, minarea, maxarea, objective, threshold):
iw, ih = image.size
err = lambda w, h: 1-(lambda x: x if x < 1 else 1/x)(iw/ih/(w/h))
err = lambda w, h: 1-(lambda x: x if x < 1 else 1/x)(iw/ih/(w/h)) # pylint: disable=unnecessary-lambda-assignment,unnecessary-direct-lambda-call
wh = max(((w, h) for w in range(mindim, maxdim+1, 64) for h in range(mindim, maxdim+1, 64)
if minarea <= w * h <= maxarea and err(w, h) <= threshold),
key= lambda wh: (wh[0]*wh[1], -err(*wh))[::1 if objective=='Maximize area' else -1],
@@ -126,7 +127,7 @@ def multicrop_pic(image: Image, mindim, maxdim, minarea, maxarea, objective, thr
return wh and center_crop(image, *wh)
def preprocess_work(process_src, process_dst, process_width, process_height, preprocess_txt_action, process_keep_original_size, process_flip, process_split, process_caption, process_caption_deepbooru=False, split_threshold=0.5, overlap_ratio=0.2, process_focal_crop=False, process_focal_crop_face_weight=0.9, process_focal_crop_entropy_weight=0.3, process_focal_crop_edges_weight=0.5, process_focal_crop_debug=False, process_multicrop=None, process_multicrop_mindim=None, process_multicrop_maxdim=None, process_multicrop_minarea=None, process_multicrop_maxarea=None, process_multicrop_objective=None, process_multicrop_threshold=None):
def preprocess_work(process_src, process_dst, process_width, process_height, preprocess_txt_action, process_keep_original_size, process_keep_channels, process_flip, process_split, process_caption, process_caption_deepbooru, process_caption_only, split_threshold, overlap_ratio, process_focal_crop, process_focal_crop_face_weight, process_focal_crop_entropy_weight, process_focal_crop_edges_weight, process_focal_crop_debug, process_multicrop, process_multicrop_mindim, process_multicrop_maxdim, process_multicrop_minarea, process_multicrop_maxarea, process_multicrop_objective, process_multicrop_threshold):
width = process_width
height = process_height
@@ -148,22 +149,24 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
params = PreprocessParams()
params.dstdir = dst
params.flip = process_flip
params.process_caption_only = process_caption_only
params.process_caption = process_caption
params.process_caption_deepbooru = process_caption_deepbooru
params.preprocess_txt_action = preprocess_txt_action
pbar = tqdm.tqdm(files)
pbar = tqdm(files)
for index, imagefile in enumerate(pbar):
params.subindex = 0
filename = os.path.join(src, imagefile)
try:
img = Image.open(filename)
img = ImageOps.exif_transpose(img)
img = img.convert("RGB")
if not process_keep_channels:
img = img.convert("RGB")
except Exception:
continue
description = f"Preprocessing [Image {index}/{len(files)}]"
description = f"Preprocessing image {index + 1}/{len(files)}"
pbar.set_description(description)
shared.state.textinfo = description
@@ -3,7 +3,7 @@ import html
import csv
from collections import namedtuple
import torch
import tqdm
from tqdm.rich import tqdm
import safetensors.torch
import numpy as np
from PIL import Image, PngImagePlugin
@@ -448,7 +448,7 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st
is_training_inpainting_model = shared.sd_model.model.conditioning_key in {'hybrid', 'concat'}
img_c = None
pbar = tqdm.tqdm(total=steps - initial_step)
pbar = tqdm(total=steps - initial_step)
try:
sd_hijack_checkpoint.add()
+9 -6
View File
@@ -1,18 +1,21 @@
import modules.scripts
from modules import sd_samplers
from modules import sd_samplers, shared
from modules.generation_parameters_copypaste import create_override_settings_dict
from modules.processing import StableDiffusionProcessingTxt2Img, process_images
from modules.shared import opts, sd_model, debug
# from modules.shared import opts, sd_model, debug
from modules.ui import plaintext_to_html
from modules.memstats import memory_stats
def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, steps: int, sampler_index: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, enable_hr: bool, denoising_strength: float, hr_scale: float, hr_upscaler: str, hr_second_pass_steps: int, hr_resize_x: int, hr_resize_y: int, override_settings_texts, *args): # pylint: disable=unused-argument
if shared.sd_model is None:
shared.log.warning('Model not loaded')
return
override_settings = create_override_settings_dict(override_settings_texts)
p = StableDiffusionProcessingTxt2Img(
sd_model=sd_model,
outpath_samples=opts.outdir_samples or opts.outdir_txt2img_samples,
outpath_grids=opts.outdir_grids or opts.outdir_txt2img_grids,
sd_model=shared.sd_model,
outpath_samples=shared.opts.outdir_samples or shared.opts.outdir_txt2img_samples,
outpath_grids=shared.opts.outdir_grids or shared.opts.outdir_txt2img_grids,
prompt=prompt,
styles=prompt_styles,
negative_prompt=negative_prompt,
@@ -47,5 +50,5 @@ def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, step
processed = process_images(p)
p.close()
generation_info_js = processed.js()
debug(f'Processed: {len(processed.images)} Memory: {memory_stats()} txt')
shared.debug(f'Processed: {len(processed.images)} Memory: {memory_stats()} txt')
return processed.images, generation_info_js, plaintext_to_html(processed.info), plaintext_to_html(processed.comments)
+6 -2
View File
@@ -975,12 +975,14 @@ def create_ui():
with gr.Row():
process_keep_original_size = gr.Checkbox(label='Keep original size', elem_id="train_process_keep_original_size")
process_keep_channels = gr.Checkbox(label='Keep original image channels', elem_id="train_process_keep_channels")
process_flip = gr.Checkbox(label='Create flipped copies', elem_id="train_process_flip")
process_split = gr.Checkbox(label='Split oversized images', elem_id="train_process_split")
process_focal_crop = gr.Checkbox(label='Auto focal point crop', elem_id="train_process_focal_crop")
process_multicrop = gr.Checkbox(label='Auto-sized crop', elem_id="train_process_multicrop")
process_caption = gr.Checkbox(label='Use BLIP for caption', elem_id="train_process_caption")
process_caption_deepbooru = gr.Checkbox(label='Use deepbooru for caption', visible=True, elem_id="train_process_caption_deepbooru")
process_caption_only = gr.Checkbox(label='Create captions only', elem_id="train_process_multicrop")
process_caption = gr.Checkbox(label='Create BLIP captions', elem_id="train_process_caption")
process_caption_deepbooru = gr.Checkbox(label='Create Deepbooru captions', visible=True, elem_id="train_process_caption_deepbooru")
with gr.Row(visible=False) as process_split_extra_row:
process_split_threshold = gr.Slider(label='Split image threshold', value=0.5, minimum=0.0, maximum=1.0, step=0.05, elem_id="train_process_split_threshold")
@@ -1142,8 +1144,10 @@ def create_ui():
process_height,
preprocess_txt_action,
process_keep_original_size,
process_keep_channels,
process_flip,
process_split,
process_caption_only,
process_caption,
process_caption_deepbooru,
process_split_threshold,