diff --git a/CHANGELOG.md b/CHANGELOG.md index 188f7d2ef..6ba764ebd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,20 @@ # Change Log for SD.Next +## Update for 2024-09-14 + +- flux avoid unet load if unchanged +- flux mark specific unet as unavailable if load failed +- xyz grid full refactor +- xyz grid multi-mode: *selectable-script* and *alwayson-script* +- xyz grid allow usage combined with other scripts +- xyz grid allow **unet** selection +- xyz grid allow passing **model args** directly: + allowed params will be checked against models call signature + example: `width=768; height=512, width=512; height=768` +- xyz grid allow passing **processing args** directly: + params are set directly on main processing object and can be known or new params + example: `steps=10, steps=20; test=unknown` + ## Update for 2024-09-13 ### Highlights for 2024-09-13 diff --git a/javascript/sdnext.css b/javascript/sdnext.css index d2b59771b..dcabfaf8a 100644 --- a/javascript/sdnext.css +++ b/javascript/sdnext.css @@ -45,7 +45,7 @@ input[type='color'] { width: 64px; height: 32px; } .gradio-dropdown ul.options { z-index: 1000; min-width: fit-content; max-height: 50vh !important; white-space: nowrap; } .gradio-dropdown ul.options li.item { padding: var(--spacing-xs); } .gradio-dropdown ul.options li.item:not(:has(.hide)) { background-color: var(--primary-500); } -.gradio-dropdown .token { padding: var(--spacing-xs) !important; } +.gradio-dropdown .token { padding: var(--spacing-xs) !important; overflow-x: hidden; } .gradio-html { color: var(--body-text-color); } .gradio-html .min { min-height: 0; } .gradio-html div.wrap { height: 100%; } diff --git a/launch.py b/launch.py index 1db8a81f6..a2dbeb740 100755 --- a/launch.py +++ b/launch.py @@ -169,6 +169,12 @@ def start_server(immediate=True, server=None): uvicorn = None if args.test: installer.log.info("Test only") + installer.log.critical('Logging: level=critical') + installer.log.error('Logging: level=error') + installer.log.warning('Logging: level=warning') + installer.log.info('Logging: level=info') + installer.log.debug('Logging: level=debug') + installer.log.trace('Logging: level=trace') server.wants_restart = False else: if args.api_only: diff --git a/modules/devices.py b/modules/devices.py index a006e1ca3..358937ebf 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -67,7 +67,7 @@ def get_gpu_info(): } elif torch.version.cuda: return { - 'device': f'{torch.cuda.get_device_name(torch.cuda.current_device())} n={torch.cuda.device_count()} arch={torch.cuda.get_arch_list()[-1]} cap={torch.cuda.get_device_capability(device)}', + 'device': f'{torch.cuda.get_device_name(torch.cuda.current_device())} n={torch.cuda.device_count()} arch={torch.cuda.get_arch_list()[-1]} capability={torch.cuda.get_device_capability(device)}', 'cuda': torch.version.cuda, 'cudnn': torch.backends.cudnn.version(), 'driver': get_driver(), diff --git a/modules/loader.py b/modules/loader.py index 0aad4d76b..e696d8024 100644 --- a/modules/loader.py +++ b/modules/loader.py @@ -39,6 +39,9 @@ timer.startup.record("torch") import transformers # pylint: disable=W0611,C0411 timer.startup.record("transformers") +import accelerate # pylint: disable=W0611,C0411 +timer.startup.record("accelerate") + import onnxruntime # pylint: disable=W0611,C0411 onnxruntime.set_default_logger_severity(3) timer.startup.record("onnx") @@ -86,6 +89,8 @@ def get_packages(): "torch": getattr(torch, "__long_version__", torch.__version__), "diffusers": diffusers.__version__, "gradio": gradio.__version__, + "transformers": transformers.__version__, + "accelerate": accelerate.__version__, } errors.log.info(f'Load packages: {get_packages()}') diff --git a/modules/model_flux.py b/modules/model_flux.py index 0fd4e39ca..270599771 100644 --- a/modules/model_flux.py +++ b/modules/model_flux.py @@ -107,17 +107,27 @@ def load_flux_bnb(checkpoint_info, diffusers_load_config): # pylint: disable=unu install('bitsandbytes', quiet=True) from diffusers import FluxTransformer2DModel quant = get_quant(repo_path) - if quant == 'fp8': - quantization_config = transformers.BitsAndBytesConfig(load_in_8bit=True) - if transformer is None: + try: + if quant == 'fp8': + quantization_config = transformers.BitsAndBytesConfig(load_in_8bit=True, bnb_4bit_compute_dtype=devices.dtype) + debug(f'Quantization: {quantization_config}') transformer = FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config) - elif quant == 'fp4': - quantization_config = transformers.BitsAndBytesConfig(load_in_4bit=True) - if transformer is None: + elif quant == 'fp4': + quantization_config = transformers.BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=devices.dtype, bnb_4bit_quant_type= 'fp4') + debug(f'Quantization: {quantization_config}') transformer = FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config) - else: - if transformer is None: + elif quant == 'nf4': + quantization_config = transformers.BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=devices.dtype, bnb_4bit_quant_type= 'nf4') + debug(f'Quantization: {quantization_config}') + transformer = FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config) + else: transformer = FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config) + except Exception as e: + shared.log.error(f"Loading FLUX: Failed to load BnB transformer: {e}") + transformer, text_encoder_2 = None, None + if debug: + from modules import errors + errors.display(e, 'FLUX:') return transformer, text_encoder_2 @@ -130,19 +140,19 @@ def load_transformer(file_path): # triggered by opts.sd_unet change "cache_dir": shared.opts.hfcache_dir, } shared.log.info(f'Loading UNet: type=FLUX file="{file_path}" offload={shared.opts.diffusers_offload_mode} quant={quant} dtype={devices.dtype}') - if 'nf4' in quant: - from modules.model_flux_nf4 import load_flux_nf4 - _transformer, _text_encoder_2 = load_flux_nf4(file_path) - if _transformer is not None: - transformer = _transformer - elif quant == 'qint8' or quant == 'qint4': + if quant == 'qint8' or quant == 'qint4': _transformer, _text_encoder_2 = load_flux_quanto(file_path) if _transformer is not None: transformer = _transformer - elif quant == 'fp8' or quant == 'fp4': + elif quant == 'fp8' or quant == 'fp4' or quant == 'nf4': _transformer, _text_encoder_2 = load_flux_bnb(file_path, diffusers_load_config) if _transformer is not None: transformer = _transformer + elif 'nf4' in quant: # TODO right now this is not working for civitai published nf4 models + from modules.model_flux_nf4 import load_flux_nf4 + _transformer, _text_encoder_2 = load_flux_nf4(file_path) + if _transformer is not None: + transformer = _transformer else: from diffusers import FluxTransformer2DModel transformer = FluxTransformer2DModel.from_single_file(file_path, **diffusers_load_config) @@ -169,7 +179,10 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch from modules import sd_unet _transformer = load_transformer(sd_unet.unet_dict[shared.opts.sd_unet]) if _transformer is not None: + sd_unet.loaded_unet = shared.opts.sd_unet transformer = _transformer + else: + sd_unet.failed_unet.append(shared.opts.sd_unet) except Exception as e: shared.log.error(f"Loading FLUX: Failed to load UNet: {e}") if debug: diff --git a/modules/model_flux_nf4.py b/modules/model_flux_nf4.py index 1644a667a..9be36babb 100644 --- a/modules/model_flux_nf4.py +++ b/modules/model_flux_nf4.py @@ -198,16 +198,23 @@ def load_flux_nf4(checkpoint_info): _replace_with_bnb_linear(transformer, "nf4") - for param_name, param in converted_state_dict.items(): - if param_name not in expected_state_dict_keys: - continue - is_param_float8_e4m3fn = hasattr(torch, "float8_e4m3fn") and param.dtype == torch.float8_e4m3fn - if torch.is_floating_point(param) and not is_param_float8_e4m3fn: - param = param.to(devices.dtype) - if not check_quantized_param(transformer, param_name): - set_module_tensor_to_device(transformer, param_name, device=0, value=param) - else: - create_quantized_param(transformer, param, param_name, target_device=0, state_dict=original_state_dict, pre_quantized=True) + try: + for param_name, param in converted_state_dict.items(): + if param_name not in expected_state_dict_keys: + continue + is_param_float8_e4m3fn = hasattr(torch, "float8_e4m3fn") and param.dtype == torch.float8_e4m3fn + if torch.is_floating_point(param) and not is_param_float8_e4m3fn: + param = param.to(devices.dtype) + if not check_quantized_param(transformer, param_name): + set_module_tensor_to_device(transformer, param_name, device=0, value=param) + else: + create_quantized_param(transformer, param, param_name, target_device=0, state_dict=original_state_dict, pre_quantized=True) + except Exception as e: + transformer, text_encoder_2 = None, None + shared.log.error(f"Loading FLUX: Failed to load UNET: {e}") + if debug: + from modules import errors + errors.display(e, 'FLUX:') del original_state_dict devices.torch_gc(force=True) diff --git a/modules/sd_unet.py b/modules/sd_unet.py index 16d942a22..a3094d88d 100644 --- a/modules/sd_unet.py +++ b/modules/sd_unet.py @@ -3,10 +3,13 @@ from modules import shared, devices, files_cache, sd_models unet_dict = {} +loaded_unet = None +failed_unet = [] debug = os.environ.get('SD_LOAD_DEBUG', None) is not None def load_unet(model): + global loaded_unet # pylint: disable=global-statement if shared.opts.sd_unet == 'None': return if shared.opts.sd_unet not in list(unet_dict): @@ -19,16 +22,21 @@ def load_unet(model): config = None config_file = 'default' try: - if "StableCascade" in model.__class__.__name__: + if shared.opts.sd_unet == loaded_unet or shared.opts.sd_unet in failed_unet: + pass + elif "StableCascade" in model.__class__.__name__: from modules.model_stablecascade import load_prior prior_unet, prior_text_encoder = load_prior(unet_dict[shared.opts.sd_unet], config_file=config_file) + loaded_unet = shared.opts.sd_unet if prior_unet is not None: model.prior_pipe.prior = None # Prevent OOM model.prior_pipe.prior = prior_unet.to(devices.device, dtype=devices.dtype_unet) if prior_text_encoder is not None: model.prior_pipe.text_encoder = None # Prevent OOM model.prior_pipe.text_encoder = prior_text_encoder.to(devices.device, dtype=devices.dtype) - if "Flux" in model.__class__.__name__: + elif "Flux" in model.__class__.__name__: + sd_models.load_diffuser() # TODO forcing reloading entire flux as loading transformers only leads to massive memory usage + """ from modules.model_flux import load_transformer transformer = load_transformer(unet_dict[shared.opts.sd_unet]) if transformer is not None: @@ -36,8 +44,10 @@ def load_unet(model): if shared.opts.diffusers_offload_mode == 'none': sd_models.move_model(transformer, devices.device) model.transformer = transformer + loaded_unet = shared.opts.sd_unet from modules.sd_models import set_diffuser_offload set_diffuser_offload(model, 'model') + """ else: if not hasattr(model, 'unet') or model.unet is None: shared.log.error('UNet not found in current model') diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py index 2800a3722..e728b5f33 100644 --- a/scripts/xyz_grid.py +++ b/scripts/xyz_grid.py @@ -1,7 +1,4 @@ -# pylint: disable=unused-argument - -import os -import re +# xyz grid that shows as selectable script import csv import random from collections import namedtuple @@ -11,477 +8,14 @@ from io import StringIO from PIL import Image import numpy as np import gradio as gr -from modules import shared, errors, scripts, images, sd_samplers, processing, sd_models, sd_vae, ipadapter +from scripts.xyz_grid_shared import str_permutations, list_to_csv_string, re_range # pylint: disable=no-name-in-module +from scripts.xyz_grid_classes import axis_options, AxisOption, SharedSettingsStackHelper # pylint: disable=no-name-in-module +from scripts.xyz_grid_draw import draw_xyz_grid # pylint: disable=no-name-in-module +from modules import shared, errors, scripts, images, processing from modules.ui_components import ToolButton import modules.ui_symbols as symbols -def apply_field(field): - def fun(p, x, xs): - shared.log.debug(f'XYZ grid apply field: {field}={x}') - setattr(p, field, x) - return fun - - -def apply_task_args(field): - def fun(p, x, xs): - shared.log.debug(f'XYZ grid apply task-arg: {field}={x}') - p.task_args[field] = x - return fun - - -def apply_setting(field): - def fun(p, x, xs): - shared.log.debug(f'XYZ grid apply setting: {field}={x}') - shared.opts.data[field] = x - return fun - - -def apply_prompt(p, x, xs): - if xs[0] not in p.prompt and xs[0] not in p.negative_prompt: - shared.log.warning(f"XYZ grid: prompt S/R did not find {xs[0]} in prompt or negative prompt.") - else: - p.prompt = p.prompt.replace(xs[0], x) - p.negative_prompt = p.negative_prompt.replace(xs[0], x) - shared.log.debug(f'XYZ grid apply prompt: "{xs[0]}"="{x}"') - - -def apply_order(p, x, xs): - token_order = [] - for token in x: - token_order.append((p.prompt.find(token), token)) - token_order.sort(key=lambda t: t[0]) - prompt_parts = [] - for _, token in token_order: - n = p.prompt.find(token) - prompt_parts.append(p.prompt[0:n]) - p.prompt = p.prompt[n + len(token):] - prompt_tmp = "" - for idx, part in enumerate(prompt_parts): - prompt_tmp += part - prompt_tmp += x[idx] - p.prompt = prompt_tmp + p.prompt - - -def apply_sampler(p, x, xs): - sampler_name = sd_samplers.samplers_map.get(x.lower(), None) - if sampler_name is None: - shared.log.warning(f"XYZ grid: unknown sampler: {x}") - else: - p.sampler_name = sampler_name - shared.log.debug(f'XYZ grid apply sampler: "{x}"') - - -def apply_hr_sampler_name(p, x, xs): - hr_sampler_name = sd_samplers.samplers_map.get(x.lower(), None) - if hr_sampler_name is None: - shared.log.warning(f"XYZ grid: unknown sampler: {x}") - else: - p.hr_sampler_name = hr_sampler_name - shared.log.debug(f'XYZ grid apply HR sampler: "{x}"') - - -def confirm_samplers(p, xs): - for x in xs: - if x.lower() not in sd_samplers.samplers_map: - shared.log.warning(f"XYZ grid: unknown sampler: {x}") - - -def apply_checkpoint(p, x, xs): - if x == shared.opts.sd_model_checkpoint: - return - info = sd_models.get_closet_checkpoint_match(x) - if info is None: - shared.log.warning(f"XYZ grid: apply checkpoint unknown checkpoint: {x}") - else: - sd_models.reload_model_weights(shared.sd_model, info) - p.override_settings['sd_model_checkpoint'] = info.name - shared.log.debug(f'XYZ grid apply checkpoint: "{x}"') - - -def apply_refiner(p, x, xs): - if x == shared.opts.sd_model_refiner: - return - if x == 'None': - return - info = sd_models.get_closet_checkpoint_match(x) - if info is None: - shared.log.warning(f"XYZ grid: apply refiner unknown checkpoint: {x}") - else: - sd_models.reload_model_weights(shared.sd_refiner, info) - p.override_settings['sd_model_refiner'] = info.name - shared.log.debug(f'XYZ grid apply refiner: "{x}"') - - -def apply_dict(p, x, xs): - if x == shared.opts.sd_model_dict: - return - info_dict = sd_models.get_closet_checkpoint_match(x) - info_ckpt = sd_models.get_closet_checkpoint_match(shared.opts.sd_model_checkpoint) - if info_dict is None or info_ckpt is None: - shared.log.warning(f"XYZ grid: apply dict unknown checkpoint: {x}") - else: - shared.opts.sd_model_dict = info_dict.name # this will trigger reload_model_weights via onchange handler - p.override_settings['sd_model_checkpoint'] = info_ckpt.name - p.override_settings['sd_model_dict'] = info_dict.name - shared.log.debug(f'XYZ grid apply model dict: "{x}"') - - -def apply_clip_skip(p, x, xs): - p.clip_skip = x - shared.opts.data["clip_skip"] = x - shared.log.debug(f'XYZ grid apply clip-skip: "{x}"') - - -def find_vae(name: str): - if name.lower() in ['auto', 'automatic']: - return sd_vae.unspecified - if name.lower() == 'none': - return None - else: - choices = [x for x in sorted(sd_vae.vae_dict, key=lambda x: len(x)) if name.lower().strip() in x.lower()] - if len(choices) == 0: - shared.log.warning(f"No VAE found for {name}; using automatic") - return sd_vae.unspecified - else: - return sd_vae.vae_dict[choices[0]] - - -def apply_vae(p, x, xs): - sd_vae.reload_vae_weights(shared.sd_model, vae_file=find_vae(x)) - shared.log.debug(f'XYZ grid apply VAE: "{x}"') - - -def list_lora(): - import sys - lora = [v for k, v in sys.modules.items() if k == 'networks'][0] - loras = [v.fullname for v in lora.available_networks.values()] - return ['None'] + loras - - -def apply_lora(p, x, xs): - if x == 'None': - return - x = os.path.basename(x) - p.prompt = p.prompt + f" " - shared.log.debug(f'XYZ grid apply LoRA: "{x}"') - - -def apply_te(p, x, xs): - shared.opts.data["sd_text_encoder"] = x - sd_models.reload_text_encoder() - shared.log.debug(f'XYZ grid apply text-encoder: "{x}"') - - -def apply_styles(p: processing.StableDiffusionProcessingTxt2Img, x: str, _): - p.styles.extend(x.split(',')) - shared.log.debug(f'XYZ grid apply style: "{x}"') - - -def apply_upscaler(p: processing.StableDiffusionProcessingTxt2Img, opt, x): - p.enable_hr = True - p.hr_force = True - p.denoising_strength = 0.0 - p.hr_upscaler = opt - shared.log.debug(f'XYZ grid apply upscaler: "{x}"') - - -def apply_context(p: processing.StableDiffusionProcessingTxt2Img, opt, x): - p.resize_mode = 5 - p.resize_context = opt - shared.log.debug(f'XYZ grid apply resize-context: "{x}"') - - -def apply_face_restore(p, opt, x): - opt = opt.lower() - if opt == 'codeformer': - is_active = True - p.face_restoration_model = 'CodeFormer' - elif opt == 'gfpgan': - is_active = True - p.face_restoration_model = 'GFPGAN' - else: - is_active = opt in ('true', 'yes', 'y', '1') - p.restore_faces = is_active - shared.log.debug(f'XYZ grid apply face-restore: "{x}"') - - -def apply_override(field): - def fun(p, x, xs): - p.override_settings[field] = x - shared.log.debug(f'XYZ grid apply override: "{field}"="{x}"') - return fun - - -def format_value_add_label(p, opt, x): - if type(x) == float: - x = round(x, 8) - return f"{opt.label}: {x}" - - -def format_value(p, opt, x): - if type(x) == float: - x = round(x, 8) - return x - - -def format_value_join_list(p, opt, x): - return ", ".join(x) - - -def do_nothing(p, x, xs): - pass - - -def format_nothing(p, opt, x): - return "" - - -def str_permutations(x): - """dummy function for specifying it in AxisOption's type when you want to get a list of permutations""" - return x - - -def list_to_csv_string(data_list): - with StringIO() as o: - csv.writer(o).writerow(data_list) - return o.getvalue().strip() - - -class AxisOption: - def __init__(self, label, tipe, apply, fmt=format_value_add_label, confirm=None, cost=0.0, choices=None): - self.label = label - self.type = tipe - self.apply = apply - self.format_value = fmt - self.confirm = confirm - self.cost = cost - self.choices = choices - - -class AxisOptionImg2Img(AxisOption): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.is_img2img = True - -class AxisOptionTxt2Img(AxisOption): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.is_img2img = False - - -axis_options = [ - AxisOption("Nothing", str, do_nothing, fmt=format_nothing), - AxisOption("Prompt S/R", str, apply_prompt, fmt=format_value), - AxisOption("Model", str, apply_checkpoint, fmt=format_value, cost=1.0, choices=lambda: sorted(sd_models.checkpoints_list)), - AxisOption("VAE", str, apply_vae, cost=0.7, choices=lambda: ['None'] + list(sd_vae.vae_dict)), - AxisOption("LoRA", str, apply_lora, cost=0.5, choices=list_lora), - AxisOption("LoRA strength", float, apply_setting('extra_networks_default_multiplier')), - AxisOption("Text encoder", str, apply_te, cost=0.7, choices=lambda: ['None', 'T5 FP4', 'T5 FP8', 'T5 FP16']), - AxisOption("Styles", str, apply_styles, choices=lambda: [s.name for s in shared.prompt_styles.styles.values()]), - AxisOption("Seed", int, apply_field("seed")), - AxisOption("Steps", int, apply_field("steps")), - AxisOption("CFG scale", float, apply_field("cfg_scale")), - AxisOption("Guidance end", float, apply_field("cfg_end")), - AxisOption("Variation seed", int, apply_field("subseed")), - AxisOption("Variation strength", float, apply_field("subseed_strength")), - AxisOption("Clip skip", float, apply_clip_skip), - AxisOption("Denoising strength", float, apply_field("denoising_strength")), - AxisOption("Prompt order", str_permutations, apply_order, fmt=format_value_join_list), - AxisOption("Model dictionary", str, apply_dict, fmt=format_value, cost=1.0, choices=lambda: ['None'] + list(sd_models.checkpoints_list)), - AxisOptionImg2Img("Image mask weight", float, apply_field("inpainting_mask_weight")), - AxisOptionTxt2Img("[Sampler] Name", str, apply_sampler, fmt=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]), - AxisOptionImg2Img("[Sampler] Name", str, apply_sampler, fmt=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers_for_img2img]), - AxisOption("[Sampler] Timestep spacing", str, apply_setting("schedulers_timestep_spacing"), choices=lambda: ['default', 'linspace', 'leading', 'trailing']), - AxisOption("[Sampler] Sigma min", float, apply_field("s_min")), - AxisOption("[Sampler] Sigma max", float, apply_field("s_max")), - AxisOption("[Sampler] Sigma tmin", float, apply_field("s_tmin")), - AxisOption("[Sampler] Sigma tmax", float, apply_field("s_tmax")), - AxisOption("[Sampler] Sigma churn", float, apply_field("s_churn")), - AxisOption("[Sampler] Sigma noise", float, apply_field("s_noise")), - AxisOption("[Sampler] Shift", float, apply_setting("schedulers_shift")), - AxisOption("[Sampler] ETA", float, apply_setting("scheduler_eta")), - AxisOption("[Sampler] Solver order", int, apply_setting("schedulers_solver_order")), - AxisOption("[Second pass] Upscaler", str, apply_field("hr_upscaler"), choices=lambda: [*shared.latent_upscale_modes, *[x.name for x in shared.sd_upscalers]]), - AxisOption("[Second pass] Sampler", str, apply_hr_sampler_name, fmt=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]), - AxisOption("[Second pass] Denoising strength", float, apply_field("denoising_strength")), - AxisOption("[Second pass] Hires steps", int, apply_field("hr_second_pass_steps")), - AxisOption("[Second pass] CFG scale", float, apply_field("image_cfg_scale")), - AxisOption("[Second pass] Guidance rescale", float, apply_field("diffusers_guidance_rescale")), - AxisOption("[Refiner] Model", str, apply_refiner, fmt=format_value, cost=1.0, choices=lambda: ['None'] + sorted(sd_models.checkpoints_list)), - AxisOption("[Refiner] Refiner start", float, apply_field("refiner_start")), - AxisOption("[Refiner] Refiner steps", float, apply_field("refiner_steps")), - AxisOption("[Postprocess] Upscaler", str, apply_upscaler, choices=lambda: [x.name for x in shared.sd_upscalers][1:]), - AxisOption("[Postprocess] Context", str, apply_context, choices=lambda: ["Add with forward", "Remove with forward", "Add with backward", "Remove with backward"]), - AxisOption("[Postprocess] Face restore", str, apply_face_restore, fmt=format_value), - AxisOption("[HDR] Mode", int, apply_field("hdr_mode")), - AxisOption("[HDR] Brightness", float, apply_field("hdr_brightness")), - AxisOption("[HDR] Color", float, apply_field("hdr_color")), - AxisOption("[HDR] Sharpen", float, apply_field("hdr_sharpen")), - AxisOption("[HDR] Clamp boundary", float, apply_field("hdr_boundary")), - AxisOption("[HDR] Clamp threshold", float, apply_field("hdr_threshold")), - AxisOption("[HDR] Maximize center shift", float, apply_field("hdr_max_center")), - AxisOption("[HDR] Maximize boundary", float, apply_field("hdr_max_boundry")), - AxisOption("[HDR] Tint color hex", str, apply_field("hdr_color_picker")), - AxisOption("[HDR] Tint ratio", float, apply_field("hdr_tint_ratio")), - AxisOption("[Token Merging] ToMe ratio", float, apply_setting('tome_ratio')), - AxisOption("[Token Merging] ToDo ratio", float, apply_setting('todo_ratio')), - AxisOption("[FreeU] 1st stage backbone factor", float, apply_setting('freeu_b1')), - AxisOption("[FreeU] 2nd stage backbone factor", float, apply_setting('freeu_b2')), - AxisOption("[FreeU] 1st stage skip factor", float, apply_setting('freeu_s1')), - AxisOption("[FreeU] 2nd stage skip factor", float, apply_setting('freeu_s2')), - AxisOption("[IP adapter] Name", str, apply_field('ip_adapter_names'), cost=1.0, choices=lambda: list(ipadapter.ADAPTERS)), - AxisOption("[IP adapter] Scale", float, apply_field('ip_adapter_scales')), - AxisOption("[IP adapter] Starts", float, apply_field('ip_adapter_starts')), - AxisOption("[IP adapter] Ends", float, apply_field('ip_adapter_ends')), - AxisOption("[HiDiffusion] T1", float, apply_override('hidiffusion_t1')), - AxisOption("[HiDiffusion] T2", float, apply_override('hidiffusion_t2')), - AxisOption("[HiDiffusion] Agression step", float, apply_field('hidiffusion_steps')), - AxisOption("[PAG] Attention scale", float, apply_field('pag_scale')), - AxisOption("[PAG] Adaptive scaling", float, apply_field('pag_adaptive')), - AxisOption("[PAG] Applied layers", str, apply_setting('pag_apply_layers')), -] - - -def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend, include_lone_images, include_sub_grids, first_axes_processed, second_axes_processed, margin_size, no_grid): - hor_texts = [[images.GridAnnotation(x)] for x in x_labels] - ver_texts = [[images.GridAnnotation(y)] for y in y_labels] - title_texts = [[images.GridAnnotation(z)] for z in z_labels] - list_size = (len(xs) * len(ys) * len(zs)) - processed_result = None - shared.state.job_count = list_size * p.n_iter - - def process_cell(x, y, z, ix, iy, iz): - nonlocal processed_result - - def index(ix, iy, iz): - return ix + iy * len(xs) + iz * len(xs) * len(ys) - - shared.state.job = 'grid' - processed: processing.Processed = cell(x, y, z, ix, iy, iz) - if processed_result is None: - processed_result = copy(processed) - if processed_result is None: - shared.log.error('XYZ grid: no processing results') - return processing.Processed(p, []) - processed_result.images = [None] * list_size - processed_result.all_prompts = [None] * list_size - processed_result.all_seeds = [None] * list_size - processed_result.infotexts = [None] * list_size - processed_result.index_of_first_image = 1 - idx = index(ix, iy, iz) - if processed is not None and processed.images: - processed_result.images[idx] = processed.images[0] - processed_result.all_prompts[idx] = processed.prompt - processed_result.all_seeds[idx] = processed.seed - processed_result.infotexts[idx] = processed.infotexts[0] - else: - cell_mode = "P" - cell_size = (processed_result.width, processed_result.height) - if processed_result.images[0] is not None: - cell_mode = processed_result.images[0].mode - cell_size = processed_result.images[0].size - processed_result.images[idx] = Image.new(cell_mode, cell_size) - - if first_axes_processed == 'x': - for ix, x in enumerate(xs): - if second_axes_processed == 'y': - for iy, y in enumerate(ys): - for iz, z in enumerate(zs): - process_cell(x, y, z, ix, iy, iz) - else: - for iz, z in enumerate(zs): - for iy, y in enumerate(ys): - process_cell(x, y, z, ix, iy, iz) - elif first_axes_processed == 'y': - for iy, y in enumerate(ys): - if second_axes_processed == 'x': - for ix, x in enumerate(xs): - for iz, z in enumerate(zs): - process_cell(x, y, z, ix, iy, iz) - else: - for iz, z in enumerate(zs): - for ix, x in enumerate(xs): - process_cell(x, y, z, ix, iy, iz) - elif first_axes_processed == 'z': - for iz, z in enumerate(zs): - if second_axes_processed == 'x': - for ix, x in enumerate(xs): - for iy, y in enumerate(ys): - process_cell(x, y, z, ix, iy, iz) - else: - for iy, y in enumerate(ys): - for ix, x in enumerate(xs): - process_cell(x, y, z, ix, iy, iz) - - if not processed_result: - shared.log.error("XYZ grid: Failed to initialize processing") - return processing.Processed(p, []) - elif not any(processed_result.images): - shared.log.error("XYZ grid: Failed to return processed image") - return processing.Processed(p, []) - - z_count = len(zs) - for i in range(z_count): - start_index = (i * len(xs) * len(ys)) + i - end_index = start_index + len(xs) * len(ys) - if (not no_grid or include_sub_grids) and images.check_grid_size(processed_result.images[start_index:end_index]): - grid = images.image_grid(processed_result.images[start_index:end_index], rows=len(ys)) - if draw_legend: - grid = images.draw_grid_annotations(grid, processed_result.images[start_index].size[0], processed_result.images[start_index].size[1], hor_texts, ver_texts, margin_size, title=title_texts[i]) - processed_result.images.insert(i, grid) - processed_result.all_prompts.insert(i, processed_result.all_prompts[start_index]) - processed_result.all_seeds.insert(i, processed_result.all_seeds[start_index]) - processed_result.infotexts.insert(i, processed_result.infotexts[start_index]) - sub_grid_size = processed_result.images[0].size - if not no_grid and images.check_grid_size(processed_result.images[:z_count]): - z_grid = images.image_grid(processed_result.images[:z_count], rows=1) - if draw_legend: - z_grid = images.draw_grid_annotations(z_grid, sub_grid_size[0], sub_grid_size[1], [[images.GridAnnotation()] for _ in z_labels], [[images.GridAnnotation()]]) - processed_result.images.insert(0, z_grid) - #processed_result.all_prompts.insert(0, processed_result.all_prompts[0]) - #processed_result.all_seeds.insert(0, processed_result.all_seeds[0]) - processed_result.infotexts.insert(0, processed_result.infotexts[0]) - return processed_result - - -class SharedSettingsStackHelper(object): - vae = None - schedulers_solver_order = None - tome_ratio = None - todo_ratio = None - sd_model_checkpoint = None - sd_model_dict = None - sd_vae_checkpoint = None - - def __enter__(self): - #Save overridden settings so they can be restored later. - self.vae = shared.opts.sd_vae - self.schedulers_solver_order = shared.opts.schedulers_solver_order - self.tome_ratio = shared.opts.tome_ratio - self.todo_ratio = shared.opts.todo_ratio - self.sd_model_checkpoint = shared.opts.sd_model_checkpoint - self.sd_model_dict = shared.opts.sd_model_dict - self.sd_vae_checkpoint = shared.opts.sd_vae - - def __exit__(self, exc_type, exc_value, tb): - #Restore overriden settings after plot generation. - shared.opts.data["sd_vae"] = self.vae - shared.opts.data["schedulers_solver_order"] = self.schedulers_solver_order - shared.opts.data["tome_ratio"] = self.tome_ratio - shared.opts.data["todo_ratio"] = self.todo_ratio - if self.sd_model_dict != shared.opts.sd_model_dict: - shared.opts.data["sd_model_dict"] = self.sd_model_dict - if self.sd_model_checkpoint != shared.opts.sd_model_checkpoint: - shared.opts.data["sd_model_checkpoint"] = self.sd_model_checkpoint - sd_models.reload_model_weights() - if self.sd_vae_checkpoint != shared.opts.sd_vae: - shared.opts.data["sd_vae"] = self.sd_vae_checkpoint - sd_vae.reload_vae_weights() - - -re_range = re.compile(r'([-+]?[0-9]*\.?[0-9]+)-([-+]?[0-9]*\.?[0-9]+):?([0-9]+)?') - class Script(scripts.Script): current_axis_options = [] @@ -699,7 +233,6 @@ class Script(scripts.Script): shared.state.xyz_plot_x = AxisInfo(x_opt, xs) shared.state.xyz_plot_y = AxisInfo(y_opt, ys) shared.state.xyz_plot_z = AxisInfo(z_opt, zs) - # If one of the axes is very slow to change between (like SD model checkpoint), then make sure it is in the outer iteration of the nested `for` loop. first_axes_processed = 'z' second_axes_processed = 'y' if x_opt.cost > y_opt.cost and x_opt.cost > z_opt.cost: @@ -803,8 +336,5 @@ class Script(scripts.Script): del processed.all_seeds[1] del processed.infotexts[1] elif no_grid: - # del processed.images[0] - # del processed.all_prompts[0] - # del processed.all_seeds[0] del processed.infotexts[0] return processed diff --git a/scripts/xyz_grid_classes.py b/scripts/xyz_grid_classes.py new file mode 100644 index 000000000..f9e807c56 --- /dev/null +++ b/scripts/xyz_grid_classes.py @@ -0,0 +1,151 @@ +from scripts.xyz_grid_shared import apply_field, apply_task_args, apply_setting, apply_prompt, apply_order, apply_sampler, apply_hr_sampler_name, confirm_samplers, apply_checkpoint, apply_refiner, apply_unet, apply_dict, apply_clip_skip, apply_vae, list_lora, apply_lora, apply_te, apply_styles, apply_upscaler, apply_context, apply_face_restore, apply_override, apply_processing, format_value_add_label, format_value, format_value_join_list, do_nothing, format_nothing, str_permutations # pylint: disable=no-name-in-module +from modules import shared, sd_samplers, ipadapter, sd_models, sd_vae, sd_unet + + +class AxisOption: + def __init__(self, label, tipe, apply, fmt=format_value_add_label, confirm=None, cost=0.0, choices=None): + self.label = label + self.type = tipe + self.apply = apply + self.format_value = fmt + self.confirm = confirm + self.cost = cost + self.choices = choices + + +class AxisOptionImg2Img(AxisOption): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.is_img2img = True + +class AxisOptionTxt2Img(AxisOption): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.is_img2img = False + + +class SharedSettingsStackHelper(object): + vae = None + schedulers_solver_order = None + tome_ratio = None + todo_ratio = None + sd_model_checkpoint = None + sd_model_refiner = None + sd_model_dict = None + sd_vae = None + sd_unet = None + sd_text_encoder = None + + def __enter__(self): + #Save overridden settings so they can be restored later. + self.vae = shared.opts.sd_vae + self.schedulers_solver_order = shared.opts.schedulers_solver_order + self.tome_ratio = shared.opts.tome_ratio + self.todo_ratio = shared.opts.todo_ratio + self.sd_model_checkpoint = shared.opts.sd_model_checkpoint + self.sd_model_refiner = shared.opts.sd_model_refiner + self.sd_model_dict = shared.opts.sd_model_dict + self.sd_vae = shared.opts.sd_vae + self.sd_unet = shared.opts.sd_unet + self.sd_text_encoder = shared.opts.sd_text_encoder + + def __exit__(self, exc_type, exc_value, tb): + #Restore overriden settings after plot generation. + shared.opts.data["sd_vae"] = self.vae + shared.opts.data["schedulers_solver_order"] = self.schedulers_solver_order + shared.opts.data["tome_ratio"] = self.tome_ratio + shared.opts.data["todo_ratio"] = self.todo_ratio + if self.sd_model_checkpoint != shared.opts.sd_model_checkpoint: + shared.opts.data["sd_model_checkpoint"] = self.sd_model_checkpoint + sd_models.reload_model_weights(op='model') + if self.sd_model_refiner != shared.opts.sd_model_refiner: + shared.opts.data["sd_model_refiner"] = self.sd_model_refiner + sd_models.reload_model_weights(op='refiner') + if self.sd_model_dict != shared.opts.sd_model_dict: + shared.opts.data["sd_model_dict"] = self.sd_model_dict + sd_models.reload_model_weights(op='dict') + if self.sd_vae != shared.opts.sd_vae: + shared.opts.data["sd_vae"] = self.sd_vae + sd_vae.reload_vae_weights() + if self.sd_text_encoder != shared.opts.sd_text_encoder: + shared.opts.data["sd_text_encoder"] = self.sd_text_encoder + sd_models.reload_text_encoder() + if self.sd_unet != shared.opts.sd_unet: + shared.opts.data["sd_unet"] = self.sd_unet + sd_unet.load_unet(shared.sd_model) + + +axis_options = [ + AxisOption("Nothing", str, do_nothing, fmt=format_nothing), + AxisOption("Prompt S/R", str, apply_prompt, fmt=format_value), + AxisOption("Model", str, apply_checkpoint, fmt=format_value, cost=1.0, choices=lambda: sorted(sd_models.checkpoints_list)), + AxisOption("UNET", str, apply_unet, cost=0.9, choices=lambda: ['None'] + list(sd_unet.unet_dict)), + AxisOption("VAE", str, apply_vae, cost=0.7, choices=lambda: ['None'] + list(sd_vae.vae_dict)), + AxisOption("LoRA", str, apply_lora, cost=0.5, choices=list_lora), + AxisOption("LoRA strength", float, apply_setting('extra_networks_default_multiplier')), + AxisOption("Text encoder", str, apply_te, cost=0.7, choices=lambda: ['None', 'T5 FP4', 'T5 FP8', 'T5 FP16']), + AxisOption("Styles", str, apply_styles, choices=lambda: [s.name for s in shared.prompt_styles.styles.values()]), + AxisOption("Seed", int, apply_field("seed")), + AxisOption("Steps", int, apply_field("steps")), + AxisOption("CFG scale", float, apply_field("cfg_scale")), + AxisOption("Guidance end", float, apply_field("cfg_end")), + AxisOption("Variation seed", int, apply_field("subseed")), + AxisOption("Variation strength", float, apply_field("subseed_strength")), + AxisOption("Clip skip", float, apply_clip_skip), + AxisOption("Denoising strength", float, apply_field("denoising_strength")), + AxisOption("Prompt order", str_permutations, apply_order, fmt=format_value_join_list), + AxisOption("Model dictionary", str, apply_dict, fmt=format_value, cost=1.0, choices=lambda: ['None'] + list(sd_models.checkpoints_list)), + AxisOption("Model args", str, apply_task_args), + AxisOption("Processing args", str, apply_processing), + AxisOptionImg2Img("Image mask weight", float, apply_field("inpainting_mask_weight")), + AxisOptionTxt2Img("[Sampler] Name", str, apply_sampler, fmt=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]), + AxisOptionImg2Img("[Sampler] Name", str, apply_sampler, fmt=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers_for_img2img]), + AxisOption("[Sampler] Timestep spacing", str, apply_setting("schedulers_timestep_spacing"), choices=lambda: ['default', 'linspace', 'leading', 'trailing']), + AxisOption("[Sampler] Sigma min", float, apply_field("s_min")), + AxisOption("[Sampler] Sigma max", float, apply_field("s_max")), + AxisOption("[Sampler] Sigma tmin", float, apply_field("s_tmin")), + AxisOption("[Sampler] Sigma tmax", float, apply_field("s_tmax")), + AxisOption("[Sampler] Sigma churn", float, apply_field("s_churn")), + AxisOption("[Sampler] Sigma noise", float, apply_field("s_noise")), + AxisOption("[Sampler] Shift", float, apply_setting("schedulers_shift")), + AxisOption("[Sampler] ETA", float, apply_setting("scheduler_eta")), + AxisOption("[Sampler] Solver order", int, apply_setting("schedulers_solver_order")), + AxisOption("[Second pass] Upscaler", str, apply_field("hr_upscaler"), choices=lambda: [*shared.latent_upscale_modes, *[x.name for x in shared.sd_upscalers]]), + AxisOption("[Second pass] Sampler", str, apply_hr_sampler_name, fmt=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]), + AxisOption("[Second pass] Denoising strength", float, apply_field("denoising_strength")), + AxisOption("[Second pass] Hires steps", int, apply_field("hr_second_pass_steps")), + AxisOption("[Second pass] CFG scale", float, apply_field("image_cfg_scale")), + AxisOption("[Second pass] Guidance rescale", float, apply_field("diffusers_guidance_rescale")), + AxisOption("[Refiner] Model", str, apply_refiner, fmt=format_value, cost=1.0, choices=lambda: ['None'] + sorted(sd_models.checkpoints_list)), + AxisOption("[Refiner] Refiner start", float, apply_field("refiner_start")), + AxisOption("[Refiner] Refiner steps", float, apply_field("refiner_steps")), + AxisOption("[Postprocess] Upscaler", str, apply_upscaler, choices=lambda: [x.name for x in shared.sd_upscalers][1:]), + AxisOption("[Postprocess] Context", str, apply_context, choices=lambda: ["Add with forward", "Remove with forward", "Add with backward", "Remove with backward"]), + AxisOption("[Postprocess] Face restore", str, apply_face_restore, fmt=format_value), + AxisOption("[HDR] Mode", int, apply_field("hdr_mode")), + AxisOption("[HDR] Brightness", float, apply_field("hdr_brightness")), + AxisOption("[HDR] Color", float, apply_field("hdr_color")), + AxisOption("[HDR] Sharpen", float, apply_field("hdr_sharpen")), + AxisOption("[HDR] Clamp boundary", float, apply_field("hdr_boundary")), + AxisOption("[HDR] Clamp threshold", float, apply_field("hdr_threshold")), + AxisOption("[HDR] Maximize center shift", float, apply_field("hdr_max_center")), + AxisOption("[HDR] Maximize boundary", float, apply_field("hdr_max_boundry")), + AxisOption("[HDR] Tint color hex", str, apply_field("hdr_color_picker")), + AxisOption("[HDR] Tint ratio", float, apply_field("hdr_tint_ratio")), + AxisOption("[Token Merging] ToMe ratio", float, apply_setting('tome_ratio')), + AxisOption("[Token Merging] ToDo ratio", float, apply_setting('todo_ratio')), + AxisOption("[FreeU] 1st stage backbone factor", float, apply_setting('freeu_b1')), + AxisOption("[FreeU] 2nd stage backbone factor", float, apply_setting('freeu_b2')), + AxisOption("[FreeU] 1st stage skip factor", float, apply_setting('freeu_s1')), + AxisOption("[FreeU] 2nd stage skip factor", float, apply_setting('freeu_s2')), + AxisOption("[IP adapter] Name", str, apply_field('ip_adapter_names'), cost=1.0, choices=lambda: list(ipadapter.ADAPTERS)), + AxisOption("[IP adapter] Scale", float, apply_field('ip_adapter_scales')), + AxisOption("[IP adapter] Starts", float, apply_field('ip_adapter_starts')), + AxisOption("[IP adapter] Ends", float, apply_field('ip_adapter_ends')), + AxisOption("[HiDiffusion] T1", float, apply_override('hidiffusion_t1')), + AxisOption("[HiDiffusion] T2", float, apply_override('hidiffusion_t2')), + AxisOption("[HiDiffusion] Agression step", float, apply_field('hidiffusion_steps')), + AxisOption("[PAG] Attention scale", float, apply_field('pag_scale')), + AxisOption("[PAG] Adaptive scaling", float, apply_field('pag_adaptive')), + AxisOption("[PAG] Applied layers", str, apply_setting('pag_apply_layers')), +] diff --git a/scripts/xyz_grid_draw.py b/scripts/xyz_grid_draw.py new file mode 100644 index 000000000..be4862e3d --- /dev/null +++ b/scripts/xyz_grid_draw.py @@ -0,0 +1,105 @@ +from copy import copy +from PIL import Image +from modules import shared, images, processing + + +def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend, include_lone_images, include_sub_grids, first_axes_processed, second_axes_processed, margin_size, no_grid): # pylint: disable=unused-argument + hor_texts = [[images.GridAnnotation(x)] for x in x_labels] + ver_texts = [[images.GridAnnotation(y)] for y in y_labels] + title_texts = [[images.GridAnnotation(z)] for z in z_labels] + list_size = (len(xs) * len(ys) * len(zs)) + processed_result = None + shared.state.job_count = list_size * p.n_iter + + def process_cell(x, y, z, ix, iy, iz): + nonlocal processed_result + + def index(ix, iy, iz): + return ix + iy * len(xs) + iz * len(xs) * len(ys) + + shared.state.job = 'grid' + processed: processing.Processed = cell(x, y, z, ix, iy, iz) + if processed_result is None: + processed_result = copy(processed) + if processed_result is None: + shared.log.error('XYZ grid: no processing results') + return processing.Processed(p, []) + processed_result.images = [None] * list_size + processed_result.all_prompts = [None] * list_size + processed_result.all_seeds = [None] * list_size + processed_result.infotexts = [None] * list_size + processed_result.index_of_first_image = 1 + idx = index(ix, iy, iz) + if processed is not None and processed.images: + processed_result.images[idx] = processed.images[0] + processed_result.all_prompts[idx] = processed.prompt + processed_result.all_seeds[idx] = processed.seed + processed_result.infotexts[idx] = processed.infotexts[0] + else: + cell_mode = "P" + cell_size = (processed_result.width, processed_result.height) + if processed_result.images[0] is not None: + cell_mode = processed_result.images[0].mode + cell_size = processed_result.images[0].size + processed_result.images[idx] = Image.new(cell_mode, cell_size) + + if first_axes_processed == 'x': + for ix, x in enumerate(xs): + if second_axes_processed == 'y': + for iy, y in enumerate(ys): + for iz, z in enumerate(zs): + process_cell(x, y, z, ix, iy, iz) + else: + for iz, z in enumerate(zs): + for iy, y in enumerate(ys): + process_cell(x, y, z, ix, iy, iz) + elif first_axes_processed == 'y': + for iy, y in enumerate(ys): + if second_axes_processed == 'x': + for ix, x in enumerate(xs): + for iz, z in enumerate(zs): + process_cell(x, y, z, ix, iy, iz) + else: + for iz, z in enumerate(zs): + for ix, x in enumerate(xs): + process_cell(x, y, z, ix, iy, iz) + elif first_axes_processed == 'z': + for iz, z in enumerate(zs): + if second_axes_processed == 'x': + for ix, x in enumerate(xs): + for iy, y in enumerate(ys): + process_cell(x, y, z, ix, iy, iz) + else: + for iy, y in enumerate(ys): + for ix, x in enumerate(xs): + process_cell(x, y, z, ix, iy, iz) + + if not processed_result: + shared.log.error("XYZ grid: Failed to initialize processing") + return processing.Processed(p, []) + elif not any(processed_result.images): + shared.log.error("XYZ grid: Failed to return processed image") + return processing.Processed(p, []) + + z_count = len(zs) + for i in range(z_count): + start_index = (i * len(xs) * len(ys)) + i + end_index = start_index + len(xs) * len(ys) + if (not no_grid or include_sub_grids) and images.check_grid_size(processed_result.images[start_index:end_index]): + grid = images.image_grid(processed_result.images[start_index:end_index], rows=len(ys)) + if draw_legend: + grid = images.draw_grid_annotations(grid, processed_result.images[start_index].size[0], processed_result.images[start_index].size[1], hor_texts, ver_texts, margin_size, title=title_texts[i]) + processed_result.images.insert(i, grid) + processed_result.all_prompts.insert(i, processed_result.all_prompts[start_index]) + processed_result.all_seeds.insert(i, processed_result.all_seeds[start_index]) + processed_result.infotexts.insert(i, processed_result.infotexts[start_index]) + sub_grid_size = processed_result.images[0].size + if not no_grid and images.check_grid_size(processed_result.images[:z_count]): + z_grid = images.image_grid(processed_result.images[:z_count], rows=1) + if draw_legend: + z_grid = images.draw_grid_annotations(z_grid, sub_grid_size[0], sub_grid_size[1], [[images.GridAnnotation()] for _ in z_labels], [[images.GridAnnotation()]]) + processed_result.images.insert(0, z_grid) + #processed_result.all_prompts.insert(0, processed_result.all_prompts[0]) + #processed_result.all_seeds.insert(0, processed_result.all_seeds[0]) + processed_result.infotexts.insert(0, processed_result.infotexts[0]) + return processed_result diff --git a/scripts/xyz_grid_on.py b/scripts/xyz_grid_on.py index ac362ba06..6f9e86daa 100644 --- a/scripts/xyz_grid_on.py +++ b/scripts/xyz_grid_on.py @@ -1,7 +1,4 @@ -# pylint: disable=unused-argument - -import os -import re +# xyz grid that shows up as alwayson script import csv import random from collections import namedtuple @@ -11,7 +8,10 @@ from io import StringIO from PIL import Image import numpy as np import gradio as gr -from modules import shared, errors, scripts, images, sd_samplers, processing, sd_models, sd_vae, ipadapter +from scripts.xyz_grid_shared import str_permutations, list_to_csv_string, re_range # pylint: disable=no-name-in-module +from scripts.xyz_grid_classes import axis_options, AxisOption, SharedSettingsStackHelper # pylint: disable=no-name-in-module +from scripts.xyz_grid_draw import draw_xyz_grid # pylint: disable=no-name-in-module +from modules import shared, errors, scripts, images, processing from modules.ui_components import ToolButton import modules.ui_symbols as symbols @@ -20,474 +20,6 @@ active = False cache = None -def apply_field(field): - def fun(p, x, xs): - shared.log.debug(f'XYZ grid apply field: {field}={x}') - setattr(p, field, x) - return fun - - -def apply_task_args(field): - def fun(p, x, xs): - shared.log.debug(f'XYZ grid apply task-arg: {field}={x}') - p.task_args[field] = x - return fun - - -def apply_setting(field): - def fun(p, x, xs): - shared.log.debug(f'XYZ grid apply setting: {field}={x}') - shared.opts.data[field] = x - return fun - - -def apply_prompt(p, x, xs): - if xs[0] not in p.prompt and xs[0] not in p.negative_prompt: - shared.log.warning(f"XYZ grid: prompt S/R did not find {xs[0]} in prompt or negative prompt.") - else: - p.prompt = p.prompt.replace(xs[0], x) - p.all_prompts = p.batch_size * [p.prompt] - p.negative_prompt = p.negative_prompt.replace(xs[0], x) - p.all_negative_prompts = p.batch_size * [p.negative_prompt] - shared.log.debug(f'XYZ grid apply prompt: "{xs[0]}"="{x}"') - - -def apply_order(p, x, xs): - token_order = [] - for token in x: - token_order.append((p.prompt.find(token), token)) - token_order.sort(key=lambda t: t[0]) - prompt_parts = [] - for _, token in token_order: - n = p.prompt.find(token) - prompt_parts.append(p.prompt[0:n]) - p.prompt = p.prompt[n + len(token):] - prompt_tmp = "" - for idx, part in enumerate(prompt_parts): - prompt_tmp += part - prompt_tmp += x[idx] - p.prompt = prompt_tmp + p.prompt - - -def apply_sampler(p, x, xs): - sampler_name = sd_samplers.samplers_map.get(x.lower(), None) - if sampler_name is None: - shared.log.warning(f"XYZ grid: unknown sampler: {x}") - else: - p.sampler_name = sampler_name - shared.log.debug(f'XYZ grid apply sampler: "{x}"') - - -def apply_hr_sampler_name(p, x, xs): - hr_sampler_name = sd_samplers.samplers_map.get(x.lower(), None) - if hr_sampler_name is None: - shared.log.warning(f"XYZ grid: unknown sampler: {x}") - else: - p.hr_sampler_name = hr_sampler_name - shared.log.debug(f'XYZ grid apply HR sampler: "{x}"') - - -def confirm_samplers(p, xs): - for x in xs: - if x.lower() not in sd_samplers.samplers_map: - shared.log.warning(f"XYZ grid: unknown sampler: {x}") - - -def apply_checkpoint(p, x, xs): - if x == shared.opts.sd_model_checkpoint: - return - info = sd_models.get_closet_checkpoint_match(x) - if info is None: - shared.log.warning(f"XYZ grid: apply checkpoint unknown checkpoint: {x}") - else: - sd_models.reload_model_weights(shared.sd_model, info) - p.override_settings['sd_model_checkpoint'] = info.name - shared.log.debug(f'XYZ grid apply checkpoint: "{x}"') - - -def apply_refiner(p, x, xs): - if x == shared.opts.sd_model_refiner: - return - if x == 'None': - return - info = sd_models.get_closet_checkpoint_match(x) - if info is None: - shared.log.warning(f"XYZ grid: apply refiner unknown checkpoint: {x}") - else: - sd_models.reload_model_weights(shared.sd_refiner, info) - p.override_settings['sd_model_refiner'] = info.name - shared.log.debug(f'XYZ grid apply refiner: "{x}"') - - -def apply_dict(p, x, xs): - if x == shared.opts.sd_model_dict: - return - info_dict = sd_models.get_closet_checkpoint_match(x) - info_ckpt = sd_models.get_closet_checkpoint_match(shared.opts.sd_model_checkpoint) - if info_dict is None or info_ckpt is None: - shared.log.warning(f"XYZ grid: apply dict unknown checkpoint: {x}") - else: - shared.opts.sd_model_dict = info_dict.name # this will trigger reload_model_weights via onchange handler - p.override_settings['sd_model_checkpoint'] = info_ckpt.name - p.override_settings['sd_model_dict'] = info_dict.name - shared.log.debug(f'XYZ grid apply model dict: "{x}"') - - -def apply_clip_skip(p, x, xs): - p.clip_skip = x - shared.opts.data["clip_skip"] = x - shared.log.debug(f'XYZ grid apply clip-skip: "{x}"') - - -def find_vae(name: str): - if name.lower() in ['auto', 'automatic']: - return sd_vae.unspecified - if name.lower() == 'none': - return None - else: - choices = [x for x in sorted(sd_vae.vae_dict, key=lambda x: len(x)) if name.lower().strip() in x.lower()] - if len(choices) == 0: - shared.log.warning(f"No VAE found for {name}; using automatic") - return sd_vae.unspecified - else: - return sd_vae.vae_dict[choices[0]] - - -def apply_vae(p, x, xs): - sd_vae.reload_vae_weights(shared.sd_model, vae_file=find_vae(x)) - shared.log.debug(f'XYZ grid apply VAE: "{x}"') - - -def list_lora(): - import sys - lora = [v for k, v in sys.modules.items() if k == 'networks'][0] - loras = [v.fullname for v in lora.available_networks.values()] - return ['None'] + loras - - -def apply_lora(p, x, xs): - if x == 'None': - return - x = os.path.basename(x) - p.prompt = p.prompt + f" " - shared.log.debug(f'XYZ grid apply LoRA: "{x}"') - - -def apply_te(p, x, xs): - shared.opts.data["sd_text_encoder"] = x - sd_models.reload_text_encoder() - shared.log.debug(f'XYZ grid apply text-encoder: "{x}"') - - -def apply_styles(p: processing.StableDiffusionProcessingTxt2Img, x: str, _): - p.styles.extend(x.split(',')) - shared.log.debug(f'XYZ grid apply style: "{x}"') - - -def apply_upscaler(p: processing.StableDiffusionProcessingTxt2Img, opt, x): - p.enable_hr = True - p.hr_force = True - p.denoising_strength = 0.0 - p.hr_upscaler = opt - shared.log.debug(f'XYZ grid apply upscaler: "{x}"') - - -def apply_context(p: processing.StableDiffusionProcessingTxt2Img, opt, x): - p.resize_mode = 5 - p.resize_context = opt - shared.log.debug(f'XYZ grid apply resize-context: "{x}"') - - -def apply_face_restore(p, opt, x): - opt = opt.lower() - if opt == 'codeformer': - is_active = True - p.face_restoration_model = 'CodeFormer' - elif opt == 'gfpgan': - is_active = True - p.face_restoration_model = 'GFPGAN' - else: - is_active = opt in ('true', 'yes', 'y', '1') - p.restore_faces = is_active - shared.log.debug(f'XYZ grid apply face-restore: "{x}"') - - -def apply_override(field): - def fun(p, x, xs): - p.override_settings[field] = x - shared.log.debug(f'XYZ grid apply override: "{field}"="{x}"') - return fun - - -def format_value_add_label(p, opt, x): - if type(x) == float: - x = round(x, 8) - return f"{opt.label}: {x}" - - -def format_value(p, opt, x): - if type(x) == float: - x = round(x, 8) - return x - - -def format_value_join_list(p, opt, x): - return ", ".join(x) - - -def do_nothing(p, x, xs): - pass - - -def format_nothing(p, opt, x): - return "" - - -def str_permutations(x): - """dummy function for specifying it in AxisOption's type when you want to get a list of permutations""" - return x - - -def list_to_csv_string(data_list): - with StringIO() as o: - csv.writer(o).writerow(data_list) - return o.getvalue().strip() - - -class AxisOption: - def __init__(self, label, tipe, apply, fmt=format_value_add_label, confirm=None, cost=0.0, choices=None): - self.label = label - self.type = tipe - self.apply = apply - self.format_value = fmt - self.confirm = confirm - self.cost = cost - self.choices = choices - - -class AxisOptionImg2Img(AxisOption): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.is_img2img = True - -class AxisOptionTxt2Img(AxisOption): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.is_img2img = False - - -axis_options = [ - AxisOption("Nothing", str, do_nothing, fmt=format_nothing), - AxisOption("Prompt S/R", str, apply_prompt, fmt=format_value), - AxisOption("Model", str, apply_checkpoint, fmt=format_value, cost=1.0, choices=lambda: sorted(sd_models.checkpoints_list)), - AxisOption("VAE", str, apply_vae, cost=0.7, choices=lambda: ['None'] + list(sd_vae.vae_dict)), - AxisOption("LoRA", str, apply_lora, cost=0.5, choices=list_lora), - AxisOption("LoRA strength", float, apply_setting('extra_networks_default_multiplier')), - AxisOption("Text encoder", str, apply_te, cost=0.7, choices=lambda: ['None', 'T5 FP4', 'T5 FP8', 'T5 FP16']), - AxisOption("Styles", str, apply_styles, choices=lambda: [s.name for s in shared.prompt_styles.styles.values()]), - AxisOption("Seed", int, apply_field("seed")), - AxisOption("Steps", int, apply_field("steps")), - AxisOption("CFG scale", float, apply_field("cfg_scale")), - AxisOption("Guidance end", float, apply_field("cfg_end")), - AxisOption("Variation seed", int, apply_field("subseed")), - AxisOption("Variation strength", float, apply_field("subseed_strength")), - AxisOption("Clip skip", float, apply_clip_skip), - AxisOption("Denoising strength", float, apply_field("denoising_strength")), - AxisOption("Prompt order", str_permutations, apply_order, fmt=format_value_join_list), - AxisOption("Model dictionary", str, apply_dict, fmt=format_value, cost=1.0, choices=lambda: ['None'] + list(sd_models.checkpoints_list)), - AxisOptionImg2Img("Image mask weight", float, apply_field("inpainting_mask_weight")), - AxisOptionTxt2Img("[Sampler] Name", str, apply_sampler, fmt=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]), - AxisOptionImg2Img("[Sampler] Name", str, apply_sampler, fmt=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers_for_img2img]), - AxisOption("[Sampler] Timestep spacing", str, apply_setting("schedulers_timestep_spacing"), choices=lambda: ['default', 'linspace', 'leading', 'trailing']), - AxisOption("[Sampler] Sigma min", float, apply_field("s_min")), - AxisOption("[Sampler] Sigma max", float, apply_field("s_max")), - AxisOption("[Sampler] Sigma tmin", float, apply_field("s_tmin")), - AxisOption("[Sampler] Sigma tmax", float, apply_field("s_tmax")), - AxisOption("[Sampler] Sigma churn", float, apply_field("s_churn")), - AxisOption("[Sampler] Sigma noise", float, apply_field("s_noise")), - AxisOption("[Sampler] Shift", float, apply_setting("schedulers_shift")), - AxisOption("[Sampler] ETA", float, apply_setting("scheduler_eta")), - AxisOption("[Sampler] Solver order", int, apply_setting("schedulers_solver_order")), - AxisOption("[Second pass] Upscaler", str, apply_field("hr_upscaler"), choices=lambda: [*shared.latent_upscale_modes, *[x.name for x in shared.sd_upscalers]]), - AxisOption("[Second pass] Sampler", str, apply_hr_sampler_name, fmt=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]), - AxisOption("[Second pass] Denoising strength", float, apply_field("denoising_strength")), - AxisOption("[Second pass] Hires steps", int, apply_field("hr_second_pass_steps")), - AxisOption("[Second pass] CFG scale", float, apply_field("image_cfg_scale")), - AxisOption("[Second pass] Guidance rescale", float, apply_field("diffusers_guidance_rescale")), - AxisOption("[Refiner] Model", str, apply_refiner, fmt=format_value, cost=1.0, choices=lambda: ['None'] + sorted(sd_models.checkpoints_list)), - AxisOption("[Refiner] Refiner start", float, apply_field("refiner_start")), - AxisOption("[Refiner] Refiner steps", float, apply_field("refiner_steps")), - AxisOption("[Postprocess] Upscaler", str, apply_upscaler, choices=lambda: [x.name for x in shared.sd_upscalers][1:]), - AxisOption("[Postprocess] Context", str, apply_context, choices=lambda: ["Add with forward", "Remove with forward", "Add with backward", "Remove with backward"]), - AxisOption("[Postprocess] Face restore", str, apply_face_restore, fmt=format_value), - AxisOption("[HDR] Mode", int, apply_field("hdr_mode")), - AxisOption("[HDR] Brightness", float, apply_field("hdr_brightness")), - AxisOption("[HDR] Color", float, apply_field("hdr_color")), - AxisOption("[HDR] Sharpen", float, apply_field("hdr_sharpen")), - AxisOption("[HDR] Clamp boundary", float, apply_field("hdr_boundary")), - AxisOption("[HDR] Clamp threshold", float, apply_field("hdr_threshold")), - AxisOption("[HDR] Maximize center shift", float, apply_field("hdr_max_center")), - AxisOption("[HDR] Maximize boundary", float, apply_field("hdr_max_boundry")), - AxisOption("[HDR] Tint color hex", str, apply_field("hdr_color_picker")), - AxisOption("[HDR] Tint ratio", float, apply_field("hdr_tint_ratio")), - AxisOption("[Token Merging] ToMe ratio", float, apply_setting('tome_ratio')), - AxisOption("[Token Merging] ToDo ratio", float, apply_setting('todo_ratio')), - AxisOption("[FreeU] 1st stage backbone factor", float, apply_setting('freeu_b1')), - AxisOption("[FreeU] 2nd stage backbone factor", float, apply_setting('freeu_b2')), - AxisOption("[FreeU] 1st stage skip factor", float, apply_setting('freeu_s1')), - AxisOption("[FreeU] 2nd stage skip factor", float, apply_setting('freeu_s2')), - AxisOption("[IP adapter] Name", str, apply_field('ip_adapter_names'), cost=1.0, choices=lambda: list(ipadapter.ADAPTERS)), - AxisOption("[IP adapter] Scale", float, apply_field('ip_adapter_scales')), - AxisOption("[IP adapter] Starts", float, apply_field('ip_adapter_starts')), - AxisOption("[IP adapter] Ends", float, apply_field('ip_adapter_ends')), - AxisOption("[HiDiffusion] T1", float, apply_override('hidiffusion_t1')), - AxisOption("[HiDiffusion] T2", float, apply_override('hidiffusion_t2')), - AxisOption("[HiDiffusion] Agression step", float, apply_field('hidiffusion_steps')), - AxisOption("[PAG] Attention scale", float, apply_field('pag_scale')), - AxisOption("[PAG] Adaptive scaling", float, apply_field('pag_adaptive')), - AxisOption("[PAG] Applied layers", str, apply_setting('pag_apply_layers')), -] - - -def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend, include_lone_images, include_sub_grids, first_axes_processed, second_axes_processed, margin_size, no_grid): - hor_texts = [[images.GridAnnotation(x)] for x in x_labels] - ver_texts = [[images.GridAnnotation(y)] for y in y_labels] - title_texts = [[images.GridAnnotation(z)] for z in z_labels] - list_size = (len(xs) * len(ys) * len(zs)) - processed_result = None - shared.state.job_count = list_size * p.n_iter - - def process_cell(x, y, z, ix, iy, iz): - nonlocal processed_result - - def index(ix, iy, iz): - return ix + iy * len(xs) + iz * len(xs) * len(ys) - - shared.state.job = 'grid' - processed: processing.Processed = cell(x, y, z, ix, iy, iz) - if processed_result is None: - processed_result = copy(processed) - if processed_result is None: - shared.log.error('XYZ grid: no processing results') - return processing.Processed(p, []) - processed_result.images = [None] * list_size - processed_result.all_prompts = [None] * list_size - processed_result.all_seeds = [None] * list_size - processed_result.infotexts = [None] * list_size - processed_result.index_of_first_image = 1 - idx = index(ix, iy, iz) - if processed is not None and processed.images: - processed_result.images[idx] = processed.images[0] - processed_result.all_prompts[idx] = processed.prompt - processed_result.all_seeds[idx] = processed.seed - processed_result.infotexts[idx] = processed.infotexts[0] - else: - cell_mode = "P" - cell_size = (processed_result.width, processed_result.height) - if processed_result.images[0] is not None: - cell_mode = processed_result.images[0].mode - cell_size = processed_result.images[0].size - processed_result.images[idx] = Image.new(cell_mode, cell_size) - - if first_axes_processed == 'x': - for ix, x in enumerate(xs): - if second_axes_processed == 'y': - for iy, y in enumerate(ys): - for iz, z in enumerate(zs): - process_cell(x, y, z, ix, iy, iz) - else: - for iz, z in enumerate(zs): - for iy, y in enumerate(ys): - process_cell(x, y, z, ix, iy, iz) - elif first_axes_processed == 'y': - for iy, y in enumerate(ys): - if second_axes_processed == 'x': - for ix, x in enumerate(xs): - for iz, z in enumerate(zs): - process_cell(x, y, z, ix, iy, iz) - else: - for iz, z in enumerate(zs): - for ix, x in enumerate(xs): - process_cell(x, y, z, ix, iy, iz) - elif first_axes_processed == 'z': - for iz, z in enumerate(zs): - if second_axes_processed == 'x': - for ix, x in enumerate(xs): - for iy, y in enumerate(ys): - process_cell(x, y, z, ix, iy, iz) - else: - for iy, y in enumerate(ys): - for ix, x in enumerate(xs): - process_cell(x, y, z, ix, iy, iz) - - if not processed_result: - shared.log.error("XYZ grid: Failed to initialize processing") - return processing.Processed(p, []) - elif not any(processed_result.images): - shared.log.error("XYZ grid: Failed to return processed image") - return processing.Processed(p, []) - - z_count = len(zs) - for i in range(z_count): - start_index = (i * len(xs) * len(ys)) + i - end_index = start_index + len(xs) * len(ys) - if (not no_grid or include_sub_grids) and images.check_grid_size(processed_result.images[start_index:end_index]): - grid = images.image_grid(processed_result.images[start_index:end_index], rows=len(ys)) - if draw_legend: - grid = images.draw_grid_annotations(grid, processed_result.images[start_index].size[0], processed_result.images[start_index].size[1], hor_texts, ver_texts, margin_size, title=title_texts[i]) - processed_result.images.insert(i, grid) - processed_result.all_prompts.insert(i, processed_result.all_prompts[start_index]) - processed_result.all_seeds.insert(i, processed_result.all_seeds[start_index]) - processed_result.infotexts.insert(i, processed_result.infotexts[start_index]) - sub_grid_size = processed_result.images[0].size - if not no_grid and images.check_grid_size(processed_result.images[:z_count]): - z_grid = images.image_grid(processed_result.images[:z_count], rows=1) - if draw_legend: - z_grid = images.draw_grid_annotations(z_grid, sub_grid_size[0], sub_grid_size[1], [[images.GridAnnotation()] for _ in z_labels], [[images.GridAnnotation()]]) - processed_result.images.insert(0, z_grid) - #processed_result.all_prompts.insert(0, processed_result.all_prompts[0]) - #processed_result.all_seeds.insert(0, processed_result.all_seeds[0]) - processed_result.infotexts.insert(0, processed_result.infotexts[0]) - return processed_result - - -class SharedSettingsStackHelper(object): - vae = None - schedulers_solver_order = None - tome_ratio = None - todo_ratio = None - sd_model_checkpoint = None - sd_model_dict = None - sd_vae_checkpoint = None - - def __enter__(self): - #Save overridden settings so they can be restored later. - self.vae = shared.opts.sd_vae - self.schedulers_solver_order = shared.opts.schedulers_solver_order - self.tome_ratio = shared.opts.tome_ratio - self.todo_ratio = shared.opts.todo_ratio - self.sd_model_checkpoint = shared.opts.sd_model_checkpoint - self.sd_model_dict = shared.opts.sd_model_dict - self.sd_vae_checkpoint = shared.opts.sd_vae - - def __exit__(self, exc_type, exc_value, tb): - #Restore overriden settings after plot generation. - shared.opts.data["sd_vae"] = self.vae - shared.opts.data["schedulers_solver_order"] = self.schedulers_solver_order - shared.opts.data["tome_ratio"] = self.tome_ratio - shared.opts.data["todo_ratio"] = self.todo_ratio - if self.sd_model_dict != shared.opts.sd_model_dict: - shared.opts.data["sd_model_dict"] = self.sd_model_dict - if self.sd_model_checkpoint != shared.opts.sd_model_checkpoint: - shared.opts.data["sd_model_checkpoint"] = self.sd_model_checkpoint - sd_models.reload_model_weights() - if self.sd_vae_checkpoint != shared.opts.sd_vae: - shared.opts.data["sd_vae"] = self.sd_vae_checkpoint - sd_vae.reload_vae_weights() - - -re_range = re.compile(r'([-+]?[0-9]*\.?[0-9]+)-([-+]?[0-9]*\.?[0-9]+):?([0-9]+)?') - class Script(scripts.Script): current_axis_options = [] @@ -823,7 +355,7 @@ class Script(scripts.Script): cache = processed return processed - def process_images(self, p, enabled, x_type, x_values, x_values_dropdown, y_type, y_values, y_values_dropdown, z_type, z_values, z_values_dropdown, csv_mode, draw_legend, no_fixed_seeds, no_grid, include_lone_images, include_sub_grids, margin_size): # pylint: disable=W0221 + def process_images(self, p, enabled, x_type, x_values, x_values_dropdown, y_type, y_values, y_values_dropdown, z_type, z_values, z_values_dropdown, csv_mode, draw_legend, no_fixed_seeds, no_grid, include_lone_images, include_sub_grids, margin_size): # pylint: disable=W0221, W0613 global cache # pylint: disable=W0603 if cache is not None and hasattr(cache, 'images'): samples = cache.images.copy() diff --git a/scripts/xyz_grid_shared.py b/scripts/xyz_grid_shared.py new file mode 100644 index 000000000..144da09f9 --- /dev/null +++ b/scripts/xyz_grid_shared.py @@ -0,0 +1,270 @@ +# pylint: disable=unused-argument + +import os +import re +import csv +from io import StringIO +from modules import shared, processing, sd_samplers, sd_models, sd_vae, sd_unet + + +re_range = re.compile(r'([-+]?[0-9]*\.?[0-9]+)-([-+]?[0-9]*\.?[0-9]+):?([0-9]+)?') + + +def apply_field(field): + def fun(p, x, xs): + shared.log.debug(f'XYZ grid apply field: {field}={x}') + setattr(p, field, x) + return fun + + +def apply_task_args(p, x, xs): + for section in x.split(';'): + k, v = section.split('=') + k, v = k.strip(), v.strip() + if v.replace('.','',1).isdigit(): + v = float(v) if '.' in v else int(v) + p.task_args[k] = v + shared.log.debug(f'XYZ grid apply task-arg: {k}={type(v)}:{v}') + + +def apply_processing(p, x, xs): + for section in x.split(';'): + k, v = section.split('=') + k, v = k.strip(), v.strip() + if v.replace('.','',1).isdigit(): + v = float(v) if '.' in v else int(v) + found = 'existing' if hasattr(p, k) else 'new' + setattr(p, k, v) + shared.log.debug(f'XYZ grid apply processing-arg: type={found} {k}={type(v)}:{v} ') + + +def apply_setting(field): + def fun(p, x, xs): + shared.log.debug(f'XYZ grid apply setting: {field}={x}') + shared.opts.data[field] = x + return fun + + +def apply_prompt(p, x, xs): + if xs[0] not in p.prompt and xs[0] not in p.negative_prompt: + shared.log.warning(f"XYZ grid: prompt S/R did not find {xs[0]} in prompt or negative prompt.") + else: + p.prompt = p.prompt.replace(xs[0], x) + for i in range(len(p.all_prompts)): + p.all_prompts[i] = p.all_prompts[i].replace(xs[0], x) + p.negative_prompt = p.negative_prompt.replace(xs[0], x) + for i in range(len(p.all_negative_prompts)): + p.all_negative_prompts[i] = p.all_negative_prompts[i].replace(xs[0], x) + shared.log.debug(f'XYZ grid apply prompt: "{xs[0]}"="{x}"') + + +def apply_order(p, x, xs): + token_order = [] + for token in x: + token_order.append((p.prompt.find(token), token)) + token_order.sort(key=lambda t: t[0]) + prompt_parts = [] + for _, token in token_order: + n = p.prompt.find(token) + prompt_parts.append(p.prompt[0:n]) + p.prompt = p.prompt[n + len(token):] + prompt_tmp = "" + for idx, part in enumerate(prompt_parts): + prompt_tmp += part + prompt_tmp += x[idx] + p.prompt = prompt_tmp + p.prompt + + +def apply_sampler(p, x, xs): + sampler_name = sd_samplers.samplers_map.get(x.lower(), None) + if sampler_name is None: + shared.log.warning(f"XYZ grid: unknown sampler: {x}") + else: + p.sampler_name = sampler_name + shared.log.debug(f'XYZ grid apply sampler: "{x}"') + + +def apply_hr_sampler_name(p, x, xs): + hr_sampler_name = sd_samplers.samplers_map.get(x.lower(), None) + if hr_sampler_name is None: + shared.log.warning(f"XYZ grid: unknown sampler: {x}") + else: + p.hr_sampler_name = hr_sampler_name + shared.log.debug(f'XYZ grid apply HR sampler: "{x}"') + + +def confirm_samplers(p, xs): + for x in xs: + if x.lower() not in sd_samplers.samplers_map: + shared.log.warning(f"XYZ grid: unknown sampler: {x}") + + +def apply_checkpoint(p, x, xs): + if x == shared.opts.sd_model_checkpoint: + return + info = sd_models.get_closet_checkpoint_match(x) + if info is None: + shared.log.warning(f"XYZ grid: apply checkpoint unknown checkpoint: {x}") + else: + sd_models.reload_model_weights(shared.sd_model, info) + p.override_settings['sd_model_checkpoint'] = info.name + shared.log.debug(f'XYZ grid apply checkpoint: "{x}"') + + +def apply_refiner(p, x, xs): + if x == shared.opts.sd_model_refiner: + return + if x == 'None': + return + info = sd_models.get_closet_checkpoint_match(x) + if info is None: + shared.log.warning(f"XYZ grid: apply refiner unknown checkpoint: {x}") + else: + sd_models.reload_model_weights(shared.sd_refiner, info) + p.override_settings['sd_model_refiner'] = info.name + shared.log.debug(f'XYZ grid apply refiner: "{x}"') + + +def apply_unet(p, x, xs): + if x == shared.opts.sd_unet: + return + if x == 'None': + return + p.override_settings['sd_unet'] = x + sd_unet.load_unet(shared.sd_model) + shared.log.debug(f'XYZ grid apply unet: "{x}"') + + +def apply_dict(p, x, xs): + if x == shared.opts.sd_model_dict: + return + info_dict = sd_models.get_closet_checkpoint_match(x) + info_ckpt = sd_models.get_closet_checkpoint_match(shared.opts.sd_model_checkpoint) + if info_dict is None or info_ckpt is None: + shared.log.warning(f"XYZ grid: apply dict unknown checkpoint: {x}") + else: + shared.opts.sd_model_dict = info_dict.name # this will trigger reload_model_weights via onchange handler + p.override_settings['sd_model_checkpoint'] = info_ckpt.name + p.override_settings['sd_model_dict'] = info_dict.name + shared.log.debug(f'XYZ grid apply model dict: "{x}"') + + +def apply_clip_skip(p, x, xs): + p.clip_skip = x + shared.opts.data["clip_skip"] = x + shared.log.debug(f'XYZ grid apply clip-skip: "{x}"') + + +def find_vae(name: str): + if name.lower() in ['auto', 'automatic']: + return sd_vae.unspecified + if name.lower() == 'none': + return None + else: + choices = [x for x in sorted(sd_vae.vae_dict, key=lambda x: len(x)) if name.lower().strip() in x.lower()] + if len(choices) == 0: + shared.log.warning(f"No VAE found for {name}; using automatic") + return sd_vae.unspecified + else: + return sd_vae.vae_dict[choices[0]] + + +def apply_vae(p, x, xs): + sd_vae.reload_vae_weights(shared.sd_model, vae_file=find_vae(x)) + shared.log.debug(f'XYZ grid apply VAE: "{x}"') + + +def list_lora(): + import sys + lora = [v for k, v in sys.modules.items() if k == 'networks'][0] + loras = [v.fullname for v in lora.available_networks.values()] + return ['None'] + loras + + +def apply_lora(p, x, xs): + if x == 'None': + return + x = os.path.basename(x) + p.prompt = p.prompt + f" " + shared.log.debug(f'XYZ grid apply LoRA: "{x}"') + + +def apply_te(p, x, xs): + shared.opts.data["sd_text_encoder"] = x + sd_models.reload_text_encoder() + shared.log.debug(f'XYZ grid apply text-encoder: "{x}"') + + +def apply_styles(p: processing.StableDiffusionProcessingTxt2Img, x: str, _): + p.styles.extend(x.split(',')) + shared.log.debug(f'XYZ grid apply style: "{x}"') + + +def apply_upscaler(p: processing.StableDiffusionProcessingTxt2Img, opt, x): + p.enable_hr = True + p.hr_force = True + p.denoising_strength = 0.0 + p.hr_upscaler = opt + shared.log.debug(f'XYZ grid apply upscaler: "{x}"') + + +def apply_context(p: processing.StableDiffusionProcessingTxt2Img, opt, x): + p.resize_mode = 5 + p.resize_context = opt + shared.log.debug(f'XYZ grid apply resize-context: "{x}"') + + +def apply_face_restore(p, opt, x): + opt = opt.lower() + if opt == 'codeformer': + is_active = True + p.face_restoration_model = 'CodeFormer' + elif opt == 'gfpgan': + is_active = True + p.face_restoration_model = 'GFPGAN' + else: + is_active = opt in ('true', 'yes', 'y', '1') + p.restore_faces = is_active + shared.log.debug(f'XYZ grid apply face-restore: "{x}"') + + +def apply_override(field): + def fun(p, x, xs): + p.override_settings[field] = x + shared.log.debug(f'XYZ grid apply override: "{field}"="{x}"') + return fun + + +def format_value_add_label(p, opt, x): + if type(x) == float: + x = round(x, 8) + return f"{opt.label}: {x}" + + +def format_value(p, opt, x): + if type(x) == float: + x = round(x, 8) + return x + + +def format_value_join_list(p, opt, x): + return ", ".join(x) + + +def do_nothing(p, x, xs): + pass + + +def format_nothing(p, opt, x): + return "" + + +def str_permutations(x): + """dummy function for specifying it in AxisOption's type when you want to get a list of permutations""" + return x + + +def list_to_csv_string(data_list): + with StringIO() as o: + csv.writer(o).writerow(data_list) + return o.getvalue().strip() diff --git a/webui.py b/webui.py index 6013aafd6..029e22de0 100644 --- a/webui.py +++ b/webui.py @@ -117,11 +117,6 @@ def initialize(): modelloader.load_upscalers() timer.startup.record("upscalers") - shared.opts.onchange("sd_vae", wrap_queued_call(lambda: modules.sd_vae.reload_vae_weights()), call=False) - shared.opts.onchange("sd_unet", wrap_queued_call(lambda: modules.sd_unet.load_unet(shared.sd_model)), call=False) - shared.opts.onchange("temp_dir", gr_tempdir.on_tmpdir_changed) - timer.startup.record("onchange") - shared.reload_hypernetworks() shared.prompt_styles.reload() @@ -170,13 +165,16 @@ def load_model(): shared.state.end() thread_model.join() thread_refiner.join() + timer.startup.record("checkpoint") shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(op='model')), call=False) shared.opts.onchange("sd_model_refiner", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(op='refiner')), call=False) shared.opts.onchange("sd_text_encoder", wrap_queued_call(lambda: modules.sd_models.reload_text_encoder()), call=False) shared.opts.onchange("sd_model_dict", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(op='dict')), call=False) shared.opts.onchange("sd_vae", wrap_queued_call(lambda: modules.sd_vae.reload_vae_weights()), call=False) + shared.opts.onchange("sd_unet", wrap_queued_call(lambda: modules.sd_unet.load_unet(shared.sd_model)), call=False) shared.opts.onchange("sd_backend", wrap_queued_call(lambda: modules.sd_models.change_backend()), call=False) - timer.startup.record("checkpoint") + shared.opts.onchange("temp_dir", gr_tempdir.on_tmpdir_changed) + timer.startup.record("onchange") def create_api(app):