From fe496f4ebcb8b3d8981d69a1c2459b09cb429723 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Fri, 5 May 2023 09:05:59 -0400 Subject: [PATCH] add train preprocess options --- cli/train/train.py | 1 - installer.py | 11 +++--- modules/img2img.py | 35 ++++++++++--------- modules/sd_models.py | 4 +-- modules/shared.py | 2 +- modules/textual_inversion/preprocess.py | 23 ++++++------ .../textual_inversion/textual_inversion.py | 4 +-- modules/txt2img.py | 15 ++++---- modules/ui.py | 8 +++-- 9 files changed, 56 insertions(+), 47 deletions(-) diff --git a/cli/train/train.py b/cli/train/train.py index b1e1d48fb..bbad21565 100755 --- a/cli/train/train.py +++ b/cli/train/train.py @@ -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 diff --git a/installer.py b/installer.py index f2b4d5b5d..2da153806 100644 --- a/installer.py +++ b/installer.py @@ -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}') diff --git a/modules/img2img.py b/modules/img2img.py index 58b9449aa..d0b5183f4 100644 --- a/modules/img2img.py +++ b/modules/img2img.py @@ -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) diff --git a/modules/sd_models.py b/modules/sd_models.py index c8790b53e..e5f4fb6a2 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -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 diff --git a/modules/shared.py b/modules/shared.py index de12c0787..66aba96fa 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -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"), diff --git a/modules/textual_inversion/preprocess.py b/modules/textual_inversion/preprocess.py index 2b0714df3..41224dd18 100644 --- a/modules/textual_inversion/preprocess.py +++ b/modules/textual_inversion/preprocess.py @@ -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 diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py index 72be41bb4..bee86ceac 100644 --- a/modules/textual_inversion/textual_inversion.py +++ b/modules/textual_inversion/textual_inversion.py @@ -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() diff --git a/modules/txt2img.py b/modules/txt2img.py index 36bf6dfb6..d7533f4b2 100644 --- a/modules/txt2img.py +++ b/modules/txt2img.py @@ -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) diff --git a/modules/ui.py b/modules/ui.py index 8d00e7db5..1be72b29f 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -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,