diff --git a/extensions-builtin/prompt-bracket-checker/javascript/prompt-bracket-checker.js b/extensions-builtin/prompt-bracket-checker/javascript/prompt-bracket-checker.js index eccfb0f9d..251a1f578 100644 --- a/extensions-builtin/prompt-bracket-checker/javascript/prompt-bracket-checker.js +++ b/extensions-builtin/prompt-bracket-checker/javascript/prompt-bracket-checker.js @@ -93,10 +93,12 @@ function checkBrackets(evt) { } var shadowRootLoaded = setInterval(function() { - var shadowTextArea = document.querySelector('gradio-app').shadowRoot.querySelectorAll('#txt2img_prompt > label > textarea'); - if(shadowTextArea.length < 1) { - return false; - } + var sahdowRoot = document.querySelector('gradio-app').shadowRoot; + if(! sahdowRoot) return false; + + var shadowTextArea = sahdowRoot.querySelectorAll('#txt2img_prompt > label > textarea'); + if(shadowTextArea.length < 1) return false; + clearInterval(shadowRootLoaded); diff --git a/javascript/hints.js b/javascript/hints.js index fa5e5ae89..e746e20d5 100644 --- a/javascript/hints.js +++ b/javascript/hints.js @@ -92,6 +92,7 @@ titles = { "Weighted sum": "Result = A * (1 - M) + B * M", "Add difference": "Result = A + (B - C) * M", + "No interpolation": "Result = A", "Initialization text": "If the number of tokens is more than the number of vectors, some may be skipped.\nLeave the textbox empty to start with zeroed out vectors", "Learning rate": "How fast should training go. Low values will take longer to train, high values may fail to converge (not generate accurate results) and/or may break the embedding (This has happened if you see Loss: nan in the training info textbox. If this happens, you need to manually restore your embedding from an older not-broken backup).\n\nYou can set a single numeric value, or multiple learning rates using the syntax:\n\n rate_1:max_steps_1, rate_2:max_steps_2, ...\n\nEG: 0.005:100, 1e-3:1000, 1e-5\n\nWill train with rate of 0.005 for first 100 steps, then 1e-3 until 1000 steps, then 1e-5 for all remaining steps.", diff --git a/javascript/progressbar.js b/javascript/progressbar.js index 18c771a23..ff6d757ba 100644 --- a/javascript/progressbar.js +++ b/javascript/progressbar.js @@ -81,8 +81,13 @@ function request(url, data, handler, errorHandler){ xhr.onreadystatechange = function () { if (xhr.readyState === 4) { if (xhr.status === 200) { - var js = JSON.parse(xhr.responseText); - handler(js) + try { + var js = JSON.parse(xhr.responseText); + handler(js) + } catch (error) { + console.error(error); + errorHandler() + } } else{ errorHandler() } @@ -155,7 +160,7 @@ function requestProgress(id_task, progressbarContainer, gallery, atEnd, onProgre } var fun = function(id_task, id_live_preview){ - request("/internal/progress", {"id_task": id_task, "id_live_preview": id_live_preview}, function(res){ + request("./internal/progress", {"id_task": id_task, "id_live_preview": id_live_preview}, function(res){ if(res.completed){ removeProgressBar() return diff --git a/javascript/ui.js b/javascript/ui.js index 7d3d57a3c..37788a3e4 100644 --- a/javascript/ui.js +++ b/javascript/ui.js @@ -172,6 +172,15 @@ function submit_img2img(){ return res } +function modelmerger(){ + var id = randomId() + requestProgress(id, gradioApp().getElementById('modelmerger_results_panel'), null, function(){}) + + var res = create_submit_args(arguments) + res[0] = id + return res +} + function ask_for_style_name(_, prompt_text, negative_prompt_text) { name_ = prompt('Style name:') diff --git a/modules/devices.py b/modules/devices.py index 206184fbf..524ec7af4 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -169,8 +169,10 @@ orig_Tensor_cumsum = torch.Tensor.cumsum def cumsum_fix(input, cumsum_func, *args, **kwargs): if input.device.type == 'mps': output_dtype = kwargs.get('dtype', input.dtype) - if any(output_dtype == broken_dtype for broken_dtype in [torch.bool, torch.int8, torch.int16, torch.int64]): + if output_dtype == torch.int64: return cumsum_func(input.cpu(), *args, **kwargs).to(input.device) + elif cumsum_needs_bool_fix and output_dtype == torch.bool or cumsum_needs_int_fix and (output_dtype == torch.int8 or output_dtype == torch.int16): + return cumsum_func(input.to(torch.int32), *args, **kwargs).to(torch.int64) return cumsum_func(input, *args, **kwargs) @@ -181,9 +183,10 @@ if has_mps(): torch.nn.functional.layer_norm = layer_norm_fix torch.Tensor.numpy = numpy_fix elif version.parse(torch.__version__) > version.parse("1.13.1"): - if not torch.Tensor([1,2]).to(torch.device("mps")).equal(torch.Tensor([1,1]).to(torch.device("mps")).cumsum(0, dtype=torch.int16)): - torch.cumsum = lambda input, *args, **kwargs: ( cumsum_fix(input, orig_cumsum, *args, **kwargs) ) - torch.Tensor.cumsum = lambda self, *args, **kwargs: ( cumsum_fix(self, orig_Tensor_cumsum, *args, **kwargs) ) + cumsum_needs_int_fix = not torch.Tensor([1,2]).to(torch.device("mps")).equal(torch.ShortTensor([1,1]).to(torch.device("mps")).cumsum(0)) + cumsum_needs_bool_fix = not torch.BoolTensor([True,True]).to(device=torch.device("mps"), dtype=torch.int64).equal(torch.BoolTensor([True,False]).to(torch.device("mps")).cumsum(0)) + torch.cumsum = lambda input, *args, **kwargs: ( cumsum_fix(input, orig_cumsum, *args, **kwargs) ) + torch.Tensor.cumsum = lambda self, *args, **kwargs: ( cumsum_fix(self, orig_Tensor_cumsum, *args, **kwargs) ) orig_narrow = torch.narrow torch.narrow = lambda *args, **kwargs: ( orig_narrow(*args, **kwargs).clone() ) diff --git a/modules/extras.py b/modules/extras.py index 22668fcda..1218f88fc 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -15,7 +15,7 @@ from typing import Callable, List, OrderedDict, Tuple from functools import partial from dataclasses import dataclass -from modules import processing, shared, images, devices, sd_models, sd_samplers +from modules import processing, shared, images, devices, sd_models, sd_samplers, sd_vae from modules.shared import opts import modules.gfpgan_model from modules.ui import plaintext_to_html @@ -251,7 +251,8 @@ def run_pnginfo(image): def create_config(ckpt_result, config_source, a, b, c): def config(x): - return sd_models.find_checkpoint_config(x) if x else None + res = sd_models.find_checkpoint_config(x) if x else None + return res if res != shared.sd_default_config else None if config_source == 0: cfg = config(a) or config(b) or config(c) @@ -274,10 +275,25 @@ def create_config(ckpt_result, config_source, a, b, c): shutil.copyfile(cfg, checkpoint_filename) -def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format, config_source): +checkpoint_dict_skip_on_merge = ["cond_stage_model.transformer.text_model.embeddings.position_ids"] + + +def to_half(tensor, enable): + if enable and tensor.dtype == torch.float: + return tensor.half() + + return tensor + + +def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format, config_source, bake_in_vae): shared.state.begin() shared.state.job = 'model-merge' + def fail(message): + shared.state.textinfo = message + shared.state.end() + return [*[gr.update() for _ in range(4)], message] + def weighted_sum(theta0, theta1, alpha): return ((1 - alpha) * theta0) + (alpha * theta1) @@ -287,57 +303,96 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam def add_difference(theta0, theta1_2_diff, alpha): return theta0 + (alpha * theta1_2_diff) - primary_model_info = sd_models.checkpoints_list[primary_model_name] - secondary_model_info = sd_models.checkpoints_list[secondary_model_name] - tertiary_model_info = sd_models.checkpoints_list.get(tertiary_model_name, None) - result_is_inpainting_model = False + def filename_weighted_sum(): + a = primary_model_info.model_name + b = secondary_model_info.model_name + Ma = round(1 - multiplier, 2) + Mb = round(multiplier, 2) + + return f"{Ma}({a}) + {Mb}({b})" + + def filename_add_difference(): + a = primary_model_info.model_name + b = secondary_model_info.model_name + c = tertiary_model_info.model_name + M = round(multiplier, 2) + + return f"{a} + {M}({b} - {c})" + + def filename_nothing(): + return primary_model_info.model_name theta_funcs = { - "Weighted sum": (None, weighted_sum), - "Add difference": (get_difference, add_difference), + "Weighted sum": (filename_weighted_sum, None, weighted_sum), + "Add difference": (filename_add_difference, get_difference, add_difference), + "No interpolation": (filename_nothing, None, None), } - theta_func1, theta_func2 = theta_funcs[interp_method] + filename_generator, theta_func1, theta_func2 = theta_funcs[interp_method] + shared.state.job_count = (1 if theta_func1 else 0) + (1 if theta_func2 else 0) - if theta_func1 and not tertiary_model_info: - shared.state.textinfo = "Failed: Interpolation method requires a tertiary model." - shared.state.end() - return ["Failed: Interpolation method requires a tertiary model."] + [gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)] + if not primary_model_name: + return fail("Failed: Merging requires a primary model.") - shared.state.textinfo = f"Loading {secondary_model_info.filename}..." - print(f"Loading {secondary_model_info.filename}...") - theta_1 = sd_models.read_state_dict(secondary_model_info.filename, map_location='cpu') + primary_model_info = sd_models.checkpoints_list[primary_model_name] + + if theta_func2 and not secondary_model_name: + return fail("Failed: Merging requires a secondary model.") + + secondary_model_info = sd_models.checkpoints_list[secondary_model_name] if theta_func2 else None + + if theta_func1 and not tertiary_model_name: + return fail(f"Failed: Interpolation method ({interp_method}) requires a tertiary model.") + + tertiary_model_info = sd_models.checkpoints_list[tertiary_model_name] if theta_func1 else None + + result_is_inpainting_model = False + + if theta_func2: + shared.state.textinfo = f"Loading B" + print(f"Loading {secondary_model_info.filename}...") + theta_1 = sd_models.read_state_dict(secondary_model_info.filename, map_location='cpu') + else: + theta_1 = None if theta_func1: + shared.state.textinfo = f"Loading C" print(f"Loading {tertiary_model_info.filename}...") theta_2 = sd_models.read_state_dict(tertiary_model_info.filename, map_location='cpu') + shared.state.textinfo = 'Merging B and C' + shared.state.sampling_steps = len(theta_1.keys()) for key in tqdm.tqdm(theta_1.keys()): + if key in checkpoint_dict_skip_on_merge: + continue + if 'model' in key: if key in theta_2: t2 = theta_2.get(key, torch.zeros_like(theta_1[key])) theta_1[key] = theta_func1(theta_1[key], t2) else: theta_1[key] = torch.zeros_like(theta_1[key]) + + shared.state.sampling_step += 1 del theta_2 + shared.state.nextjob() + shared.state.textinfo = f"Loading {primary_model_info.filename}..." print(f"Loading {primary_model_info.filename}...") theta_0 = sd_models.read_state_dict(primary_model_info.filename, map_location='cpu') print("Merging...") - - chckpoint_dict_skip_on_merge = ["cond_stage_model.transformer.text_model.embeddings.position_ids"] - + shared.state.textinfo = 'Merging A and B' + shared.state.sampling_steps = len(theta_0.keys()) for key in tqdm.tqdm(theta_0.keys()): - if 'model' in key and key in theta_1: + if theta_1 and 'model' in key and key in theta_1: - if key in chckpoint_dict_skip_on_merge: + if key in checkpoint_dict_skip_on_merge: continue a = theta_0[key] b = theta_1[key] - shared.state.textinfo = f'Merging layer {key}' # this enables merging an inpainting model (A) with another one (B); # where normal model would have 4 channels, for latenst space, inpainting model would # have another 4 channels for unmasked picture's latent space, plus one channel for mask, for a total of 9 @@ -352,36 +407,39 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam else: theta_0[key] = theta_func2(a, b, multiplier) - if save_as_half: - theta_0[key] = theta_0[key].half() + theta_0[key] = to_half(theta_0[key], save_as_half) - # I believe this part should be discarded, but I'll leave it for now until I am sure - for key in theta_1.keys(): - if 'model' in key and key not in theta_0: + shared.state.sampling_step += 1 - if key in chckpoint_dict_skip_on_merge: - continue - - theta_0[key] = theta_1[key] - if save_as_half: - theta_0[key] = theta_0[key].half() del theta_1 + bake_in_vae_filename = sd_vae.vae_dict.get(bake_in_vae, None) + if bake_in_vae_filename is not None: + print(f"Baking in VAE from {bake_in_vae_filename}") + shared.state.textinfo = 'Baking in VAE' + vae_dict = sd_vae.load_vae_dict(bake_in_vae_filename, map_location='cpu') + + for key in vae_dict.keys(): + theta_0_key = 'first_stage_model.' + key + if theta_0_key in theta_0: + theta_0[theta_0_key] = to_half(vae_dict[key], save_as_half) + + del vae_dict + + if save_as_half and not theta_func2: + for key in theta_0.keys(): + theta_0[key] = to_half(theta_0[key], save_as_half) + ckpt_dir = shared.cmd_opts.ckpt_dir or sd_models.model_path - filename = \ - primary_model_info.model_name + '_' + str(round(1-multiplier, 2)) + '-' + \ - secondary_model_info.model_name + '_' + str(round(multiplier, 2)) + '-' + \ - interp_method.replace(" ", "_") + \ - '-merged.' + \ - ("inpainting." if result_is_inpainting_model else "") + \ - checkpoint_format - - filename = filename if custom_name == '' else (custom_name + '.' + checkpoint_format) + filename = filename_generator() if custom_name == '' else custom_name + filename += ".inpainting" if result_is_inpainting_model else "" + filename += "." + checkpoint_format output_modelname = os.path.join(ckpt_dir, filename) - shared.state.textinfo = f"Saving to {output_modelname}..." + shared.state.nextjob() + shared.state.textinfo = "Saving" print(f"Saving to {output_modelname}...") _, extension = os.path.splitext(output_modelname) @@ -394,8 +452,8 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam create_config(output_modelname, config_source, primary_model_info, secondary_model_info, tertiary_model_info) - print("Checkpoint saved.") - shared.state.textinfo = "Checkpoint saved to " + output_modelname + print(f"Checkpoint saved to {output_modelname}.") + shared.state.textinfo = "Checkpoint saved" shared.state.end() - return ["Checkpoint saved to " + output_modelname] + [gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)] + return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)], "Checkpoint saved to " + output_modelname] diff --git a/modules/progress.py b/modules/progress.py index 3327b8830..c69ecf3d1 100644 --- a/modules/progress.py +++ b/modules/progress.py @@ -67,10 +67,13 @@ def progressapi(req: ProgressRequest): progress = 0 - if shared.state.job_count > 0: - progress += shared.state.job_no / shared.state.job_count - if shared.state.sampling_steps > 0: - progress += 1 / shared.state.job_count * shared.state.sampling_step / shared.state.sampling_steps + job_count, job_no = shared.state.job_count, shared.state.job_no + sampling_steps, sampling_step = shared.state.sampling_steps, shared.state.sampling_step + + if job_count > 0: + progress += job_no / job_count + if sampling_steps > 0 and job_count > 0: + progress += 1 / job_count * sampling_step / sampling_steps progress = min(progress, 1) diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py index 6872eab12..4af457dfd 100644 --- a/modules/sd_hijack.py +++ b/modules/sd_hijack.py @@ -69,6 +69,13 @@ def undo_optimizations(): ldm.modules.diffusionmodules.model.AttnBlock.forward = diffusionmodules_model_AttnBlock_forward +def fix_checkpoint(): + """checkpoints are now added and removed in embedding/hypernet code, since torch doesn't want + checkpoints to be added when not training (there's a warning)""" + + pass + + class StableDiffusionModelHijack: fixes = None comments = [] diff --git a/modules/sd_models.py b/modules/sd_models.py index 6a6ec2e09..3ea1dea7f 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -41,14 +41,16 @@ class CheckpointInfo: if name.startswith("\\") or name.startswith("/"): name = name[1:] - self.title = name + self.name = name self.model_name = os.path.splitext(name.replace("/", "_").replace("\\", "_"))[0] self.hash = model_hash(filename) - self.sha256 = hashes.sha256_from_cache(self.filename, "checkpoint/" + self.title) + self.sha256 = hashes.sha256_from_cache(self.filename, "checkpoint/" + name) self.shorthash = self.sha256[0:10] if self.sha256 else None - self.ids = [self.hash, self.model_name, self.title, f'{name} [{self.hash}]'] + ([self.shorthash, self.sha256] if self.shorthash else []) + self.title = name if self.shorthash is None else f'{name} [{self.shorthash}]' + + self.ids = [self.hash, self.model_name, self.title, name, f'{name} [{self.hash}]'] + ([self.shorthash, self.sha256, f'{self.name} [{self.shorthash}]'] if self.shorthash else []) def register(self): checkpoints_list[self.title] = self @@ -56,13 +58,15 @@ class CheckpointInfo: checkpoint_alisases[id] = self def calculate_shorthash(self): - self.sha256 = hashes.sha256(self.filename, "checkpoint/" + self.title) + self.sha256 = hashes.sha256(self.filename, "checkpoint/" + self.name) self.shorthash = self.sha256[0:10] if self.shorthash not in self.ids: self.ids += [self.shorthash, self.sha256] self.register() + self.title = f'{self.name} [{self.shorthash}]' + return self.shorthash @@ -225,7 +229,10 @@ def read_state_dict(checkpoint_file, print_global_state=False, map_location=None def load_model_weights(model, checkpoint_info: CheckpointInfo): + title = checkpoint_info.title sd_model_hash = checkpoint_info.calculate_shorthash() + if checkpoint_info.title != title: + shared.opts.data["sd_model_checkpoint"] = checkpoint_info.title cache_enabled = shared.opts.sd_checkpoint_cache > 0 diff --git a/modules/sd_vae.py b/modules/sd_vae.py index da1bf15c4..4ce238b86 100644 --- a/modules/sd_vae.py +++ b/modules/sd_vae.py @@ -120,6 +120,12 @@ def resolve_vae(checkpoint_file): return None, None +def load_vae_dict(filename, map_location): + vae_ckpt = sd_models.read_state_dict(filename, map_location=map_location) + vae_dict_1 = {k: v for k, v in vae_ckpt.items() if k[0:4] != "loss" and k not in vae_ignore_keys} + return vae_dict_1 + + def load_vae(model, vae_file=None, vae_source="from unknown source"): global vae_dict, loaded_vae_file # save_settings = False @@ -137,8 +143,7 @@ def load_vae(model, vae_file=None, vae_source="from unknown source"): print(f"Loading VAE weights {vae_source}: {vae_file}") store_base_vae(model) - vae_ckpt = sd_models.read_state_dict(vae_file, map_location=shared.weight_load_location) - vae_dict_1 = {k: v for k, v in vae_ckpt.items() if k[0:4] != "loss" and k not in vae_ignore_keys} + vae_dict_1 = load_vae_dict(vae_file, map_location=shared.weight_load_location) _load_vae_dict(model, vae_dict_1) if cache_enabled: diff --git a/modules/shared.py b/modules/shared.py index 66f4bfbce..b32540b13 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -20,10 +20,11 @@ from modules.paths import models_path, script_path, sd_path demo = None +sd_default_config = os.path.join(script_path, "configs/v1-inference.yaml") sd_model_file = os.path.join(script_path, 'model.ckpt') default_sd_model_file = sd_model_file parser = argparse.ArgumentParser() -parser.add_argument("--config", type=str, default=os.path.join(script_path, "configs/v1-inference.yaml"), help="path to config which constructs model",) +parser.add_argument("--config", type=str, default=sd_default_config, help="path to config which constructs model",) parser.add_argument("--ckpt", type=str, default=sd_model_file, help="path to checkpoint of stable diffusion model; if specified, this checkpoint will be added to the list of checkpoints and loaded",) parser.add_argument("--ckpt-dir", type=str, default=None, help="Path to directory with stable diffusion checkpoints") parser.add_argument("--vae-dir", type=str, default=None, help="Path to directory with VAE files") @@ -453,7 +454,7 @@ options_templates.update(options_section(('ui', "User interface"), { "js_modal_lightbox_initially_zoomed": OptionInfo(True, "Show images zoomed in by default in full page image viewer"), "show_progress_in_title": OptionInfo(True, "Show generation progress in window title."), "samplers_in_dropdown": OptionInfo(True, "Use dropdown for sampler selection instead of radio group"), - "dimensions_and_batch_together": OptionInfo(True, "Show Witdth/Height and Batch sliders in same row"), + "dimensions_and_batch_together": OptionInfo(True, "Show Width/Height and Batch sliders in same row"), 'quicksettings': OptionInfo("sd_model_checkpoint", "Quicksettings list"), 'ui_reorder': OptionInfo(", ".join(ui_reorder_categories), "txt2img/img2img UI item order"), 'localization': OptionInfo("None", "Localization (requires restart)", gr.Dropdown, lambda: {"choices": ["None"] + list(localization.localizations.keys())}, refresh=lambda: localization.list_localizations(cmd_opts.localizations_dir)), diff --git a/modules/textual_inversion/preprocess.py b/modules/textual_inversion/preprocess.py index 64abff4da..c0ac11d39 100644 --- a/modules/textual_inversion/preprocess.py +++ b/modules/textual_inversion/preprocess.py @@ -12,7 +12,7 @@ from modules.shared import opts, cmd_opts from modules.textual_inversion import autocrop -def preprocess(id_task, process_src, process_dst, process_width, process_height, preprocess_txt_action, 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): +def preprocess(id_task, process_src, process_dst, process_width, process_height, preprocess_txt_action, 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): try: if process_caption: shared.interrogator.load() @@ -20,7 +20,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_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) + preprocess_work(process_src, process_dst, process_width, process_height, preprocess_txt_action, 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) finally: @@ -109,8 +109,30 @@ def split_pic(image, inverse_xy, width, height, overlap_ratio): splitted = image.crop((0, y, to_w, y + to_h)) yield splitted +# not using torchvision.transforms.CenterCrop because it doesn't allow float regions +def center_crop(image: Image, w: int, h: int): + iw, ih = image.size + if ih / h < iw / w: + sw = w * ih / h + box = (iw - sw) / 2, 0, iw - (iw - sw) / 2, ih + else: + sh = h * iw / w + box = 0, (ih - sh) / 2, iw, ih - (ih - sh) / 2 + return image.resize((w, h), Image.Resampling.LANCZOS, box) -def preprocess_work(process_src, process_dst, process_width, process_height, preprocess_txt_action, 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): + +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)) + 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], + default=None + ) + return wh and center_crop(image, *wh) + + +def preprocess_work(process_src, process_dst, process_width, process_height, preprocess_txt_action, 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): width = process_width height = process_height src = os.path.abspath(process_src) @@ -194,6 +216,14 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre save_pic(focal, index, params, existing_caption=existing_caption) process_default_resize = False + if process_multicrop: + cropped = multicrop_pic(img, process_multicrop_mindim, process_multicrop_maxdim, process_multicrop_minarea, process_multicrop_maxarea, process_multicrop_objective, process_multicrop_threshold) + if cropped is not None: + save_pic(cropped, index, params, existing_caption=existing_caption) + else: + print(f"skipped {img.width}x{img.height} image {filename} (can't find suitable size within error threshold)") + process_default_resize = False + if process_default_resize: img = images.resize_image(1, img, width, height) save_pic(img, index, params, existing_caption=existing_caption) diff --git a/modules/ui.py b/modules/ui.py index 28e226b51..92e53f43a 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -20,7 +20,7 @@ import numpy as np from PIL import Image, PngImagePlugin from modules.call_queue import wrap_gradio_gpu_call, wrap_queued_call, wrap_gradio_call -from modules import sd_hijack, sd_models, localization, script_callbacks, ui_extensions, deepbooru +from modules import sd_hijack, sd_models, localization, script_callbacks, ui_extensions, deepbooru, sd_vae from modules.ui_components import FormRow, FormGroup, ToolButton, FormHTML from modules.paths import script_path @@ -439,7 +439,7 @@ def apply_setting(key, value): opts.data_labels[key].onchange() opts.save(shared.config_filename) - return value + return getattr(opts, key) def update_generation_info(generation_info, html_info, img_index): @@ -597,6 +597,16 @@ def ordered_ui_categories(): yield category +def get_value_for_setting(key): + value = getattr(opts, key) + + info = opts.data_labels[key] + args = info.component_args() if callable(info.component_args) else info.component_args or {} + args = {k: v for k, v in args.items() if k not in {'precision'}} + + return gr.update(value=value, **args) + + def create_ui(): import modules.img2img import modules.txt2img @@ -1185,7 +1195,7 @@ def create_ui(): with gr.Column(variant='compact'): gr.HTML(value="
A merger of the two checkpoints will be generated in your checkpoint directory.
") - with FormRow(): + with FormRow(elem_id="modelmerger_models"): primary_model_name = gr.Dropdown(modules.sd_models.checkpoint_tiles(), elem_id="modelmerger_primary_model_name", label="Primary model (A)") create_refresh_button(primary_model_name, modules.sd_models.list_models, lambda: {"choices": modules.sd_models.checkpoint_tiles()}, "refresh_checkpoint_A") @@ -1197,19 +1207,27 @@ def create_ui(): custom_name = gr.Textbox(label="Custom Name (Optional)", elem_id="modelmerger_custom_name") interp_amount = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Multiplier (M) - set to 0 to get model A', value=0.3, elem_id="modelmerger_interp_amount") - interp_method = gr.Radio(choices=["Weighted sum", "Add difference"], value="Weighted sum", label="Interpolation Method", elem_id="modelmerger_interp_method") + interp_method = gr.Radio(choices=["No interpolation", "Weighted sum", "Add difference"], value="Weighted sum", label="Interpolation Method", elem_id="modelmerger_interp_method") with FormRow(): checkpoint_format = gr.Radio(choices=["ckpt", "safetensors"], value="ckpt", label="Checkpoint format", elem_id="modelmerger_checkpoint_format") save_as_half = gr.Checkbox(value=False, label="Save as float16", elem_id="modelmerger_save_as_half") - config_source = gr.Radio(choices=["A, B or C", "B", "C", "Don't"], value="A, B or C", label="Copy config from", type="index", elem_id="modelmerger_config_method") + with FormRow(): + with gr.Column(): + config_source = gr.Radio(choices=["A, B or C", "B", "C", "Don't"], value="A, B or C", label="Copy config from", type="index", elem_id="modelmerger_config_method") + + with gr.Column(): + with FormRow(): + bake_in_vae = gr.Dropdown(choices=["None"] + list(sd_vae.vae_dict), value="None", label="Bake in VAE", elem_id="modelmerger_bake_in_vae") + create_refresh_button(bake_in_vae, sd_vae.refresh_vae_list, lambda: {"choices": ["None"] + list(sd_vae.vae_dict)}, "modelmerger_refresh_bake_in_vae") with gr.Row(): modelmerger_merge = gr.Button(elem_id="modelmerger_merge", value="Merge", variant='primary') - with gr.Column(variant='panel'): - submit_result = gr.Textbox(elem_id="modelmerger_result", show_label=False) + with gr.Column(variant='compact', elem_id="modelmerger_results_container"): + with gr.Group(elem_id="modelmerger_results_panel"): + modelmerger_result = gr.HTML(elem_id="modelmerger_result", show_label=False) with gr.Blocks(analytics_enabled=False) as train_interface: with gr.Row().style(equal_height=False): @@ -1260,6 +1278,7 @@ def create_ui(): 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") @@ -1272,7 +1291,19 @@ def create_ui(): process_focal_crop_entropy_weight = gr.Slider(label='Focal point entropy weight', value=0.15, minimum=0.0, maximum=1.0, step=0.05, elem_id="train_process_focal_crop_entropy_weight") process_focal_crop_edges_weight = gr.Slider(label='Focal point edges weight', value=0.5, minimum=0.0, maximum=1.0, step=0.05, elem_id="train_process_focal_crop_edges_weight") process_focal_crop_debug = gr.Checkbox(label='Create debug image', elem_id="train_process_focal_crop_debug") - + + with gr.Column(visible=False) as process_multicrop_col: + gr.Markdown('Each image is center-cropped with an automatically chosen width and height.') + with gr.Row(): + process_multicrop_mindim = gr.Slider(minimum=64, maximum=2048, step=8, label="Dimension lower bound", value=384, elem_id="train_process_multicrop_mindim") + process_multicrop_maxdim = gr.Slider(minimum=64, maximum=2048, step=8, label="Dimension upper bound", value=768, elem_id="train_process_multicrop_maxdim") + with gr.Row(): + process_multicrop_minarea = gr.Slider(minimum=64*64, maximum=2048*2048, step=1, label="Area lower bound", value=64*64, elem_id="train_process_multicrop_minarea") + process_multicrop_maxarea = gr.Slider(minimum=64*64, maximum=2048*2048, step=1, label="Area upper bound", value=640*640, elem_id="train_process_multicrop_maxarea") + with gr.Row(): + process_multicrop_objective = gr.Radio(["Maximize area", "Minimize error"], value="Maximize area", label="Resizing objective", elem_id="train_process_multicrop_objective") + process_multicrop_threshold = gr.Slider(minimum=0, maximum=1, step=0.01, label="Error threshold", value=0.1, elem_id="train_process_multicrop_threshold") + with gr.Row(): with gr.Column(scale=3): gr.HTML(value="") @@ -1294,6 +1325,12 @@ def create_ui(): outputs=[process_focal_crop_row], ) + process_multicrop.change( + fn=lambda show: gr_show(show), + inputs=[process_multicrop], + outputs=[process_multicrop_col], + ) + def get_textual_inversion_template_names(): return sorted([x for x in textual_inversion.textual_inversion_templates]) @@ -1413,6 +1450,13 @@ def create_ui(): 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, ], outputs=[ ti_output, @@ -1566,7 +1610,7 @@ def create_ui(): opts.save(shared.config_filename) - return gr.update(value=value), opts.dumpjson() + return get_value_for_setting(key), opts.dumpjson() with gr.Blocks(analytics_enabled=False) as settings_interface: with gr.Row(): @@ -1738,7 +1782,7 @@ def create_ui(): component_keys = [k for k in opts.data_labels.keys() if k in component_dict] def get_settings_values(): - return [getattr(opts, key) for key in component_keys] + return [get_value_for_setting(key) for key in component_keys] demo.load( fn=get_settings_values, @@ -1753,12 +1797,15 @@ def create_ui(): print("Error loading/saving model file:", file=sys.stderr) print(traceback.format_exc(), file=sys.stderr) modules.sd_models.list_models() # to remove the potentially missing models from the list - return [f"Error merging checkpoints: {e}"] + [gr.Dropdown.update(choices=modules.sd_models.checkpoint_tiles()) for _ in range(4)] + return [*[gr.Dropdown.update(choices=modules.sd_models.checkpoint_tiles()) for _ in range(4)], f"Error merging checkpoints: {e}"] return results + modelmerger_merge.click(fn=lambda: '', inputs=[], outputs=[modelmerger_result]) modelmerger_merge.click( - fn=modelmerger, + fn=wrap_gradio_gpu_call(modelmerger, extra_outputs=lambda: [gr.update() for _ in range(4)]), + _js='modelmerger', inputs=[ + dummy_component, primary_model_name, secondary_model_name, tertiary_model_name, @@ -1768,13 +1815,14 @@ def create_ui(): custom_name, checkpoint_format, config_source, + bake_in_vae, ], outputs=[ - submit_result, primary_model_name, secondary_model_name, tertiary_model_name, component_dict['sd_model_checkpoint'], + modelmerger_result, ] ) diff --git a/style.css b/style.css index 7deedac58..ec0b4c5a6 100644 --- a/style.css +++ b/style.css @@ -635,6 +635,16 @@ canvas[key="mask"] { margin: 0.6em 0em 0.55em 0; } +#modelmerger_results_container{ + margin-top: 1em; + overflow: visible; +} + +#modelmerger_models{ + gap: 0; +} + + #quicksettings .gr-button-tool{ margin: 0; } diff --git a/webui.sh b/webui.sh index 6e07778ff..1edf921dd 100755 --- a/webui.sh +++ b/webui.sh @@ -165,7 +165,7 @@ else printf "\n%s\n" "${delimiter}" printf "Launching launch.py..." printf "\n%s\n" "${delimiter}" - gpu_info=$(lspci | grep VGA) + gpu_info=$(lspci 2>/dev/null | grep VGA) if echo "$gpu_info" | grep -q "AMD" then if [[ -z "${TORCH_COMMAND}" ]]