diff --git a/CHANGELOG.md b/CHANGELOG.md index 4441e92f4..5381b08b5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,7 @@ Improvements: - gguf transformer loader (prototype) - OpenVINO: add accuracy option - ZLUDA: guess GPU arch +- Major model load refactor Fixes: - fix send-to-control diff --git a/modules/extras.py b/modules/extras.py index e22360f8a..162491580 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -188,7 +188,7 @@ def run_modelmerger(id_task, **kwargs): # pylint: disable=unused-argument _, extension = os.path.splitext(output_modelname) if os.path.exists(output_modelname) and not kwargs.get("overwrite", False): - return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)], f"Model alredy exists: {output_modelname}"] + return [*[gr.Dropdown.update(choices=sd_models.checkpoint_titles()) for _ in range(4)], f"Model alredy exists: {output_modelname}"] if extension.lower() == ".safetensors": safetensors.torch.save_file(theta_0, output_modelname, metadata=metadata) else: @@ -202,7 +202,7 @@ def run_modelmerger(id_task, **kwargs): # pylint: disable=unused-argument created_model.calculate_shorthash() devices.torch_gc(force=True) shared.state.end() - return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)], f"Model saved to {output_modelname}"] + return [*[gr.Dropdown.update(choices=sd_models.checkpoint_titles()) for _ in range(4)], f"Model saved to {output_modelname}"] def run_modelconvert(model, checkpoint_formats, precision, conv_type, custom_name, unet_conv, text_encoder_conv, diff --git a/modules/face/faceid.py b/modules/face/faceid.py index 754ce59a3..e2d5efccb 100644 --- a/modules/face/faceid.py +++ b/modules/face/faceid.py @@ -6,9 +6,10 @@ import numpy as np import diffusers import huggingface_hub as hf from PIL import Image -from modules import processing, shared, devices, extra_networks, sd_models, sd_hijack_freeu, script_callbacks, ipadapter +from modules import processing, shared, devices, extra_networks, sd_hijack_freeu, script_callbacks, ipadapter, token_merge from modules.sd_hijack_hypertile import context_hypertile_vae, context_hypertile_unet + FACEID_MODELS = { "FaceID Base": "h94/IP-Adapter-FaceID/ip-adapter-faceid_sd15.bin", "FaceID Plus v1": "h94/IP-Adapter-FaceID/ip-adapter-faceid-plus_sd15.bin", @@ -69,7 +70,7 @@ def face_id( shared.prompt_styles.apply_styles_to_extra(p) if shared.opts.cuda_compile_backend == 'none': - sd_models.apply_token_merging(p.sd_model) + token_merge.apply_token_merging(p.sd_model) sd_hijack_freeu.apply_freeu(p, not shared.native) script_callbacks.before_process_callback(p) @@ -246,7 +247,7 @@ def face_id( if faceid_model is not None and original_load_ip_adapter is not None: faceid_model.__class__.load_ip_adapter = original_load_ip_adapter if shared.opts.cuda_compile_backend == 'none': - sd_models.remove_token_merging(p.sd_model) + token_merge.remove_token_merging(p.sd_model) script_callbacks.after_process_callback(p) return processed_images diff --git a/modules/loader.py b/modules/loader.py index a2970abfd..05e5ec394 100644 --- a/modules/loader.py +++ b/modules/loader.py @@ -44,6 +44,8 @@ if ".dev" in torch.__version__ or "+git" in torch.__version__: timer.startup.record("torch") import transformers # pylint: disable=W0611,C0411 +from transformers import logging as transformers_logging # pylint: disable=W0611,C0411 +transformers_logging.set_verbosity_error() timer.startup.record("transformers") import accelerate # pylint: disable=W0611,C0411 diff --git a/modules/model_flux.py b/modules/model_flux.py index 38207f73b..9bbc24f83 100644 --- a/modules/model_flux.py +++ b/modules/model_flux.py @@ -122,10 +122,12 @@ def quant_flux_bnb(checkpoint_info, transformer, text_encoder_2): bnb_4bit_quant_type=shared.opts.bnb_quantization_type, bnb_4bit_compute_dtype=devices.dtype ) - if 'Model' in shared.opts.bnb_quantization and transformer is None: + if ('Model' in shared.opts.bnb_quantization) and (transformer is None): transformer = diffusers.FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, quantization_config=bnb_config, torch_dtype=devices.dtype) shared.log.debug(f'Quantization: module=transformer type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}') - if 'Text Encoder' in shared.opts.bnb_quantization and text_encoder_2 is None: + if ('Text Encoder' in shared.opts.bnb_quantization) and (text_encoder_2 is None): + if repo_id == 'sayakpaul/flux.1-dev-nf4': + repo_id = 'black-forest-labs/FLUX.1-dev' # workaround since sayakpaul model is missing model_index.json text_encoder_2 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_2", cache_dir=cache_dir, quantization_config=bnb_config, torch_dtype=devices.dtype) shared.log.debug(f'Quantization: module=t5 type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}') except Exception as e: diff --git a/modules/model_quant.py b/modules/model_quant.py index d54d6ff6d..1348662de 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -23,6 +23,7 @@ def load_bnb(msg='', silent=False): bnb = None if not silent: raise + return None def load_quanto(msg='', silent=False): @@ -42,6 +43,7 @@ def load_quanto(msg='', silent=False): quanto = None if not silent: raise + return None def get_quant(name): diff --git a/modules/model_sd3.py b/modules/model_sd3.py index 639f6e4eb..da99e6c4b 100644 --- a/modules/model_sd3.py +++ b/modules/model_sd3.py @@ -85,6 +85,9 @@ def load_missing(kwargs, fn, cache_dir): if 'text_encoder_3' not in kwargs and 'text_encoder_3' not in keys: kwargs['text_encoder_3'] = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_3", variant='fp16', cache_dir=cache_dir, torch_dtype=devices.dtype) shared.log.debug(f'Load model: type=SD3 missing=te3 repo="{repo_id}"') + if 'vae' not in kwargs and 'vae' not in keys: + kwargs['vae'] = diffusers.AutoencoderKL.from_pretrained(repo_id, subfolder='vae', cache_dir=cache_dir, torch_dtype=devices.dtype) + shared.log.debug(f'Load model: type=SD3 missing=vae repo="{repo_id}"') # if 'transformer' not in kwargs and 'transformer' not in keys: # kwargs['transformer'] = diffusers.SD3Transformer2DModel.from_pretrained(default_repo_id, subfolder="transformer", cache_dir=cache_dir, torch_dtype=devices.dtype) return kwargs @@ -120,7 +123,8 @@ def load_sd3(checkpoint_info, cache_dir=None, config=None): kwargs = {} kwargs = load_overrides(kwargs, cache_dir) - kwargs = load_quants(kwargs, repo_id, cache_dir) + if fn is None or not os.path.exists(fn): + kwargs = load_quants(kwargs, repo_id, cache_dir) loader = diffusers.StableDiffusion3Pipeline.from_pretrained if fn is not None and os.path.exists(fn): diff --git a/modules/onnx_impl/ui.py b/modules/onnx_impl/ui.py index f73e477c4..49af8d98b 100644 --- a/modules/onnx_impl/ui.py +++ b/modules/onnx_impl/ui.py @@ -15,7 +15,7 @@ def create_ui(): from modules.ui_common import create_refresh_button from modules.ui_components import DropdownMulti from modules.shared import log, opts, cmd_opts, refresh_checkpoints - from modules.sd_models import checkpoint_tiles, get_closet_checkpoint_match + from modules.sd_models import checkpoint_titles, get_closet_checkpoint_match from modules.paths import sd_configs_path from .execution_providers import ExecutionProvider, install_execution_provider from .utils import check_diffusers_cache @@ -46,7 +46,7 @@ def create_ui(): with gr.TabItem("Manage cache", id="manage_cache"): cache_state_dirname = gr.Textbox(value=None, visible=False) with gr.Row(): - model_dropdown = gr.Dropdown(label="Model", value="Please select model", choices=checkpoint_tiles()) + model_dropdown = gr.Dropdown(label="Model", value="Please select model", choices=checkpoint_titles()) create_refresh_button(model_dropdown, refresh_checkpoints, {}, "onnx_cache_refresh_diffusers_model") with gr.Row(): def remove_cache_onnx_converted(dirname: str): diff --git a/modules/processing.py b/modules/processing.py index 04350ee39..99d0cb351 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -4,7 +4,7 @@ import time from contextlib import nullcontext import numpy as np from PIL import Image, ImageOps -from modules import shared, devices, errors, images, scripts, memstats, lowvram, script_callbacks, extra_networks, detailer, sd_hijack_freeu, sd_models, sd_vae, processing_helpers, timer, face_restoration +from modules import shared, devices, errors, images, scripts, memstats, lowvram, script_callbacks, extra_networks, detailer, sd_hijack_freeu, sd_models, sd_checkpoint, sd_vae, processing_helpers, timer, face_restoration, token_merge from modules.sd_hijack_hypertile import context_hypertile_vae, context_hypertile_unet from modules.processing_class import StableDiffusionProcessing, StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, StableDiffusionProcessingControl # pylint: disable=unused-import from modules.processing_info import create_infotext @@ -46,7 +46,8 @@ class Processed: self.width = p.width if hasattr(p, 'width') else (self.images[0].width if len(self.images) > 0 else 0) self.height = p.height if hasattr(p, 'height') else (self.images[0].height if len(self.images) > 0 else 0) self.sampler_name = p.sampler_name or '' - self.cfg_scale = p.cfg_scale or 0 + self.cfg_scale = p.cfg_scale if p.cfg_scale > 1 else None + self.cfg_end = p.cfg_end if p.cfg_end < 0 else None self.image_cfg_scale = p.image_cfg_scale or 0 self.steps = p.steps or 0 self.batch_size = max(1, p.batch_size) @@ -96,6 +97,7 @@ class Processed: "height": self.height, "sampler_name": self.sampler_name, "cfg_scale": self.cfg_scale, + "cfg_end": self.cfg_end, "steps": self.steps, "batch_size": self.batch_size, "detailer": self.detailer, @@ -136,11 +138,11 @@ def process_images(p: StableDiffusionProcessing) -> Processed: processed = None try: # if no checkpoint override or the override checkpoint can't be found, remove override entry and load opts checkpoint - if p.override_settings.get('sd_model_checkpoint', None) is not None and sd_models.checkpoint_aliases.get(p.override_settings.get('sd_model_checkpoint')) is None: + if p.override_settings.get('sd_model_checkpoint', None) is not None and sd_checkpoint.checkpoint_aliases.get(p.override_settings.get('sd_model_checkpoint')) is None: shared.log.warning(f"Override not found: checkpoint={p.override_settings.get('sd_model_checkpoint', None)}") p.override_settings.pop('sd_model_checkpoint', None) sd_models.reload_model_weights() - if p.override_settings.get('sd_model_refiner', None) is not None and sd_models.checkpoint_aliases.get(p.override_settings.get('sd_model_refiner')) is None: + if p.override_settings.get('sd_model_refiner', None) is not None and sd_checkpoint.checkpoint_aliases.get(p.override_settings.get('sd_model_refiner')) is None: shared.log.warning(f"Override not found: refiner={p.override_settings.get('sd_model_refiner', None)}") p.override_settings.pop('sd_model_refiner', None) sd_models.reload_model_weights() @@ -162,7 +164,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed: shared.prompt_styles.apply_styles_to_extra(p) shared.prompt_styles.extract_comments(p) if shared.opts.cuda_compile_backend == 'none': - sd_models.apply_token_merging(p.sd_model) + token_merge.apply_token_merging(p.sd_model) sd_hijack_freeu.apply_freeu(p, not shared.native) if p.width is not None: @@ -205,7 +207,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed: finally: pag.unapply() if shared.opts.cuda_compile_backend == 'none': - sd_models.remove_token_merging(p.sd_model) + token_merge.remove_token_merging(p.sd_model) script_callbacks.after_process_callback(p) diff --git a/modules/processing_info.py b/modules/processing_info.py index 29513167d..e798211b1 100644 --- a/modules/processing_info.py +++ b/modules/processing_info.py @@ -41,11 +41,12 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts=None, all_seeds=No # basic "Steps": p.steps, "Seed": all_seeds[index], - "Sampler": p.sampler_name, - "CFG scale": p.cfg_scale, + "Sampler": p.sampler_name if p.sampler_name != 'Default' else None, + "CFG scale": p.cfg_scale if p.cfg_scale > 1.0 else None, + "CFG end": p.cfg_end if p.cfg_end < 1.0 else None, "Size": f"{p.width}x{p.height}" if hasattr(p, 'width') and hasattr(p, 'height') else None, "Batch": f'{p.n_iter}x{p.batch_size}' if p.n_iter > 1 or p.batch_size > 1 else None, - "Parser": shared.opts.prompt_attention, + "Parser": shared.opts.prompt_attention.split()[0], "Model": None if (not shared.opts.add_model_name_to_info) or (not shared.sd_model.sd_checkpoint_info.model_name) else shared.sd_model.sd_checkpoint_info.model_name.replace(',', '').replace(':', ''), "Model hash": getattr(p, 'sd_model_hash', None if (not shared.opts.add_model_hash_to_info) or (not shared.sd_model.sd_model_hash) else shared.sd_model.sd_model_hash), "VAE": (None if not shared.opts.add_model_name_to_info or sd_vae.loaded_vae_file is None else os.path.splitext(os.path.basename(sd_vae.loaded_vae_file))[0]) if p.full_quality else 'TAESD', diff --git a/modules/processing_original.py b/modules/processing_original.py index 852eb9a37..649023aae 100644 --- a/modules/processing_original.py +++ b/modules/processing_original.py @@ -1,7 +1,7 @@ import torch import numpy as np from PIL import Image -from modules import shared, devices, processing, images, sd_models, sd_vae, sd_samplers, processing_helpers, prompt_parser +from modules import shared, devices, processing, images, sd_vae, sd_samplers, processing_helpers, prompt_parser, token_merge from modules.sd_hijack_hypertile import hypertile_set @@ -135,10 +135,10 @@ def sample_txt2img(p: processing.StableDiffusionProcessingTxt2Img, conditioning, p.sampler.initialize(p) samples = samples[:, :, p.truncate_y//2:samples.shape[2]-(p.truncate_y+1)//2, p.truncate_x//2:samples.shape[3]-(p.truncate_x+1)//2] noise = create_random_tensors(samples.shape[1:], seeds=seeds, subseeds=subseeds, subseed_strength=subseed_strength, p=p) - sd_models.apply_token_merging(p.sd_model) + token_merge.apply_token_merging(p.sd_model) hypertile_set(p, hr=True) samples = p.sampler.sample_img2img(p, samples, noise, conditioning, unconditional_conditioning, steps=p.hr_second_pass_steps or p.steps, image_conditioning=image_conditioning) - sd_models.apply_token_merging(p.sd_model) + token_merge.apply_token_merging(p.sd_model) else: p.ops.append('upscale') x = None diff --git a/modules/sd_checkpoint.py b/modules/sd_checkpoint.py new file mode 100644 index 000000000..45834a1f0 --- /dev/null +++ b/modules/sd_checkpoint.py @@ -0,0 +1,382 @@ +import os +import re +import time +import json +import collections +from modules import shared, paths, modelloader, hashes, sd_hijack_accelerate + + +checkpoints_list = {} +checkpoint_aliases = {} +checkpoints_loaded = collections.OrderedDict() +model_dir = "Stable-diffusion" +model_path = os.path.abspath(os.path.join(paths.models_path, model_dir)) +sd_metadata_file = os.path.join(paths.data_path, "metadata.json") +sd_metadata = None +sd_metadata_pending = 0 +sd_metadata_timer = 0 + + +class CheckpointInfo: + def __init__(self, filename, sha=None): + self.name = None + self.hash = sha + self.filename = filename + self.type = '' + relname = filename + app_path = os.path.abspath(paths.script_path) + + def rel(fn, path): + try: + return os.path.relpath(fn, path) + except Exception: + return fn + + if relname.startswith('..'): + relname = os.path.abspath(relname) + if relname.startswith(shared.opts.ckpt_dir): + relname = rel(filename, shared.opts.ckpt_dir) + elif relname.startswith(shared.opts.diffusers_dir): + relname = rel(filename, shared.opts.diffusers_dir) + elif relname.startswith(model_path): + relname = rel(filename, model_path) + elif relname.startswith(paths.script_path): + relname = rel(filename, paths.script_path) + elif relname.startswith(app_path): + relname = rel(filename, app_path) + else: + relname = os.path.abspath(relname) + relname, ext = os.path.splitext(relname) + ext = ext.lower()[1:] + + if os.path.isfile(filename): # ckpt or safetensor + self.name = relname + self.filename = filename + self.sha256 = hashes.sha256_from_cache(self.filename, f"checkpoint/{relname}") + self.type = ext + if 'nf4' in filename: + self.type = 'transformer' + else: # maybe a diffuser + if self.hash is None: + repo = [r for r in modelloader.diffuser_repos if self.filename == r['name']] + else: + repo = [r for r in modelloader.diffuser_repos if self.hash == r['hash']] + if len(repo) == 0: + self.name = filename + self.filename = filename + self.sha256 = None + self.type = 'unknown' + else: + self.name = os.path.join(os.path.basename(shared.opts.diffusers_dir), repo[0]['name']) + self.filename = repo[0]['path'] + self.sha256 = repo[0]['hash'] + self.type = 'diffusers' + + self.shorthash = self.sha256[0:10] if self.sha256 else None + self.title = self.name if self.shorthash is None else f'{self.name} [{self.shorthash}]' + self.path = self.filename + self.model_name = os.path.basename(self.name) + self.metadata = read_metadata_from_safetensors(filename) + # shared.log.debug(f'Checkpoint: type={self.type} name={self.name} filename={self.filename} hash={self.shorthash} title={self.title}') + + def register(self): + checkpoints_list[self.title] = self + for i in [self.name, self.filename, self.shorthash, self.title]: + if i is not None: + checkpoint_aliases[i] = self + + def calculate_shorthash(self): + self.sha256 = hashes.sha256(self.filename, f"checkpoint/{self.name}") + if self.sha256 is None: + return None + self.shorthash = self.sha256[0:10] + if self.title in checkpoints_list: + checkpoints_list.pop(self.title) + self.title = f'{self.name} [{self.shorthash}]' + self.register() + return self.shorthash + + +def setup_model(): + list_models() + sd_hijack_accelerate.hijack_hfhub() + # sd_hijack_accelerate.hijack_torch_conv() + if not shared.native: + enable_midas_autodownload() + + +def checkpoint_titles(use_short=False): # pylint: disable=unused-argument + def convert(name): + return int(name) if name.isdigit() else name.lower() + def alphanumeric_key(key): + return [convert(c) for c in re.split('([0-9]+)', key)] + return sorted([x.title for x in checkpoints_list.values()], key=alphanumeric_key) + + +def list_models(): + t0 = time.time() + global checkpoints_list # pylint: disable=global-statement + checkpoints_list.clear() + checkpoint_aliases.clear() + ext_filter = [".safetensors"] if shared.opts.sd_disable_ckpt or shared.native else [".ckpt", ".safetensors"] + model_list = list(modelloader.load_models(model_path=model_path, model_url=None, command_path=shared.opts.ckpt_dir, ext_filter=ext_filter, download_name=None, ext_blacklist=[".vae.ckpt", ".vae.safetensors"])) + for filename in sorted(model_list, key=str.lower): + checkpoint_info = CheckpointInfo(filename) + if checkpoint_info.name is not None: + checkpoint_info.register() + if shared.native: + for repo in modelloader.load_diffusers_models(clear=True): + checkpoint_info = CheckpointInfo(repo['name'], sha=repo['hash']) + if checkpoint_info.name is not None: + checkpoint_info.register() + if shared.cmd_opts.ckpt is not None: + if not os.path.exists(shared.cmd_opts.ckpt) and not shared.native: + if shared.cmd_opts.ckpt.lower() != "none": + shared.log.warning(f'Load model: path="{shared.cmd_opts.ckpt}" not found') + else: + checkpoint_info = CheckpointInfo(shared.cmd_opts.ckpt) + if checkpoint_info.name is not None: + checkpoint_info.register() + shared.opts.data['sd_model_checkpoint'] = checkpoint_info.title + elif shared.cmd_opts.ckpt != shared.default_sd_model_file and shared.cmd_opts.ckpt is not None: + shared.log.warning(f'Load model: path="{shared.cmd_opts.ckpt}" not found') + shared.log.info(f'Available Models: path="{shared.opts.ckpt_dir}" items={len(checkpoints_list)} time={time.time()-t0:.2f}') + checkpoints_list = dict(sorted(checkpoints_list.items(), key=lambda cp: cp[1].filename)) + +def update_model_hashes(): + txt = [] + lst = [ckpt for ckpt in checkpoints_list.values() if ckpt.hash is None] + # shared.log.info(f'Models list: short hash missing for {len(lst)} out of {len(checkpoints_list)} models') + for ckpt in lst: + ckpt.hash = model_hash(ckpt.filename) + # txt.append(f'Calculated short hash: {ckpt.title} {ckpt.hash}') + # txt.append(f'Updated short hashes for {len(lst)} out of {len(checkpoints_list)} models') + lst = [ckpt for ckpt in checkpoints_list.values() if ckpt.sha256 is None or ckpt.shorthash is None] + shared.log.info(f'Models list: hash missing={len(lst)} total={len(checkpoints_list)}') + for ckpt in lst: + ckpt.sha256 = hashes.sha256(ckpt.filename, f"checkpoint/{ckpt.name}") + ckpt.shorthash = ckpt.sha256[0:10] if ckpt.sha256 is not None else None + if ckpt.sha256 is not None: + txt.append(f'Hash: {ckpt.title} {ckpt.shorthash}') + txt.append(f'Updated hashes for {len(lst)} out of {len(checkpoints_list)} models') + txt = '
'.join(txt) + return txt + + +def get_closet_checkpoint_match(s: str): + if s.startswith('https://huggingface.co/'): + s = s.replace('https://huggingface.co/', '') + if s.startswith('huggingface/'): + model_name = s.replace('huggingface/', '') + checkpoint_info = CheckpointInfo(model_name) # create a virutal model info + checkpoint_info.type = 'huggingface' + return checkpoint_info + + # alias search + checkpoint_info = checkpoint_aliases.get(s, None) + if checkpoint_info is not None: + return checkpoint_info + + # models search + found = sorted([info for info in checkpoints_list.values() if os.path.basename(info.title).lower().startswith(s.lower())], key=lambda x: len(x.title)) + if found and len(found) == 1: + return found[0] + + # reference search + """ + found = sorted([info for info in shared.reference_models.values() if os.path.basename(info['path']).lower().startswith(s.lower())], key=lambda x: len(x['path'])) + if found and len(found) == 1: + checkpoint_info = CheckpointInfo(found[0]['path']) # create a virutal model info + checkpoint_info.type = 'huggingface' + return checkpoint_info + """ + + # huggingface search + if shared.opts.sd_checkpoint_autodownload and s.count('/') == 1: + modelloader.hf_login() + found = modelloader.find_diffuser(s, full=True) + shared.log.info(f'HF search: model="{s}" results={found}') + if found is not None and len(found) == 1 and found[0] == s: + checkpoint_info = CheckpointInfo(s) + checkpoint_info.type = 'huggingface' + return checkpoint_info + + # civitai search + if shared.opts.sd_checkpoint_autodownload and s.startswith("https://civitai.com/api/download/models"): + fn = modelloader.download_civit_model_thread(model_name=None, model_url=s, model_path='', model_type='Model', token=None) + if fn is not None: + checkpoint_info = CheckpointInfo(fn) + return checkpoint_info + + return None + + +def model_hash(filename): + """old hash that only looks at a small part of the file and is prone to collisions""" + try: + with open(filename, "rb") as file: + import hashlib + # t0 = time.time() + m = hashlib.sha256() + file.seek(0x100000) + m.update(file.read(0x10000)) + shorthash = m.hexdigest()[0:8] + # t1 = time.time() + # shared.log.debug(f'Calculating short hash: {filename} hash={shorthash} time={(t1-t0):.2f}') + return shorthash + except FileNotFoundError: + return 'NOFILE' + except Exception: + return 'NOHASH' + + +def select_checkpoint(op='model'): + if op == 'dict': + model_checkpoint = shared.opts.sd_model_dict + elif op == 'refiner': + model_checkpoint = shared.opts.data.get('sd_model_refiner', None) + else: + model_checkpoint = shared.opts.sd_model_checkpoint + if model_checkpoint is None or model_checkpoint == 'None': + return None + checkpoint_info = get_closet_checkpoint_match(model_checkpoint) + if checkpoint_info is not None: + shared.log.info(f'Load {op}: select="{checkpoint_info.title if checkpoint_info is not None else None}"') + return checkpoint_info + if len(checkpoints_list) == 0: + shared.log.warning("Cannot generate without a checkpoint") + shared.log.info("Set system paths to use existing folders") + shared.log.info(" or use --models-dir to specify base folder with all models") + shared.log.info(" or use --ckpt-dir to specify folder with sd models") + shared.log.info(" or use --ckpt to force using specific model") + return None + # checkpoint_info = next(iter(checkpoints_list.values())) + if model_checkpoint is not None: + if model_checkpoint != 'model.safetensors' and model_checkpoint != 'stabilityai/stable-diffusion-xl-base-1.0': + shared.log.info(f'Load {op}: search="{model_checkpoint}" not found') + else: + shared.log.info("Selecting first available checkpoint") + # shared.log.warning(f"Loading fallback checkpoint: {checkpoint_info.title}") + # shared.opts.data['sd_model_checkpoint'] = checkpoint_info.title + else: + shared.log.info(f'Load {op}: select="{checkpoint_info.title if checkpoint_info is not None else None}"') + return checkpoint_info + + +def read_metadata_from_safetensors(filename): + global sd_metadata # pylint: disable=global-statement + if sd_metadata is None: + sd_metadata = shared.readfile(sd_metadata_file, lock=True) if os.path.isfile(sd_metadata_file) else {} + res = sd_metadata.get(filename, None) + if res is not None: + return res + if not filename.endswith(".safetensors"): + return {} + if shared.cmd_opts.no_metadata: + return {} + res = {} + # try: + t0 = time.time() + with open(filename, mode="rb") as file: + try: + metadata_len = file.read(8) + metadata_len = int.from_bytes(metadata_len, "little") + json_start = file.read(2) + if metadata_len <= 2 or json_start not in (b'{"', b"{'"): + shared.log.error(f'Model metadata invalid: file="{filename}"') + json_data = json_start + file.read(metadata_len-2) + json_obj = json.loads(json_data) + for k, v in json_obj.get("__metadata__", {}).items(): + if v.startswith("data:"): + v = 'data' + if k == 'format' and v == 'pt': + continue + large = True if len(v) > 2048 else False + if large and k == 'ss_datasets': + continue + if large and k == 'workflow': + continue + if large and k == 'prompt': + continue + if large and k == 'ss_bucket_info': + continue + if v[0:1] == '{': + try: + v = json.loads(v) + if large and k == 'ss_tag_frequency': + v = { i: len(j) for i, j in v.items() } + if large and k == 'sd_merge_models': + scrub_dict(v, ['sd_merge_recipe']) + except Exception: + pass + res[k] = v + except Exception as e: + shared.log.error(f'Model metadata: file="{filename}" {e}') + sd_metadata[filename] = res + global sd_metadata_pending # pylint: disable=global-statement + sd_metadata_pending += 1 + t1 = time.time() + global sd_metadata_timer # pylint: disable=global-statement + sd_metadata_timer += (t1 - t0) + # except Exception as e: + # shared.log.error(f"Error reading metadata from: {filename} {e}") + return res + + +def enable_midas_autodownload(): + """ + Gives the ldm.modules.midas.api.load_model function automatic downloading. + + When the 512-depth-ema model, and other future models like it, is loaded, + it calls midas.api.load_model to load the associated midas depth model. + This function applies a wrapper to download the model to the correct + location automatically. + """ + from urllib import request + import ldm.modules.midas.api + midas_path = os.path.join(paths.models_path, 'midas') + for k, v in ldm.modules.midas.api.ISL_PATHS.items(): + file_name = os.path.basename(v) + ldm.modules.midas.api.ISL_PATHS[k] = os.path.join(midas_path, file_name) + midas_urls = { + "dpt_large": "https://github.com/intel-isl/DPT/releases/download/1_0/dpt_large-midas-2f21e586.pt", + "dpt_hybrid": "https://github.com/intel-isl/DPT/releases/download/1_0/dpt_hybrid-midas-501f0c75.pt", + "midas_v21": "https://github.com/AlexeyAB/MiDaS/releases/download/midas_dpt/midas_v21-f6b98070.pt", + "midas_v21_small": "https://github.com/AlexeyAB/MiDaS/releases/download/midas_dpt/midas_v21_small-70d6b9c8.pt", + } + ldm.modules.midas.api.load_model_inner = ldm.modules.midas.api.load_model + + def load_model_wrapper(model_type): + path = ldm.modules.midas.api.ISL_PATHS[model_type] + if not os.path.exists(path): + if not os.path.exists(midas_path): + os.mkdir(midas_path) + shared.log.info(f"Downloading midas model weights for {model_type} to {path}") + request.urlretrieve(midas_urls[model_type], path) + shared.log.info(f"{model_type} downloaded") + return ldm.modules.midas.api.load_model_inner(model_type) + + ldm.modules.midas.api.load_model = load_model_wrapper + + +def scrub_dict(dict_obj, keys): + for key in list(dict_obj.keys()): + if not isinstance(dict_obj, dict): + continue + if key in keys: + dict_obj.pop(key, None) + elif isinstance(dict_obj[key], dict): + scrub_dict(dict_obj[key], keys) + elif isinstance(dict_obj[key], list): + for item in dict_obj[key]: + scrub_dict(item, keys) + + +def write_metadata(): + global sd_metadata_pending # pylint: disable=global-statement + if sd_metadata_pending == 0: + shared.log.debug(f'Model metadata: file="{sd_metadata_file}" no changes') + return + shared.writefile(sd_metadata, sd_metadata_file) + shared.log.info(f'Model metadata saved: file="{sd_metadata_file}" items={sd_metadata_pending} time={sd_metadata_timer:.2f}') + sd_metadata_pending = 0 diff --git a/modules/sd_detect.py b/modules/sd_detect.py new file mode 100644 index 000000000..7144a7be7 --- /dev/null +++ b/modules/sd_detect.py @@ -0,0 +1,150 @@ +import os +import torch +import diffusers +from modules import shared, shared_items, devices, errors + + +debug_load = os.environ.get('SD_LOAD_DEBUG', None) + + +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 + pipeline = None + if guess == 'Autodetect': + try: + guess = 'Stable Diffusion XL' if 'XL' in f.upper() else 'Stable Diffusion' + # guess by size + if os.path.isfile(f) and f.endswith('.safetensors'): + size = round(os.path.getsize(f) / 1024 / 1024) + if (size > 0 and size < 128): + warn(f'Model size smaller than expected: {f} size={size} MB') + elif (size >= 316 and size <= 324) or (size >= 156 and size <= 164): # 320 or 160 + warn(f'Model detected as VAE model, but attempting to load as model: {op}={f} size={size} MB') + guess = 'VAE' + elif (size >= 4970 and size <= 4976): # 4973 + guess = 'Stable Diffusion 2' # SD v2 but could be eps or v-prediction + # elif size < 0: # unknown + # guess = 'Stable Diffusion 2B' + elif (size >= 5791 and size <= 5799): # 5795 + if op == 'model': + warn(f'Model detected as SD-XL refiner model, but attempting to load a base model: {op}={f} size={size} MB') + guess = 'Stable Diffusion XL Refiner' + elif (size >= 6611 and size <= 7220): # 6617, HassakuXL is 6776, monkrenRealisticINT_v10 is 7217 + guess = 'Stable Diffusion XL' + elif (size >= 3361 and size <= 3369): # 3368 + guess = 'Stable Diffusion Upscale' + elif (size >= 4891 and size <= 4899): # 4897 + guess = 'Stable Diffusion XL Inpaint' + elif (size >= 9791 and size <= 9799): # 9794 + guess = 'Stable Diffusion XL Instruct' + elif (size > 3138 and size < 3142): #3140 + guess = 'Stable Diffusion XL' + elif (size > 5692 and size < 5698) or (size > 4134 and size < 4138) or (size > 10362 and size < 10366) or (size > 15028 and size < 15228): + guess = 'Stable Diffusion 3' + elif (size > 18414 and size < 18420): # sd35-large aio + guess = 'Stable Diffusion 3' + elif (size > 20000 and size < 40000): + guess = 'FLUX' + # guess by name + """ + if 'LCM_' in f.upper() or 'LCM-' in f.upper() or '_LCM' in f.upper() or '-LCM' in f.upper(): + if shared.backend == shared.Backend.ORIGINAL: + warn(f'Model detected as LCM model, but attempting to load using backend=original: {op}={f} size={size} MB') + guess = 'Latent Consistency Model' + """ + if 'instaflow' in f.lower(): + guess = 'InstaFlow' + if 'segmoe' in f.lower(): + guess = 'SegMoE' + if 'hunyuandit' in f.lower(): + guess = 'HunyuanDiT' + if 'pixart-xl' in f.lower(): + guess = 'PixArt-Alpha' + if 'stable-diffusion-3' in f.lower(): + guess = 'Stable Diffusion 3' + if 'stable-cascade' in f.lower() or 'stablecascade' in f.lower() or 'wuerstchen3' in f.lower() or ('sotediffusion' in f.lower() and "v2" in f.lower()): + if devices.dtype == torch.float16: + warn('Stable Cascade does not support Float16') + guess = 'Stable Cascade' + if 'pixart-sigma' in f.lower(): + guess = 'PixArt-Sigma' + if 'lumina-next' in f.lower(): + guess = 'Lumina-Next' + if 'kolors' in f.lower(): + guess = 'Kolors' + if 'auraflow' in f.lower(): + guess = 'AuraFlow' + if 'cogview' in f.lower(): + guess = 'CogView' + if 'meissonic' in f.lower(): + guess = 'Meissonic' + pipeline = 'custom' + if 'omnigen' in f.lower(): + guess = 'OmniGen' + pipeline = 'custom' + if 'flux' in f.lower(): + guess = 'FLUX' + if size > 11000 and size < 20000: + warn(f'Model detected as FLUX UNET model, but attempting to load a base model: {op}={f} size={size} MB') + # switch for specific variant + if guess == 'Stable Diffusion' and 'inpaint' in f.lower(): + guess = 'Stable Diffusion Inpaint' + elif guess == 'Stable Diffusion' and 'instruct' in f.lower(): + guess = 'Stable Diffusion Instruct' + if guess == 'Stable Diffusion XL' and 'inpaint' in f.lower(): + guess = 'Stable Diffusion XL Inpaint' + elif guess == 'Stable Diffusion XL' and 'instruct' in f.lower(): + guess = 'Stable Diffusion XL Instruct' + # get actual pipeline + pipeline = shared_items.get_pipelines().get(guess, None) if pipeline is None else pipeline + if not quiet: + shared.log.info(f'Autodetect {op}: detect="{guess}" class={getattr(pipeline, "__name__", None)} file="{f}" size={size}MB') + except Exception as e: + shared.log.error(f'Autodetect {op}: file="{f}" {e}') + if debug_load: + errors.display(e, f'Load {op}: {f}') + return None, None + else: + try: + size = round(os.path.getsize(f) / 1024 / 1024) + pipeline = shared_items.get_pipelines().get(guess, None) if pipeline is None else pipeline + if not quiet: + shared.log.info(f'Load {op}: detect="{guess}" class={getattr(pipeline, "__name__", None)} file="{f}" size={size}MB') + except Exception as e: + shared.log.error(f'Load {op}: detect="{guess}" file="{f}" {e}') + + if pipeline is None: + shared.log.warning(f'Load {op}: detect="{guess}" file="{f}" size={size} not recognized') + pipeline = diffusers.StableDiffusionPipeline + return pipeline, guess + + +def get_load_config(model_file, model_type, config_type='yaml'): + if config_type == 'yaml': + yaml = os.path.splitext(model_file)[0] + '.yaml' + if os.path.exists(yaml): + return yaml + if model_type == 'Stable Diffusion': + return 'configs/v1-inference.yaml' + if model_type == 'Stable Diffusion XL': + return 'configs/sd_xl_base.yaml' + if model_type == 'Stable Diffusion XL Refiner': + return 'configs/sd_xl_refiner.yaml' + if model_type == 'Stable Diffusion 2': + return None # dont know if its eps or v so let diffusers sort it out + # return 'configs/v2-inference-512-base.yaml' + # return 'configs/v2-inference-768-v.yaml' + elif config_type == 'json': + if not shared.opts.diffuser_cache_config: + return None + if model_type == 'Stable Diffusion': + return 'configs/sd15' + if model_type == 'Stable Diffusion XL': + return 'configs/sdxl' + if model_type == 'Stable Diffusion 3': + return 'configs/sd3' + if model_type == 'FLUX': + return 'configs/flux' + return None diff --git a/modules/sd_models.py b/modules/sd_models.py index bc293f5fc..9662a005d 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -1,16 +1,11 @@ -import re import io import sys import json -import time import copy import inspect import logging import contextlib -import collections import os.path -from os import mkdir -from urllib import request from enum import Enum import diffusers import diffusers.loaders.single_file_utils @@ -18,20 +13,16 @@ from rich import progress # pylint: disable=redefined-builtin import torch import safetensors.torch from omegaconf import OmegaConf -from transformers import logging as transformers_logging from ldm.util import instantiate_from_config -from modules import paths, shared, shared_items, shared_state, modelloader, devices, script_callbacks, sd_vae, sd_unet, errors, hashes, sd_models_config, sd_models_compile, sd_hijack_accelerate +from modules import paths, shared, shared_state, modelloader, devices, script_callbacks, sd_vae, sd_unet, errors, sd_models_config, sd_models_compile, sd_hijack_accelerate, sd_detect from modules.timer import Timer from modules.memstats import memory_stats from modules.modeldata import model_data +from modules.sd_checkpoint import CheckpointInfo, select_checkpoint, list_models, checkpoints_list, checkpoint_titles, get_closet_checkpoint_match, update_model_hashes, setup_model, write_metadata, read_metadata_from_safetensors # pylint: disable=unused-import -transformers_logging.set_verbosity_error() model_dir = "Stable-diffusion" model_path = os.path.abspath(os.path.join(paths.models_path, model_dir)) -checkpoints_list = {} -checkpoint_aliases = {} -checkpoints_loaded = collections.OrderedDict() sd_metadata_file = os.path.join(paths.data_path, "metadata.json") sd_metadata = None sd_metadata_pending = 0 @@ -42,368 +33,11 @@ debug_process = shared.log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is diffusers_version = int(diffusers.__version__.split('.')[1]) -class CheckpointInfo: - def __init__(self, filename, sha=None): - self.name = None - self.hash = sha - self.filename = filename - self.type = '' - relname = filename - app_path = os.path.abspath(paths.script_path) - - def rel(fn, path): - try: - return os.path.relpath(fn, path) - except Exception: - return fn - - if relname.startswith('..'): - relname = os.path.abspath(relname) - if relname.startswith(shared.opts.ckpt_dir): - relname = rel(filename, shared.opts.ckpt_dir) - elif relname.startswith(shared.opts.diffusers_dir): - relname = rel(filename, shared.opts.diffusers_dir) - elif relname.startswith(model_path): - relname = rel(filename, model_path) - elif relname.startswith(paths.script_path): - relname = rel(filename, paths.script_path) - elif relname.startswith(app_path): - relname = rel(filename, app_path) - else: - relname = os.path.abspath(relname) - relname, ext = os.path.splitext(relname) - ext = ext.lower()[1:] - - if os.path.isfile(filename): # ckpt or safetensor - self.name = relname - self.filename = filename - self.sha256 = hashes.sha256_from_cache(self.filename, f"checkpoint/{relname}") - self.type = ext - if 'nf4' in filename: - self.type = 'transformer' - else: # maybe a diffuser - if self.hash is None: - repo = [r for r in modelloader.diffuser_repos if self.filename == r['name']] - else: - repo = [r for r in modelloader.diffuser_repos if self.hash == r['hash']] - if len(repo) == 0: - self.name = filename - self.filename = filename - self.sha256 = None - self.type = 'unknown' - else: - self.name = os.path.join(os.path.basename(shared.opts.diffusers_dir), repo[0]['name']) - self.filename = repo[0]['path'] - self.sha256 = repo[0]['hash'] - self.type = 'diffusers' - - self.shorthash = self.sha256[0:10] if self.sha256 else None - self.title = self.name if self.shorthash is None else f'{self.name} [{self.shorthash}]' - self.path = self.filename - self.model_name = os.path.basename(self.name) - self.metadata = read_metadata_from_safetensors(filename) - # shared.log.debug(f'Checkpoint: type={self.type} name={self.name} filename={self.filename} hash={self.shorthash} title={self.title}') - - def register(self): - checkpoints_list[self.title] = self - for i in [self.name, self.filename, self.shorthash, self.title]: - if i is not None: - checkpoint_aliases[i] = self - - def calculate_shorthash(self): - self.sha256 = hashes.sha256(self.filename, f"checkpoint/{self.name}") - if self.sha256 is None: - return None - self.shorthash = self.sha256[0:10] - if self.title in checkpoints_list: - checkpoints_list.pop(self.title) - self.title = f'{self.name} [{self.shorthash}]' - self.register() - return self.shorthash - - class NoWatermark: def apply_watermark(self, img): return img -def setup_model(): - list_models() - sd_hijack_accelerate.hijack_hfhub() - # sd_hijack_accelerate.hijack_torch_conv() - if not shared.native: - enable_midas_autodownload() - - -def checkpoint_tiles(use_short=False): # pylint: disable=unused-argument - def convert(name): - return int(name) if name.isdigit() else name.lower() - def alphanumeric_key(key): - return [convert(c) for c in re.split('([0-9]+)', key)] - return sorted([x.title for x in checkpoints_list.values()], key=alphanumeric_key) - - -def list_models(): - t0 = time.time() - global checkpoints_list # pylint: disable=global-statement - checkpoints_list.clear() - checkpoint_aliases.clear() - ext_filter = [".safetensors"] if shared.opts.sd_disable_ckpt or shared.native else [".ckpt", ".safetensors"] - model_list = list(modelloader.load_models(model_path=model_path, model_url=None, command_path=shared.opts.ckpt_dir, ext_filter=ext_filter, download_name=None, ext_blacklist=[".vae.ckpt", ".vae.safetensors"])) - for filename in sorted(model_list, key=str.lower): - checkpoint_info = CheckpointInfo(filename) - if checkpoint_info.name is not None: - checkpoint_info.register() - if shared.native: - for repo in modelloader.load_diffusers_models(clear=True): - checkpoint_info = CheckpointInfo(repo['name'], sha=repo['hash']) - if checkpoint_info.name is not None: - checkpoint_info.register() - if shared.cmd_opts.ckpt is not None: - if not os.path.exists(shared.cmd_opts.ckpt) and not shared.native: - if shared.cmd_opts.ckpt.lower() != "none": - shared.log.warning(f'Load model: path="{shared.cmd_opts.ckpt}" not found') - else: - checkpoint_info = CheckpointInfo(shared.cmd_opts.ckpt) - if checkpoint_info.name is not None: - checkpoint_info.register() - shared.opts.data['sd_model_checkpoint'] = checkpoint_info.title - elif shared.cmd_opts.ckpt != shared.default_sd_model_file and shared.cmd_opts.ckpt is not None: - shared.log.warning(f'Load model: path="{shared.cmd_opts.ckpt}" not found') - shared.log.info(f'Available Models: path="{shared.opts.ckpt_dir}" items={len(checkpoints_list)} time={time.time()-t0:.2f}') - checkpoints_list = dict(sorted(checkpoints_list.items(), key=lambda cp: cp[1].filename)) - - -def update_model_hashes(): - txt = [] - lst = [ckpt for ckpt in checkpoints_list.values() if ckpt.hash is None] - # shared.log.info(f'Models list: short hash missing for {len(lst)} out of {len(checkpoints_list)} models') - for ckpt in lst: - ckpt.hash = model_hash(ckpt.filename) - # txt.append(f'Calculated short hash: {ckpt.title} {ckpt.hash}') - # txt.append(f'Updated short hashes for {len(lst)} out of {len(checkpoints_list)} models') - lst = [ckpt for ckpt in checkpoints_list.values() if ckpt.sha256 is None or ckpt.shorthash is None] - shared.log.info(f'Models list: hash missing={len(lst)} total={len(checkpoints_list)}') - for ckpt in lst: - ckpt.sha256 = hashes.sha256(ckpt.filename, f"checkpoint/{ckpt.name}") - ckpt.shorthash = ckpt.sha256[0:10] if ckpt.sha256 is not None else None - if ckpt.sha256 is not None: - txt.append(f'Hash: {ckpt.title} {ckpt.shorthash}') - txt.append(f'Updated hashes for {len(lst)} out of {len(checkpoints_list)} models') - txt = '
'.join(txt) - return txt - - -def get_closet_checkpoint_match(s: str): - if s.startswith('https://huggingface.co/'): - s = s.replace('https://huggingface.co/', '') - if s.startswith('huggingface/'): - model_name = s.replace('huggingface/', '') - checkpoint_info = CheckpointInfo(model_name) # create a virutal model info - checkpoint_info.type = 'huggingface' - return checkpoint_info - - # alias search - checkpoint_info = checkpoint_aliases.get(s, None) - if checkpoint_info is not None: - return checkpoint_info - - # models search - found = sorted([info for info in checkpoints_list.values() if os.path.basename(info.title).lower().startswith(s.lower())], key=lambda x: len(x.title)) - if found and len(found) == 1: - return found[0] - - # reference search - """ - found = sorted([info for info in shared.reference_models.values() if os.path.basename(info['path']).lower().startswith(s.lower())], key=lambda x: len(x['path'])) - if found and len(found) == 1: - checkpoint_info = CheckpointInfo(found[0]['path']) # create a virutal model info - checkpoint_info.type = 'huggingface' - return checkpoint_info - """ - - # huggingface search - if shared.opts.sd_checkpoint_autodownload and s.count('/') == 1: - modelloader.hf_login() - found = modelloader.find_diffuser(s, full=True) - shared.log.info(f'HF search: model="{s}" results={found}') - if found is not None and len(found) == 1 and found[0] == s: - checkpoint_info = CheckpointInfo(s) - checkpoint_info.type = 'huggingface' - return checkpoint_info - - # civitai search - if shared.opts.sd_checkpoint_autodownload and s.startswith("https://civitai.com/api/download/models"): - fn = modelloader.download_civit_model_thread(model_name=None, model_url=s, model_path='', model_type='Model', token=None) - if fn is not None: - checkpoint_info = CheckpointInfo(fn) - return checkpoint_info - - return None - - -def model_hash(filename): - """old hash that only looks at a small part of the file and is prone to collisions""" - try: - with open(filename, "rb") as file: - import hashlib - # t0 = time.time() - m = hashlib.sha256() - file.seek(0x100000) - m.update(file.read(0x10000)) - shorthash = m.hexdigest()[0:8] - # t1 = time.time() - # shared.log.debug(f'Calculating short hash: {filename} hash={shorthash} time={(t1-t0):.2f}') - return shorthash - except FileNotFoundError: - return 'NOFILE' - except Exception: - return 'NOHASH' - - -def select_checkpoint(op='model'): - if op == 'dict': - model_checkpoint = shared.opts.sd_model_dict - elif op == 'refiner': - model_checkpoint = shared.opts.data.get('sd_model_refiner', None) - else: - model_checkpoint = shared.opts.sd_model_checkpoint - if model_checkpoint is None or model_checkpoint == 'None': - return None - checkpoint_info = get_closet_checkpoint_match(model_checkpoint) - if checkpoint_info is not None: - shared.log.info(f'Load {op}: select="{checkpoint_info.title if checkpoint_info is not None else None}"') - return checkpoint_info - if len(checkpoints_list) == 0: - shared.log.warning("Cannot generate without a checkpoint") - shared.log.info("Set system paths to use existing folders") - shared.log.info(" or use --models-dir to specify base folder with all models") - shared.log.info(" or use --ckpt-dir to specify folder with sd models") - shared.log.info(" or use --ckpt to force using specific model") - return None - # checkpoint_info = next(iter(checkpoints_list.values())) - if model_checkpoint is not None: - if model_checkpoint != 'model.safetensors' and model_checkpoint != 'stabilityai/stable-diffusion-xl-base-1.0': - shared.log.info(f'Load {op}: search="{model_checkpoint}" not found') - else: - shared.log.info("Selecting first available checkpoint") - # shared.log.warning(f"Loading fallback checkpoint: {checkpoint_info.title}") - # shared.opts.data['sd_model_checkpoint'] = checkpoint_info.title - else: - shared.log.info(f'Load {op}: select="{checkpoint_info.title if checkpoint_info is not None else None}"') - return checkpoint_info - - -checkpoint_dict_replacements = { - 'cond_stage_model.transformer.embeddings.': 'cond_stage_model.transformer.text_model.embeddings.', - 'cond_stage_model.transformer.encoder.': 'cond_stage_model.transformer.text_model.encoder.', - 'cond_stage_model.transformer.final_layer_norm.': 'cond_stage_model.transformer.text_model.final_layer_norm.', -} - - -def transform_checkpoint_dict_key(k): - for text, replacement in checkpoint_dict_replacements.items(): - if k.startswith(text): - k = replacement + k[len(text):] - return k - - -def get_state_dict_from_checkpoint(pl_sd): - pl_sd = pl_sd.pop("state_dict", pl_sd) - pl_sd.pop("state_dict", None) - sd = {} - for k, v in pl_sd.items(): - new_key = transform_checkpoint_dict_key(k) - if new_key is not None: - sd[new_key] = v - pl_sd.clear() - pl_sd.update(sd) - return pl_sd - - -def write_metadata(): - global sd_metadata_pending # pylint: disable=global-statement - if sd_metadata_pending == 0: - shared.log.debug(f'Model metadata: file="{sd_metadata_file}" no changes') - return - shared.writefile(sd_metadata, sd_metadata_file) - shared.log.info(f'Model metadata saved: file="{sd_metadata_file}" items={sd_metadata_pending} time={sd_metadata_timer:.2f}') - sd_metadata_pending = 0 - - -def scrub_dict(dict_obj, keys): - for key in list(dict_obj.keys()): - if not isinstance(dict_obj, dict): - continue - if key in keys: - dict_obj.pop(key, None) - elif isinstance(dict_obj[key], dict): - scrub_dict(dict_obj[key], keys) - elif isinstance(dict_obj[key], list): - for item in dict_obj[key]: - scrub_dict(item, keys) - - -def read_metadata_from_safetensors(filename): - global sd_metadata # pylint: disable=global-statement - if sd_metadata is None: - sd_metadata = shared.readfile(sd_metadata_file, lock=True) if os.path.isfile(sd_metadata_file) else {} - res = sd_metadata.get(filename, None) - if res is not None: - return res - if not filename.endswith(".safetensors"): - return {} - if shared.cmd_opts.no_metadata: - return {} - res = {} - # try: - t0 = time.time() - with open(filename, mode="rb") as file: - try: - metadata_len = file.read(8) - metadata_len = int.from_bytes(metadata_len, "little") - json_start = file.read(2) - if metadata_len <= 2 or json_start not in (b'{"', b"{'"): - shared.log.error(f'Model metadata invalid: file="{filename}"') - json_data = json_start + file.read(metadata_len-2) - json_obj = json.loads(json_data) - for k, v in json_obj.get("__metadata__", {}).items(): - if v.startswith("data:"): - v = 'data' - if k == 'format' and v == 'pt': - continue - large = True if len(v) > 2048 else False - if large and k == 'ss_datasets': - continue - if large and k == 'workflow': - continue - if large and k == 'prompt': - continue - if large and k == 'ss_bucket_info': - continue - if v[0:1] == '{': - try: - v = json.loads(v) - if large and k == 'ss_tag_frequency': - v = { i: len(j) for i, j in v.items() } - if large and k == 'sd_merge_models': - scrub_dict(v, ['sd_merge_recipe']) - except Exception: - pass - res[k] = v - except Exception as e: - shared.log.error(f'Model metadata: file="{filename}" {e}') - sd_metadata[filename] = res - global sd_metadata_pending # pylint: disable=global-statement - sd_metadata_pending += 1 - t1 = time.time() - global sd_metadata_timer # pylint: disable=global-statement - sd_metadata_timer += (t1 - t0) - # except Exception as e: - # shared.log.error(f"Error reading metadata from: {filename} {e}") - return res - - def read_state_dict(checkpoint_file, map_location=None, what:str='model'): # pylint: disable=unused-argument if not os.path.isfile(checkpoint_file): shared.log.error(f'Load dict: path="{checkpoint_file}" not a file') @@ -449,26 +83,55 @@ def get_safetensor_keys(filename): return keys +def get_state_dict_from_checkpoint(pl_sd): + checkpoint_dict_replacements = { + 'cond_stage_model.transformer.embeddings.': 'cond_stage_model.transformer.text_model.embeddings.', + 'cond_stage_model.transformer.encoder.': 'cond_stage_model.transformer.text_model.encoder.', + 'cond_stage_model.transformer.final_layer_norm.': 'cond_stage_model.transformer.text_model.final_layer_norm.', + } + + def transform_checkpoint_dict_key(k): + for text, replacement in checkpoint_dict_replacements.items(): + if k.startswith(text): + k = replacement + k[len(text):] + return k + + pl_sd = pl_sd.pop("state_dict", pl_sd) + pl_sd.pop("state_dict", None) + sd = {} + for k, v in pl_sd.items(): + new_key = transform_checkpoint_dict_key(k) + if new_key is not None: + sd[new_key] = v + pl_sd.clear() + pl_sd.update(sd) + return pl_sd + + def get_checkpoint_state_dict(checkpoint_info: CheckpointInfo, timer): if not os.path.isfile(checkpoint_info.filename): return None + """ if checkpoint_info in checkpoints_loaded: shared.log.info("Load model: cache") checkpoints_loaded.move_to_end(checkpoint_info, last=True) # FIFO -> LRU cache return checkpoints_loaded[checkpoint_info] + """ res = read_state_dict(checkpoint_info.filename, what='model') + """ if shared.opts.sd_checkpoint_cache > 0 and not shared.native: # cache newly loaded model checkpoints_loaded[checkpoint_info] = res # clean up cache if limit is reached while len(checkpoints_loaded) > shared.opts.sd_checkpoint_cache: checkpoints_loaded.popitem(last=False) + """ timer.record("load") return res def load_model_weights(model: torch.nn.Module, checkpoint_info: CheckpointInfo, state_dict, timer): - _pipeline, _model_type = detect_pipeline(checkpoint_info.path, 'model') + _pipeline, _model_type = sd_detect.detect_pipeline(checkpoint_info.path, 'model') shared.log.debug(f'Load model: memory={memory_stats()}') timer.record("hash") if model_data.sd_dict == 'None': @@ -520,41 +183,6 @@ def load_model_weights(model: torch.nn.Module, checkpoint_info: CheckpointInfo, return True -def enable_midas_autodownload(): - """ - Gives the ldm.modules.midas.api.load_model function automatic downloading. - - When the 512-depth-ema model, and other future models like it, is loaded, - it calls midas.api.load_model to load the associated midas depth model. - This function applies a wrapper to download the model to the correct - location automatically. - """ - import ldm.modules.midas.api - midas_path = os.path.join(paths.models_path, 'midas') - for k, v in ldm.modules.midas.api.ISL_PATHS.items(): - file_name = os.path.basename(v) - ldm.modules.midas.api.ISL_PATHS[k] = os.path.join(midas_path, file_name) - midas_urls = { - "dpt_large": "https://github.com/intel-isl/DPT/releases/download/1_0/dpt_large-midas-2f21e586.pt", - "dpt_hybrid": "https://github.com/intel-isl/DPT/releases/download/1_0/dpt_hybrid-midas-501f0c75.pt", - "midas_v21": "https://github.com/AlexeyAB/MiDaS/releases/download/midas_dpt/midas_v21-f6b98070.pt", - "midas_v21_small": "https://github.com/AlexeyAB/MiDaS/releases/download/midas_dpt/midas_v21_small-70d6b9c8.pt", - } - ldm.modules.midas.api.load_model_inner = ldm.modules.midas.api.load_model - - def load_model_wrapper(model_type): - path = ldm.modules.midas.api.ISL_PATHS[model_type] - if not os.path.exists(path): - if not os.path.exists(midas_path): - mkdir(midas_path) - shared.log.info(f"Downloading midas model weights for {model_type} to {path}") - request.urlretrieve(midas_urls[model_type], path) - shared.log.info(f"{model_type} downloaded") - return ldm.modules.midas.api.load_model_inner(model_type) - - ldm.modules.midas.api.load_model = load_model_wrapper - - def repair_config(sd_config): if "use_ema" not in sd_config.model.params: sd_config.model.params.use_ema = False @@ -580,7 +208,6 @@ def change_backend(): unload_model_weights() shared.backend = shared.Backend.ORIGINAL if shared.opts.sd_backend == 'original' else shared.Backend.DIFFUSERS shared.native = shared.backend == shared.Backend.DIFFUSERS - checkpoints_loaded.clear() from modules.sd_samplers import list_samplers list_samplers() list_models() @@ -588,118 +215,6 @@ def change_backend(): refresh_vae_list() -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 - pipeline = None - if guess == 'Autodetect': - try: - guess = 'Stable Diffusion XL' if 'XL' in f.upper() else 'Stable Diffusion' - # guess by size - if os.path.isfile(f) and f.endswith('.safetensors'): - size = round(os.path.getsize(f) / 1024 / 1024) - if (size > 0 and size < 128): - warn(f'Model size smaller than expected: {f} size={size} MB') - elif (size >= 316 and size <= 324) or (size >= 156 and size <= 164): # 320 or 160 - warn(f'Model detected as VAE model, but attempting to load as model: {op}={f} size={size} MB') - guess = 'VAE' - elif (size >= 4970 and size <= 4976): # 4973 - guess = 'Stable Diffusion 2' # SD v2 but could be eps or v-prediction - # elif size < 0: # unknown - # guess = 'Stable Diffusion 2B' - elif (size >= 5791 and size <= 5799): # 5795 - if op == 'model': - warn(f'Model detected as SD-XL refiner model, but attempting to load a base model: {op}={f} size={size} MB') - guess = 'Stable Diffusion XL Refiner' - elif (size >= 6611 and size <= 7220): # 6617, HassakuXL is 6776, monkrenRealisticINT_v10 is 7217 - guess = 'Stable Diffusion XL' - elif (size >= 3361 and size <= 3369): # 3368 - guess = 'Stable Diffusion Upscale' - elif (size >= 4891 and size <= 4899): # 4897 - guess = 'Stable Diffusion XL Inpaint' - elif (size >= 9791 and size <= 9799): # 9794 - guess = 'Stable Diffusion XL Instruct' - elif (size > 3138 and size < 3142): #3140 - guess = 'Stable Diffusion XL' - elif (size > 5692 and size < 5698) or (size > 4134 and size < 4138) or (size > 10362 and size < 10366) or (size > 15028 and size < 15228): - guess = 'Stable Diffusion 3' - elif (size > 20000 and size < 40000): - guess = 'FLUX' - # guess by name - """ - if 'LCM_' in f.upper() or 'LCM-' in f.upper() or '_LCM' in f.upper() or '-LCM' in f.upper(): - if shared.backend == shared.Backend.ORIGINAL: - warn(f'Model detected as LCM model, but attempting to load using backend=original: {op}={f} size={size} MB') - guess = 'Latent Consistency Model' - """ - if 'instaflow' in f.lower(): - guess = 'InstaFlow' - if 'segmoe' in f.lower(): - guess = 'SegMoE' - if 'hunyuandit' in f.lower(): - guess = 'HunyuanDiT' - if 'pixart-xl' in f.lower(): - guess = 'PixArt-Alpha' - if 'stable-diffusion-3' in f.lower(): - guess = 'Stable Diffusion 3' - if 'stable-cascade' in f.lower() or 'stablecascade' in f.lower() or 'wuerstchen3' in f.lower() or ('sotediffusion' in f.lower() and "v2" in f.lower()): - if devices.dtype == torch.float16: - warn('Stable Cascade does not support Float16') - guess = 'Stable Cascade' - if 'pixart-sigma' in f.lower(): - guess = 'PixArt-Sigma' - if 'lumina-next' in f.lower(): - guess = 'Lumina-Next' - if 'kolors' in f.lower(): - guess = 'Kolors' - if 'auraflow' in f.lower(): - guess = 'AuraFlow' - if 'cogview' in f.lower(): - guess = 'CogView' - if 'meissonic' in f.lower(): - guess = 'Meissonic' - pipeline = 'custom' - if 'omnigen' in f.lower(): - guess = 'OmniGen' - pipeline = 'custom' - if 'flux' in f.lower(): - guess = 'FLUX' - if size > 11000 and size < 20000: - warn(f'Model detected as FLUX UNET model, but attempting to load a base model: {op}={f} size={size} MB') - # switch for specific variant - if guess == 'Stable Diffusion' and 'inpaint' in f.lower(): - guess = 'Stable Diffusion Inpaint' - elif guess == 'Stable Diffusion' and 'instruct' in f.lower(): - guess = 'Stable Diffusion Instruct' - if guess == 'Stable Diffusion XL' and 'inpaint' in f.lower(): - guess = 'Stable Diffusion XL Inpaint' - elif guess == 'Stable Diffusion XL' and 'instruct' in f.lower(): - guess = 'Stable Diffusion XL Instruct' - # get actual pipeline - pipeline = shared_items.get_pipelines().get(guess, None) if pipeline is None else pipeline - if not quiet: - shared.log.info(f'Autodetect {op}: detect="{guess}" class={getattr(pipeline, "__name__", None)} file="{f}" size={size}MB') - except Exception as e: - shared.log.error(f'Autodetect {op}: file="{f}" {e}') - if debug_load: - errors.display(e, f'Load {op}: {f}') - return None, None - else: - try: - size = round(os.path.getsize(f) / 1024 / 1024) - pipeline = shared_items.get_pipelines().get(guess, None) if pipeline is None else pipeline - if not quiet: - shared.log.info(f'Load {op}: detect="{guess}" class={getattr(pipeline, "__name__", None)} file="{f}" size={size}MB') - except Exception as e: - shared.log.error(f'Load {op}: detect="{guess}" file="{f}" {e}') - - if pipeline is None: - shared.log.warning(f'Load {op}: detect="{guess}" file="{f}" size={size} not recognized') - pipeline = diffusers.StableDiffusionPipeline - return pipeline, guess - - def copy_diffuser_options(new_pipe, orig_pipe): new_pipe.sd_checkpoint_info = getattr(orig_pipe, 'sd_checkpoint_info', None) new_pipe.sd_model_checkpoint = getattr(orig_pipe, 'sd_model_checkpoint', None) @@ -997,35 +512,6 @@ def move_base(model, device): return R -def get_load_config(model_file, model_type, config_type='yaml'): - if config_type == 'yaml': - yaml = os.path.splitext(model_file)[0] + '.yaml' - if os.path.exists(yaml): - return yaml - if model_type == 'Stable Diffusion': - return 'configs/v1-inference.yaml' - if model_type == 'Stable Diffusion XL': - return 'configs/sd_xl_base.yaml' - if model_type == 'Stable Diffusion XL Refiner': - return 'configs/sd_xl_refiner.yaml' - if model_type == 'Stable Diffusion 2': - return None # dont know if its eps or v so let diffusers sort it out - # return 'configs/v2-inference-512-base.yaml' - # return 'configs/v2-inference-768-v.yaml' - elif config_type == 'json': - if not shared.opts.diffuser_cache_config: - return None - if model_type == 'Stable Diffusion': - return 'configs/sd15' - if model_type == 'Stable Diffusion XL': - return 'configs/sdxl' - if model_type == 'Stable Diffusion 3': - return 'configs/sd3' - if model_type == 'FLUX': - return 'configs/flux' - return None - - def patch_diffuser_config(sd_model, model_file): def load_config(fn, k): model_file = os.path.splitext(fn)[0] @@ -1216,7 +702,7 @@ def load_diffuser_file(model_type, pipeline, checkpoint_info, diffusers_load_con if shared.opts.diffusers_force_zeros: diffusers_load_config['force_zeros_for_empty_prompt '] = shared.opts.diffusers_force_zeros else: - model_config = get_load_config(checkpoint_info.path, model_type, config_type='json') + model_config = sd_detect.get_load_config(checkpoint_info.path, model_type, config_type='json') if model_config is not None: if debug_load: shared.log.debug(f'Load {op}: config="{model_config}"') @@ -1307,7 +793,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No return # detect pipeline - pipeline, model_type = detect_pipeline(checkpoint_info.path, op) + pipeline, model_type = sd_detect.detect_pipeline(checkpoint_info.path, op) # preload vae so it can be used as param vae = None @@ -1782,7 +1268,7 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, timer=None, shared.log.info(f"Model loaded in {timer.summary()}") current_checkpoint_info = None devices.torch_gc(force=True) - shared.log.info(f'Model load finished: {memory_stats()} cached={len(checkpoints_loaded.keys())}') + shared.log.info(f'Model load finished: {memory_stats()}') def reload_text_encoder(initial=False): @@ -1842,7 +1328,7 @@ def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model', state_dict = get_checkpoint_state_dict(checkpoint_info, timer) if not shared.native else None checkpoint_config = sd_models_config.find_checkpoint_config(state_dict, checkpoint_info) timer.record("config") - if sd_model is None or checkpoint_config != getattr(sd_model, 'used_config', None): + if sd_model is None or checkpoint_config != getattr(sd_model, 'used_config', None) or force: sd_model = None if not shared.native: load_model(checkpoint_info, already_loaded_state_dict=state_dict, timer=timer, op=op) @@ -1936,82 +1422,6 @@ def unload_model_weights(op='model'): shared.log.debug(f'Unload weights {op}: {memory_stats()}') -def apply_token_merging(sd_model): - current_tome = getattr(sd_model, 'applied_tome', 0) - current_todo = getattr(sd_model, 'applied_todo', 0) - - if shared.opts.token_merging_method == 'ToMe' and shared.opts.tome_ratio > 0: - if current_tome == shared.opts.tome_ratio: - return - if shared.opts.hypertile_unet_enabled and not shared.cmd_opts.experimental: - shared.log.warning('Token merging not supported with HyperTile for UNet') - return - try: - import installer - installer.install('tomesd', 'tomesd', ignore=False) - import tomesd - tomesd.apply_patch( - sd_model, - ratio=shared.opts.tome_ratio, - use_rand=False, # can cause issues with some samplers - merge_attn=True, - merge_crossattn=False, - merge_mlp=False - ) - shared.log.info(f'Applying ToMe: ratio={shared.opts.tome_ratio}') - sd_model.applied_tome = shared.opts.tome_ratio - except Exception: - shared.log.warning(f'Token merging not supported: pipeline={sd_model.__class__.__name__}') - else: - sd_model.applied_tome = 0 - - if shared.opts.token_merging_method == 'ToDo' and shared.opts.todo_ratio > 0: - if current_todo == shared.opts.todo_ratio: - return - if shared.opts.hypertile_unet_enabled and not shared.cmd_opts.experimental: - shared.log.warning('Token merging not supported with HyperTile for UNet') - return - try: - from modules.todo.todo_utils import patch_attention_proc - token_merge_args = { - "ratio": shared.opts.todo_ratio, - "merge_tokens": "keys/values", - "merge_method": "downsample", - "downsample_method": "nearest", - "downsample_factor": 2, - "timestep_threshold_switch": 0.0, - "timestep_threshold_stop": 0.0, - "downsample_factor_level_2": 1, - "ratio_level_2": 0.0, - } - patch_attention_proc(sd_model.unet, token_merge_args=token_merge_args) - shared.log.info(f'Applying ToDo: ratio={shared.opts.todo_ratio}') - sd_model.applied_todo = shared.opts.todo_ratio - except Exception: - shared.log.warning(f'Token merging not supported: pipeline={sd_model.__class__.__name__}') - else: - sd_model.applied_todo = 0 - - -def remove_token_merging(sd_model): - current_tome = getattr(sd_model, 'applied_tome', 0) - current_todo = getattr(sd_model, 'applied_todo', 0) - try: - if current_tome > 0: - import tomesd - tomesd.remove_patch(sd_model) - sd_model.applied_tome = 0 - except Exception: - pass - try: - if current_todo > 0: - from modules.todo.todo_utils import remove_patch - remove_patch(sd_model) - sd_model.applied_todo = 0 - except Exception: - pass - - def path_to_repo(fn: str = ''): if isinstance(fn, CheckpointInfo): fn = fn.name diff --git a/modules/sd_vae.py b/modules/sd_vae.py index 52ba77bba..f266f8c38 100644 --- a/modules/sd_vae.py +++ b/modules/sd_vae.py @@ -2,7 +2,7 @@ import os import glob from copy import deepcopy import torch -from modules import shared, errors, paths, devices, script_callbacks, sd_models +from modules import shared, errors, paths, devices, script_callbacks, sd_models, sd_detect vae_ignore_keys = {"model_ema.decay", "model_ema.num_updates"} @@ -206,8 +206,8 @@ def load_vae_diffusers(model_file, vae_file=None, vae_source="unknown-source"): diffusers_load_config['variant'] = shared.opts.diffusers_vae_load_variant if shared.opts.diffusers_vae_upcast != 'default': diffusers_load_config['force_upcast'] = True if shared.opts.diffusers_vae_upcast == 'true' else False - _pipeline, model_type = sd_models.detect_pipeline(model_file, 'vae') - vae_config = sd_models.get_load_config(model_file, model_type, config_type='json') + _pipeline, model_type = sd_detect.detect_pipeline(model_file, 'vae') + vae_config = sd_detect.get_load_config(model_file, model_type, config_type='json') if vae_config is not None: diffusers_load_config['config'] = os.path.join(vae_config, 'vae') shared.log.info(f'Load module: type=VAE model="{vae_file}" source={vae_source} config={diffusers_load_config}') diff --git a/modules/shared.py b/modules/shared.py index a3a9a5482..ae94f26cf 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -280,12 +280,13 @@ def options_section(section_identifier, options_dict): return options_dict -def list_checkpoint_tiles(): +def list_checkpoint_titles(): import modules.sd_models # pylint: disable=W0621 - return modules.sd_models.checkpoint_tiles() + return modules.sd_models.checkpoint_titles() -default_checkpoint = list_checkpoint_tiles()[0] if len(list_checkpoint_tiles()) > 0 else "model.safetensors" +list_checkpoint_tiles = list_checkpoint_titles # alias for legacy typo +default_checkpoint = list_checkpoint_titles()[0] if len(list_checkpoint_titles()) > 0 else "model.safetensors" def is_url(string): @@ -427,12 +428,12 @@ startup_offload_mode, startup_cross_attention, startup_sdp_options = get_default options_templates.update(options_section(('sd', "Execution & Models"), { "sd_backend": OptionInfo(default_backend, "Execution backend", gr.Radio, {"choices": ["diffusers", "original"] }), - "sd_model_checkpoint": OptionInfo(default_checkpoint, "Base model", DropdownEditable, lambda: {"choices": list_checkpoint_tiles()}, refresh=refresh_checkpoints), - "sd_model_refiner": OptionInfo('None', "Refiner model", gr.Dropdown, lambda: {"choices": ['None'] + list_checkpoint_tiles()}, refresh=refresh_checkpoints), + "sd_model_checkpoint": OptionInfo(default_checkpoint, "Base model", DropdownEditable, lambda: {"choices": list_checkpoint_titles()}, refresh=refresh_checkpoints), + "sd_model_refiner": OptionInfo('None', "Refiner model", gr.Dropdown, lambda: {"choices": ['None'] + list_checkpoint_titles()}, refresh=refresh_checkpoints), "sd_vae": OptionInfo("Automatic", "VAE model", gr.Dropdown, lambda: {"choices": shared_items.sd_vae_items()}, refresh=shared_items.refresh_vae_list), "sd_unet": OptionInfo("None", "UNET model", gr.Dropdown, lambda: {"choices": shared_items.sd_unet_items()}, refresh=shared_items.refresh_unet_list), "sd_text_encoder": OptionInfo('None', "Text encoder model", gr.Dropdown, lambda: {"choices": shared_items.sd_te_items()}, refresh=shared_items.refresh_te_list), - "sd_model_dict": OptionInfo('None', "Use separate base dict", gr.Dropdown, lambda: {"choices": ['None'] + list_checkpoint_tiles()}, refresh=refresh_checkpoints), + "sd_model_dict": OptionInfo('None', "Use separate base dict", gr.Dropdown, lambda: {"choices": ['None'] + list_checkpoint_titles()}, refresh=refresh_checkpoints), "sd_checkpoint_autoload": OptionInfo(True, "Model autoload on start"), "sd_checkpoint_autodownload": OptionInfo(True, "Model auto-download on demand"), "sd_textencoder_cache": OptionInfo(True, "Cache text encoder results"), diff --git a/modules/token_merge.py b/modules/token_merge.py new file mode 100644 index 000000000..f97c1fc8e --- /dev/null +++ b/modules/token_merge.py @@ -0,0 +1,77 @@ +from modules import shared + + +def apply_token_merging(sd_model): + current_tome = getattr(sd_model, 'applied_tome', 0) + current_todo = getattr(sd_model, 'applied_todo', 0) + + if shared.opts.token_merging_method == 'ToMe' and shared.opts.tome_ratio > 0: + if current_tome == shared.opts.tome_ratio: + return + if shared.opts.hypertile_unet_enabled and not shared.cmd_opts.experimental: + shared.log.warning('Token merging not supported with HyperTile for UNet') + return + try: + import installer + installer.install('tomesd', 'tomesd', ignore=False) + import tomesd + tomesd.apply_patch( + sd_model, + ratio=shared.opts.tome_ratio, + use_rand=False, # can cause issues with some samplers + merge_attn=True, + merge_crossattn=False, + merge_mlp=False + ) + shared.log.info(f'Applying ToMe: ratio={shared.opts.tome_ratio}') + sd_model.applied_tome = shared.opts.tome_ratio + except Exception: + shared.log.warning(f'Token merging not supported: pipeline={sd_model.__class__.__name__}') + else: + sd_model.applied_tome = 0 + + if shared.opts.token_merging_method == 'ToDo' and shared.opts.todo_ratio > 0: + if current_todo == shared.opts.todo_ratio: + return + if shared.opts.hypertile_unet_enabled and not shared.cmd_opts.experimental: + shared.log.warning('Token merging not supported with HyperTile for UNet') + return + try: + from modules.todo.todo_utils import patch_attention_proc + token_merge_args = { + "ratio": shared.opts.todo_ratio, + "merge_tokens": "keys/values", + "merge_method": "downsample", + "downsample_method": "nearest", + "downsample_factor": 2, + "timestep_threshold_switch": 0.0, + "timestep_threshold_stop": 0.0, + "downsample_factor_level_2": 1, + "ratio_level_2": 0.0, + } + patch_attention_proc(sd_model.unet, token_merge_args=token_merge_args) + shared.log.info(f'Applying ToDo: ratio={shared.opts.todo_ratio}') + sd_model.applied_todo = shared.opts.todo_ratio + except Exception: + shared.log.warning(f'Token merging not supported: pipeline={sd_model.__class__.__name__}') + else: + sd_model.applied_todo = 0 + + +def remove_token_merging(sd_model): + current_tome = getattr(sd_model, 'applied_tome', 0) + current_todo = getattr(sd_model, 'applied_todo', 0) + try: + if current_tome > 0: + import tomesd + tomesd.remove_patch(sd_model) + sd_model.applied_tome = 0 + except Exception: + pass + try: + if current_todo > 0: + from modules.todo.todo_utils import remove_patch + remove_patch(sd_model) + sd_model.applied_todo = 0 + except Exception: + pass diff --git a/modules/ui_control.py b/modules/ui_control.py index 388e8ede9..4d7c59bee 100644 --- a/modules/ui_control.py +++ b/modules/ui_control.py @@ -612,6 +612,7 @@ def create_ui(_blocks: gr.Blocks=None): (mask_controls[6], "Mask auto"), # advanced (cfg_scale, "CFG scale"), + (cfg_end, "CFG end"), (clip_skip, "Clip skip"), (image_cfg_scale, "Image CFG scale"), (diffusers_guidance_rescale, "CFG rescale"), diff --git a/modules/ui_img2img.py b/modules/ui_img2img.py index d46ea4dd3..44f48c3c6 100644 --- a/modules/ui_img2img.py +++ b/modules/ui_img2img.py @@ -263,6 +263,7 @@ def create_ui(): (refiner_start, "Refiner start"), # advanced (cfg_scale, "CFG scale"), + (cfg_end, "CFG end"), (image_cfg_scale, "Image CFG scale"), (clip_skip, "Clip skip"), (diffusers_guidance_rescale, "CFG rescale"), diff --git a/modules/ui_models.py b/modules/ui_models.py index 051ca39a7..e9be428b4 100644 --- a/modules/ui_models.py +++ b/modules/ui_models.py @@ -59,8 +59,8 @@ def create_ui(): with gr.Tab(label="Convert"): with gr.Row(): - model_name = gr.Dropdown(sd_models.checkpoint_tiles(), label="Original model") - create_refresh_button(model_name, sd_models.list_models, lambda: {"choices": sd_models.checkpoint_tiles()}, "refresh_checkpoint_Z") + model_name = gr.Dropdown(sd_models.checkpoint_titles(), label="Original model") + create_refresh_button(model_name, sd_models.list_models, lambda: {"choices": sd_models.checkpoint_titles()}, "refresh_checkpoint_Z") with gr.Row(): custom_name = gr.Textbox(label="Output model name") with gr.Row(): @@ -98,7 +98,7 @@ def create_ui(): with gr.Tab(label="Merge"): def sd_model_choices(): - return ['None'] + sd_models.checkpoint_tiles() + return ['None'] + sd_models.checkpoint_titles() with gr.Row(equal_height=False): with gr.Column(variant='compact'): @@ -213,10 +213,10 @@ def create_ui(): del kwargs['dummy_component'] if kwargs.get("custom_name", None) is None: log.error('Merge: no output model specified') - return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)], "No output model specified"] + return [*[gr.Dropdown.update(choices=sd_models.checkpoint_titles()) for _ in range(4)], "No output model specified"] elif kwargs.get("primary_model_name", None) is None or kwargs.get("secondary_model_name", None) is None: log.error('Merge: no models selected') - return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)], "No models selected"] + return [*[gr.Dropdown.update(choices=sd_models.checkpoint_titles()) for _ in range(4)], "No models selected"] else: log.debug(f'Merge start: {kwargs}') try: @@ -224,7 +224,7 @@ def create_ui(): except Exception as e: modules.errors.display(e, 'Merge') sd_models.list_models() # to remove the potentially missing models from the list - return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)], f"Error merging checkpoints: {e}"] + return [*[gr.Dropdown.update(choices=sd_models.checkpoint_titles()) for _ in range(4)], f"Error merging checkpoints: {e}"] return results def tertiary(mode): diff --git a/modules/ui_txt2img.py b/modules/ui_txt2img.py index 1ae3a8dad..ece8829c5 100644 --- a/modules/ui_txt2img.py +++ b/modules/ui_txt2img.py @@ -116,6 +116,7 @@ def create_ui(): (subseed_strength, "Variation strength"), # advanced (cfg_scale, "CFG scale"), + (cfg_end, "CFG end"), (clip_skip, "Clip skip"), (image_cfg_scale, "Image CFG scale"), (diffusers_guidance_rescale, "CFG rescale"), diff --git a/scripts/x_adapter.py b/scripts/x_adapter.py index 7c1341701..553a20d30 100644 --- a/scripts/x_adapter.py +++ b/scripts/x_adapter.py @@ -22,7 +22,7 @@ class Script(scripts.Script): with gr.Row(): gr.HTML('  X-Adapter
') with gr.Row(): - model = gr.Dropdown(label='Adapter model', choices=['None'] + sd_models.checkpoint_tiles(), value='None') + model = gr.Dropdown(label='Adapter model', choices=['None'] + sd_models.checkpoint_titles(), value='None') sampler = gr.Dropdown(label='Adapter sampler', choices=[s.name for s in sd_samplers.samplers], value='Default') with gr.Row(): width = gr.Slider(label='Adapter width', minimum=64, maximum=2048, step=8, value=1024) @@ -34,7 +34,7 @@ class Script(scripts.Script): lora = gr.Textbox('', label='Adapter LoRA', default='') return model, sampler, width, height, start, scale, lora - def run(self, p: processing.StableDiffusionProcessing, model, sampler, width, height, start, scale, lora): # pylint: disable=arguments-differ + def run(self, p: processing.StableDiffusionProcessing, model, sampler, width, height, start, scale, lora): # pylint: disable=arguments-differ, unused-argument from modules.xadapter.xadapter_hijacks import PositionNet diffusers.models.embeddings.PositionNet = PositionNet # patch diffusers==0.26 from diffusers==0.20 from modules.xadapter.adapter import Adapter_XL diff --git a/scripts/xyz_grid_classes.py b/scripts/xyz_grid_classes.py index 8a78d2c40..335d66186 100644 --- a/scripts/xyz_grid_classes.py +++ b/scripts/xyz_grid_classes.py @@ -99,7 +99,7 @@ axis_options = [ AxisOption("[Param] Height", int, apply_field("height")), AxisOption("[Param] Seed", int, apply_seed), AxisOption("[Param] Steps", int, apply_field("steps")), - AxisOption("[Param] CFG scale", float, apply_field("cfg_scale")), + AxisOption("[Param] Guidance scale", float, apply_field("cfg_scale")), AxisOption("[Param] Guidance end", float, apply_field("cfg_end")), AxisOption("[Param] Variation seed", int, apply_field("subseed")), AxisOption("[Param] Variation strength", float, apply_field("subseed_strength")), @@ -125,7 +125,7 @@ axis_options = [ AxisOption("[Refine] Sampler", str, apply_hr_sampler_name, fmt=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]), AxisOption("[Refine] Denoising strength", float, apply_field("denoising_strength")), AxisOption("[Refine] Hires steps", int, apply_field("hr_second_pass_steps")), - AxisOption("[Refine] CFG scale", float, apply_field("image_cfg_scale")), + AxisOption("[Refine] Guidance scale", float, apply_field("image_cfg_scale")), AxisOption("[Refine] Guidance rescale", float, apply_field("diffusers_guidance_rescale")), AxisOption("[Refine] Refiner start", float, apply_field("refiner_start")), AxisOption("[Refine] Refiner steps", float, apply_field("refiner_steps")),