From 3f1c236ce4ac98945dc387602dbb0bb25e771fe0 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Thu, 13 Jun 2024 20:13:18 -0400 Subject: [PATCH] force apply vae config on model load --- CHANGELOG.md | 4 ++++ modules/sd_models.py | 11 +++++++---- modules/sd_vae.py | 23 +++++++++++++++++++++-- 3 files changed, 32 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e47aeda32..372bd7317 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ # Change Log for SD.Next +## Update for 2024-06-14 + +- force apply vae config on model load + ## Update for 2024-06-13 ### Highlights for 2024-06-13 diff --git a/modules/sd_models.py b/modules/sd_models.py index 2fb4a0292..535820686 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -545,7 +545,7 @@ def change_backend(): refresh_vae_list() -def detect_pipeline(f: str, op: str = 'model', warning=True): +def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False): guess = shared.opts.diffusers_pipeline warn = shared.log.warning if warning else lambda *args, **kwargs: None size = 0 @@ -642,7 +642,8 @@ def detect_pipeline(f: str, op: str = 'model', warning=True): guess = 'Stable Diffusion XL Instruct' # get actual pipeline pipeline = shared_items.get_pipelines().get(guess, None) - shared.log.info(f'Autodetect: {op}="{guess}" class={pipeline.__name__} file="{f}" size={size}MB') + if not quiet: + shared.log.info(f'Autodetect: {op}="{guess}" class={pipeline.__name__} file="{f}" size={size}MB') except Exception as e: shared.log.error(f'Error detecting diffusers pipeline: model={f} {e}') return None, None @@ -650,7 +651,8 @@ def detect_pipeline(f: str, op: str = 'model', warning=True): try: size = round(os.path.getsize(f) / 1024 / 1024) pipeline = shared_items.get_pipelines().get(guess, None) - shared.log.info(f'Diffusers: {op}="{guess}" class={pipeline.__name__} file="{f}" size={size}MB') + if not quiet: + shared.log.info(f'Diffusers: {op}="{guess}" class={pipeline.__name__} file="{f}" size={size}MB') except Exception as e: shared.log.error(f'Error loading diffusers pipeline: model={f} {e}') @@ -1157,7 +1159,8 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No timer.record("embeddings") set_diffuser_options(sd_model, vae, op) - + if op == 'model': + sd_vae.apply_vae_config(shared.sd_model.sd_checkpoint_info.filename, vae_file, sd_model) if op == 'refiner' and shared.opts.diffusers_move_refiner: shared.log.debug('Moving refiner model to CPU') move_model(sd_model, devices.cpu) diff --git a/modules/sd_vae.py b/modules/sd_vae.py index 2e27393e2..94a6c6b49 100644 --- a/modules/sd_vae.py +++ b/modules/sd_vae.py @@ -155,8 +155,6 @@ def load_vae(model, vae_file=None, vae_source="unknown-source"): except Exception as e: shared.log.error(f"Loading VAE failed: model={vae_file} source={vae_source} {e}") restore_base_vae(model) - # If vae used is not in dict, update it - # It will be removed on refresh though vae_opt = get_filename(vae_file) if vae_opt not in vae_dict: vae_dict[vae_opt] = vae_file @@ -165,6 +163,26 @@ def load_vae(model, vae_file=None, vae_source="unknown-source"): loaded_vae_file = vae_file +def apply_vae_config(model_file, vae_file, sd_model): + def get_vae_config(): + config_file = os.path.join(paths.sd_configs_path, os.path.splitext(os.path.basename(model_file))[0] + '_vae.json') + if config_file is not None and os.path.exists(config_file): + return shared.readfile(config_file) + config_file = os.path.join(paths.sd_configs_path, os.path.splitext(os.path.basename(vae_file))[0] + '.json') if vae_file else None + if config_file is not None and os.path.exists(config_file): + return shared.readfile(config_file) + config_file = os.path.join(paths.sd_configs_path, shared.sd_model_type, 'vae', 'config.json') + if config_file is not None and os.path.exists(config_file): + return shared.readfile(config_file) + return {} + + if hasattr(sd_model, 'vae') and hasattr(sd_model.vae, 'config'): + config = get_vae_config() + for k, v in config.items(): + if k in sd_model.vae.config and not k.startswith('_'): + sd_model.vae.config[k] = v + + def load_vae_diffusers(model_file, vae_file=None, vae_source="unknown-source"): if vae_file is None: return None @@ -262,6 +280,7 @@ def reload_vae_weights(sd_model=None, vae_file=unspecified): vae = load_vae_diffusers(shared.sd_model.sd_checkpoint_info.filename, vae_file, vae_source) if vae is not None: sd_models.set_diffuser_options(sd_model, vae=vae, op='vae') + apply_vae_config(shared.sd_model.sd_checkpoint_info.filename, vae_file, sd_model) if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram: sd_models.move_model(sd_model, devices.device)