From 8241e33868e945c73a55836fb3790e2181131034 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 3 Jul 2023 16:48:03 -0400 Subject: [PATCH] major diffusers update --- CHANGELOG.md | 5 ++ DIFFUSERS.md | 77 +++++++++++++++---- TODO.md | 1 + extensions-builtin/LDSR/ldsr_model_arch.py | 13 ++-- .../Lora/extra_networks_lora.py | 2 +- extensions-builtin/Lora/lora.py | 19 +++-- javascript/extraNetworks.js | 1 + javascript/style.css | 2 +- modules/cmd_args.py | 2 +- modules/lora_diffusers.py | 34 ++++++++ modules/modelloader.py | 20 ++--- modules/processing.py | 10 ++- modules/sd_hijack.py | 2 - modules/sd_models.py | 46 ++++++----- modules/sd_samplers_kdiffusion.py | 2 +- modules/shared.py | 62 ++------------- .../textual_inversion/textual_inversion.py | 48 +++++++----- modules/ui.py | 18 ++++- modules/ui_extra_networks.py | 24 +++--- modules/ui_extra_networks_checkpoints.py | 4 +- modules/ui_extra_networks_hypernets.py | 2 +- modules/ui_models.py | 21 ++--- webui.py | 2 +- 23 files changed, 247 insertions(+), 170 deletions(-) create mode 100644 modules/lora_diffusers.py diff --git a/CHANGELOG.md b/CHANGELOG.md index e60a29277..91ef21dc7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,10 @@ # Change Log for SD.Next +## Update for ... + +- add settings -> extra networks -> do not automatically build extra network pages + speeds up app start if you have a lot of extra networks and you want to build them manually when needed + ## Update for 07/01/2023 Small quality-of-life updates and bugfixes: diff --git a/DIFFUSERS.md b/DIFFUSERS.md index 52616aa2d..e43898400 100644 --- a/DIFFUSERS.md +++ b/DIFFUSERS.md @@ -9,6 +9,7 @@ initial support merged into `dev` branch default sd 1.5 model will be downloaded automatically to `models/Diffusers` on first startup, disable **controlnet** and **multi-diffusion** extensions as right now they are not compatible with diffusers +lora support is not compatible with setting `Use LyCoris handler for all Lora types`, make sure its disabled to update repo, do not use `--upgrade` flag, use manual `git pull` instead @@ -16,38 +17,80 @@ to update repo, do not use `--upgrade` flag, use manual `git pull` instead ### Standard +goal is to test standard workflows (so not diffusers) to ensure there are no regressions +so diffusers code can be merged into `master` and we can continue with development there + - run with `webui --debug --backend original` -- goal is to test standard workflows (so not diffusers) to ensure there are no regressions - so diffusers code can be merged into `master` and we can continue with development there ### Diffusers -- sd 1.5 and sd 2.1 model -- model downloader: tabs -> models -> hf hub -- txt2img, img2img, inpaint, outpaint, process -- hires fix, restore faces, etc? +whats implemented so far? -### Experimental - don't test yet +- simple model downloader for huggingface models: tabs -> models -> hf hub +- use huggingface models +- extra networks ui +- use safetensor models with diffusers backend +- standard workflows: + - txt2img, img2img, inpaint, outpaint, process + - hires fix, restore faces, etc? +- textual inversion + yes, this applies to standard embedddings, don't need ones from huggingface +- lora + yes, this applies to standard loras, don't need ones from huggingface + but seems that diffuser lora support is somewhat limited, so quite a few loras may not work + you should see which lora loads without issues in console log +- system info tab with updated information +- kandinsky model + works for me -- cuda model compile using `reduce overhead` model with or without `fullgraph` -- kandinsky model +### Experimental + +- cuda model compile + in settings -> compute settings + diffusers recommend `reduce overhead`, but other methods are available as well + it seems that fullgraph is possible (with sufficient vram) when using diffusers +- deepfloyd + in theory it should work, but its 20gb model so cant test it just yet + note that access is gated, so you'll need to download using your huggingface credentials + (you can still do it from sdnext ui, just need access token) ## Todo -- lora -- embedding -- safetensors models -- cleanup logging -- controlnet extension -- multidiffusion extension - sdxl model ## Limitations -- extra networks +even if extensions are not supported, runtime errors are never nice +will need to handle in the code before we get out of alpha + - controlnet -- multi-diffusion + `sd_model.model?.diffusion_model?` +- multi-diffusion + `sd_model.first_stage_model?.encoder?` +- lycoris + `lyco_patch_lora` ## Issues - TBD + +## Notes for HF + +- removed `quicksettings` alternative completely +- added simple model downloader in ui: *tabs -> models -> huggingface* +- redone **textual inversion** support, core is now in `modules/textual_inversion/textual_inversion.py:load_diffusers_embedding()` + the point is that sdnext pre-loads all compatible embeddings on model load so they are available in prompt context +- added support for diffuser models in **safetensors/ckpt** format + btw, when i use: `diffusers.StableDiffusionPipeline.from_ckpt` + first time it downloads something - what is that? + > Downloading (…)lve/main/config.json: 4.55k + > Downloading pytorch_model.bin: 1.22G + and in general, loading safetensors model is quite slow, is that expected? + for example, 2sec vs 18sec +- in `modules/modelloader.py:download_diffusers_model()` i get unknown property for `hf.model_info(hub_id).cardData` + can you double-check if this is linter issue or actual problem? +- redone **lora** support, core is now in `modules/lora_diffusers.py` +- question on `pipe.load_lora_weights` + does it support loading multiple loras? i don't see any notes on that in docs + also, lora strength is specified using `cross_attention_kwargs={"scale": x}` during pipeline execution + which means if there are multiple loras, they all have the same strength? diff --git a/TODO.md b/TODO.md index 19ce7d49f..ef0e08985 100644 --- a/TODO.md +++ b/TODO.md @@ -41,6 +41,7 @@ Tech that can be integrated as part of the core workflow... - [DataComp CLiP](https://github.com/mlfoundations/open_clip/blob/main/docs/datacomp_models.md) - [ClipSeg](https://github.com/timojl/clipseg) - [DragGAN](https://github.com/XingangPan/DragGAN) +- [LamaCleaner]([Title](https://github.com/Sanster/lama-cleaner)) - `TensorRT` ## Random diff --git a/extensions-builtin/LDSR/ldsr_model_arch.py b/extensions-builtin/LDSR/ldsr_model_arch.py index 41d97d071..9411e7374 100644 --- a/extensions-builtin/LDSR/ldsr_model_arch.py +++ b/extensions-builtin/LDSR/ldsr_model_arch.py @@ -23,10 +23,10 @@ class LDSR: global cached_ldsr_model if shared.opts.ldsr_cached and cached_ldsr_model is not None: - print("Loading model from cache") + shared.log.info("LDSR Loading model from cache") model: torch.nn.Module = cached_ldsr_model else: - print(f"Loading model from {self.modelPath}") + shared.log.info(f"LDSR Loading model from {self.modelPath}") _, extension = os.path.splitext(self.modelPath) if extension.lower() == ".safetensors": pl_sd = safetensors.torch.load_file(self.modelPath, device="cpu") @@ -126,11 +126,10 @@ class LDSR: height_downsampled_pre = int(np.ceil(hd)) if down_sample_rate != 1: - print( - f'Downsampling from [{width_og}, {height_og}] to [{width_downsampled_pre}, {height_downsampled_pre}]') + shared.log.info(f'LDSR Downsampling from [{width_og}, {height_og}] to [{width_downsampled_pre}, {height_downsampled_pre}]') im_og = im_og.resize((width_downsampled_pre, height_downsampled_pre), Image.LANCZOS) else: - print(f"Down sample rate is 1 from {target_scale} / 4 (Not downsampling)") + shared.log.info(f"LDSR Downsample rate is 1 from {target_scale} / 4 (Not downsampling)") # pad width and height to multiples of 64, pads with the edge values of image to avoid artifacts pad_w, pad_h = np.max(((2, 2), np.ceil(np.array(im_og.size) / 64).astype(int)), axis=0) * 64 - im_og.size @@ -183,7 +182,7 @@ def convsample_ddim(model, cond, steps, shape, eta=1.0, callback=None, normals_s ddim = DDIMSampler(model) bs = shape[0] shape = shape[1:] - print(f"Sampling with eta = {eta}; steps: {steps}") + shared.log.info(f"LDSR Sampling with eta = {eta}; steps: {steps}") samples, intermediates = ddim.sample(steps, batch_size=bs, shape=shape, conditioning=cond, callback=callback, normals_sequence=normals_sequence, quantize_x0=quantize_x0, eta=eta, mask=mask, x0=x0, temperature=temperature, verbose=False, @@ -206,7 +205,7 @@ def make_convolutional_sample(batch, model, custom_steps=None, eta=1.0, quantize if custom_shape is not None: z = torch.randn(custom_shape) - print(f"Generating {custom_shape[0]} samples of shape {custom_shape[1:]}") + shared.log.info(f"LDSR Generating {custom_shape[0]} samples of shape {custom_shape[1:]}") z0 = None diff --git a/extensions-builtin/Lora/extra_networks_lora.py b/extensions-builtin/Lora/extra_networks_lora.py index b5fea4d2e..bee0477ed 100644 --- a/extensions-builtin/Lora/extra_networks_lora.py +++ b/extensions-builtin/Lora/extra_networks_lora.py @@ -1,5 +1,5 @@ -from modules import extra_networks, shared import lora +from modules import extra_networks, shared class ExtraNetworkLora(extra_networks.ExtraNetwork): diff --git a/extensions-builtin/Lora/lora.py b/extensions-builtin/Lora/lora.py index 5273c9f82..088be31c1 100644 --- a/extensions-builtin/Lora/lora.py +++ b/extensions-builtin/Lora/lora.py @@ -1,8 +1,7 @@ import os import re -import torch from typing import Union - +import torch from modules import shared, devices, sd_models, errors, scripts, sd_hijack, hashes metadata_tags_order = {"ss_sd_model_name": 1, "ss_resolution": 2, "ss_clip_skip": 3, "ss_num_train_images": 10, "ss_tag_frequency": 20} @@ -127,7 +126,6 @@ class LoraModule: self.multiplier = 1.0 self.modules = {} self.mtime = None - self.mentioned_name = None """the text that was used to add lora to prompt - can be either name or an alias""" @@ -154,6 +152,13 @@ def assign_lora_names_to_compvis_modules(sd_model): sd_model.lora_layer_mapping = lora_layer_mapping +def load_diffuser_lora(name, lora_on_disk, multiplier): + lora = LoraModule(name, lora_on_disk) + lora.mtime = os.path.getmtime(lora_on_disk.filename) + from modules.lora_diffusers import load_diffusers_lora + load_diffusers_lora(name, lora_on_disk, multiplier) + return lora + def load_lora(name, lora_on_disk): lora = LoraModule(name, lora_on_disk) @@ -205,7 +210,6 @@ def load_lora(name, lora_on_disk): else: print(f'Lora layer {key_diffusers} matched a layer with unsupported type: {type(sd_module).__name__}') continue - raise AssertionError(f"Lora layer {key_diffusers} matched a layer with unsupported type: {type(sd_module).__name__}") with torch.no_grad(): module.weight.copy_(weight) @@ -243,14 +247,17 @@ def load_loras(names, multipliers=None): failed_to_load_loras = [] for i, name in enumerate(names): - lora = already_loaded.get(name, None) + lora = already_loaded.get(name, None) if shared.backend == shared.Backend.ORIGINAL else None lora_on_disk = loras_on_disk[i] if lora_on_disk is not None: if lora is None or os.path.getmtime(lora_on_disk.filename) > lora.mtime: try: - lora = load_lora(name, lora_on_disk) + if shared.backend == shared.Backend.DIFFUSERS: + lora = load_diffuser_lora(name, lora_on_disk, multipliers[i] if multipliers else 1.0) + else: + lora = load_lora(name, lora_on_disk) except Exception as e: errors.display(e, f"loading Lora {lora_on_disk.filename}") continue diff --git a/javascript/extraNetworks.js b/javascript/extraNetworks.js index f2e763362..39f75c91c 100644 --- a/javascript/extraNetworks.js +++ b/javascript/extraNetworks.js @@ -19,6 +19,7 @@ function setupExtraNetworksForTab(tabname) { searchTerm = search.value.toLowerCase(); gradioApp().querySelectorAll(`#${tabname}_extra_tabs div.card`).forEach((elem) => { text = `${elem.querySelector('.name').textContent.toLowerCase()} ${elem.querySelector('.search_term').textContent.toLowerCase()}`; + text = text.replace('models--', 'Diffusers') elem.style.display = text.indexOf(searchTerm) == -1 ? 'none' : ''; }); }); diff --git a/javascript/style.css b/javascript/style.css index 19faaa021..eef8d75ee 100644 --- a/javascript/style.css +++ b/javascript/style.css @@ -540,7 +540,7 @@ table.settings-value-table td{ .extra-networks-tab { padding: 0 !important; } .extra-network-subdirs { background: var(--input-background-fill); } .extra-networks-page { display: flex } -.extra-networks .custom-button { min-width: 60px; width: 100%; background: none; justify-content: left; padding: 2px 8px 2px 8px; box-shadow: none; } +.extra-networks .custom-button { min-width: 80px; max-width: 240px; width: 100%; background: none; justify-content: left; text-align: left; padding: 2px 8px 2px 8px; box-shadow: none; } .extra-networks .custom-button:hover { background: var(--button-primary-background-fill) } .extra-network-cards { display: flex; flex-wrap: wrap; height: 50vh; overflow-y: scroll; overflow-x: hidden; scroll-snap-type: y mandatory; width: -webkit-fill-available; } .extra-network-cards .card { height: fit-content; margin: 0.5em; position: relative; scroll-snap-align: start; scroll-margin-top: 0; } diff --git a/modules/cmd_args.py b/modules/cmd_args.py index 7d85e014c..fe37814e9 100644 --- a/modules/cmd_args.py +++ b/modules/cmd_args.py @@ -45,7 +45,7 @@ group.add_argument('--use-directml', default = False, action='store_true', help group.add_argument("--use-cuda", default=False, action='store_true', help="Force use nVidia CUDA backend, default: %(default)s") group.add_argument("--use-rocm", default=False, action='store_true', help="Force use AMD ROCm backend, default: %(default)s") group.add_argument('--subpath', type=str, help='Customize the URL subpath for usage with reverse proxy') -group.add_argument('--backend', type=str, choices=[None, 'original', 'diffusers'], default=None, required=False, help='force backend type') +group.add_argument('--backend', type=str, choices=['original', 'diffusers'], required=False, help='force model pipeline type') # removed args are added here as hidden in fixed format for compatbility reasons diff --git a/modules/lora_diffusers.py b/modules/lora_diffusers.py new file mode 100644 index 000000000..a6a5cbad7 --- /dev/null +++ b/modules/lora_diffusers.py @@ -0,0 +1,34 @@ +import diffusers +from modules import shared + +lora_state = { # TODO this is ugly but diffusers + 'multiplier': 1.0, + 'active': False, + 'loaded': 0, +} + +def unload_diffusers_lora(): + try: + pipe = shared.sd_model + lora_state['active'] = False + lora_state['loaded'] = 0 + pipe._remove_text_encoder_monkey_patch() # pylint: disable=W0212 + proc_cls_name = next(iter(pipe.unet.attn_processors.values())).__class__.__name__ + non_lora_proc_cls = getattr(diffusers.models.attention_processor, proc_cls_name[len("LORA"):]) + pipe.unet.set_attn_processor(non_lora_proc_cls()) + # shared.log.debug('Diffusers LoRA unloaded') + except Exception: + pass + + +def load_diffusers_lora(name, lora, strength = 1.0): + try: + pipe = shared.sd_model + pipe.load_lora_weights(lora.filename, cache_dir=shared.opts.diffusers_dir, local_files_only=True) + lora_state['active'] = True + lora_state['loaded'] += 1 + lora_state['multiplier'] = strength + # pipe.unet.load_attn_procs("pcuenq/pokemon-lora") + shared.log.info(f"Diffusers LoRA loaded: {name} {lora_state['multiplier']}") + except Exception as e: + shared.log.error(f"Diffusers LoRA loading failed: {name} {e}") diff --git a/modules/modelloader.py b/modules/modelloader.py index 758dc7e34..6de1efe6a 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -12,7 +12,7 @@ from modules.paths import script_path, models_path diffuser_repos = [] -def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config: Dict[str, str] = None): +def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config: Dict[str, str] = None, token = None): from diffusers import DiffusionPipeline import huggingface_hub as hf @@ -21,14 +21,16 @@ def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config "force_download": False, "resume_download": True, "cache_dir": shared.opts.diffusers_dir, + # "use_auth_token": True, } - if cache_dir is not None: download_config["cache_dir"] = cache_dir - + shared.log.debug(f"Diffusers downloading: {hub_id} to {cache_dir}") + if token is not None and len(token) > 2: + shared.log.debug(f"Diffusers authentication: {token}") + hf.login(token) pipeline_dir = DiffusionPipeline.download(hub_id, **download_config) model_info_dict = hf.model_info(hub_id).cardData # TODO hfhub card-data? - # some checkpoints need to be downloaded as "hidden" as they just serve as pre- or post-pipelines of other pipelines if model_info_dict is not None and "prior" in model_info_dict: download_dir = DiffusionPipeline.download(model_info_dict["prior"], **download_config) @@ -36,10 +38,8 @@ def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config # mark prior as hidden with open(os.path.join(download_dir, "hidden"), "w", encoding="utf-8") as f: f.write("True") - with open(os.path.join(pipeline_dir, "model_info.json"), "w", encoding="utf-8") as json_file: json.dump(model_info_dict, json_file) - return pipeline_dir @@ -61,7 +61,7 @@ def load_diffusers_models(model_path: str, command_path: str = None): output.append(str(r.repo_id)) except Exception as e: shared.log.error(f"Error listing diffusers: {place} {e}") - shared.log.debug(f'Scanning diffusers cache: {len(output)} {model_path} {command_path}') + shared.log.debug(f'Scanning diffusers cache: {model_path} {command_path} {len(output)}') return output @@ -105,7 +105,7 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None for place in places: for full_path in shared.walk_files(place, allowed_extensions=ext_filter): if os.path.islink(full_path) and not os.path.exists(full_path): - print(f"Skipping broken symlink: {full_path}") + shared.log.error(f"Skipping broken symlink: {full_path}") continue if ext_blacklist is not None and any(full_path.endswith(x) for x in ext_blacklist): continue @@ -172,13 +172,13 @@ def move_files(src_path: str, dest_path: str, ext_filter: str = None): if ext_filter is not None: if ext_filter not in file: continue - print(f"Moving {file} from {src_path} to {dest_path}.") + shared.log.warning(f"Moving {file} from {src_path} to {dest_path}.") try: shutil.move(fullpath, dest_path) except Exception: pass if len(os.listdir(src_path)) == 0: - print(f"Removing empty folder: {src_path}") + shared.log.info(f"Removing empty folder: {src_path}") shutil.rmtree(src_path, True) except Exception: pass diff --git a/modules/processing.py b/modules/processing.py index e8808a910..68c18c977 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -26,6 +26,7 @@ import modules.images as images import modules.styles import modules.sd_models as sd_models import modules.sd_vae as sd_vae +from modules.lora_diffusers import lora_state, unload_diffusers_lora opt_C = 4 @@ -697,6 +698,10 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: # TODO(Patrick): For wrapped pipelines this is currently a no-op shared.sd_model.scheduler = scheduler.sampler + cross_attention_kwargs={} + if lora_state['active']: + cross_attention_kwargs['scale'] = lora_state['multiplier'] + task_specific_kwargs={} if sd_models.get_diffusers_task(shared.sd_model) == sd_models.DiffusersTaskType.TEXT_2_IMAGE: task_specific_kwargs = {"height": p.height, "width": p.width} elif sd_models.get_diffusers_task(shared.sd_model) == sd_models.DiffusersTaskType.IMAGE_2_IMAGE: @@ -704,7 +709,6 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: elif sd_models.get_diffusers_task(shared.sd_model) == sd_models.DiffusersTaskType.INPAINTING: # TODO(PVP): change out to latents once possible with `diffusers` task_specific_kwargs = {"image": p.init_images[0], "mask_image": p.image_mask, "strength": p.denoising_strength} - output = shared.sd_model( prompt=prompts, negative_prompt=negative_prompts, @@ -712,9 +716,13 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: guidance_scale=p.cfg_scale, generator=generator, output_type="np", + cross_attention_kwargs=cross_attention_kwargs, **task_specific_kwargs ) x_samples_ddim = output.images + if lora_state['active']: + unload_diffusers_lora() + else: raise ValueError(f"Unknown backend {backend}") diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py index 5e397ebb9..5464cee44 100644 --- a/modules/sd_hijack.py +++ b/modules/sd_hijack.py @@ -174,14 +174,12 @@ class StableDiffusionModelHijack: sd_hijack_unet.hijack_ddpm_edit() if opts.cuda_compile and opts.cuda_compile_mode == 'ipex': - import logging if shared.cmd_opts.use_ipex: shared.log.info("Model compile enabled: IPEX Optimize Graph Mode") else: shared.log.warning("Model compile skipped: IPEX Method is for Intel GPU's with OneAPI") elif opts.cuda_compile and opts.cuda_compile_mode != 'none' and shared.backend == shared.Backend.ORIGINAL: try: - import logging import torch._dynamo as dynamo # pylint: disable=unused-import # torch._dynamo.config.log_level = logging.WARNING if opts.cuda_compile_verbose else logging.CRITICAL # pylint: disable=protected-access torch._dynamo.config.verbose = opts.cuda_compile_verbose # pylint: disable=protected-access diff --git a/modules/sd_models.py b/modules/sd_models.py index b72f0d7dd..2181b09d7 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -41,8 +41,10 @@ class CheckpointInfo: self.name = None self.hash = None self.filename = filename + self.type = '' abspath = os.path.abspath(filename) - if shared.backend == shared.Backend.ORIGINAL: + + if os.path.isfile(abspath): # ckpt or safetensor if shared.opts.ckpt_dir is not None and abspath.startswith(shared.opts.ckpt_dir): name = abspath.replace(shared.opts.ckpt_dir, '') elif abspath.startswith(model_path): @@ -54,7 +56,9 @@ class CheckpointInfo: self.name = name self.hash = model_hash(self.filename) self.sha256 = hashes.sha256_from_cache(self.filename, f"checkpoint/{name}") - elif shared.backend == shared.Backend.DIFFUSERS: + self.path = abspath + self.type = abspath.split('.')[-1].lower() + else: # maybe a diffuser repo = [r for r in modelloader.diffuser_repos if filename == r['filename']] if len(repo) == 0: error_message = f'Cannot find diffuser model: {filename}' @@ -64,6 +68,7 @@ class CheckpointInfo: self.hash = repo[0]['hash'][:8] self.sha256 = repo[0]['hash'] self.path = repo[0]['path'] + self.type = 'diffusers' if os.path.isfile(repo[0]['model_info']): file_path = repo[0]['model_info'] @@ -71,8 +76,6 @@ class CheckpointInfo: self.model_info = json.load(json_file) else: self.model_info = None - else: - raise ValueError(f'Unknown backend: {shared.backend}') self.name_for_extra = os.path.splitext(os.path.basename(filename))[0] self.model_name = os.path.splitext(name.replace("/", "_").replace("\\", "_"))[0] @@ -123,13 +126,10 @@ def checkpoint_tiles(): def list_models(): checkpoints_list.clear() checkpoint_aliases.clear() - if shared.backend == shared.Backend.ORIGINAL: - ext_filter=[".safetensors"] if shared.opts.sd_disable_ckpt else [".ckpt", ".safetensors"] - model_list = modelloader.load_models(model_path=os.path.join(models_path, 'Stable-diffusion'), model_url=None, command_path=shared.opts.ckpt_dir, ext_filter=ext_filter, download_name=None, ext_blacklist=[".vae.ckpt", ".vae.safetensors"]) - else: - global model_path # pylint: disable=global-statement - model_path = os.path.join(models_path, 'Diffusers') - model_list = modelloader.load_diffusers_models(model_path=model_path, command_path=shared.opts.diffusers_dir) + ext_filter=[".safetensors"] if shared.opts.sd_disable_ckpt else [".ckpt", ".safetensors"] + model_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"]) + if shared.backend == shared.Backend.DIFFUSERS: + model_list += modelloader.load_diffusers_models(model_path=os.path.join(models_path, 'Diffusers'), command_path=shared.opts.diffusers_dir) for filename in sorted(model_list, key=str.lower): checkpoint_info = CheckpointInfo(filename) @@ -158,8 +158,8 @@ def list_models(): model_list = modelloader.load_models(model_path=model_path, model_url=model_url, command_path=shared.opts.ckpt_dir, ext_filter=[".ckpt", ".safetensors"], download_name="v1-5-pruned-emaonly.safetensors", ext_blacklist=[".vae.ckpt", ".vae.safetensors"]) else: default_model_id = "runwayml/stable-diffusion-v1-5" - modelloader.download_diffusers_model(default_model_id, model_path) - model_list = modelloader.load_diffusers_models(model_path=model_path, command_path=shared.opts.diffusers_dir) + modelloader.download_diffusers_model(default_model_id, os.path.join(models_path, 'Diffusers')) + model_list = modelloader.load_diffusers_models(model_path=os.path.join(models_path, 'Diffusers'), command_path=shared.opts.diffusers_dir) for filename in sorted(model_list, key=str.lower): checkpoint_info = CheckpointInfo(filename) @@ -549,19 +549,22 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No try: if shared.cmd_opts.ckpt is not None and model_data.initial: # initial load model_name = modelloader.find_diffuser(shared.cmd_opts.ckpt) - if model_name is not None: shared.log.info(f'Loading diffuser model: {model_name}') model_file = modelloader.download_diffusers_model(hub_id=model_name) sd_model = diffusers.DiffusionPipeline.from_pretrained(model_file, **diffusers_load_config) - list_models() # rescan for downloaded model checkpoint_info = CheckpointInfo(model_name) if sd_model is None: checkpoint_info = checkpoint_info or select_checkpoint() shared.log.info(f'Loading diffuser model: {checkpoint_info.filename}') - sd_model = diffusers.DiffusionPipeline.from_pretrained(checkpoint_info.path, **diffusers_load_config) + if not os.path.isfile(checkpoint_info.path): + sd_model = diffusers.DiffusionPipeline.from_pretrained(checkpoint_info.path, **diffusers_load_config) + else: + diffusers_load_config["local_files_only "] = True + diffusers_load_config["extract_ema"] = True + sd_model = diffusers.StableDiffusionPipeline.from_ckpt(checkpoint_info.path, **diffusers_load_config) if "StableDiffusion" in sd_model.__class__.__name__: sd_model.scheduler = diffusers.UniPCMultistepScheduler.from_config(sd_model.scheduler.config) @@ -570,9 +573,9 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No sd_model.scheduler.name = 'DDIM' # Prior pipelines - if checkpoint_info.model_info is not None and "prior" in checkpoint_info.model_info: + if hasattr(checkpoint_info, 'model_info') and checkpoint_info.model_info is not None and "prior" in checkpoint_info.model_info: prior_id = checkpoint_info.model_info["prior"] - shared.log.info(f"Loading prior {prior_id} for {checkpoint_info.filename}") + shared.log.info(f"Loading diffuser prior: {checkpoint_info.filename} {prior_id}") prior = diffusers.DiffusionPipeline.from_pretrained(prior_id, **diffusers_load_config) sd_model = PriorPipeline(prior=prior, main=sd_model) # wrap sd_model @@ -586,7 +589,6 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No sd_model.to(devices.device) sd_model.unet.to(memory_format=torch.channels_last) import torch._dynamo as dynamo # pylint: disable=unused-import - # torch._dynamo.config.log_level = logging.WARNING if shared.opts.cuda_compile_verbose else logging.CRITICAL # pylint: disable=protected-access torch._dynamo.config.verbose = shared.opts.cuda_compile_verbose # pylint: disable=protected-access torch._dynamo.config.suppress_errors = shared.opts.cuda_compile_errors # pylint: disable=protected-access sd_model.unet = torch.compile(sd_model.unet, mode=shared.opts.cuda_compile_mode, fullgraph=shared.opts.cuda_compile_fullgraph) # pylint: disable=attribute-defined-outside-init @@ -602,6 +604,12 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No shared.log.error("Failed to load diffusers model") errors.display(e, "loading Diffusers model") shared.sd_model = sd_model + + from modules.textual_inversion import textual_inversion + embedding_db = textual_inversion.EmbeddingDatabase() + embedding_db.add_embedding_dir(shared.opts.embeddings_dir) + embedding_db.load_textual_inversion_embeddings(force_reload=True) + timer.record("load") shared.log.info(f"Model loaded in {timer.summary()}") devices.torch_gc(force=True) diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py index 2393b2ec2..f15d3c09b 100644 --- a/modules/sd_samplers_kdiffusion.py +++ b/modules/sd_samplers_kdiffusion.py @@ -330,7 +330,7 @@ class KDiffusionSampler: try: return BrownianTreeNoiseSampler(x.to("cpu"), sigma_min, sigma_max, seed=current_iter_seeds, transform=lambda x: x.to("cpu"), transform_last=lambda x: x.to(shared.device)) # pylint: disable=E1123 except Exception: - print("ERROR Please apply this patch to repositories/k-diffusion/k_diffusion/sampling.py: https://github.com/crowsonkb/k-diffusion/pull/68/files") + shared.log.error("Apply this patch to repositories/k-diffusion/k_diffusion/sampling.py: https://github.com/crowsonkb/k-diffusion/pull/68/files") return None else: return BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=current_iter_seeds) diff --git a/modules/shared.py b/modules/shared.py index 3b31df4d2..5f480f562 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -6,12 +6,10 @@ import datetime import urllib.request from urllib.parse import urlparse from enum import Enum -import tempfile import gradio as gr import tqdm import requests -import diffusers -from modules import errors, ui_components, shared_items, cmd_args, modelloader +from modules import errors, ui_components, shared_items, cmd_args from modules.paths_internal import models_path, script_path, data_path, sd_configs_path, sd_default_config, sd_model_file, default_sd_model_file, extensions_dir, extensions_builtin_dir # pylint: disable=W0611 import modules.interrogate import modules.memmon @@ -235,48 +233,6 @@ def list_checkpoint_tiles(): default_checkpoint = list_checkpoint_tiles()[0] if len(list_checkpoint_tiles()) > 0 else "model.ckpt" -def load_diffusers_lora(lora_repo: str): - pipe = sys.modules[__name__].sd_model - if lora_repo == "": - pipe._remove_text_encoder_monkey_patch() # pylint: disable=W0212 - proc_cls_name = next(iter(pipe.unet.attn_processors.values())).__class__.__name__ - non_lora_proc_cls = getattr(diffusers.models.attention_processor, proc_cls_name[len("LORA"):]) - pipe.unet.set_attn_processor(non_lora_proc_cls()) - return "" - elif len(lora_repo.split('/')) == 2: - lora_dir = os.path.dirname(opts.data["diffusers_dir"]) - cache_dir = os.path.join(lora_dir, "Diffusers_LoRA") - pipe.load_lora_weights(lora_repo, cache_dir=cache_dir) - print(f"Loaded {lora_repo}") - return f"{lora_repo} is loaded. Pass empty text field to remove LoRA or pass new LoRA id." - else: - print(f"{lora_repo} is not a valid LoRA identifier.") - return "" - - -def load_diffusers_text_inv(text_inv_repo: str): - pipe = sys.modules[__name__].sd_model - if text_inv_repo == "": - pipe.tokenizer = pipe.tokenizer.__class__.from_pretrained(pipe.tokenizer.name_or_path) - pipe.text_encoder.resize_token_embeddings(len(pipe.tokenizer)) - return "" - elif is_url(text_inv_repo): - with tempfile.TemporaryDirectory() as temp_dir: - os.system(f"wget -P {temp_dir} {text_inv_repo}") - temp_file_path = os.path.join(temp_dir, text_inv_repo.split('/')[-1]) - pipe.load_textual_inversion(temp_file_path) - text_inv_repo = '/'.join(text_inv_repo.split('/')[-2:]) - print(f"Loaded Civit.ai Textual Inv: {text_inv_repo}") - elif len(text_inv_repo.split('/')) == 2: - text_inv_dir = os.path.dirname(opts.data["diffusers_dir"]) - cache_dir = os.path.join(text_inv_dir, "Diffusers_Text_Inv") - pipe.load_textual_inversion(text_inv_repo, cache_dir=cache_dir) - print(f"Loaded {text_inv_repo}") - text_inv_tokens = pipe.tokenizer.added_tokens_encoder.keys() - text_inv_tokens = [t for t in text_inv_tokens if not (len(t.split("_")) > 1 and t.split("_")[-1].isdigit())] - return f"{', '.join(text_inv_tokens)} loaded. Pass empty text field to remove all or add new textual inversion id." - - def refresh_checkpoints(): import modules.sd_models # pylint: disable=W0621 return modules.sd_models.list_models() @@ -349,7 +305,7 @@ options_templates.update(options_section(('sd', "Stable Diffusion"), { "prompt_mean_norm": OptionInfo(True, "Prompt attention mean normalization"), "comma_padding_backtrack": OptionInfo(20, "Prompt padding for long prompts", gr.Slider, {"minimum": 0, "maximum": 74, "step": 1 }), "sd_disable_ckpt": OptionInfo(False, "Disallow usage of checkpoints in ckpt format"), - "sd_backend": OptionInfo("Original", "Stable Diffusion backend (experimental)", gr.Radio, lambda: {"choices": ["Original", "Diffusers"] }), + "sd_backend": OptionInfo("original", "Stable Diffusion backend (experimental)", gr.Radio, lambda: {"choices": ["original", "diffusers"] }), })) options_templates.update(options_section(('optimizations', "Optimizations"), { @@ -483,7 +439,6 @@ options_templates.update(options_section(('ui', "User Interface"), { "ui_tab_reorder": OptionInfo("From Text, From Image, Process Image", "UI tabs order"), "ui_scripts_reorder": OptionInfo("Enable Dynamic Thresholding, ControlNet", "UI scripts order"), "ui_reorder": OptionInfo(", ".join(ui_reorder_categories), "txt2img/img2img UI item order"), - "ui_extra_networks_tab_reorder": OptionInfo("Checkpoints, Lora, LyCORIS, Textual Inversion, Hypernetworks", "Extra networks tab order"), })) options_templates.update(options_section(('live-preview', "Live Previews"), { @@ -573,11 +528,13 @@ options_templates.update(options_section(('upscaling', "Upscaling"), { })) options_templates.update(options_section(('extra_networks', "Extra Networks"), { + "ui_extra_networks_tab_reorder": OptionInfo("Checkpoints, Lora, LyCORIS, Textual Inversion, Hypernetworks", "Extra networks tab order"), "extra_networks_card_cover": OptionInfo("inline", "UI position", gr.Radio, lambda: {"choices": ["cover", "inline", "sidebar"]}), "extra_networks_sidebar_width": OptionInfo(35, "UI sidebar width (%)", gr.Slider, {"minimum": 10, "maximum": 80, "step": 1}), "extra_networks_card_size": OptionInfo(200, "UI card size (px)", gr.Slider, {"minimum": 20, "maximum": 2000, "step": 1}), "extra_networks_card_square": OptionInfo(False, "UI disable variable aspect ratio"), "extra_networks_card_fit": OptionInfo("cover", "UI image contain method", gr.Radio, lambda: {"choices": ["contain", "cover", "fill"]}), + "extra_network_skip_indexing": OptionInfo(False, "Do not automatically build extra network pages", gr.Checkbox), "lyco_patch_lora": OptionInfo(False, "Use LyCoris handler for all Lora types", gr.Checkbox), "lora_disable": OptionInfo(False, "Disable built-in Lora handler", gr.Checkbox, { "visible": True }, onchange=lora_disable), "lora_functional": OptionInfo(False, "Use Kohya method for handling multiple Loras", gr.Checkbox), @@ -734,13 +691,10 @@ opts = Options() config_filename = cmd_opts.config opts.load(config_filename) cmd_opts = cmd_args.compatibility_args(opts, cmd_opts) -if cmd_opts.backend == 'diffusers': - log.info('Overriding backend to Diffusers') - opts.data['sd_backend'] = 'Diffusers' -if cmd_opts.backend == 'original': - log.info('Overriding backend to Diffusers') - opts.data['sd_backend'] = 'Original' -backend = Backend.DIFFUSERS if opts.sd_backend == 'Diffusers' else Backend.ORIGINAL +if cmd_opts.backend: + opts.data['sd_backend'] = cmd_opts.backend.lower() +backend = Backend.DIFFUSERS if opts.sd_backend == 'diffusers' else Backend.ORIGINAL +log.info(f'Pipeline: {cmd_opts.backend.lower()}') prompt_styles = modules.styles.StyleDatabase(opts.styles_dir) cmd_opts.disable_extension_access = (cmd_opts.share or cmd_opts.listen or (cmd_opts.server_name or False)) and not cmd_opts.insecure diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py index 2fdcb03d9..27170eb40 100644 --- a/modules/textual_inversion/textual_inversion.py +++ b/modules/textual_inversion/textual_inversion.py @@ -113,33 +113,52 @@ class EmbeddingDatabase: def register_embedding(self, embedding, model): self.word_embeddings[embedding.name] = embedding - ids = model.cond_stage_model.tokenize([embedding.name])[0] - first_id = ids[0] if first_id not in self.ids_lookup: self.ids_lookup[first_id] = [] - self.ids_lookup[first_id] = sorted(self.ids_lookup[first_id] + [(ids, embedding)], key=lambda x: len(x[0]), reverse=True) - return embedding def get_expected_shape(self): if shared.sd_model is None: shared.log.error('Model not loaded') return 0 + if shared.backend == shared.Backend.DIFFUSERS: + return 0 vec = shared.sd_model.cond_stage_model.encode_embedding_init_text(",", 1) return vec.shape[1] + def load_diffusers_embedding(self, filename: str, path: str): + fn, ext = os.path.splitext(filename) + if ext.lower() != ".pt" and ext.lower() != ".safetensors": + return + pipe = shared.sd_model + if filename == "": + pipe.tokenizer = pipe.tokenizer.__class__.from_pretrained(pipe.tokenizer.name_or_path) + pipe.text_encoder.resize_token_embeddings(len(pipe.tokenizer)) + return + name = os.path.basename(fn) + embedding = Embedding(vec=None, name=name) + try: + pipe.load_textual_inversion(path, cache_dir=shared.opts.data["diffusers_dir"], local_files_only=True) + self.word_embeddings[name] = embedding + except Exception: + self.skipped_embeddings[name] = embedding + text_inv_tokens = pipe.tokenizer.added_tokens_encoder.keys() + text_inv_tokens = [t for t in text_inv_tokens if not (len(t.split("_")) > 1 and t.split("_")[-1].isdigit())] + def load_from_file(self, path, filename): name, ext = os.path.splitext(filename) ext = ext.upper() + if shared.backend == shared.Backend.DIFFUSERS: + self.load_diffusers_embedding(filename, path) + return if ext in ['.PNG', '.WEBP', '.JXL', '.AVIF']: _, second_ext = os.path.splitext(name) if second_ext.upper() == '.PREVIEW': return - embed_image = Image.open(path) if hasattr(embed_image, 'text') and 'sd-ti-embedding' in embed_image.text: data = embedding_from_b64(embed_image.text['sd-ti-embedding']) @@ -206,15 +225,12 @@ class EmbeddingDatabase: continue def load_textual_inversion_embeddings(self, force_reload=False): - if shared.backend == shared.Backend.DIFFUSERS: # TODO Diffusers - return if not force_reload: need_reload = False for embdir in self.embedding_dirs.values(): if embdir.has_changed(): need_reload = True break - if not need_reload: return @@ -241,32 +257,25 @@ class EmbeddingDatabase: def find_embedding_at_position(self, tokens, offset): token = tokens[offset] possible_matches = self.ids_lookup.get(token, None) - if possible_matches is None: return None, None - for ids, embedding in possible_matches: if tokens[offset:offset + len(ids)] == ids: return embedding, len(ids) - return None, None def create_embedding(name, num_vectors_per_token, overwrite_old, init_text='*'): cond_model = shared.sd_model.cond_stage_model - with devices.autocast(): cond_model([""]) # will send cond model to GPU if lowvram/medvram is active - #cond_model expects at least some text, so we provide '*' as backup. embedded = cond_model.encode_embedding_init_text(init_text or '*', num_vectors_per_token) vec = torch.zeros((num_vectors_per_token, embedded.shape[1]), device=devices.device) - #Only copy if we provided an init_text, otherwise keep vectors as zeros if init_text: for i in range(num_vectors_per_token): vec[i] = embedded[i * int(embedded.shape[0]) // num_vectors_per_token] - # Remove illegal characters from name. name = "".join( x for x in name if (x.isalnum() or x in "._- ")) fn = os.path.join(shared.opts.embeddings_dir, f"{name}.pt") @@ -299,29 +308,33 @@ def write_loss(log_directory, filename, step, epoch_len, values): **values, }) + def tensorboard_setup(log_directory): os.makedirs(os.path.join(log_directory, "tensorboard"), exist_ok=True) return SummaryWriter( log_dir=os.path.join(log_directory, "tensorboard"), flush_secs=shared.opts.training_tensorboard_flush_every) + def tensorboard_add(tensorboard_writer, loss, global_step, step, learn_rate, epoch_num): tensorboard_add_scaler(tensorboard_writer, "Loss/train", loss, global_step) tensorboard_add_scaler(tensorboard_writer, f"Loss/train/epoch-{epoch_num}", loss, step) tensorboard_add_scaler(tensorboard_writer, "Learn rate/train", learn_rate, global_step) tensorboard_add_scaler(tensorboard_writer, f"Learn rate/train/epoch-{epoch_num}", learn_rate, step) + def tensorboard_add_scaler(tensorboard_writer, tag, value, step): tensorboard_writer.add_scalar(tag=tag, scalar_value=value, global_step=step) + def tensorboard_add_image(tensorboard_writer, tag, pil_image, step): # Convert a pil image to a torch tensor img_tensor = torch.as_tensor(np.array(pil_image, copy=True)) img_tensor = img_tensor.view(pil_image.size[1], pil_image.size[0], len(pil_image.getbands())) img_tensor = img_tensor.permute((2, 0, 1)) - tensorboard_writer.add_image(tag, img_tensor, global_step=step) + def validate_train_inputs(model_name, learn_rate, batch_size, gradient_step, data_root, template_file, template_filename, steps, save_model_every, create_image_every, log_directory, name="embedding"): assert model_name, f"{name} not selected" assert learn_rate, "Learning rate is empty or 0" @@ -383,15 +396,12 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st images_embeds_dir = None hijack = sd_hijack.model_hijack - embedding = hijack.embedding_db.word_embeddings[embedding_name] checkpoint = sd_models.select_checkpoint() - initial_step = embedding.step or 0 if initial_step >= steps: shared.state.textinfo = "Model has already been trained beyond specified max steps" return embedding, filename - scheduler = LearnRateScheduler(learn_rate, steps, initial_step) clip_grad = torch.nn.utils.clip_grad_value_ if clip_grad_mode == "value" else \ torch.nn.utils.clip_grad_norm_ if clip_grad_mode == "norm" else \ diff --git a/modules/ui.py b/modules/ui.py index 2a27df235..42378732b 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -321,7 +321,7 @@ def create_override_settings_dropdown(tabname, row): # pylint: disable=unused-ar return dropdown -def create_ui(): +def create_ui(startup_timer): import modules.img2img # pylint: disable=redefined-outer-name import modules.txt2img # pylint: disable=redefined-outer-name reload_javascript() @@ -334,7 +334,7 @@ def create_ui(): txt_prompt_img = gr.File(label="", elem_id="txt2img_prompt_image", file_count="single", type="binary", visible=False) with FormRow(variant='compact', elem_id="txt2img_extra_networks", visible=False) as extra_networks_ui: from modules import ui_extra_networks - extra_networks_ui = ui_extra_networks.create_ui(extra_networks_ui, extra_networks_button, 'txt2img') + extra_networks_ui = ui_extra_networks.create_ui(extra_networks_ui, extra_networks_button, 'txt2img', skip_indexing=opts.extra_network_skip_indexing) with gr.Row().style(equal_height=False, elem_id="txt2img_interface"): with gr.Column(variant='compact', elem_id="txt2img_settings"): for category in ordered_ui_categories(): @@ -492,9 +492,10 @@ def create_ui(): ui_extra_networks.setup_ui(extra_networks_ui, txt2img_gallery) + startup_timer.record("ui-txt2img") + modules.scripts.scripts_current = modules.scripts.scripts_img2img modules.scripts.scripts_img2img.initialize_scripts(is_img2img=True) - with gr.Blocks(analytics_enabled=False) as img2img_interface: img2img_prompt, img2img_prompt_styles, img2img_negative_prompt, submit, img2img_interrogate, img2img_deepbooru, img2img_prompt_style_apply, img2img_save_style, img2img_paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button = create_toprow(is_img2img=True) @@ -502,7 +503,7 @@ def create_ui(): with FormRow(variant='compact', elem_id="img2img_extra_networks", visible=False) as extra_networks_ui: from modules import ui_extra_networks - extra_networks_ui_img2img = ui_extra_networks.create_ui(extra_networks_ui, extra_networks_button, 'img2img') + extra_networks_ui_img2img = ui_extra_networks.create_ui(extra_networks_ui, extra_networks_button, 'img2img', skip_indexing=opts.extra_network_skip_indexing) with FormRow().style(equal_height=False, elem_id="img2img_interface"): with gr.Column(variant='compact', elem_id="img2img_settings"): @@ -849,16 +850,21 @@ def create_ui(): paste_button=img2img_paste, tabname="img2img", source_text_component=img2img_prompt, source_image_component=None, )) + startup_timer.record("ui-img2img") + modules.scripts.scripts_current = None with gr.Blocks(analytics_enabled=False) as extras_interface: ui_postprocessing.create_ui() + startup_timer.record("ui-extras") with gr.Blocks(analytics_enabled=False) as train_interface: ui_train.create_ui(txt2img_preview_params = [txt2img_prompt, txt2img_negative_prompt, steps, sampler_index, cfg_scale, seed, width, height]) + startup_timer.record("ui-train") with gr.Blocks(analytics_enabled=False) as models_interface: ui_models.create_ui() + startup_timer.record("ui-models") def create_setting_component(key, is_quicksettings=False): def fun(): @@ -1049,6 +1055,7 @@ def create_ui(): outputs=[dummy_component] ) + startup_timer.record("ui-settings") interfaces = [ (txt2img_interface, "From Text", "txt2img"), @@ -1061,6 +1068,7 @@ def create_ui(): interfaces += [(settings_interface, "Settings", "settings")] extensions_interface = ui_extensions.create_ui() interfaces += [(extensions_interface, "Extensions", "extensions")] + startup_timer.record("ui-extensions") modules.shared.tab_names = [] for _interface, label, _ifid in interfaces: @@ -1136,6 +1144,8 @@ def create_ui(): queue=False, ) + startup_timer.record("ui-defaults") + loadsave.dump_defaults() demo.ui_loadsave = loadsave diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py index ba4e7662e..abd0cd80b 100644 --- a/modules/ui_extra_networks.py +++ b/modules/ui_extra_networks.py @@ -143,7 +143,10 @@ class ExtraNetworksPage: shared.log.info(f"Extra network created thumbnails: {self.name} {created}") self.missing_thumbs.clear() - def create_html(self, tabname): + def create_html(self, tabname, skip = False): + self_name_id = self.name.replace(" ", "_") + if skip: + return f"
Extra network page not ready
Click refresh to try again
" items_html = '' subdirs = {} allowed_folders = [os.path.abspath(x) for x in self.allowed_directories_for_previews()] @@ -151,7 +154,9 @@ class ExtraNetworksPage: for root, dirs, _files in os.walk(parentdir, followlinks=True): for dirname in dirs: x = os.path.join(root, dirname) - if not os.path.isdir(x): + if shared.opts.diffusers_dir in x: + subdirs[os.path.basename(shared.opts.diffusers_dir)] = 1 + if (not os.path.isdir(x)) or ('models--' in x): continue subdir = os.path.abspath(x)[len(parentdir):].replace("\\", "/") while subdir.startswith("/"): @@ -172,22 +177,15 @@ class ExtraNetworksPage: self.metadata[item["name"]] = item.get("metadata", {}) self.info[item["name"]] = self.find_info(item['filename']) items_html += self.create_html_for_item(item, tabname) - # if items_html == '': - # dirs = "".join([f"
  • {x}
    • " for x in self.allowed_directories_for_previews()]) - # items_html = f'
    No models found: {dirs}
    ' - self_name_id = self.name.replace(" ", "_") if len(subdirs_html) > 0 or len(items_html) > 0: - res = f""" -
    {subdirs_html}
    -
    {items_html}
    - """ + res = f"
    {subdirs_html}
    {items_html}
    " else: return '' threading.Thread(target=self.create_thumb).start() return res except Exception as e: shared.log.error(f'Extra networks page error: {e}') - return '' + return f"
    Extra network error
    {e}
    " def list_items(self): raise NotImplementedError @@ -290,7 +288,7 @@ def sort_extra_pages(pages): return sorted(pages, key=lambda x: tab_scores[x.name]) -def create_ui(container, button, tabname): +def create_ui(container, button, tabname, skip_indexing = False): ui = ExtraNetworksUi() ui.pages = [] ui.stored_extra_pages = sort_extra_pages(extra_pages) @@ -308,7 +306,7 @@ def create_ui(container, button, tabname): ui.description_target_filename = gr.Textbox('Description save filename', elem_id=tabname+"_description_filename", visible=False) for page in ui.stored_extra_pages: - page_html = page.create_html(ui.tabname) + page_html = page.create_html(ui.tabname, skip_indexing) if len(page_html) > 0: with gr.Tab(page.title, id=page.title.lower().replace(" ", "_"), elem_classes="extra-networks-tab"): page_elem = gr.HTML(page_html, elem_id=tabname+page.name+"_extra_page", elem_classes="extra-networks-page") diff --git a/modules/ui_extra_networks_checkpoints.py b/modules/ui_extra_networks_checkpoints.py index fe65d99df..37bee332b 100644 --- a/modules/ui_extra_networks_checkpoints.py +++ b/modules/ui_extra_networks_checkpoints.py @@ -21,10 +21,10 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage): "filename": path, "preview": self.find_preview(path), "description": self.find_description(path), - "search_term": self.search_terms_from_path(checkpoint.filename) + " " + (checkpoint.sha256 or ""), + "search_term": f'{self.search_terms_from_path(checkpoint.filename)} {(checkpoint.sha256 or "")} /{checkpoint.type}/', "onclick": '"' + html.escape(f"""return selectCheckpoint({json.dumps(name)})""") + '"', "local_preview": f"{path}.{shared.opts.samples_format}", } def allowed_directories_for_previews(self): - return [v for v in [shared.opts.ckpt_dir, sd_models.model_path] if v is not None] + return [v for v in [shared.opts.ckpt_dir, shared.opts.diffusers_dir, sd_models.model_path] if v is not None] diff --git a/modules/ui_extra_networks_hypernets.py b/modules/ui_extra_networks_hypernets.py index d29863212..01cb22c6c 100644 --- a/modules/ui_extra_networks_hypernets.py +++ b/modules/ui_extra_networks_hypernets.py @@ -19,7 +19,7 @@ class ExtraNetworksPageHypernetworks(ui_extra_networks.ExtraNetworksPage): "preview": self.find_preview(path), "description": self.find_description(path), "search_term": self.search_terms_from_path(path), - "prompt": json.dumps(f""), + "prompt": json.dumps(f""), "local_preview": f"{path}.preview.{shared.opts.samples_format}", } diff --git a/modules/ui_models.py b/modules/ui_models.py index 369a3baec..b7c078a96 100644 --- a/modules/ui_models.py +++ b/modules/ui_models.py @@ -187,11 +187,11 @@ def create_ui(): def hf_select(evt: gr.SelectData): return data[evt.index[0]][0] - def hf_download_model(hub_id: str): + def hf_download_model(hub_id: str, token): from modules.shared import log, opts from modules.modelloader import download_diffusers_model try: - download_diffusers_model(hub_id, cache_dir=opts.diffusers_dir) + download_diffusers_model(hub_id, cache_dir=opts.diffusers_dir, token=token) except Exception as e: log.error(f"Diffuser model downloaded error: model={hub_id} {e}") return f"Diffuser model downloaded error: model={hub_id} {e}" @@ -200,12 +200,14 @@ def create_ui(): log.info(f"Diffuser model downloaded: model={hub_id}") return f'Diffuser model downloaded: model={hub_id}' - with gr.Row(): - hf_search_text = gr.Textbox('', label = 'Seach models', placeholder='search huggingface models') - - with gr.Row(): - hf_selected = gr.Textbox('', label = 'Select model', placeholder='select model from search results or enter model name manually') - with gr.Row(): + with gr.Column(scale=6): + with gr.Row(): + hf_search_text = gr.Textbox('', label = 'Seach models', placeholder='search huggingface models') + with gr.Row(): + hf_selected = gr.Textbox('', label = 'Select model', placeholder='select model from search results or enter model name manually') + with gr.Row(): + hf_token = gr.Textbox('', label = 'Huggingface token', placeholder='optional access token for private or gated models') + with gr.Column(scale=1): hf_download_model_btn = gr.Button(value="Download model", variant='primary') with gr.Row(): @@ -214,10 +216,9 @@ def create_ui(): hf_search_text.submit(fn=hf_search, inputs=[hf_search_text], outputs=[hf_results]) hf_results.select(hf_select, inputs=None, outputs=[hf_selected]) - hf_download_model_btn.click(fn=hf_download_model, inputs=[hf_selected], outputs=[models_outcome]) + hf_download_model_btn.click(fn=hf_download_model, inputs=[hf_selected, hf_token], outputs=[models_outcome]) # TODO load_diffusers_lora - # TODO load_diffusers_text_inv with gr.Tab(label="CivitAI"): pass diff --git a/webui.py b/webui.py index d95c290b8..b78919c47 100644 --- a/webui.py +++ b/webui.py @@ -229,7 +229,7 @@ def start_ui(): log.debug('Creating UI') modules.script_callbacks.before_ui_callback() startup_timer.record("before-ui") - shared.demo = modules.ui.create_ui() + shared.demo = modules.ui.create_ui(startup_timer) startup_timer.record("ui") if cmd_opts.disable_queue: log.info('Server queues disabled')