repair img2img

Add missing k-diffusion sigma params to infotext

.eslintrc.js

@@ -87,5 +87,9 @@ module.exports = {
         modalNextImage: "readonly",
         // token-counters.js
         setupTokenCounters: "readonly",
+        // localStorage.js
+        localSet: "readonly",
+        localGet: "readonly",
+        localRemove: "readonly"
     }
 };

.github/workflows/on_pull_request.yaml

@@ -1,4 +1,4 @@
-name: Run Linting/Formatting on Pull Requests
+name: Linter
 
 on:
   - push
@@ -6,7 +6,9 @@ on:
 
 jobs:
   lint-python:
+    name: ruff
     runs-on: ubuntu-latest
+    if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name != github.event.pull_request.base.repo.full_name
     steps:
       - name: Checkout Code
         uses: actions/checkout@v3
@@ -18,11 +20,13 @@ jobs:
           #     not to have GHA download an (at the time of writing) 4 GB cache
           #     of PyTorch and other dependencies.
       - name: Install Ruff
-        run: pip install ruff==0.0.265
+        run: pip install ruff==0.0.272
       - name: Run Ruff
         run: ruff .
   lint-js:
+    name: eslint
     runs-on: ubuntu-latest
+    if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name != github.event.pull_request.base.repo.full_name
     steps:
       - name: Checkout Code
         uses: actions/checkout@v3

.github/workflows/run_tests.yaml

@@ -1,4 +1,4 @@
-name: Run basic features tests on CPU with empty SD model
+name: Tests
 
 on:
   - push
@@ -6,7 +6,9 @@ on:
 
 jobs:
   test:
+    name: tests on CPU with empty model
     runs-on: ubuntu-latest
+    if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name != github.event.pull_request.base.repo.full_name
     steps:
       - name: Checkout Code
         uses: actions/checkout@v3
@@ -39,10 +41,11 @@ jobs:
           --skip-prepare-environment
           --skip-torch-cuda-test
           --test-server
+          --do-not-download-clip
           --no-half
           --disable-opt-split-attention
           --use-cpu all
-          --add-stop-route
+          --api-server-stop
           2>&1 | tee output.txt &
       - name: Run tests
         run: |
@@ -50,7 +53,7 @@ jobs:
           python -m pytest -vv --junitxml=test/results.xml --cov . --cov-report=xml --verify-base-url test
       - name: Kill test server
         if: always()
-        run: curl -vv -XPOST http://127.0.0.1:7860/_stop && sleep 10
+        run: curl -vv -XPOST http://127.0.0.1:7860/sdapi/v1/server-stop && sleep 10
       - name: Show coverage
         run: |
           python -m coverage combine .coverage*

.github/workflows/warns_merge_master.yml

@@ -0,0 +1,19 @@
+name: Pull requests can't target master branch
+
+"on":
+  pull_request:
+    types:
+      - opened
+      - synchronize
+      - reopened
+    branches:
+      - master
+
+jobs:
+  check:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Warning marge into master
+        run: |
+          echo -e "::warning::This pull request directly merge into \"master\" branch, normally development happens on \"dev\" branch."
+          exit 1

CHANGELOG.md

@@ -1,3 +1,98 @@
+## 1.5.1
+
+### Minor:
+ * support parsing text encoder blocks in some new LoRAs
+ * delete scale checker script due to user demand
+
+### Extensions and API:
+ * add postprocess_batch_list script callback
+
+### Bug Fixes:
+ * fix TI training for SD1
+ * fix reload altclip model error
+ * prepend the pythonpath instead of overriding it
+ * fix typo in SD_WEBUI_RESTARTING
+ * if txt2img/img2img raises an exception, finally call state.end()
+ * fix composable diffusion weight parsing
+ * restyle Startup profile for black users
+ * fix webui not launching with --nowebui
+ * catch exception for non git extensions
+ * fix some options missing from /sdapi/v1/options
+ * fix for extension update status always saying "unknown"
+ * fix display of extra network cards that have `<>` in the name
+ * update lora extension to work with python 3.8
+
+
+## 1.5.0
+
+### Features:
+ * SD XL support
+ * user metadata system for custom networks
+ * extended Lora metadata editor: set activation text, default weight, view tags, training info
+ * Lora extension rework to include other types of networks (all that were previously handled by LyCORIS extension)
+ * show github stars for extenstions
+ * img2img batch mode can read extra stuff from png info
+ * img2img batch works with subdirectories
+ * hotkeys to move prompt elements: alt+left/right
+ * restyle time taken/VRAM display
+ * add textual inversion hashes to infotext
+ * optimization: cache git extension repo information
+ * move generate button next to the generated picture for mobile clients
+ * hide cards for networks of incompatible Stable Diffusion version in Lora extra networks interface
+ * skip installing packages with pip if they all are already installed - startup speedup of about 2 seconds
+
+### Minor:
+ * checkbox to check/uncheck all extensions in the Installed tab
+ * add gradio user to infotext and to filename patterns
+ * allow gif for extra network previews
+ * add options to change colors in grid
+ * use natural sort for items in extra networks
+ * Mac: use empty_cache() from torch 2 to clear VRAM
+ * added automatic support for installing the right libraries for Navi3 (AMD)
+ * add option SWIN_torch_compile to accelerate SwinIR upscale
+ * suppress printing TI embedding info at start to console by default
+ * speedup extra networks listing
+ * added `[none]` filename token.
+ * removed thumbs extra networks view mode (use settings tab to change width/height/scale to get thumbs)
+ * add always_discard_next_to_last_sigma option to XYZ plot
+ * automatically switch to 32-bit float VAE if the generated picture has NaNs without the need for `--no-half-vae` commandline flag.
+ 
+### Extensions and API:
+ * api endpoints: /sdapi/v1/server-kill, /sdapi/v1/server-restart, /sdapi/v1/server-stop
+ * allow Script to have custom metaclass
+ * add model exists status check /sdapi/v1/options
+ * rename --add-stop-route to --api-server-stop
+ * add `before_hr` script callback
+ * add callback `after_extra_networks_activate`
+ * disable rich exception output in console for API by default, use WEBUI_RICH_EXCEPTIONS env var to enable
+ * return http 404 when thumb file not found
+ * allow replacing extensions index with environment variable
+ 
+### Bug Fixes:
+ * fix for catch errors when retrieving extension index #11290
+ * fix very slow loading speed of .safetensors files when reading from network drives
+ * API cache cleanup
+ * fix UnicodeEncodeError when writing to file CLIP Interrogator batch mode
+ * fix warning of 'has_mps' deprecated from PyTorch
+ * fix problem with extra network saving images as previews losing generation info
+ * fix throwing exception when trying to resize image with I;16 mode
+ * fix for #11534: canvas zoom and pan extension hijacking shortcut keys
+ * fixed launch script to be runnable from any directory
+ * don't add "Seed Resize: -1x-1" to API image metadata
+ * correctly remove end parenthesis with ctrl+up/down
+ * fixing --subpath on newer gradio version
+ * fix: check fill size none zero when resize  (fixes #11425)
+ * use submit and blur for quick settings textbox
+ * save img2img batch with images.save_image()
+ * prevent running preload.py for disabled extensions
+ * fix: previously, model name was added together with directory name to infotext and to [model_name] filename pattern; directory name is now not included
+
+
+## 1.4.1
+
+### Bug Fixes:
+ * add queue lock for refresh-checkpoints
+
 ## 1.4.0
 
 ### Features:

README.md

@@ -88,7 +88,7 @@ A browser interface based on Gradio library for Stable Diffusion.
 - [Alt-Diffusion](https://arxiv.org/abs/2211.06679) support - see [wiki](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#alt-diffusion) for instructions
 - Now without any bad letters!
 - Load checkpoints in safetensors format
-- Eased resolution restriction: generated image's domension must be a multiple of 8 rather than 64
+- Eased resolution restriction: generated image's dimension must be a multiple of 8 rather than 64
 - Now with a license!
 - Reorder elements in the UI from settings screen
 
@@ -115,7 +115,7 @@ Alternatively, use online services (like Google Colab):
 1. Install the dependencies:
 ```bash
 # Debian-based:
-sudo apt install wget git python3 python3-venv
+sudo apt install wget git python3 python3-venv libgl1 libglib2.0-0
 # Red Hat-based:
 sudo dnf install wget git python3
 # Arch-based:
@@ -123,7 +123,7 @@ sudo pacman -S wget git python3
 ```
 2. Navigate to the directory you would like the webui to be installed and execute the following command:
 ```bash
-bash <(wget -qO- https://raw.githubusercontent.com/AUTOMATIC1111/stable-diffusion-webui/master/webui.sh)
+wget -q https://raw.githubusercontent.com/AUTOMATIC1111/stable-diffusion-webui/master/webui.sh
 ```
 3. Run `webui.sh`.
 4. Check `webui-user.sh` for options.
@@ -135,8 +135,11 @@ Find the instructions [here](https://github.com/AUTOMATIC1111/stable-diffusion-w
 Here's how to add code to this repo: [Contributing](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Contributing)
 
 ## Documentation
+
 The documentation was moved from this README over to the project's [wiki](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki).
 
+For the purposes of getting Google and other search engines to crawl the wiki, here's a link to the (not for humans) [crawlable wiki](https://github-wiki-see.page/m/AUTOMATIC1111/stable-diffusion-webui/wiki).
+
 ## Credits
 Licenses for borrowed code can be found in `Settings -> Licenses` screen, and also in `html/licenses.html` file.
 
@@ -165,5 +168,7 @@ Licenses for borrowed code can be found in `Settings -> Licenses` screen, and al
 - Security advice - RyotaK
 - UniPC sampler - Wenliang Zhao - https://github.com/wl-zhao/UniPC
 - TAESD - Ollin Boer Bohan - https://github.com/madebyollin/taesd
+- LyCORIS - KohakuBlueleaf
+- Restart sampling - lambertae - https://github.com/Newbeeer/diffusion_restart_sampling
 - Initial Gradio script - posted on 4chan by an Anonymous user. Thank you Anonymous user.
 - (You)

extensions-builtin/LDSR/ldsr_model_arch.py

@@ -12,7 +12,7 @@ import safetensors.torch
 
 from ldm.models.diffusion.ddim import DDIMSampler
 from ldm.util import instantiate_from_config, ismap
-from modules import shared, sd_hijack
+from modules import shared, sd_hijack, devices
 
 cached_ldsr_model: torch.nn.Module = None
 
@@ -112,8 +112,7 @@ class LDSR:
 
 
         gc.collect()
-        if torch.cuda.is_available:
+        devices.torch_gc()
-            torch.cuda.empty_cache()
 
         im_og = image
         width_og, height_og = im_og.size
@@ -150,8 +149,7 @@ class LDSR:
 
         del model
         gc.collect()
-        if torch.cuda.is_available:
+        devices.torch_gc()
-            torch.cuda.empty_cache()
 
         return a
 

extensions-builtin/LDSR/scripts/ldsr_model.py

@@ -1,7 +1,6 @@
 import os
 
-from basicsr.utils.download_util import load_file_from_url
+from modules.modelloader import load_file_from_url
-
 from modules.upscaler import Upscaler, UpscalerData
 from ldsr_model_arch import LDSR
 from modules import shared, script_callbacks, errors
@@ -43,20 +42,17 @@ class UpscalerLDSR(Upscaler):
         if local_safetensors_path is not None and os.path.exists(local_safetensors_path):
             model = local_safetensors_path
         else:
-            model = local_ckpt_path if local_ckpt_path is not None else load_file_from_url(url=self.model_url, model_dir=self.model_download_path, file_name="model.ckpt", progress=True)
+            model = local_ckpt_path or load_file_from_url(self.model_url, model_dir=self.model_download_path, file_name="model.ckpt")
 
-        yaml = local_yaml_path if local_yaml_path is not None else load_file_from_url(url=self.yaml_url, model_dir=self.model_download_path, file_name="project.yaml", progress=True)
+        yaml = local_yaml_path or load_file_from_url(self.yaml_url, model_dir=self.model_download_path, file_name="project.yaml")
 
-        try:
         return LDSR(model, yaml)
-        except Exception:
-            errors.report("Error importing LDSR", exc_info=True)
-        return None
 
     def do_upscale(self, img, path):
+        try:
             ldsr = self.load_model(path)
-        if ldsr is None:
+        except Exception:
-            print("NO LDSR!")
+            errors.report(f"Failed loading LDSR model {path}", exc_info=True)
             return img
         ddim_steps = shared.opts.ldsr_steps
         return ldsr.super_resolution(img, ddim_steps, self.scale)

extensions-builtin/Lora/extra_networks_lora.py

@@ -1,5 +1,5 @@
 from modules import extra_networks, shared
-import lora
+import networks
 
 
 class ExtraNetworkLora(extra_networks.ExtraNetwork):
@@ -9,24 +9,38 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
     def activate(self, p, params_list):
         additional = shared.opts.sd_lora
 
-        if additional != "None" and additional in lora.available_loras and not any(x for x in params_list if x.items[0] == additional):
+        if additional != "None" and additional in networks.available_networks and not any(x for x in params_list if x.items[0] == additional):
             p.all_prompts = [x + f"<lora:{additional}:{shared.opts.extra_networks_default_multiplier}>" for x in p.all_prompts]
             params_list.append(extra_networks.ExtraNetworkParams(items=[additional, shared.opts.extra_networks_default_multiplier]))
 
         names = []
-        multipliers = []
+        te_multipliers = []
+        unet_multipliers = []
+        dyn_dims = []
         for params in params_list:
             assert params.items
 
-            names.append(params.items[0])
+            names.append(params.positional[0])
-            multipliers.append(float(params.items[1]) if len(params.items) > 1 else 1.0)
 
-        lora.load_loras(names, multipliers)
+            te_multiplier = float(params.positional[1]) if len(params.positional) > 1 else 1.0
+            te_multiplier = float(params.named.get("te", te_multiplier))
+
+            unet_multiplier = float(params.positional[2]) if len(params.positional) > 2 else te_multiplier
+            unet_multiplier = float(params.named.get("unet", unet_multiplier))
+
+            dyn_dim = int(params.positional[3]) if len(params.positional) > 3 else None
+            dyn_dim = int(params.named["dyn"]) if "dyn" in params.named else dyn_dim
+
+            te_multipliers.append(te_multiplier)
+            unet_multipliers.append(unet_multiplier)
+            dyn_dims.append(dyn_dim)
+
+        networks.load_networks(names, te_multipliers, unet_multipliers, dyn_dims)
 
         if shared.opts.lora_add_hashes_to_infotext:
-            lora_hashes = []
+            network_hashes = []
-            for item in lora.loaded_loras:
+            for item in networks.loaded_networks:
-                shorthash = item.lora_on_disk.shorthash
+                shorthash = item.network_on_disk.shorthash
                 if not shorthash:
                     continue
 
@@ -36,10 +50,10 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
 
                 alias = alias.replace(":", "").replace(",", "")
 
-                lora_hashes.append(f"{alias}: {shorthash}")
+                network_hashes.append(f"{alias}: {shorthash}")
 
-            if lora_hashes:
+            if network_hashes:
-                p.extra_generation_params["Lora hashes"] = ", ".join(lora_hashes)
+                p.extra_generation_params["Lora hashes"] = ", ".join(network_hashes)
 
     def deactivate(self, p):
         pass

extensions-builtin/Lora/lora.py

@@ -1,506 +1,9 @@
-import os
+import networks
-import re
-import torch
-from typing import Union
 
-from modules import shared, devices, sd_models, errors, scripts, sd_hijack, hashes
+list_available_loras = networks.list_available_networks
 
-metadata_tags_order = {"ss_sd_model_name": 1, "ss_resolution": 2, "ss_clip_skip": 3, "ss_num_train_images": 10, "ss_tag_frequency": 20}
+available_loras = networks.available_networks
-
+available_lora_aliases = networks.available_network_aliases
-re_digits = re.compile(r"\d+")
+available_lora_hash_lookup = networks.available_network_hash_lookup
-re_x_proj = re.compile(r"(.*)_([qkv]_proj)$")
+forbidden_lora_aliases = networks.forbidden_network_aliases
-re_compiled = {}
+loaded_loras = networks.loaded_networks
-
-suffix_conversion = {
-    "attentions": {},
-    "resnets": {
-        "conv1": "in_layers_2",
-        "conv2": "out_layers_3",
-        "time_emb_proj": "emb_layers_1",
-        "conv_shortcut": "skip_connection",
-    }
-}
-
-
-def convert_diffusers_name_to_compvis(key, is_sd2):
-    def match(match_list, regex_text):
-        regex = re_compiled.get(regex_text)
-        if regex is None:
-            regex = re.compile(regex_text)
-            re_compiled[regex_text] = regex
-
-        r = re.match(regex, key)
-        if not r:
-            return False
-
-        match_list.clear()
-        match_list.extend([int(x) if re.match(re_digits, x) else x for x in r.groups()])
-        return True
-
-    m = []
-
-    if match(m, r"lora_unet_down_blocks_(\d+)_(attentions|resnets)_(\d+)_(.+)"):
-        suffix = suffix_conversion.get(m[1], {}).get(m[3], m[3])
-        return f"diffusion_model_input_blocks_{1 + m[0] * 3 + m[2]}_{1 if m[1] == 'attentions' else 0}_{suffix}"
-
-    if match(m, r"lora_unet_mid_block_(attentions|resnets)_(\d+)_(.+)"):
-        suffix = suffix_conversion.get(m[0], {}).get(m[2], m[2])
-        return f"diffusion_model_middle_block_{1 if m[0] == 'attentions' else m[1] * 2}_{suffix}"
-
-    if match(m, r"lora_unet_up_blocks_(\d+)_(attentions|resnets)_(\d+)_(.+)"):
-        suffix = suffix_conversion.get(m[1], {}).get(m[3], m[3])
-        return f"diffusion_model_output_blocks_{m[0] * 3 + m[2]}_{1 if m[1] == 'attentions' else 0}_{suffix}"
-
-    if match(m, r"lora_unet_down_blocks_(\d+)_downsamplers_0_conv"):
-        return f"diffusion_model_input_blocks_{3 + m[0] * 3}_0_op"
-
-    if match(m, r"lora_unet_up_blocks_(\d+)_upsamplers_0_conv"):
-        return f"diffusion_model_output_blocks_{2 + m[0] * 3}_{2 if m[0]>0 else 1}_conv"
-
-    if match(m, r"lora_te_text_model_encoder_layers_(\d+)_(.+)"):
-        if is_sd2:
-            if 'mlp_fc1' in m[1]:
-                return f"model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc1', 'mlp_c_fc')}"
-            elif 'mlp_fc2' in m[1]:
-                return f"model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc2', 'mlp_c_proj')}"
-            else:
-                return f"model_transformer_resblocks_{m[0]}_{m[1].replace('self_attn', 'attn')}"
-
-        return f"transformer_text_model_encoder_layers_{m[0]}_{m[1]}"
-
-    return key
-
-
-class LoraOnDisk:
-    def __init__(self, name, filename):
-        self.name = name
-        self.filename = filename
-        self.metadata = {}
-        self.is_safetensors = os.path.splitext(filename)[1].lower() == ".safetensors"
-
-        if self.is_safetensors:
-            try:
-                self.metadata = sd_models.read_metadata_from_safetensors(filename)
-            except Exception as e:
-                errors.display(e, f"reading lora {filename}")
-
-        if self.metadata:
-            m = {}
-            for k, v in sorted(self.metadata.items(), key=lambda x: metadata_tags_order.get(x[0], 999)):
-                m[k] = v
-
-            self.metadata = m
-
-        self.ssmd_cover_images = self.metadata.pop('ssmd_cover_images', None)  # those are cover images and they are too big to display in UI as text
-        self.alias = self.metadata.get('ss_output_name', self.name)
-
-        self.hash = None
-        self.shorthash = None
-        self.set_hash(
-            self.metadata.get('sshs_model_hash') or
-            hashes.sha256_from_cache(self.filename, "lora/" + self.name, use_addnet_hash=self.is_safetensors) or
-            ''
-        )
-
-    def set_hash(self, v):
-        self.hash = v
-        self.shorthash = self.hash[0:12]
-
-        if self.shorthash:
-            available_lora_hash_lookup[self.shorthash] = self
-
-    def read_hash(self):
-        if not self.hash:
-            self.set_hash(hashes.sha256(self.filename, "lora/" + self.name, use_addnet_hash=self.is_safetensors) or '')
-
-    def get_alias(self):
-        if shared.opts.lora_preferred_name == "Filename" or self.alias.lower() in forbidden_lora_aliases:
-            return self.name
-        else:
-            return self.alias
-
-
-class LoraModule:
-    def __init__(self, name, lora_on_disk: LoraOnDisk):
-        self.name = name
-        self.lora_on_disk = lora_on_disk
-        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"""
-
-
-class LoraUpDownModule:
-    def __init__(self):
-        self.up = None
-        self.down = None
-        self.alpha = None
-
-
-def assign_lora_names_to_compvis_modules(sd_model):
-    lora_layer_mapping = {}
-
-    for name, module in shared.sd_model.cond_stage_model.wrapped.named_modules():
-        lora_name = name.replace(".", "_")
-        lora_layer_mapping[lora_name] = module
-        module.lora_layer_name = lora_name
-
-    for name, module in shared.sd_model.model.named_modules():
-        lora_name = name.replace(".", "_")
-        lora_layer_mapping[lora_name] = module
-        module.lora_layer_name = lora_name
-
-    sd_model.lora_layer_mapping = lora_layer_mapping
-
-
-def load_lora(name, lora_on_disk):
-    lora = LoraModule(name, lora_on_disk)
-    lora.mtime = os.path.getmtime(lora_on_disk.filename)
-
-    sd = sd_models.read_state_dict(lora_on_disk.filename)
-
-    # this should not be needed but is here as an emergency fix for an unknown error people are experiencing in 1.2.0
-    if not hasattr(shared.sd_model, 'lora_layer_mapping'):
-        assign_lora_names_to_compvis_modules(shared.sd_model)
-
-    keys_failed_to_match = {}
-    is_sd2 = 'model_transformer_resblocks' in shared.sd_model.lora_layer_mapping
-
-    for key_diffusers, weight in sd.items():
-        key_diffusers_without_lora_parts, lora_key = key_diffusers.split(".", 1)
-        key = convert_diffusers_name_to_compvis(key_diffusers_without_lora_parts, is_sd2)
-
-        sd_module = shared.sd_model.lora_layer_mapping.get(key, None)
-
-        if sd_module is None:
-            m = re_x_proj.match(key)
-            if m:
-                sd_module = shared.sd_model.lora_layer_mapping.get(m.group(1), None)
-
-        if sd_module is None:
-            keys_failed_to_match[key_diffusers] = key
-            continue
-
-        lora_module = lora.modules.get(key, None)
-        if lora_module is None:
-            lora_module = LoraUpDownModule()
-            lora.modules[key] = lora_module
-
-        if lora_key == "alpha":
-            lora_module.alpha = weight.item()
-            continue
-
-        if type(sd_module) == torch.nn.Linear:
-            module = torch.nn.Linear(weight.shape[1], weight.shape[0], bias=False)
-        elif type(sd_module) == torch.nn.modules.linear.NonDynamicallyQuantizableLinear:
-            module = torch.nn.Linear(weight.shape[1], weight.shape[0], bias=False)
-        elif type(sd_module) == torch.nn.MultiheadAttention:
-            module = torch.nn.Linear(weight.shape[1], weight.shape[0], bias=False)
-        elif type(sd_module) == torch.nn.Conv2d and weight.shape[2:] == (1, 1):
-            module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], (1, 1), bias=False)
-        elif type(sd_module) == torch.nn.Conv2d and weight.shape[2:] == (3, 3):
-            module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], (3, 3), bias=False)
-        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)
-
-        module.to(device=devices.cpu, dtype=devices.dtype)
-
-        if lora_key == "lora_up.weight":
-            lora_module.up = module
-        elif lora_key == "lora_down.weight":
-            lora_module.down = module
-        else:
-            raise AssertionError(f"Bad Lora layer name: {key_diffusers} - must end in lora_up.weight, lora_down.weight or alpha")
-
-    if keys_failed_to_match:
-        print(f"Failed to match keys when loading Lora {lora_on_disk.filename}: {keys_failed_to_match}")
-
-    return lora
-
-
-def load_loras(names, multipliers=None):
-    already_loaded = {}
-
-    for lora in loaded_loras:
-        if lora.name in names:
-            already_loaded[lora.name] = lora
-
-    loaded_loras.clear()
-
-    loras_on_disk = [available_lora_aliases.get(name, None) for name in names]
-    if any(x is None for x in loras_on_disk):
-        list_available_loras()
-
-        loras_on_disk = [available_lora_aliases.get(name, None) for name in names]
-
-    failed_to_load_loras = []
-
-    for i, name in enumerate(names):
-        lora = already_loaded.get(name, 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)
-                except Exception as e:
-                    errors.display(e, f"loading Lora {lora_on_disk.filename}")
-                    continue
-
-            lora.mentioned_name = name
-
-            lora_on_disk.read_hash()
-
-        if lora is None:
-            failed_to_load_loras.append(name)
-            print(f"Couldn't find Lora with name {name}")
-            continue
-
-        lora.multiplier = multipliers[i] if multipliers else 1.0
-        loaded_loras.append(lora)
-
-    if failed_to_load_loras:
-        sd_hijack.model_hijack.comments.append("Failed to find Loras: " + ", ".join(failed_to_load_loras))
-
-
-def lora_calc_updown(lora, module, target):
-    with torch.no_grad():
-        up = module.up.weight.to(target.device, dtype=target.dtype)
-        down = module.down.weight.to(target.device, dtype=target.dtype)
-
-        if up.shape[2:] == (1, 1) and down.shape[2:] == (1, 1):
-            updown = (up.squeeze(2).squeeze(2) @ down.squeeze(2).squeeze(2)).unsqueeze(2).unsqueeze(3)
-        elif up.shape[2:] == (3, 3) or down.shape[2:] == (3, 3):
-            updown = torch.nn.functional.conv2d(down.permute(1, 0, 2, 3), up).permute(1, 0, 2, 3)
-        else:
-            updown = up @ down
-
-        updown = updown * lora.multiplier * (module.alpha / module.up.weight.shape[1] if module.alpha else 1.0)
-
-        return updown
-
-
-def lora_restore_weights_from_backup(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.MultiheadAttention]):
-    weights_backup = getattr(self, "lora_weights_backup", None)
-
-    if weights_backup is None:
-        return
-
-    if isinstance(self, torch.nn.MultiheadAttention):
-        self.in_proj_weight.copy_(weights_backup[0])
-        self.out_proj.weight.copy_(weights_backup[1])
-    else:
-        self.weight.copy_(weights_backup)
-
-
-def lora_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.MultiheadAttention]):
-    """
-    Applies the currently selected set of Loras to the weights of torch layer self.
-    If weights already have this particular set of loras applied, does nothing.
-    If not, restores orginal weights from backup and alters weights according to loras.
-    """
-
-    lora_layer_name = getattr(self, 'lora_layer_name', None)
-    if lora_layer_name is None:
-        return
-
-    current_names = getattr(self, "lora_current_names", ())
-    wanted_names = tuple((x.name, x.multiplier) for x in loaded_loras)
-
-    weights_backup = getattr(self, "lora_weights_backup", None)
-    if weights_backup is None:
-        if isinstance(self, torch.nn.MultiheadAttention):
-            weights_backup = (self.in_proj_weight.to(devices.cpu, copy=True), self.out_proj.weight.to(devices.cpu, copy=True))
-        else:
-            weights_backup = self.weight.to(devices.cpu, copy=True)
-
-        self.lora_weights_backup = weights_backup
-
-    if current_names != wanted_names:
-        lora_restore_weights_from_backup(self)
-
-        for lora in loaded_loras:
-            module = lora.modules.get(lora_layer_name, None)
-            if module is not None and hasattr(self, 'weight'):
-                self.weight += lora_calc_updown(lora, module, self.weight)
-                continue
-
-            module_q = lora.modules.get(lora_layer_name + "_q_proj", None)
-            module_k = lora.modules.get(lora_layer_name + "_k_proj", None)
-            module_v = lora.modules.get(lora_layer_name + "_v_proj", None)
-            module_out = lora.modules.get(lora_layer_name + "_out_proj", None)
-
-            if isinstance(self, torch.nn.MultiheadAttention) and module_q and module_k and module_v and module_out:
-                updown_q = lora_calc_updown(lora, module_q, self.in_proj_weight)
-                updown_k = lora_calc_updown(lora, module_k, self.in_proj_weight)
-                updown_v = lora_calc_updown(lora, module_v, self.in_proj_weight)
-                updown_qkv = torch.vstack([updown_q, updown_k, updown_v])
-
-                self.in_proj_weight += updown_qkv
-                self.out_proj.weight += lora_calc_updown(lora, module_out, self.out_proj.weight)
-                continue
-
-            if module is None:
-                continue
-
-            print(f'failed to calculate lora weights for layer {lora_layer_name}')
-
-        self.lora_current_names = wanted_names
-
-
-def lora_forward(module, input, original_forward):
-    """
-    Old way of applying Lora by executing operations during layer's forward.
-    Stacking many loras this way results in big performance degradation.
-    """
-
-    if len(loaded_loras) == 0:
-        return original_forward(module, input)
-
-    input = devices.cond_cast_unet(input)
-
-    lora_restore_weights_from_backup(module)
-    lora_reset_cached_weight(module)
-
-    res = original_forward(module, input)
-
-    lora_layer_name = getattr(module, 'lora_layer_name', None)
-    for lora in loaded_loras:
-        module = lora.modules.get(lora_layer_name, None)
-        if module is None:
-            continue
-
-        module.up.to(device=devices.device)
-        module.down.to(device=devices.device)
-
-        res = res + module.up(module.down(input)) * lora.multiplier * (module.alpha / module.up.weight.shape[1] if module.alpha else 1.0)
-
-    return res
-
-
-def lora_reset_cached_weight(self: Union[torch.nn.Conv2d, torch.nn.Linear]):
-    self.lora_current_names = ()
-    self.lora_weights_backup = None
-
-
-def lora_Linear_forward(self, input):
-    if shared.opts.lora_functional:
-        return lora_forward(self, input, torch.nn.Linear_forward_before_lora)
-
-    lora_apply_weights(self)
-
-    return torch.nn.Linear_forward_before_lora(self, input)
-
-
-def lora_Linear_load_state_dict(self, *args, **kwargs):
-    lora_reset_cached_weight(self)
-
-    return torch.nn.Linear_load_state_dict_before_lora(self, *args, **kwargs)
-
-
-def lora_Conv2d_forward(self, input):
-    if shared.opts.lora_functional:
-        return lora_forward(self, input, torch.nn.Conv2d_forward_before_lora)
-
-    lora_apply_weights(self)
-
-    return torch.nn.Conv2d_forward_before_lora(self, input)
-
-
-def lora_Conv2d_load_state_dict(self, *args, **kwargs):
-    lora_reset_cached_weight(self)
-
-    return torch.nn.Conv2d_load_state_dict_before_lora(self, *args, **kwargs)
-
-
-def lora_MultiheadAttention_forward(self, *args, **kwargs):
-    lora_apply_weights(self)
-
-    return torch.nn.MultiheadAttention_forward_before_lora(self, *args, **kwargs)
-
-
-def lora_MultiheadAttention_load_state_dict(self, *args, **kwargs):
-    lora_reset_cached_weight(self)
-
-    return torch.nn.MultiheadAttention_load_state_dict_before_lora(self, *args, **kwargs)
-
-
-def list_available_loras():
-    available_loras.clear()
-    available_lora_aliases.clear()
-    forbidden_lora_aliases.clear()
-    available_lora_hash_lookup.clear()
-    forbidden_lora_aliases.update({"none": 1, "Addams": 1})
-
-    os.makedirs(shared.cmd_opts.lora_dir, exist_ok=True)
-
-    candidates = list(shared.walk_files(shared.cmd_opts.lora_dir, allowed_extensions=[".pt", ".ckpt", ".safetensors"]))
-    for filename in sorted(candidates, key=str.lower):
-        if os.path.isdir(filename):
-            continue
-
-        name = os.path.splitext(os.path.basename(filename))[0]
-        try:
-            entry = LoraOnDisk(name, filename)
-        except OSError:  # should catch FileNotFoundError and PermissionError etc.
-            errors.report(f"Failed to load LoRA {name} from {filename}", exc_info=True)
-            continue
-
-        available_loras[name] = entry
-
-        if entry.alias in available_lora_aliases:
-            forbidden_lora_aliases[entry.alias.lower()] = 1
-
-        available_lora_aliases[name] = entry
-        available_lora_aliases[entry.alias] = entry
-
-
-re_lora_name = re.compile(r"(.*)\s*\([0-9a-fA-F]+\)")
-
-
-def infotext_pasted(infotext, params):
-    if "AddNet Module 1" in [x[1] for x in scripts.scripts_txt2img.infotext_fields]:
-        return  # if the other extension is active, it will handle those fields, no need to do anything
-
-    added = []
-
-    for k in params:
-        if not k.startswith("AddNet Model "):
-            continue
-
-        num = k[13:]
-
-        if params.get("AddNet Module " + num) != "LoRA":
-            continue
-
-        name = params.get("AddNet Model " + num)
-        if name is None:
-            continue
-
-        m = re_lora_name.match(name)
-        if m:
-            name = m.group(1)
-
-        multiplier = params.get("AddNet Weight A " + num, "1.0")
-
-        added.append(f"<lora:{name}:{multiplier}>")
-
-    if added:
-        params["Prompt"] += "\n" + "".join(added)
-
-
-available_loras = {}
-available_lora_aliases = {}
-available_lora_hash_lookup = {}
-forbidden_lora_aliases = {}
-loaded_loras = []
-
-list_available_loras()

extensions-builtin/Lora/lyco_helpers.py

@@ -0,0 +1,21 @@
+import torch
+
+
+def make_weight_cp(t, wa, wb):
+    temp = torch.einsum('i j k l, j r -> i r k l', t, wb)
+    return torch.einsum('i j k l, i r -> r j k l', temp, wa)
+
+
+def rebuild_conventional(up, down, shape, dyn_dim=None):
+    up = up.reshape(up.size(0), -1)
+    down = down.reshape(down.size(0), -1)
+    if dyn_dim is not None:
+        up = up[:, :dyn_dim]
+        down = down[:dyn_dim, :]
+    return (up @ down).reshape(shape)
+
+
+def rebuild_cp_decomposition(up, down, mid):
+    up = up.reshape(up.size(0), -1)
+    down = down.reshape(down.size(0), -1)
+    return torch.einsum('n m k l, i n, m j -> i j k l', mid, up, down)

extensions-builtin/Lora/network.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+import os
+from collections import namedtuple
+import enum
+
+from modules import sd_models, cache, errors, hashes, shared
+
+NetworkWeights = namedtuple('NetworkWeights', ['network_key', 'sd_key', 'w', 'sd_module'])
+
+metadata_tags_order = {"ss_sd_model_name": 1, "ss_resolution": 2, "ss_clip_skip": 3, "ss_num_train_images": 10, "ss_tag_frequency": 20}
+
+
+class SdVersion(enum.Enum):
+    Unknown = 1
+    SD1 = 2
+    SD2 = 3
+    SDXL = 4
+
+
+class NetworkOnDisk:
+    def __init__(self, name, filename):
+        self.name = name
+        self.filename = filename
+        self.metadata = {}
+        self.is_safetensors = os.path.splitext(filename)[1].lower() == ".safetensors"
+
+        def read_metadata():
+            metadata = sd_models.read_metadata_from_safetensors(filename)
+            metadata.pop('ssmd_cover_images', None)  # those are cover images, and they are too big to display in UI as text
+
+            return metadata
+
+        if self.is_safetensors:
+            try:
+                self.metadata = cache.cached_data_for_file('safetensors-metadata', "lora/" + self.name, filename, read_metadata)
+            except Exception as e:
+                errors.display(e, f"reading lora {filename}")
+
+        if self.metadata:
+            m = {}
+            for k, v in sorted(self.metadata.items(), key=lambda x: metadata_tags_order.get(x[0], 999)):
+                m[k] = v
+
+            self.metadata = m
+
+        self.alias = self.metadata.get('ss_output_name', self.name)
+
+        self.hash = None
+        self.shorthash = None
+        self.set_hash(
+            self.metadata.get('sshs_model_hash') or
+            hashes.sha256_from_cache(self.filename, "lora/" + self.name, use_addnet_hash=self.is_safetensors) or
+            ''
+        )
+
+        self.sd_version = self.detect_version()
+
+    def detect_version(self):
+        if str(self.metadata.get('ss_base_model_version', "")).startswith("sdxl_"):
+            return SdVersion.SDXL
+        elif str(self.metadata.get('ss_v2', "")) == "True":
+            return SdVersion.SD2
+        elif len(self.metadata):
+            return SdVersion.SD1
+
+        return SdVersion.Unknown
+
+    def set_hash(self, v):
+        self.hash = v
+        self.shorthash = self.hash[0:12]
+
+        if self.shorthash:
+            import networks
+            networks.available_network_hash_lookup[self.shorthash] = self
+
+    def read_hash(self):
+        if not self.hash:
+            self.set_hash(hashes.sha256(self.filename, "lora/" + self.name, use_addnet_hash=self.is_safetensors) or '')
+
+    def get_alias(self):
+        import networks
+        if shared.opts.lora_preferred_name == "Filename" or self.alias.lower() in networks.forbidden_network_aliases:
+            return self.name
+        else:
+            return self.alias
+
+
+class Network:  # LoraModule
+    def __init__(self, name, network_on_disk: NetworkOnDisk):
+        self.name = name
+        self.network_on_disk = network_on_disk
+        self.te_multiplier = 1.0
+        self.unet_multiplier = 1.0
+        self.dyn_dim = None
+        self.modules = {}
+        self.mtime = None
+
+        self.mentioned_name = None
+        """the text that was used to add the network to prompt - can be either name or an alias"""
+
+
+class ModuleType:
+    def create_module(self, net: Network, weights: NetworkWeights) -> Network | None:
+        return None
+
+
+class NetworkModule:
+    def __init__(self, net: Network, weights: NetworkWeights):
+        self.network = net
+        self.network_key = weights.network_key
+        self.sd_key = weights.sd_key
+        self.sd_module = weights.sd_module
+
+        if hasattr(self.sd_module, 'weight'):
+            self.shape = self.sd_module.weight.shape
+
+        self.dim = None
+        self.bias = weights.w.get("bias")
+        self.alpha = weights.w["alpha"].item() if "alpha" in weights.w else None
+        self.scale = weights.w["scale"].item() if "scale" in weights.w else None
+
+    def multiplier(self):
+        if 'transformer' in self.sd_key[:20]:
+            return self.network.te_multiplier
+        else:
+            return self.network.unet_multiplier
+
+    def calc_scale(self):
+        if self.scale is not None:
+            return self.scale
+        if self.dim is not None and self.alpha is not None:
+            return self.alpha / self.dim
+
+        return 1.0
+
+    def finalize_updown(self, updown, orig_weight, output_shape):
+        if self.bias is not None:
+            updown = updown.reshape(self.bias.shape)
+            updown += self.bias.to(orig_weight.device, dtype=orig_weight.dtype)
+            updown = updown.reshape(output_shape)
+
+        if len(output_shape) == 4:
+            updown = updown.reshape(output_shape)
+
+        if orig_weight.size().numel() == updown.size().numel():
+            updown = updown.reshape(orig_weight.shape)
+
+        return updown * self.calc_scale() * self.multiplier()
+
+    def calc_updown(self, target):
+        raise NotImplementedError()
+
+    def forward(self, x, y):
+        raise NotImplementedError()
+

extensions-builtin/Lora/network_full.py

@@ -0,0 +1,22 @@
+import network
+
+
+class ModuleTypeFull(network.ModuleType):
+    def create_module(self, net: network.Network, weights: network.NetworkWeights):
+        if all(x in weights.w for x in ["diff"]):
+            return NetworkModuleFull(net, weights)
+
+        return None
+
+
+class NetworkModuleFull(network.NetworkModule):
+    def __init__(self,  net: network.Network, weights: network.NetworkWeights):
+        super().__init__(net, weights)
+
+        self.weight = weights.w.get("diff")
+
+    def calc_updown(self, orig_weight):
+        output_shape = self.weight.shape
+        updown = self.weight.to(orig_weight.device, dtype=orig_weight.dtype)
+
+        return self.finalize_updown(updown, orig_weight, output_shape)

extensions-builtin/Lora/network_hada.py

@@ -0,0 +1,55 @@
+import lyco_helpers
+import network
+
+
+class ModuleTypeHada(network.ModuleType):
+    def create_module(self, net: network.Network, weights: network.NetworkWeights):
+        if all(x in weights.w for x in ["hada_w1_a", "hada_w1_b", "hada_w2_a", "hada_w2_b"]):
+            return NetworkModuleHada(net, weights)
+
+        return None
+
+
+class NetworkModuleHada(network.NetworkModule):
+    def __init__(self,  net: network.Network, weights: network.NetworkWeights):
+        super().__init__(net, weights)
+
+        if hasattr(self.sd_module, 'weight'):
+            self.shape = self.sd_module.weight.shape
+
+        self.w1a = weights.w["hada_w1_a"]
+        self.w1b = weights.w["hada_w1_b"]
+        self.dim = self.w1b.shape[0]
+        self.w2a = weights.w["hada_w2_a"]
+        self.w2b = weights.w["hada_w2_b"]
+
+        self.t1 = weights.w.get("hada_t1")
+        self.t2 = weights.w.get("hada_t2")
+
+    def calc_updown(self, orig_weight):
+        w1a = self.w1a.to(orig_weight.device, dtype=orig_weight.dtype)
+        w1b = self.w1b.to(orig_weight.device, dtype=orig_weight.dtype)
+        w2a = self.w2a.to(orig_weight.device, dtype=orig_weight.dtype)
+        w2b = self.w2b.to(orig_weight.device, dtype=orig_weight.dtype)
+
+        output_shape = [w1a.size(0), w1b.size(1)]
+
+        if self.t1 is not None:
+            output_shape = [w1a.size(1), w1b.size(1)]
+            t1 = self.t1.to(orig_weight.device, dtype=orig_weight.dtype)
+            updown1 = lyco_helpers.make_weight_cp(t1, w1a, w1b)
+            output_shape += t1.shape[2:]
+        else:
+            if len(w1b.shape) == 4:
+                output_shape += w1b.shape[2:]
+            updown1 = lyco_helpers.rebuild_conventional(w1a, w1b, output_shape)
+
+        if self.t2 is not None:
+            t2 = self.t2.to(orig_weight.device, dtype=orig_weight.dtype)
+            updown2 = lyco_helpers.make_weight_cp(t2, w2a, w2b)
+        else:
+            updown2 = lyco_helpers.rebuild_conventional(w2a, w2b, output_shape)
+
+        updown = updown1 * updown2
+
+        return self.finalize_updown(updown, orig_weight, output_shape)

extensions-builtin/Lora/network_ia3.py

@@ -0,0 +1,30 @@
+import network
+
+
+class ModuleTypeIa3(network.ModuleType):
+    def create_module(self, net: network.Network, weights: network.NetworkWeights):
+        if all(x in weights.w for x in ["weight"]):
+            return NetworkModuleIa3(net, weights)
+
+        return None
+
+
+class NetworkModuleIa3(network.NetworkModule):
+    def __init__(self,  net: network.Network, weights: network.NetworkWeights):
+        super().__init__(net, weights)
+
+        self.w = weights.w["weight"]
+        self.on_input = weights.w["on_input"].item()
+
+    def calc_updown(self, orig_weight):
+        w = self.w.to(orig_weight.device, dtype=orig_weight.dtype)
+
+        output_shape = [w.size(0), orig_weight.size(1)]
+        if self.on_input:
+            output_shape.reverse()
+        else:
+            w = w.reshape(-1, 1)
+
+        updown = orig_weight * w
+
+        return self.finalize_updown(updown, orig_weight, output_shape)

extensions-builtin/Lora/network_lokr.py

@@ -0,0 +1,64 @@
+import torch
+
+import lyco_helpers
+import network
+
+
+class ModuleTypeLokr(network.ModuleType):
+    def create_module(self, net: network.Network, weights: network.NetworkWeights):
+        has_1 = "lokr_w1" in weights.w or ("lokr_w1_a" in weights.w and "lokr_w1_b" in weights.w)
+        has_2 = "lokr_w2" in weights.w or ("lokr_w2_a" in weights.w and "lokr_w2_b" in weights.w)
+        if has_1 and has_2:
+            return NetworkModuleLokr(net, weights)
+
+        return None
+
+
+def make_kron(orig_shape, w1, w2):
+    if len(w2.shape) == 4:
+        w1 = w1.unsqueeze(2).unsqueeze(2)
+    w2 = w2.contiguous()
+    return torch.kron(w1, w2).reshape(orig_shape)
+
+
+class NetworkModuleLokr(network.NetworkModule):
+    def __init__(self,  net: network.Network, weights: network.NetworkWeights):
+        super().__init__(net, weights)
+
+        self.w1 = weights.w.get("lokr_w1")
+        self.w1a = weights.w.get("lokr_w1_a")
+        self.w1b = weights.w.get("lokr_w1_b")
+        self.dim = self.w1b.shape[0] if self.w1b is not None else self.dim
+        self.w2 = weights.w.get("lokr_w2")
+        self.w2a = weights.w.get("lokr_w2_a")
+        self.w2b = weights.w.get("lokr_w2_b")
+        self.dim = self.w2b.shape[0] if self.w2b is not None else self.dim
+        self.t2 = weights.w.get("lokr_t2")
+
+    def calc_updown(self, orig_weight):
+        if self.w1 is not None:
+            w1 = self.w1.to(orig_weight.device, dtype=orig_weight.dtype)
+        else:
+            w1a = self.w1a.to(orig_weight.device, dtype=orig_weight.dtype)
+            w1b = self.w1b.to(orig_weight.device, dtype=orig_weight.dtype)
+            w1 = w1a @ w1b
+
+        if self.w2 is not None:
+            w2 = self.w2.to(orig_weight.device, dtype=orig_weight.dtype)
+        elif self.t2 is None:
+            w2a = self.w2a.to(orig_weight.device, dtype=orig_weight.dtype)
+            w2b = self.w2b.to(orig_weight.device, dtype=orig_weight.dtype)
+            w2 = w2a @ w2b
+        else:
+            t2 = self.t2.to(orig_weight.device, dtype=orig_weight.dtype)
+            w2a = self.w2a.to(orig_weight.device, dtype=orig_weight.dtype)
+            w2b = self.w2b.to(orig_weight.device, dtype=orig_weight.dtype)
+            w2 = lyco_helpers.make_weight_cp(t2, w2a, w2b)
+
+        output_shape = [w1.size(0) * w2.size(0), w1.size(1) * w2.size(1)]
+        if len(orig_weight.shape) == 4:
+            output_shape = orig_weight.shape
+
+        updown = make_kron(output_shape, w1, w2)
+
+        return self.finalize_updown(updown, orig_weight, output_shape)

extensions-builtin/Lora/network_lora.py

@@ -0,0 +1,86 @@
+import torch
+
+import lyco_helpers
+import network
+from modules import devices
+
+
+class ModuleTypeLora(network.ModuleType):
+    def create_module(self, net: network.Network, weights: network.NetworkWeights):
+        if all(x in weights.w for x in ["lora_up.weight", "lora_down.weight"]):
+            return NetworkModuleLora(net, weights)
+
+        return None
+
+
+class NetworkModuleLora(network.NetworkModule):
+    def __init__(self,  net: network.Network, weights: network.NetworkWeights):
+        super().__init__(net, weights)
+
+        self.up_model = self.create_module(weights.w, "lora_up.weight")
+        self.down_model = self.create_module(weights.w, "lora_down.weight")
+        self.mid_model = self.create_module(weights.w, "lora_mid.weight", none_ok=True)
+
+        self.dim = weights.w["lora_down.weight"].shape[0]
+
+    def create_module(self, weights, key, none_ok=False):
+        weight = weights.get(key)
+
+        if weight is None and none_ok:
+            return None
+
+        is_linear = type(self.sd_module) in [torch.nn.Linear, torch.nn.modules.linear.NonDynamicallyQuantizableLinear, torch.nn.MultiheadAttention]
+        is_conv = type(self.sd_module) in [torch.nn.Conv2d]
+
+        if is_linear:
+            weight = weight.reshape(weight.shape[0], -1)
+            module = torch.nn.Linear(weight.shape[1], weight.shape[0], bias=False)
+        elif is_conv and key == "lora_down.weight" or key == "dyn_up":
+            if len(weight.shape) == 2:
+                weight = weight.reshape(weight.shape[0], -1, 1, 1)
+
+            if weight.shape[2] != 1 or weight.shape[3] != 1:
+                module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], self.sd_module.kernel_size, self.sd_module.stride, self.sd_module.padding, bias=False)
+            else:
+                module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], (1, 1), bias=False)
+        elif is_conv and key == "lora_mid.weight":
+            module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], self.sd_module.kernel_size, self.sd_module.stride, self.sd_module.padding, bias=False)
+        elif is_conv and key == "lora_up.weight" or key == "dyn_down":
+            module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], (1, 1), bias=False)
+        else:
+            raise AssertionError(f'Lora layer {self.network_key} matched a layer with unsupported type: {type(self.sd_module).__name__}')
+
+        with torch.no_grad():
+            if weight.shape != module.weight.shape:
+                weight = weight.reshape(module.weight.shape)
+            module.weight.copy_(weight)
+
+        module.to(device=devices.cpu, dtype=devices.dtype)
+        module.weight.requires_grad_(False)
+
+        return module
+
+    def calc_updown(self, orig_weight):
+        up = self.up_model.weight.to(orig_weight.device, dtype=orig_weight.dtype)
+        down = self.down_model.weight.to(orig_weight.device, dtype=orig_weight.dtype)
+
+        output_shape = [up.size(0), down.size(1)]
+        if self.mid_model is not None:
+            # cp-decomposition
+            mid = self.mid_model.weight.to(orig_weight.device, dtype=orig_weight.dtype)
+            updown = lyco_helpers.rebuild_cp_decomposition(up, down, mid)
+            output_shape += mid.shape[2:]
+        else:
+            if len(down.shape) == 4:
+                output_shape += down.shape[2:]
+            updown = lyco_helpers.rebuild_conventional(up, down, output_shape, self.network.dyn_dim)
+
+        return self.finalize_updown(updown, orig_weight, output_shape)
+
+    def forward(self, x, y):
+        self.up_model.to(device=devices.device)
+        self.down_model.to(device=devices.device)
+
+        return y + self.up_model(self.down_model(x)) * self.multiplier() * self.calc_scale()
+
+

extensions-builtin/Lora/networks.py

@@ -0,0 +1,468 @@
+import os
+import re
+
+import network
+import network_lora
+import network_hada
+import network_ia3
+import network_lokr
+import network_full
+
+import torch
+from typing import Union
+
+from modules import shared, devices, sd_models, errors, scripts, sd_hijack
+
+module_types = [
+    network_lora.ModuleTypeLora(),
+    network_hada.ModuleTypeHada(),
+    network_ia3.ModuleTypeIa3(),
+    network_lokr.ModuleTypeLokr(),
+    network_full.ModuleTypeFull(),
+]
+
+
+re_digits = re.compile(r"\d+")
+re_x_proj = re.compile(r"(.*)_([qkv]_proj)$")
+re_compiled = {}
+
+suffix_conversion = {
+    "attentions": {},
+    "resnets": {
+        "conv1": "in_layers_2",
+        "conv2": "out_layers_3",
+        "time_emb_proj": "emb_layers_1",
+        "conv_shortcut": "skip_connection",
+    }
+}
+
+
+def convert_diffusers_name_to_compvis(key, is_sd2):
+    def match(match_list, regex_text):
+        regex = re_compiled.get(regex_text)
+        if regex is None:
+            regex = re.compile(regex_text)
+            re_compiled[regex_text] = regex
+
+        r = re.match(regex, key)
+        if not r:
+            return False
+
+        match_list.clear()
+        match_list.extend([int(x) if re.match(re_digits, x) else x for x in r.groups()])
+        return True
+
+    m = []
+
+    if match(m, r"lora_unet_conv_in(.*)"):
+        return f'diffusion_model_input_blocks_0_0{m[0]}'
+
+    if match(m, r"lora_unet_conv_out(.*)"):
+        return f'diffusion_model_out_2{m[0]}'
+
+    if match(m, r"lora_unet_time_embedding_linear_(\d+)(.*)"):
+        return f"diffusion_model_time_embed_{m[0] * 2 - 2}{m[1]}"
+
+    if match(m, r"lora_unet_down_blocks_(\d+)_(attentions|resnets)_(\d+)_(.+)"):
+        suffix = suffix_conversion.get(m[1], {}).get(m[3], m[3])
+        return f"diffusion_model_input_blocks_{1 + m[0] * 3 + m[2]}_{1 if m[1] == 'attentions' else 0}_{suffix}"
+
+    if match(m, r"lora_unet_mid_block_(attentions|resnets)_(\d+)_(.+)"):
+        suffix = suffix_conversion.get(m[0], {}).get(m[2], m[2])
+        return f"diffusion_model_middle_block_{1 if m[0] == 'attentions' else m[1] * 2}_{suffix}"
+
+    if match(m, r"lora_unet_up_blocks_(\d+)_(attentions|resnets)_(\d+)_(.+)"):
+        suffix = suffix_conversion.get(m[1], {}).get(m[3], m[3])
+        return f"diffusion_model_output_blocks_{m[0] * 3 + m[2]}_{1 if m[1] == 'attentions' else 0}_{suffix}"
+
+    if match(m, r"lora_unet_down_blocks_(\d+)_downsamplers_0_conv"):
+        return f"diffusion_model_input_blocks_{3 + m[0] * 3}_0_op"
+
+    if match(m, r"lora_unet_up_blocks_(\d+)_upsamplers_0_conv"):
+        return f"diffusion_model_output_blocks_{2 + m[0] * 3}_{2 if m[0]>0 else 1}_conv"
+
+    if match(m, r"lora_te_text_model_encoder_layers_(\d+)_(.+)"):
+        if is_sd2:
+            if 'mlp_fc1' in m[1]:
+                return f"model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc1', 'mlp_c_fc')}"
+            elif 'mlp_fc2' in m[1]:
+                return f"model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc2', 'mlp_c_proj')}"
+            else:
+                return f"model_transformer_resblocks_{m[0]}_{m[1].replace('self_attn', 'attn')}"
+
+        return f"transformer_text_model_encoder_layers_{m[0]}_{m[1]}"
+
+    if match(m, r"lora_te2_text_model_encoder_layers_(\d+)_(.+)"):
+        if 'mlp_fc1' in m[1]:
+            return f"1_model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc1', 'mlp_c_fc')}"
+        elif 'mlp_fc2' in m[1]:
+            return f"1_model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc2', 'mlp_c_proj')}"
+        else:
+            return f"1_model_transformer_resblocks_{m[0]}_{m[1].replace('self_attn', 'attn')}"
+
+    return key
+
+
+def assign_network_names_to_compvis_modules(sd_model):
+    network_layer_mapping = {}
+
+    if shared.sd_model.is_sdxl:
+        for i, embedder in enumerate(shared.sd_model.conditioner.embedders):
+            if not hasattr(embedder, 'wrapped'):
+                continue
+
+            for name, module in embedder.wrapped.named_modules():
+                network_name = f'{i}_{name.replace(".", "_")}'
+                network_layer_mapping[network_name] = module
+                module.network_layer_name = network_name
+    else:
+        for name, module in shared.sd_model.cond_stage_model.wrapped.named_modules():
+            network_name = name.replace(".", "_")
+            network_layer_mapping[network_name] = module
+            module.network_layer_name = network_name
+
+    for name, module in shared.sd_model.model.named_modules():
+        network_name = name.replace(".", "_")
+        network_layer_mapping[network_name] = module
+        module.network_layer_name = network_name
+
+    sd_model.network_layer_mapping = network_layer_mapping
+
+
+def load_network(name, network_on_disk):
+    net = network.Network(name, network_on_disk)
+    net.mtime = os.path.getmtime(network_on_disk.filename)
+
+    sd = sd_models.read_state_dict(network_on_disk.filename)
+
+    # this should not be needed but is here as an emergency fix for an unknown error people are experiencing in 1.2.0
+    if not hasattr(shared.sd_model, 'network_layer_mapping'):
+        assign_network_names_to_compvis_modules(shared.sd_model)
+
+    keys_failed_to_match = {}
+    is_sd2 = 'model_transformer_resblocks' in shared.sd_model.network_layer_mapping
+
+    matched_networks = {}
+
+    for key_network, weight in sd.items():
+        key_network_without_network_parts, network_part = key_network.split(".", 1)
+
+        key = convert_diffusers_name_to_compvis(key_network_without_network_parts, is_sd2)
+        sd_module = shared.sd_model.network_layer_mapping.get(key, None)
+
+        if sd_module is None:
+            m = re_x_proj.match(key)
+            if m:
+                sd_module = shared.sd_model.network_layer_mapping.get(m.group(1), None)
+
+        # SDXL loras seem to already have correct compvis keys, so only need to replace "lora_unet" with "diffusion_model"
+        if sd_module is None and "lora_unet" in key_network_without_network_parts:
+            key = key_network_without_network_parts.replace("lora_unet", "diffusion_model")
+            sd_module = shared.sd_model.network_layer_mapping.get(key, None)
+        elif sd_module is None and "lora_te1_text_model" in key_network_without_network_parts:
+            key = key_network_without_network_parts.replace("lora_te1_text_model", "0_transformer_text_model")
+            sd_module = shared.sd_model.network_layer_mapping.get(key, None)
+
+            # some SD1 Loras also have correct compvis keys
+            if sd_module is None:
+                key = key_network_without_network_parts.replace("lora_te1_text_model", "transformer_text_model")
+                sd_module = shared.sd_model.network_layer_mapping.get(key, None)
+
+        if sd_module is None:
+            keys_failed_to_match[key_network] = key
+            continue
+
+        if key not in matched_networks:
+            matched_networks[key] = network.NetworkWeights(network_key=key_network, sd_key=key, w={}, sd_module=sd_module)
+
+        matched_networks[key].w[network_part] = weight
+
+    for key, weights in matched_networks.items():
+        net_module = None
+        for nettype in module_types:
+            net_module = nettype.create_module(net, weights)
+            if net_module is not None:
+                break
+
+        if net_module is None:
+            raise AssertionError(f"Could not find a module type (out of {', '.join([x.__class__.__name__ for x in module_types])}) that would accept those keys: {', '.join(weights.w)}")
+
+        net.modules[key] = net_module
+
+    if keys_failed_to_match:
+        print(f"Failed to match keys when loading network {network_on_disk.filename}: {keys_failed_to_match}")
+
+    return net
+
+
+def load_networks(names, te_multipliers=None, unet_multipliers=None, dyn_dims=None):
+    already_loaded = {}
+
+    for net in loaded_networks:
+        if net.name in names:
+            already_loaded[net.name] = net
+
+    loaded_networks.clear()
+
+    networks_on_disk = [available_network_aliases.get(name, None) for name in names]
+    if any(x is None for x in networks_on_disk):
+        list_available_networks()
+
+        networks_on_disk = [available_network_aliases.get(name, None) for name in names]
+
+    failed_to_load_networks = []
+
+    for i, name in enumerate(names):
+        net = already_loaded.get(name, None)
+
+        network_on_disk = networks_on_disk[i]
+
+        if network_on_disk is not None:
+            if net is None or os.path.getmtime(network_on_disk.filename) > net.mtime:
+                try:
+                    net = load_network(name, network_on_disk)
+                except Exception as e:
+                    errors.display(e, f"loading network {network_on_disk.filename}")
+                    continue
+
+            net.mentioned_name = name
+
+            network_on_disk.read_hash()
+
+        if net is None:
+            failed_to_load_networks.append(name)
+            print(f"Couldn't find network with name {name}")
+            continue
+
+        net.te_multiplier = te_multipliers[i] if te_multipliers else 1.0
+        net.unet_multiplier = unet_multipliers[i] if unet_multipliers else 1.0
+        net.dyn_dim = dyn_dims[i] if dyn_dims else 1.0
+        loaded_networks.append(net)
+
+    if failed_to_load_networks:
+        sd_hijack.model_hijack.comments.append("Failed to find networks: " + ", ".join(failed_to_load_networks))
+
+
+def network_restore_weights_from_backup(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.MultiheadAttention]):
+    weights_backup = getattr(self, "network_weights_backup", None)
+
+    if weights_backup is None:
+        return
+
+    if isinstance(self, torch.nn.MultiheadAttention):
+        self.in_proj_weight.copy_(weights_backup[0])
+        self.out_proj.weight.copy_(weights_backup[1])
+    else:
+        self.weight.copy_(weights_backup)
+
+
+def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.MultiheadAttention]):
+    """
+    Applies the currently selected set of networks to the weights of torch layer self.
+    If weights already have this particular set of networks applied, does nothing.
+    If not, restores orginal weights from backup and alters weights according to networks.
+    """
+
+    network_layer_name = getattr(self, 'network_layer_name', None)
+    if network_layer_name is None:
+        return
+
+    current_names = getattr(self, "network_current_names", ())
+    wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in loaded_networks)
+
+    weights_backup = getattr(self, "network_weights_backup", None)
+    if weights_backup is None:
+        if isinstance(self, torch.nn.MultiheadAttention):
+            weights_backup = (self.in_proj_weight.to(devices.cpu, copy=True), self.out_proj.weight.to(devices.cpu, copy=True))
+        else:
+            weights_backup = self.weight.to(devices.cpu, copy=True)
+
+        self.network_weights_backup = weights_backup
+
+    if current_names != wanted_names:
+        network_restore_weights_from_backup(self)
+
+        for net in loaded_networks:
+            module = net.modules.get(network_layer_name, None)
+            if module is not None and hasattr(self, 'weight'):
+                with torch.no_grad():
+                    updown = module.calc_updown(self.weight)
+
+                    if len(self.weight.shape) == 4 and self.weight.shape[1] == 9:
+                        # inpainting model. zero pad updown to make channel[1]  4 to 9
+                        updown = torch.nn.functional.pad(updown, (0, 0, 0, 0, 0, 5))
+
+                    self.weight += updown
+                    continue
+
+            module_q = net.modules.get(network_layer_name + "_q_proj", None)
+            module_k = net.modules.get(network_layer_name + "_k_proj", None)
+            module_v = net.modules.get(network_layer_name + "_v_proj", None)
+            module_out = net.modules.get(network_layer_name + "_out_proj", None)
+
+            if isinstance(self, torch.nn.MultiheadAttention) and module_q and module_k and module_v and module_out:
+                with torch.no_grad():
+                    updown_q = module_q.calc_updown(self.in_proj_weight)
+                    updown_k = module_k.calc_updown(self.in_proj_weight)
+                    updown_v = module_v.calc_updown(self.in_proj_weight)
+                    updown_qkv = torch.vstack([updown_q, updown_k, updown_v])
+                    updown_out = module_out.calc_updown(self.out_proj.weight)
+
+                    self.in_proj_weight += updown_qkv
+                    self.out_proj.weight += updown_out
+                    continue
+
+            if module is None:
+                continue
+
+            print(f'failed to calculate network weights for layer {network_layer_name}')
+
+        self.network_current_names = wanted_names
+
+
+def network_forward(module, input, original_forward):
+    """
+    Old way of applying Lora by executing operations during layer's forward.
+    Stacking many loras this way results in big performance degradation.
+    """
+
+    if len(loaded_networks) == 0:
+        return original_forward(module, input)
+
+    input = devices.cond_cast_unet(input)
+
+    network_restore_weights_from_backup(module)
+    network_reset_cached_weight(module)
+
+    y = original_forward(module, input)
+
+    network_layer_name = getattr(module, 'network_layer_name', None)
+    for lora in loaded_networks:
+        module = lora.modules.get(network_layer_name, None)
+        if module is None:
+            continue
+
+        y = module.forward(y, input)
+
+    return y
+
+
+def network_reset_cached_weight(self: Union[torch.nn.Conv2d, torch.nn.Linear]):
+    self.network_current_names = ()
+    self.network_weights_backup = None
+
+
+def network_Linear_forward(self, input):
+    if shared.opts.lora_functional:
+        return network_forward(self, input, torch.nn.Linear_forward_before_network)
+
+    network_apply_weights(self)
+
+    return torch.nn.Linear_forward_before_network(self, input)
+
+
+def network_Linear_load_state_dict(self, *args, **kwargs):
+    network_reset_cached_weight(self)
+
+    return torch.nn.Linear_load_state_dict_before_network(self, *args, **kwargs)
+
+
+def network_Conv2d_forward(self, input):
+    if shared.opts.lora_functional:
+        return network_forward(self, input, torch.nn.Conv2d_forward_before_network)
+
+    network_apply_weights(self)
+
+    return torch.nn.Conv2d_forward_before_network(self, input)
+
+
+def network_Conv2d_load_state_dict(self, *args, **kwargs):
+    network_reset_cached_weight(self)
+
+    return torch.nn.Conv2d_load_state_dict_before_network(self, *args, **kwargs)
+
+
+def network_MultiheadAttention_forward(self, *args, **kwargs):
+    network_apply_weights(self)
+
+    return torch.nn.MultiheadAttention_forward_before_network(self, *args, **kwargs)
+
+
+def network_MultiheadAttention_load_state_dict(self, *args, **kwargs):
+    network_reset_cached_weight(self)
+
+    return torch.nn.MultiheadAttention_load_state_dict_before_network(self, *args, **kwargs)
+
+
+def list_available_networks():
+    available_networks.clear()
+    available_network_aliases.clear()
+    forbidden_network_aliases.clear()
+    available_network_hash_lookup.clear()
+    forbidden_network_aliases.update({"none": 1, "Addams": 1})
+
+    os.makedirs(shared.cmd_opts.lora_dir, exist_ok=True)
+
+    candidates = list(shared.walk_files(shared.cmd_opts.lora_dir, allowed_extensions=[".pt", ".ckpt", ".safetensors"]))
+    candidates += list(shared.walk_files(shared.cmd_opts.lyco_dir_backcompat, allowed_extensions=[".pt", ".ckpt", ".safetensors"]))
+    for filename in candidates:
+        if os.path.isdir(filename):
+            continue
+
+        name = os.path.splitext(os.path.basename(filename))[0]
+        try:
+            entry = network.NetworkOnDisk(name, filename)
+        except OSError:  # should catch FileNotFoundError and PermissionError etc.
+            errors.report(f"Failed to load network {name} from {filename}", exc_info=True)
+            continue
+
+        available_networks[name] = entry
+
+        if entry.alias in available_network_aliases:
+            forbidden_network_aliases[entry.alias.lower()] = 1
+
+        available_network_aliases[name] = entry
+        available_network_aliases[entry.alias] = entry
+
+
+re_network_name = re.compile(r"(.*)\s*\([0-9a-fA-F]+\)")
+
+
+def infotext_pasted(infotext, params):
+    if "AddNet Module 1" in [x[1] for x in scripts.scripts_txt2img.infotext_fields]:
+        return  # if the other extension is active, it will handle those fields, no need to do anything
+
+    added = []
+
+    for k in params:
+        if not k.startswith("AddNet Model "):
+            continue
+
+        num = k[13:]
+
+        if params.get("AddNet Module " + num) != "LoRA":
+            continue
+
+        name = params.get("AddNet Model " + num)
+        if name is None:
+            continue
+
+        m = re_network_name.match(name)
+        if m:
+            name = m.group(1)
+
+        multiplier = params.get("AddNet Weight A " + num, "1.0")
+
+        added.append(f"<lora:{name}:{multiplier}>")
+
+    if added:
+        params["Prompt"] += "\n" + "".join(added)
+
+
+available_networks = {}
+available_network_aliases = {}
+loaded_networks = []
+available_network_hash_lookup = {}
+forbidden_network_aliases = {}
+
+list_available_networks()

extensions-builtin/Lora/preload.py

@@ -4,3 +4,4 @@ from modules import paths
 
 def preload(parser):
     parser.add_argument("--lora-dir", type=str, help="Path to directory with Lora networks.", default=os.path.join(paths.models_path, 'Lora'))
+    parser.add_argument("--lyco-dir-backcompat", type=str, help="Path to directory with LyCORIS networks (for backawards compatibility; can also use --lyco-dir).", default=os.path.join(paths.models_path, 'LyCORIS'))

extensions-builtin/Lora/scripts/lora_script.py

@@ -4,69 +4,76 @@ import torch
 import gradio as gr
 from fastapi import FastAPI
 
-import lora
+import network
+import networks
+import lora  # noqa:F401
 import extra_networks_lora
 import ui_extra_networks_lora
 from modules import script_callbacks, ui_extra_networks, extra_networks, shared
 
 def unload():
-    torch.nn.Linear.forward = torch.nn.Linear_forward_before_lora
+    torch.nn.Linear.forward = torch.nn.Linear_forward_before_network
-    torch.nn.Linear._load_from_state_dict = torch.nn.Linear_load_state_dict_before_lora
+    torch.nn.Linear._load_from_state_dict = torch.nn.Linear_load_state_dict_before_network
-    torch.nn.Conv2d.forward = torch.nn.Conv2d_forward_before_lora
+    torch.nn.Conv2d.forward = torch.nn.Conv2d_forward_before_network
-    torch.nn.Conv2d._load_from_state_dict = torch.nn.Conv2d_load_state_dict_before_lora
+    torch.nn.Conv2d._load_from_state_dict = torch.nn.Conv2d_load_state_dict_before_network
-    torch.nn.MultiheadAttention.forward = torch.nn.MultiheadAttention_forward_before_lora
+    torch.nn.MultiheadAttention.forward = torch.nn.MultiheadAttention_forward_before_network
-    torch.nn.MultiheadAttention._load_from_state_dict = torch.nn.MultiheadAttention_load_state_dict_before_lora
+    torch.nn.MultiheadAttention._load_from_state_dict = torch.nn.MultiheadAttention_load_state_dict_before_network
 
 
 def before_ui():
     ui_extra_networks.register_page(ui_extra_networks_lora.ExtraNetworksPageLora())
-    extra_networks.register_extra_network(extra_networks_lora.ExtraNetworkLora())
+
+    extra_network = extra_networks_lora.ExtraNetworkLora()
+    extra_networks.register_extra_network(extra_network)
+    extra_networks.register_extra_network_alias(extra_network, "lyco")
 
 
-if not hasattr(torch.nn, 'Linear_forward_before_lora'):
+if not hasattr(torch.nn, 'Linear_forward_before_network'):
-    torch.nn.Linear_forward_before_lora = torch.nn.Linear.forward
+    torch.nn.Linear_forward_before_network = torch.nn.Linear.forward
 
-if not hasattr(torch.nn, 'Linear_load_state_dict_before_lora'):
+if not hasattr(torch.nn, 'Linear_load_state_dict_before_network'):
-    torch.nn.Linear_load_state_dict_before_lora = torch.nn.Linear._load_from_state_dict
+    torch.nn.Linear_load_state_dict_before_network = torch.nn.Linear._load_from_state_dict
 
-if not hasattr(torch.nn, 'Conv2d_forward_before_lora'):
+if not hasattr(torch.nn, 'Conv2d_forward_before_network'):
-    torch.nn.Conv2d_forward_before_lora = torch.nn.Conv2d.forward
+    torch.nn.Conv2d_forward_before_network = torch.nn.Conv2d.forward
 
-if not hasattr(torch.nn, 'Conv2d_load_state_dict_before_lora'):
+if not hasattr(torch.nn, 'Conv2d_load_state_dict_before_network'):
-    torch.nn.Conv2d_load_state_dict_before_lora = torch.nn.Conv2d._load_from_state_dict
+    torch.nn.Conv2d_load_state_dict_before_network = torch.nn.Conv2d._load_from_state_dict
 
-if not hasattr(torch.nn, 'MultiheadAttention_forward_before_lora'):
+if not hasattr(torch.nn, 'MultiheadAttention_forward_before_network'):
-    torch.nn.MultiheadAttention_forward_before_lora = torch.nn.MultiheadAttention.forward
+    torch.nn.MultiheadAttention_forward_before_network = torch.nn.MultiheadAttention.forward
 
-if not hasattr(torch.nn, 'MultiheadAttention_load_state_dict_before_lora'):
+if not hasattr(torch.nn, 'MultiheadAttention_load_state_dict_before_network'):
-    torch.nn.MultiheadAttention_load_state_dict_before_lora = torch.nn.MultiheadAttention._load_from_state_dict
+    torch.nn.MultiheadAttention_load_state_dict_before_network = torch.nn.MultiheadAttention._load_from_state_dict
 
-torch.nn.Linear.forward = lora.lora_Linear_forward
+torch.nn.Linear.forward = networks.network_Linear_forward
-torch.nn.Linear._load_from_state_dict = lora.lora_Linear_load_state_dict
+torch.nn.Linear._load_from_state_dict = networks.network_Linear_load_state_dict
-torch.nn.Conv2d.forward = lora.lora_Conv2d_forward
+torch.nn.Conv2d.forward = networks.network_Conv2d_forward
-torch.nn.Conv2d._load_from_state_dict = lora.lora_Conv2d_load_state_dict
+torch.nn.Conv2d._load_from_state_dict = networks.network_Conv2d_load_state_dict
-torch.nn.MultiheadAttention.forward = lora.lora_MultiheadAttention_forward
+torch.nn.MultiheadAttention.forward = networks.network_MultiheadAttention_forward
-torch.nn.MultiheadAttention._load_from_state_dict = lora.lora_MultiheadAttention_load_state_dict
+torch.nn.MultiheadAttention._load_from_state_dict = networks.network_MultiheadAttention_load_state_dict
 
-script_callbacks.on_model_loaded(lora.assign_lora_names_to_compvis_modules)
+script_callbacks.on_model_loaded(networks.assign_network_names_to_compvis_modules)
 script_callbacks.on_script_unloaded(unload)
 script_callbacks.on_before_ui(before_ui)
-script_callbacks.on_infotext_pasted(lora.infotext_pasted)
+script_callbacks.on_infotext_pasted(networks.infotext_pasted)
 
 
 shared.options_templates.update(shared.options_section(('extra_networks', "Extra Networks"), {
-    "sd_lora": shared.OptionInfo("None", "Add Lora to prompt", gr.Dropdown, lambda: {"choices": ["None", *lora.available_loras]}, refresh=lora.list_available_loras),
+    "sd_lora": shared.OptionInfo("None", "Add network to prompt", gr.Dropdown, lambda: {"choices": ["None", *networks.available_networks]}, refresh=networks.list_available_networks),
     "lora_preferred_name": shared.OptionInfo("Alias from file", "When adding to prompt, refer to Lora by", gr.Radio, {"choices": ["Alias from file", "Filename"]}),
     "lora_add_hashes_to_infotext": shared.OptionInfo(True, "Add Lora hashes to infotext"),
+    "lora_show_all": shared.OptionInfo(False, "Always show all networks on the Lora page").info("otherwise, those detected as for incompatible version of Stable Diffusion will be hidden"),
+    "lora_hide_unknown_for_versions": shared.OptionInfo([], "Hide networks of unknown versions for model versions", gr.CheckboxGroup, {"choices": ["SD1", "SD2", "SDXL"]}),
 }))
 
 
 shared.options_templates.update(shared.options_section(('compatibility', "Compatibility"), {
-    "lora_functional": shared.OptionInfo(False, "Lora: use old method that takes longer when you have multiple Loras active and produces same results as kohya-ss/sd-webui-additional-networks extension"),
+    "lora_functional": shared.OptionInfo(False, "Lora/Networks: use old method that takes longer when you have multiple Loras active and produces same results as kohya-ss/sd-webui-additional-networks extension"),
 }))
 
 
-def create_lora_json(obj: lora.LoraOnDisk):
+def create_lora_json(obj: network.NetworkOnDisk):
     return {
         "name": obj.name,
         "alias": obj.alias,
@@ -75,17 +82,17 @@ def create_lora_json(obj: lora.LoraOnDisk):
     }
 
 
-def api_loras(_: gr.Blocks, app: FastAPI):
+def api_networks(_: gr.Blocks, app: FastAPI):
     @app.get("/sdapi/v1/loras")
     async def get_loras():
-        return [create_lora_json(obj) for obj in lora.available_loras.values()]
+        return [create_lora_json(obj) for obj in networks.available_networks.values()]
 
     @app.post("/sdapi/v1/refresh-loras")
     async def refresh_loras():
-        return lora.list_available_loras()
+        return networks.list_available_networks()
 
 
-script_callbacks.on_app_started(api_loras)
+script_callbacks.on_app_started(api_networks)
 
 re_lora = re.compile("<lora:([^:]+):")
 
@@ -98,19 +105,19 @@ def infotext_pasted(infotext, d):
     hashes = [x.strip().split(':', 1) for x in hashes.split(",")]
     hashes = {x[0].strip().replace(",", ""): x[1].strip() for x in hashes}
 
-    def lora_replacement(m):
+    def network_replacement(m):
         alias = m.group(1)
         shorthash = hashes.get(alias)
         if shorthash is None:
             return m.group(0)
 
-        lora_on_disk = lora.available_lora_hash_lookup.get(shorthash)
+        network_on_disk = networks.available_network_hash_lookup.get(shorthash)
-        if lora_on_disk is None:
+        if network_on_disk is None:
             return m.group(0)
 
-        return f'<lora:{lora_on_disk.get_alias()}:'
+        return f'<lora:{network_on_disk.get_alias()}:'
 
-    d["Prompt"] = re.sub(re_lora, lora_replacement, d["Prompt"])
+    d["Prompt"] = re.sub(re_lora, network_replacement, d["Prompt"])
 
 
 script_callbacks.on_infotext_pasted(infotext_pasted)

extensions-builtin/Lora/ui_edit_user_metadata.py

@@ -0,0 +1,216 @@
+import datetime
+import html
+import random
+
+import gradio as gr
+import re
+
+from modules import ui_extra_networks_user_metadata
+
+
+def is_non_comma_tagset(tags):
+    average_tag_length = sum(len(x) for x in tags.keys()) / len(tags)
+
+    return average_tag_length >= 16
+
+
+re_word = re.compile(r"[-_\w']+")
+re_comma = re.compile(r" *, *")
+
+
+def build_tags(metadata):
+    tags = {}
+
+    for _, tags_dict in metadata.get("ss_tag_frequency", {}).items():
+        for tag, tag_count in tags_dict.items():
+            tag = tag.strip()
+            tags[tag] = tags.get(tag, 0) + int(tag_count)
+
+    if tags and is_non_comma_tagset(tags):
+        new_tags = {}
+
+        for text, text_count in tags.items():
+            for word in re.findall(re_word, text):
+                if len(word) < 3:
+                    continue
+
+                new_tags[word] = new_tags.get(word, 0) + text_count
+
+        tags = new_tags
+
+    ordered_tags = sorted(tags.keys(), key=tags.get, reverse=True)
+
+    return [(tag, tags[tag]) for tag in ordered_tags]
+
+
+class LoraUserMetadataEditor(ui_extra_networks_user_metadata.UserMetadataEditor):
+    def __init__(self, ui, tabname, page):
+        super().__init__(ui, tabname, page)
+
+        self.select_sd_version = None
+
+        self.taginfo = None
+        self.edit_activation_text = None
+        self.slider_preferred_weight = None
+        self.edit_notes = None
+
+    def save_lora_user_metadata(self, name, desc, sd_version, activation_text, preferred_weight, notes):
+        user_metadata = self.get_user_metadata(name)
+        user_metadata["description"] = desc
+        user_metadata["sd version"] = sd_version
+        user_metadata["activation text"] = activation_text
+        user_metadata["preferred weight"] = preferred_weight
+        user_metadata["notes"] = notes
+
+        self.write_user_metadata(name, user_metadata)
+
+    def get_metadata_table(self, name):
+        table = super().get_metadata_table(name)
+        item = self.page.items.get(name, {})
+        metadata = item.get("metadata") or {}
+
+        keys = {
+            'ss_sd_model_name': "Model:",
+            'ss_clip_skip': "Clip skip:",
+            'ss_network_module': "Kohya module:",
+        }
+
+        for key, label in keys.items():
+            value = metadata.get(key, None)
+            if value is not None and str(value) != "None":
+                table.append((label, html.escape(value)))
+
+        ss_training_started_at = metadata.get('ss_training_started_at')
+        if ss_training_started_at:
+            table.append(("Date trained:", datetime.datetime.utcfromtimestamp(float(ss_training_started_at)).strftime('%Y-%m-%d %H:%M')))
+
+        ss_bucket_info = metadata.get("ss_bucket_info")
+        if ss_bucket_info and "buckets" in ss_bucket_info:
+            resolutions = {}
+            for _, bucket in ss_bucket_info["buckets"].items():
+                resolution = bucket["resolution"]
+                resolution = f'{resolution[1]}x{resolution[0]}'
+
+                resolutions[resolution] = resolutions.get(resolution, 0) + int(bucket["count"])
+
+            resolutions_list = sorted(resolutions.keys(), key=resolutions.get, reverse=True)
+            resolutions_text = html.escape(", ".join(resolutions_list[0:4]))
+            if len(resolutions) > 4:
+                resolutions_text += ", ..."
+                resolutions_text = f"<span title='{html.escape(', '.join(resolutions_list))}'>{resolutions_text}</span>"
+
+            table.append(('Resolutions:' if len(resolutions_list) > 1 else 'Resolution:', resolutions_text))
+
+        image_count = 0
+        for _, params in metadata.get("ss_dataset_dirs", {}).items():
+            image_count += int(params.get("img_count", 0))
+
+        if image_count:
+            table.append(("Dataset size:", image_count))
+
+        return table
+
+    def put_values_into_components(self, name):
+        user_metadata = self.get_user_metadata(name)
+        values = super().put_values_into_components(name)
+
+        item = self.page.items.get(name, {})
+        metadata = item.get("metadata") or {}
+
+        tags = build_tags(metadata)
+        gradio_tags = [(tag, str(count)) for tag, count in tags[0:24]]
+
+        return [
+            *values[0:5],
+            item.get("sd_version", "Unknown"),
+            gr.HighlightedText.update(value=gradio_tags, visible=True if tags else False),
+            user_metadata.get('activation text', ''),
+            float(user_metadata.get('preferred weight', 0.0)),
+            gr.update(visible=True if tags else False),
+            gr.update(value=self.generate_random_prompt_from_tags(tags), visible=True if tags else False),
+        ]
+
+    def generate_random_prompt(self, name):
+        item = self.page.items.get(name, {})
+        metadata = item.get("metadata") or {}
+        tags = build_tags(metadata)
+
+        return self.generate_random_prompt_from_tags(tags)
+
+    def generate_random_prompt_from_tags(self, tags):
+        max_count = None
+        res = []
+        for tag, count in tags:
+            if not max_count:
+                max_count = count
+
+            v = random.random() * max_count
+            if count > v:
+                res.append(tag)
+
+        return ", ".join(sorted(res))
+
+    def create_extra_default_items_in_left_column(self):
+
+        # this would be a lot better as gr.Radio but I can't make it work
+        self.select_sd_version = gr.Dropdown(['SD1', 'SD2', 'SDXL', 'Unknown'], value='Unknown', label='Stable Diffusion version', interactive=True)
+
+    def create_editor(self):
+        self.create_default_editor_elems()
+
+        self.taginfo = gr.HighlightedText(label="Training dataset tags")
+        self.edit_activation_text = gr.Text(label='Activation text', info="Will be added to prompt along with Lora")
+        self.slider_preferred_weight = gr.Slider(label='Preferred weight', info="Set to 0 to disable", minimum=0.0, maximum=2.0, step=0.01)
+
+        with gr.Row() as row_random_prompt:
+            with gr.Column(scale=8):
+                random_prompt = gr.Textbox(label='Random prompt', lines=4, max_lines=4, interactive=False)
+
+            with gr.Column(scale=1, min_width=120):
+                generate_random_prompt = gr.Button('Generate', size="lg", scale=1)
+
+        self.edit_notes = gr.TextArea(label='Notes', lines=4)
+
+        generate_random_prompt.click(fn=self.generate_random_prompt, inputs=[self.edit_name_input], outputs=[random_prompt], show_progress=False)
+
+        def select_tag(activation_text, evt: gr.SelectData):
+            tag = evt.value[0]
+
+            words = re.split(re_comma, activation_text)
+            if tag in words:
+                words = [x for x in words if x != tag and x.strip()]
+                return ", ".join(words)
+
+            return activation_text + ", " + tag if activation_text else tag
+
+        self.taginfo.select(fn=select_tag, inputs=[self.edit_activation_text], outputs=[self.edit_activation_text], show_progress=False)
+
+        self.create_default_buttons()
+
+        viewed_components = [
+            self.edit_name,
+            self.edit_description,
+            self.html_filedata,
+            self.html_preview,
+            self.edit_notes,
+            self.select_sd_version,
+            self.taginfo,
+            self.edit_activation_text,
+            self.slider_preferred_weight,
+            row_random_prompt,
+            random_prompt,
+        ]
+
+        self.button_edit\
+            .click(fn=self.put_values_into_components, inputs=[self.edit_name_input], outputs=viewed_components)\
+            .then(fn=lambda: gr.update(visible=True), inputs=[], outputs=[self.box])
+
+        edited_components = [
+            self.edit_description,
+            self.select_sd_version,
+            self.edit_activation_text,
+            self.slider_preferred_weight,
+            self.edit_notes,
+        ]
+
+        self.setup_save_handler(self.button_save, self.save_lora_user_metadata, edited_components)

extensions-builtin/Lora/ui_extra_networks_lora.py

@@ -1,8 +1,11 @@
-import json
 import os
-import lora
+
+import network
+import networks
 
 from modules import shared, ui_extra_networks
+from modules.ui_extra_networks import quote_js
+from ui_edit_user_metadata import LoraUserMetadataEditor
 
 
 class ExtraNetworksPageLora(ui_extra_networks.ExtraNetworksPage):
@@ -10,27 +13,66 @@ class ExtraNetworksPageLora(ui_extra_networks.ExtraNetworksPage):
         super().__init__('Lora')
 
     def refresh(self):
-        lora.list_available_loras()
+        networks.list_available_networks()
+
+    def create_item(self, name, index=None, enable_filter=True):
+        lora_on_disk = networks.available_networks.get(name)
 
-    def list_items(self):
-        for index, (name, lora_on_disk) in enumerate(lora.available_loras.items()):
         path, ext = os.path.splitext(lora_on_disk.filename)
 
         alias = lora_on_disk.get_alias()
 
-            yield {
+        item = {
             "name": name,
-                "filename": path,
+            "filename": lora_on_disk.filename,
             "preview": self.find_preview(path),
             "description": self.find_description(path),
             "search_term": self.search_terms_from_path(lora_on_disk.filename),
-                "prompt": json.dumps(f"<lora:{alias}:") + " + opts.extra_networks_default_multiplier + " + json.dumps(">"),
             "local_preview": f"{path}.{shared.opts.samples_format}",
-                "metadata": json.dumps(lora_on_disk.metadata, indent=4) if lora_on_disk.metadata else None,
+            "metadata": lora_on_disk.metadata,
             "sort_keys": {'default': index, **self.get_sort_keys(lora_on_disk.filename)},
-
+            "sd_version": lora_on_disk.sd_version.name,
         }
 
-    def allowed_directories_for_previews(self):
+        self.read_user_metadata(item)
-        return [shared.cmd_opts.lora_dir]
+        activation_text = item["user_metadata"].get("activation text")
+        preferred_weight = item["user_metadata"].get("preferred weight", 0.0)
+        item["prompt"] = quote_js(f"<lora:{alias}:") + " + " + (str(preferred_weight) if preferred_weight else "opts.extra_networks_default_multiplier") + " + " + quote_js(">")
 
+        if activation_text:
+            item["prompt"] += " + " + quote_js(" " + activation_text)
+
+        sd_version = item["user_metadata"].get("sd version")
+        if sd_version in network.SdVersion.__members__:
+            item["sd_version"] = sd_version
+            sd_version = network.SdVersion[sd_version]
+        else:
+            sd_version = lora_on_disk.sd_version
+
+        if shared.opts.lora_show_all or not enable_filter:
+            pass
+        elif sd_version == network.SdVersion.Unknown:
+            model_version = network.SdVersion.SDXL if shared.sd_model.is_sdxl else network.SdVersion.SD2 if shared.sd_model.is_sd2 else network.SdVersion.SD1
+            if model_version.name in shared.opts.lora_hide_unknown_for_versions:
+                return None
+        elif shared.sd_model.is_sdxl and sd_version != network.SdVersion.SDXL:
+            return None
+        elif shared.sd_model.is_sd2 and sd_version != network.SdVersion.SD2:
+            return None
+        elif shared.sd_model.is_sd1 and sd_version != network.SdVersion.SD1:
+            return None
+
+        return item
+
+    def list_items(self):
+        for index, name in enumerate(networks.available_networks):
+            item = self.create_item(name, index)
+
+            if item is not None:
+                yield item
+
+    def allowed_directories_for_previews(self):
+        return [shared.cmd_opts.lora_dir, shared.cmd_opts.lyco_dir_backcompat]
+
+    def create_user_metadata_editor(self, ui, tabname):
+        return LoraUserMetadataEditor(ui, tabname, self)

extensions-builtin/ScuNET/scripts/scunet_model.py

@@ -1,4 +1,3 @@
-import os.path
 import sys
 
 import PIL.Image
@@ -6,12 +5,11 @@ import numpy as np
 import torch
 from tqdm import tqdm
 
-from basicsr.utils.download_util import load_file_from_url
-
 import modules.upscaler
 from modules import devices, modelloader, script_callbacks, errors
-from scunet_model_arch import SCUNet as net
+from scunet_model_arch import SCUNet
 
+from modules.modelloader import load_file_from_url
 from modules.shared import opts
 
 
@@ -28,7 +26,7 @@ class UpscalerScuNET(modules.upscaler.Upscaler):
         scalers = []
         add_model2 = True
         for file in model_paths:
-            if "http" in file:
+            if file.startswith("http"):
                 name = self.model_name
             else:
                 name = modelloader.friendly_name(file)
@@ -87,11 +85,12 @@ class UpscalerScuNET(modules.upscaler.Upscaler):
 
     def do_upscale(self, img: PIL.Image.Image, selected_file):
 
-        torch.cuda.empty_cache()
+        devices.torch_gc()
 
+        try:
             model = self.load_model(selected_file)
-        if model is None:
+        except Exception as e:
-            print(f"ScuNET: Unable to load model from {selected_file}", file=sys.stderr)
+            print(f"ScuNET: Unable to load model from {selected_file}: {e}", file=sys.stderr)
             return img
 
         device = devices.get_device_for('scunet')
@@ -111,7 +110,7 @@ class UpscalerScuNET(modules.upscaler.Upscaler):
         torch_output = torch_output[:, :h * 1, :w * 1] # remove padding, if any
         np_output: np.ndarray = torch_output.float().cpu().clamp_(0, 1).numpy()
         del torch_img, torch_output
-        torch.cuda.empty_cache()
+        devices.torch_gc()
 
         output = np_output.transpose((1, 2, 0))  # CHW to HWC
         output = output[:, :, ::-1]  # BGR to RGB
@@ -119,15 +118,12 @@ class UpscalerScuNET(modules.upscaler.Upscaler):
 
     def load_model(self, path: str):
         device = devices.get_device_for('scunet')
-        if "http" in path:
+        if path.startswith("http"):
-            filename = load_file_from_url(url=self.model_url, model_dir=self.model_download_path, file_name="%s.pth" % self.name, progress=True)
+            # TODO: this doesn't use `path` at all?
+            filename = load_file_from_url(self.model_url, model_dir=self.model_download_path, file_name=f"{self.name}.pth")
         else:
             filename = path
-        if not os.path.exists(os.path.join(self.model_path, filename)) or filename is None:
+        model = SCUNet(in_nc=3, config=[4, 4, 4, 4, 4, 4, 4], dim=64)
-            print(f"ScuNET: Unable to load model from {filename}", file=sys.stderr)
-            return None
-
-        model = net(in_nc=3, config=[4, 4, 4, 4, 4, 4, 4], dim=64)
         model.load_state_dict(torch.load(filename), strict=True)
         model.eval()
         for _, v in model.named_parameters():

extensions-builtin/SwinIR/scripts/swinir_model.py

@@ -1,34 +1,35 @@
-import os
+import sys
+import platform
 
 import numpy as np
 import torch
 from PIL import Image
-from basicsr.utils.download_util import load_file_from_url
 from tqdm import tqdm
 
 from modules import modelloader, devices, script_callbacks, shared
 from modules.shared import opts, state
-from swinir_model_arch import SwinIR as net
+from swinir_model_arch import SwinIR
-from swinir_model_arch_v2 import Swin2SR as net2
+from swinir_model_arch_v2 import Swin2SR
 from modules.upscaler import Upscaler, UpscalerData
 
+SWINIR_MODEL_URL = "https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN.pth"
 
 device_swinir = devices.get_device_for('swinir')
 
 
 class UpscalerSwinIR(Upscaler):
     def __init__(self, dirname):
+        self._cached_model = None           # keep the model when SWIN_torch_compile is on to prevent re-compile every runs
+        self._cached_model_config = None    # to clear '_cached_model' when changing model (v1/v2) or settings
         self.name = "SwinIR"
-        self.model_url = "https://github.com/JingyunLiang/SwinIR/releases/download/v0.0" \
+        self.model_url = SWINIR_MODEL_URL
-                         "/003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR" \
-                         "-L_x4_GAN.pth "
         self.model_name = "SwinIR 4x"
         self.user_path = dirname
         super().__init__()
         scalers = []
         model_files = self.find_models(ext_filter=[".pt", ".pth"])
         for model in model_files:
-            if "http" in model:
+            if model.startswith("http"):
                 name = self.model_name
             else:
                 name = modelloader.friendly_name(model)
@@ -37,27 +38,39 @@ class UpscalerSwinIR(Upscaler):
         self.scalers = scalers
 
     def do_upscale(self, img, model_file):
+        use_compile = hasattr(opts, 'SWIN_torch_compile') and opts.SWIN_torch_compile \
+            and int(torch.__version__.split('.')[0]) >= 2 and platform.system() != "Windows"
+        current_config = (model_file, opts.SWIN_tile)
+
+        if use_compile and self._cached_model_config == current_config:
+            model = self._cached_model
+        else:
+            self._cached_model = None
+            try:
                 model = self.load_model(model_file)
-        if model is None:
+            except Exception as e:
+                print(f"Failed loading SwinIR model {model_file}: {e}", file=sys.stderr)
                 return img
             model = model.to(device_swinir, dtype=devices.dtype)
+            if use_compile:
+                model = torch.compile(model)
+                self._cached_model = model
+                self._cached_model_config = current_config
         img = upscale(img, model)
-        try:
+        devices.torch_gc()
-            torch.cuda.empty_cache()
-        except Exception:
-            pass
         return img
 
     def load_model(self, path, scale=4):
-        if "http" in path:
+        if path.startswith("http"):
-            dl_name = "%s%s" % (self.model_name.replace(" ", "_"), ".pth")
+            filename = modelloader.load_file_from_url(
-            filename = load_file_from_url(url=path, model_dir=self.model_download_path, file_name=dl_name, progress=True)
+                url=path,
+                model_dir=self.model_download_path,
+                file_name=f"{self.model_name.replace(' ', '_')}.pth",
+            )
         else:
             filename = path
-        if filename is None or not os.path.exists(filename):
-            return None
         if filename.endswith(".v2.pth"):
-            model = net2(
+            model = Swin2SR(
                 upscale=scale,
                 in_chans=3,
                 img_size=64,
@@ -72,7 +85,7 @@ class UpscalerSwinIR(Upscaler):
             )
             params = None
         else:
-            model = net(
+            model = SwinIR(
                 upscale=scale,
                 in_chans=3,
                 img_size=64,
@@ -172,6 +185,8 @@ def on_ui_settings():
 
     shared.opts.add_option("SWIN_tile", shared.OptionInfo(192, "Tile size for all SwinIR.", gr.Slider, {"minimum": 16, "maximum": 512, "step": 16}, section=('upscaling', "Upscaling")))
     shared.opts.add_option("SWIN_tile_overlap", shared.OptionInfo(8, "Tile overlap, in pixels for SwinIR. Low values = visible seam.", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}, section=('upscaling', "Upscaling")))
+    if int(torch.__version__.split('.')[0]) >= 2 and platform.system() != "Windows":    # torch.compile() require pytorch 2.0 or above, and not on Windows
+        shared.opts.add_option("SWIN_torch_compile", shared.OptionInfo(False, "Use torch.compile to accelerate SwinIR.", gr.Checkbox, {"interactive": True}, section=('upscaling', "Upscaling")).info("Takes longer on first run"))
 
 
 script_callbacks.on_ui_settings(on_ui_settings)

extensions-builtin/canvas-zoom-and-pan/javascript/zoom.js

@@ -4,12 +4,12 @@ onUiLoaded(async() => {
         inpaint: "#img2maskimg",
         inpaintSketch: "#inpaint_sketch",
         rangeGroup: "#img2img_column_size",
-        sketch: "#img2img_sketch",
+        sketch: "#img2img_sketch"
     };
     const tabNameToElementId = {
         "Inpaint sketch": elementIDs.inpaintSketch,
         "Inpaint": elementIDs.inpaint,
-        "Sketch": elementIDs.sketch,
+        "Sketch": elementIDs.sketch
     };
 
     // Helper functions
@@ -42,43 +42,110 @@ onUiLoaded(async() => {
         }
     }
 
-    // Check is hotkey valid
+    // Function for defining the "Ctrl", "Shift" and "Alt" keys
-    function isSingleLetter(value) {
+    function isModifierKey(event, key) {
+        switch (key) {
+        case "Ctrl":
+            return event.ctrlKey;
+        case "Shift":
+            return event.shiftKey;
+        case "Alt":
+            return event.altKey;
+        default:
+            return false;
+        }
+    }
+
+    // Check if hotkey is valid
+    function isValidHotkey(value) {
+        const specialKeys = ["Ctrl", "Alt", "Shift", "Disable"];
         return (
-            typeof value === "string" && value.length === 1 && /[a-z]/i.test(value)
+            (typeof value === "string" &&
+                value.length === 1 &&
+                /[a-z]/i.test(value)) ||
+            specialKeys.includes(value)
         );
     }
 
-    // Create hotkeyConfig from opts
+    // Normalize hotkey
+    function normalizeHotkey(hotkey) {
+        return hotkey.length === 1 ? "Key" + hotkey.toUpperCase() : hotkey;
+    }
+
+    // Format hotkey for display
+    function formatHotkeyForDisplay(hotkey) {
+        return hotkey.startsWith("Key") ? hotkey.slice(3) : hotkey;
+    }
+
+    // Create hotkey configuration with the provided options
     function createHotkeyConfig(defaultHotkeysConfig, hotkeysConfigOpts) {
-        const result = {};
+        const result = {}; // Resulting hotkey configuration
-        const usedKeys = new Set();
+        const usedKeys = new Set(); // Set of used hotkeys
 
+        // Iterate through defaultHotkeysConfig keys
         for (const key in defaultHotkeysConfig) {
-            if (typeof hotkeysConfigOpts[key] === "boolean") {
+            const userValue = hotkeysConfigOpts[key]; // User-provided hotkey value
-                result[key] = hotkeysConfigOpts[key];
+            const defaultValue = defaultHotkeysConfig[key]; // Default hotkey value
-                continue;
+
-            }
+            // Apply appropriate value for undefined, boolean, or object userValue
             if (
-                hotkeysConfigOpts[key] &&
+                userValue === undefined ||
-                isSingleLetter(hotkeysConfigOpts[key]) &&
+                typeof userValue === "boolean" ||
-                !usedKeys.has(hotkeysConfigOpts[key].toUpperCase())
+                typeof userValue === "object" ||
+                userValue === "disable"
             ) {
-                // If the property passed the test and has not yet been used, add 'Key' before it and save it
+                result[key] =
-                result[key] = "Key" + hotkeysConfigOpts[key].toUpperCase();
+                    userValue === undefined ? defaultValue : userValue;
-                usedKeys.add(hotkeysConfigOpts[key].toUpperCase());
+            } else if (isValidHotkey(userValue)) {
+                const normalizedUserValue = normalizeHotkey(userValue);
+
+                // Check for conflicting hotkeys
+                if (!usedKeys.has(normalizedUserValue)) {
+                    usedKeys.add(normalizedUserValue);
+                    result[key] = normalizedUserValue;
                 } else {
-                // If the property does not pass the test or has already been used, we keep the default value
                     console.error(
-                    `Hotkey: ${hotkeysConfigOpts[key]} for ${key} is repeated and conflicts with another hotkey or is not 1 letter. The default hotkey is used: ${defaultHotkeysConfig[key][3]}`
+                        `Hotkey: ${formatHotkeyForDisplay(
+                            userValue
+                        )} for ${key} is repeated and conflicts with another hotkey. The default hotkey is used: ${formatHotkeyForDisplay(
+                            defaultValue
+                        )}`
                     );
-                result[key] = defaultHotkeysConfig[key];
+                    result[key] = defaultValue;
+                }
+            } else {
+                console.error(
+                    `Hotkey: ${formatHotkeyForDisplay(
+                        userValue
+                    )} for ${key} is not valid. The default hotkey is used: ${formatHotkeyForDisplay(
+                        defaultValue
+                    )}`
+                );
+                result[key] = defaultValue;
             }
         }
 
         return result;
     }
 
+    // Disables functions in the config object based on the provided list of function names
+    function disableFunctions(config, disabledFunctions) {
+        // Bind the hasOwnProperty method to the functionMap object to avoid errors
+        const hasOwnProperty =
+            Object.prototype.hasOwnProperty.bind(functionMap);
+
+        // Loop through the disabledFunctions array and disable the corresponding functions in the config object
+        disabledFunctions.forEach(funcName => {
+            if (hasOwnProperty(funcName)) {
+                const key = functionMap[funcName];
+                config[key] = "disable";
+            }
+        });
+
+        // Return the updated config object
+        return config;
+    }
+
     /**
      * The restoreImgRedMask function displays a red mask around an image to indicate the aspect ratio.
      * If the image display property is set to 'none', the mask breaks. To fix this, the function
@@ -100,7 +167,9 @@ onUiLoaded(async() => {
         imageARPreview.style.transform = "";
         if (parseFloat(mainTab.style.width) > 865) {
             const transformString = mainTab.style.transform;
-            const scaleMatch = transformString.match(/scale\(([-+]?[0-9]*\.?[0-9]+)\)/);
+            const scaleMatch = transformString.match(
+                /scale\(([-+]?[0-9]*\.?[0-9]+)\)/
+            );
             let zoom = 1; // default zoom
 
             if (scaleMatch && scaleMatch[1]) {
@@ -124,31 +193,53 @@ onUiLoaded(async() => {
 
     // Default config
     const defaultHotkeysConfig = {
+        canvas_hotkey_zoom: "Alt",
+        canvas_hotkey_adjust: "Ctrl",
         canvas_hotkey_reset: "KeyR",
         canvas_hotkey_fullscreen: "KeyS",
         canvas_hotkey_move: "KeyF",
         canvas_hotkey_overlap: "KeyO",
+        canvas_disabled_functions: [],
         canvas_show_tooltip: true,
-        canvas_swap_controls: false
+        canvas_blur_prompt: false
     };
-    // swap the actions for ctr + wheel and shift + wheel
+
-    const hotkeysConfig = createHotkeyConfig(
+    const functionMap = {
+        "Zoom": "canvas_hotkey_zoom",
+        "Adjust brush size": "canvas_hotkey_adjust",
+        "Moving canvas": "canvas_hotkey_move",
+        "Fullscreen": "canvas_hotkey_fullscreen",
+        "Reset Zoom": "canvas_hotkey_reset",
+        "Overlap": "canvas_hotkey_overlap"
+    };
+
+    // Loading the configuration from opts
+    const preHotkeysConfig = createHotkeyConfig(
         defaultHotkeysConfig,
         hotkeysConfigOpts
     );
 
+    // Disable functions that are not needed by the user
+    const hotkeysConfig = disableFunctions(
+        preHotkeysConfig,
+        preHotkeysConfig.canvas_disabled_functions
+    );
+
     let isMoving = false;
     let mouseX, mouseY;
     let activeElement;
 
-    const elements = Object.fromEntries(Object.keys(elementIDs).map((id) => [
+    const elements = Object.fromEntries(
+        Object.keys(elementIDs).map(id => [
             id,
-        gradioApp().querySelector(elementIDs[id]),
+            gradioApp().querySelector(elementIDs[id])
-    ]));
+        ])
+    );
     const elemData = {};
 
     // Apply functionality to the range inputs. Restore redmask and correct for long images.
-    const rangeInputs = elements.rangeGroup ? Array.from(elements.rangeGroup.querySelectorAll("input")) :
+    const rangeInputs = elements.rangeGroup ?
+        Array.from(elements.rangeGroup.querySelectorAll("input")) :
         [
             gradioApp().querySelector("#img2img_width input[type='range']"),
             gradioApp().querySelector("#img2img_height input[type='range']")
@@ -180,38 +271,56 @@ onUiLoaded(async() => {
             const toolTipElemnt =
                 targetElement.querySelector(".image-container");
             const tooltip = document.createElement("div");
-            tooltip.className = "tooltip";
+            tooltip.className = "canvas-tooltip";
 
             // Creating an item of information
             const info = document.createElement("i");
-            info.className = "tooltip-info";
+            info.className = "canvas-tooltip-info";
             info.textContent = "";
 
             // Create a container for the contents of the tooltip
             const tooltipContent = document.createElement("div");
-            tooltipContent.className = "tooltip-content";
+            tooltipContent.className = "canvas-tooltip-content";
 
-            // Add info about hotkeys
+            // Define an array with hotkey information and their actions
-            const zoomKey = hotkeysConfig.canvas_swap_controls ? "Ctrl" : "Shift";
+            const hotkeysInfo = [
-            const adjustKey = hotkeysConfig.canvas_swap_controls ? "Shift" : "Ctrl";
-
-            const hotkeys = [
-                {key: `${zoomKey} + wheel`, action: "Zoom canvas"},
-                {key: `${adjustKey} + wheel`, action: "Adjust brush size"},
                 {
-                    key: hotkeysConfig.canvas_hotkey_reset.charAt(hotkeysConfig.canvas_hotkey_reset.length - 1),
+                    configKey: "canvas_hotkey_zoom",
-                    action: "Reset zoom"
+                    action: "Zoom canvas",
+                    keySuffix: " + wheel"
                 },
                 {
-                    key: hotkeysConfig.canvas_hotkey_fullscreen.charAt(hotkeysConfig.canvas_hotkey_fullscreen.length - 1),
+                    configKey: "canvas_hotkey_adjust",
+                    action: "Adjust brush size",
+                    keySuffix: " + wheel"
+                },
+                {configKey: "canvas_hotkey_reset", action: "Reset zoom"},
+                {
+                    configKey: "canvas_hotkey_fullscreen",
                     action: "Fullscreen mode"
                 },
-                {
+                {configKey: "canvas_hotkey_move", action: "Move canvas"},
-                    key: hotkeysConfig.canvas_hotkey_move.charAt(hotkeysConfig.canvas_hotkey_move.length - 1),
+                {configKey: "canvas_hotkey_overlap", action: "Overlap"}
-                    action: "Move canvas"
-                }
             ];
+
+            // Create hotkeys array with disabled property based on the config values
+            const hotkeys = hotkeysInfo.map(info => {
+                const configValue = hotkeysConfig[info.configKey];
+                const key = info.keySuffix ?
+                    `${configValue}${info.keySuffix}` :
+                    configValue.charAt(configValue.length - 1);
+                return {
+                    key,
+                    action: info.action,
+                    disabled: configValue === "disable"
+                };
+            });
+
             for (const hotkey of hotkeys) {
+                if (hotkey.disabled) {
+                    continue;
+                }
+
                 const p = document.createElement("p");
                 p.innerHTML = `<b>${hotkey.key}</b> - ${hotkey.action}`;
                 tooltipContent.appendChild(p);
@@ -346,10 +455,7 @@ onUiLoaded(async() => {
 
         // Change the zoom level based on user interaction
         function changeZoomLevel(operation, e) {
-            if (
+            if (isModifierKey(e, hotkeysConfig.canvas_hotkey_zoom)) {
-                (!hotkeysConfig.canvas_swap_controls && e.shiftKey) ||
-                (hotkeysConfig.canvas_swap_controls && e.ctrlKey)
-            ) {
                 e.preventDefault();
 
                 let zoomPosX, zoomPosY;
@@ -503,6 +609,19 @@ onUiLoaded(async() => {
 
         // Handle keydown events
         function handleKeyDown(event) {
+            // Disable key locks to make pasting from the buffer work correctly
+            if ((event.ctrlKey && event.code === 'KeyV') || (event.ctrlKey && event.code === 'KeyC') || event.code === "F5") {
+                return;
+            }
+
+            // before activating shortcut, ensure user is not actively typing in an input field
+            if (!hotkeysConfig.canvas_blur_prompt) {
+                if (event.target.nodeName === 'TEXTAREA' || event.target.nodeName === 'INPUT') {
+                    return;
+                }
+            }
+
+
             const hotkeyActions = {
                 [hotkeysConfig.canvas_hotkey_reset]: resetZoom,
                 [hotkeysConfig.canvas_hotkey_overlap]: toggleOverlap,
@@ -514,6 +633,13 @@ onUiLoaded(async() => {
                 event.preventDefault();
                 action(event);
             }
+
+            if (
+                isModifierKey(event, hotkeysConfig.canvas_hotkey_zoom) ||
+                isModifierKey(event, hotkeysConfig.canvas_hotkey_adjust)
+            ) {
+                event.preventDefault();
+            }
         }
 
         // Get Mouse position
@@ -564,11 +690,7 @@ onUiLoaded(async() => {
             changeZoomLevel(operation, e);
 
             // Handle brush size adjustment with ctrl key pressed
-            if (
+            if (isModifierKey(e, hotkeysConfig.canvas_hotkey_adjust)) {
-                (hotkeysConfig.canvas_swap_controls && e.shiftKey) ||
-                (!hotkeysConfig.canvas_swap_controls &&
-                    (e.ctrlKey || e.metaKey))
-            ) {
                 e.preventDefault();
 
                 // Increase or decrease brush size based on scroll direction
@@ -578,6 +700,20 @@ onUiLoaded(async() => {
 
         // Handle the move event for pan functionality. Updates the panX and panY variables and applies the new transform to the target element.
         function handleMoveKeyDown(e) {
+
+            // Disable key locks to make pasting from the buffer work correctly
+            if ((e.ctrlKey && e.code === 'KeyV') || (e.ctrlKey && event.code === 'KeyC') || e.code === "F5") {
+                return;
+            }
+
+            // before activating shortcut, ensure user is not actively typing in an input field
+            if (!hotkeysConfig.canvas_blur_prompt) {
+                if (e.target.nodeName === 'TEXTAREA' || e.target.nodeName === 'INPUT') {
+                    return;
+                }
+            }
+
+
             if (e.code === hotkeysConfig.canvas_hotkey_move) {
                 if (!e.ctrlKey && !e.metaKey && isKeyDownHandlerAttached) {
                     e.preventDefault();

extensions-builtin/canvas-zoom-and-pan/scripts/hotkey_config.py

@@ -1,10 +1,14 @@
+import gradio as gr
 from modules import shared
 
 shared.options_templates.update(shared.options_section(('canvas_hotkey', "Canvas Hotkeys"), {
-    "canvas_hotkey_move": shared.OptionInfo("F", "Moving the canvas"),
+    "canvas_hotkey_zoom": shared.OptionInfo("Alt", "Zoom canvas", gr.Radio, {"choices": ["Shift","Ctrl", "Alt"]}).info("If you choose 'Shift' you cannot scroll horizontally, 'Alt' can cause a little trouble in firefox"),
+    "canvas_hotkey_adjust": shared.OptionInfo("Ctrl", "Adjust brush size", gr.Radio, {"choices": ["Shift","Ctrl", "Alt"]}).info("If you choose 'Shift' you cannot scroll horizontally, 'Alt' can cause a little trouble in firefox"),
+    "canvas_hotkey_move": shared.OptionInfo("F", "Moving the canvas").info("To work correctly in firefox, turn off 'Automatically search the page text when typing' in the browser settings"),
     "canvas_hotkey_fullscreen": shared.OptionInfo("S", "Fullscreen Mode, maximizes the picture so that it fits into the screen and stretches it to its full width "),
     "canvas_hotkey_reset": shared.OptionInfo("R", "Reset zoom and canvas positon"),
-    "canvas_hotkey_overlap": shared.OptionInfo("O", "Toggle overlap ( Technical button, neededs for testing )"),
+    "canvas_hotkey_overlap": shared.OptionInfo("O", "Toggle overlap").info("Technical button, neededs for testing"),
     "canvas_show_tooltip": shared.OptionInfo(True, "Enable tooltip on the canvas"),
-    "canvas_swap_controls": shared.OptionInfo(False, "Swap hotkey combinations for Zoom and Adjust brush resize"),
+    "canvas_blur_prompt": shared.OptionInfo(False, "Take the focus off the prompt when working with a canvas"),
+    "canvas_disabled_functions": shared.OptionInfo(["Overlap"], "Disable function that you don't use", gr.CheckboxGroup, {"choices": ["Zoom","Adjust brush size", "Moving canvas","Fullscreen","Reset Zoom","Overlap"]}),
 }))

extensions-builtin/canvas-zoom-and-pan/style.css

@@ -1,4 +1,4 @@
-.tooltip-info {
+.canvas-tooltip-info {
   position: absolute;
   top: 10px;
   left: 10px;
@@ -15,7 +15,7 @@
   z-index: 100;
 }
 
-.tooltip-info::after {
+.canvas-tooltip-info::after {
   content: '';
   display: block;
   width: 2px;
@@ -24,7 +24,7 @@
   margin-top: 2px; 
 }
 
-.tooltip-info::before {
+.canvas-tooltip-info::before {
   content: '';
   display: block;
   width: 2px;
@@ -32,7 +32,7 @@
   background-color: white;
 }
 
-.tooltip-content {
+.canvas-tooltip-content {
   display: none;
   background-color: #f9f9f9; 
   color: #333; 
@@ -50,7 +50,7 @@
   z-index: 100;
 }
 
-.tooltip:hover .tooltip-content {
+.canvas-tooltip:hover .canvas-tooltip-content {
   display: block;
   animation: fadeIn 0.5s; 
   opacity: 1;

extensions-builtin/extra-options-section/scripts/extra_options_section.py

@@ -43,6 +43,6 @@ class ExtraOptionsSection(scripts.Script):
 
 
 shared.options_templates.update(shared.options_section(('ui', "User interface"), {
-    "extra_options": shared.OptionInfo([], "Options in main UI", ui_components.DropdownMulti, lambda: {"choices": list(shared.opts.data_labels.keys())}).js("info", "settingsHintsShowQuicksettings").info("setting entries that also appear in txt2img/img2img interfaces").needs_restart(),
+    "extra_options": shared.OptionInfo([], "Options in main UI", ui_components.DropdownMulti, lambda: {"choices": list(shared.opts.data_labels.keys())}).js("info", "settingsHintsShowQuicksettings").info("setting entries that also appear in txt2img/img2img interfaces").needs_reload_ui(),
-    "extra_options_accordion": shared.OptionInfo(False, "Place options in main UI into an accordion")
+    "extra_options_accordion": shared.OptionInfo(False, "Place options in main UI into an accordion").needs_restart()
 }))

extensions-builtin/mobile/javascript/mobile.js

@@ -0,0 +1,26 @@
+var isSetupForMobile = false;
+
+function isMobile() {
+    for (var tab of ["txt2img", "img2img"]) {
+        var imageTab = gradioApp().getElementById(tab + '_results');
+        if (imageTab && imageTab.offsetParent && imageTab.offsetLeft == 0) {
+            return true;
+        }
+    }
+
+    return false;
+}
+
+function reportWindowSize() {
+    var currentlyMobile = isMobile();
+    if (currentlyMobile == isSetupForMobile) return;
+    isSetupForMobile = currentlyMobile;
+
+    for (var tab of ["txt2img", "img2img"]) {
+        var button = gradioApp().getElementById(tab + '_generate_box');
+        var target = gradioApp().getElementById(currentlyMobile ? tab + '_results' : tab + '_actions_column');
+        target.insertBefore(button, target.firstElementChild);
+    }
+}
+
+window.addEventListener("resize", reportWindowSize);

html/extra-networks-card.html

@@ -1,11 +1,11 @@
-<div class='card' style={style} onclick={card_clicked} {sort_keys}>
+<div class='card' style={style} onclick={card_clicked} data-name="{name}" {sort_keys}>
 	{background_image}
+	<div class="button-row">
 		{metadata_button}
+		{edit_button}
+	</div>
 	<div class='actions'>
 		<div class='additional'>
-			<ul>
-				<a href="#" title="replace preview image with currently selected in gallery" onclick={save_card_preview}>replace preview</a>
-			</ul>
 			<span style="display:none" class='search_term{search_only}'>{search_term}</span>
 		</div>
 		<span class='name'>{name}</span>

html/image-update.svg

@@ -1,7 +0,0 @@
-<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
-  <filter id='shadow' color-interpolation-filters="sRGB">
-    <feDropShadow flood-color="black" dx="0" dy="0" flood-opacity="0.9" stdDeviation="0.5"/>
-    <feDropShadow flood-color="black" dx="0" dy="0" flood-opacity="0.9" stdDeviation="0.5"/>
-  </filter>
-  <path style="filter:url(#shadow);" fill="#FFFFFF" d="M13.18 19C13.35 19.72 13.64 20.39 14.03 21H5C3.9 21 3 20.11 3 19V5C3 3.9 3.9 3 5 3H19C20.11 3 21 3.9 21 5V11.18C20.5 11.07 20 11 19.5 11C19.33 11 19.17 11 19 11.03V5H5V19H13.18M11.21 15.83L9.25 13.47L6.5 17H13.03C13.14 15.54 13.73 14.22 14.64 13.19L13.96 12.29L11.21 15.83M19 13.5V12L16.75 14.25L19 16.5V15C20.38 15 21.5 16.12 21.5 17.5C21.5 17.9 21.41 18.28 21.24 18.62L22.33 19.71C22.75 19.08 23 18.32 23 17.5C23 15.29 21.21 13.5 19 13.5M19 20C17.62 20 16.5 18.88 16.5 17.5C16.5 17.1 16.59 16.72 16.76 16.38L15.67 15.29C15.25 15.92 15 16.68 15 17.5C15 19.71 16.79 21.5 19 21.5V23L21.25 20.75L19 18.5V20Z" />
-</svg>

javascript/edit-attention.js

@@ -100,11 +100,12 @@ function keyupEditAttention(event) {
     if (String(weight).length == 1) weight += ".0";
 
     if (closeCharacter == ')' && weight == 1) {
-        text = text.slice(0, selectionStart - 1) + text.slice(selectionStart, selectionEnd) + text.slice(selectionEnd + 5);
+        var endParenPos = text.substring(selectionEnd).indexOf(')');
+        text = text.slice(0, selectionStart - 1) + text.slice(selectionStart, selectionEnd) + text.slice(selectionEnd + endParenPos + 1);
         selectionStart--;
         selectionEnd--;
     } else {
-        text = text.slice(0, selectionEnd + 1) + weight + text.slice(selectionEnd + 1 + end - 1);
+        text = text.slice(0, selectionEnd + 1) + weight + text.slice(selectionEnd + end);
     }
 
     target.focus();

javascript/edit-order.js

@@ -0,0 +1,41 @@
+/* alt+left/right moves text in prompt */
+
+function keyupEditOrder(event) {
+    if (!opts.keyedit_move) return;
+
+    let target = event.originalTarget || event.composedPath()[0];
+    if (!target.matches("*:is([id*='_toprow'] [id*='_prompt'], .prompt) textarea")) return;
+    if (!event.altKey) return;
+
+    let isLeft = event.key == "ArrowLeft";
+    let isRight = event.key == "ArrowRight";
+    if (!isLeft && !isRight) return;
+    event.preventDefault();
+
+    let selectionStart = target.selectionStart;
+    let selectionEnd = target.selectionEnd;
+    let text = target.value;
+    let items = text.split(",");
+    let indexStart = (text.slice(0, selectionStart).match(/,/g) || []).length;
+    let indexEnd = (text.slice(0, selectionEnd).match(/,/g) || []).length;
+    let range = indexEnd - indexStart + 1;
+
+    if (isLeft && indexStart > 0) {
+        items.splice(indexStart - 1, 0, ...items.splice(indexStart, range));
+        target.value = items.join();
+        target.selectionStart = items.slice(0, indexStart - 1).join().length + (indexStart == 1 ? 0 : 1);
+        target.selectionEnd = items.slice(0, indexEnd).join().length;
+    } else if (isRight && indexEnd < items.length - 1) {
+        items.splice(indexStart + 1, 0, ...items.splice(indexStart, range));
+        target.value = items.join();
+        target.selectionStart = items.slice(0, indexStart + 1).join().length + 1;
+        target.selectionEnd = items.slice(0, indexEnd + 2).join().length;
+    }
+
+    event.preventDefault();
+    updateInput(target);
+}
+
+addEventListener('keydown', (event) => {
+    keyupEditOrder(event);
+});

javascript/extensions.js

@@ -72,3 +72,21 @@ function config_state_confirm_restore(_, config_state_name, config_restore_type)
     }
     return [confirmed, config_state_name, config_restore_type];
 }
+
+function toggle_all_extensions(event) {
+    gradioApp().querySelectorAll('#extensions .extension_toggle').forEach(function(checkbox_el) {
+        checkbox_el.checked = event.target.checked;
+    });
+}
+
+function toggle_extension() {
+    let all_extensions_toggled = true;
+    for (const checkbox_el of gradioApp().querySelectorAll('#extensions .extension_toggle')) {
+        if (!checkbox_el.checked) {
+            all_extensions_toggled = false;
+            break;
+        }
+    }
+
+    gradioApp().querySelector('#extensions .all_extensions_toggle').checked = all_extensions_toggled;
+}

javascript/extraNetworks.js

@@ -1,20 +1,38 @@
+function toggleCss(key, css, enable) {
+    var style = document.getElementById(key);
+    if (enable && !style) {
+        style = document.createElement('style');
+        style.id = key;
+        style.type = 'text/css';
+        document.head.appendChild(style);
+    }
+    if (style && !enable) {
+        document.head.removeChild(style);
+    }
+    if (style) {
+        style.innerHTML == '';
+        style.appendChild(document.createTextNode(css));
+    }
+}
+
 function setupExtraNetworksForTab(tabname) {
     gradioApp().querySelector('#' + tabname + '_extra_tabs').classList.add('extra-networks');
 
     var tabs = gradioApp().querySelector('#' + tabname + '_extra_tabs > div');
-    var search = gradioApp().querySelector('#' + tabname + '_extra_search textarea');
+    var searchDiv = gradioApp().getElementById(tabname + '_extra_search');
+    var search = searchDiv.querySelector('textarea');
     var sort = gradioApp().getElementById(tabname + '_extra_sort');
     var sortOrder = gradioApp().getElementById(tabname + '_extra_sortorder');
     var refresh = gradioApp().getElementById(tabname + '_extra_refresh');
+    var showDirsDiv = gradioApp().getElementById(tabname + '_extra_show_dirs');
+    var showDirs = gradioApp().querySelector('#' + tabname + '_extra_show_dirs input');
 
-    search.classList.add('search');
-    sort.classList.add('sort');
-    sortOrder.classList.add('sortorder');
     sort.dataset.sortkey = 'sortDefault';
-    tabs.appendChild(search);
+    tabs.appendChild(searchDiv);
     tabs.appendChild(sort);
     tabs.appendChild(sortOrder);
     tabs.appendChild(refresh);
+    tabs.appendChild(showDirsDiv);
 
     var applyFilter = function() {
         var searchTerm = search.value.toLowerCase();
@@ -80,6 +98,15 @@ function setupExtraNetworksForTab(tabname) {
     });
 
     extraNetworksApplyFilter[tabname] = applyFilter;
+
+    var showDirsUpdate = function() {
+        var css = '#' + tabname + '_extra_tabs .extra-network-subdirs { display: none; }';
+        toggleCss(tabname + '_extra_show_dirs_style', css, !showDirs.checked);
+        localSet('extra-networks-show-dirs', showDirs.checked ? 1 : 0);
+    };
+    showDirs.checked = localGet('extra-networks-show-dirs', 1) == 1;
+    showDirs.addEventListener("change", showDirsUpdate);
+    showDirsUpdate();
 }
 
 function applyExtraNetworkFilter(tabname) {
@@ -113,7 +140,7 @@ function setupExtraNetworks() {
 
 onUiLoaded(setupExtraNetworks);
 
-var re_extranet = /<([^:]+:[^:]+):[\d.]+>/;
+var re_extranet = /<([^:]+:[^:]+):[\d.]+>(.*)/;
 var re_extranet_g = /\s+<([^:]+:[^:]+):[\d.]+>/g;
 
 function tryToRemoveExtraNetworkFromPrompt(textarea, text) {
@@ -121,15 +148,22 @@ function tryToRemoveExtraNetworkFromPrompt(textarea, text) {
     var replaced = false;
     var newTextareaText;
     if (m) {
+        var extraTextAfterNet = m[2];
         var partToSearch = m[1];
-        newTextareaText = textarea.value.replaceAll(re_extranet_g, function(found) {
+        var foundAtPosition = -1;
+        newTextareaText = textarea.value.replaceAll(re_extranet_g, function(found, net, pos) {
             m = found.match(re_extranet);
             if (m[1] == partToSearch) {
                 replaced = true;
+                foundAtPosition = pos;
                 return "";
             }
             return found;
         });
+
+        if (foundAtPosition >= 0 && newTextareaText.substr(foundAtPosition, extraTextAfterNet.length) == extraTextAfterNet) {
+            newTextareaText = newTextareaText.substr(0, foundAtPosition) + newTextareaText.substr(foundAtPosition + extraTextAfterNet.length);
+        }
     } else {
         newTextareaText = textarea.value.replaceAll(new RegExp(text, "g"), function(found) {
             if (found == text) {
@@ -172,7 +206,7 @@ function saveCardPreview(event, tabname, filename) {
 }
 
 function extraNetworksSearchButton(tabs_id, event) {
-    var searchTextarea = gradioApp().querySelector("#" + tabs_id + ' > div > textarea');
+    var searchTextarea = gradioApp().querySelector("#" + tabs_id + ' > label > textarea');
     var button = event.target;
     var text = button.classList.contains("search-all") ? "" : button.textContent.trim();
 
@@ -182,19 +216,20 @@ function extraNetworksSearchButton(tabs_id, event) {
 
 var globalPopup = null;
 var globalPopupInner = null;
+function closePopup() {
+    if (!globalPopup) return;
+
+    globalPopup.style.display = "none";
+}
 function popup(contents) {
     if (!globalPopup) {
         globalPopup = document.createElement('div');
-        globalPopup.onclick = function() {
+        globalPopup.onclick = closePopup;
-            globalPopup.style.display = "none";
-        };
         globalPopup.classList.add('global-popup');
 
         var close = document.createElement('div');
         close.classList.add('global-popup-close');
-        close.onclick = function() {
+        close.onclick = closePopup;
-            globalPopup.style.display = "none";
-        };
         close.title = "Close";
         globalPopup.appendChild(close);
 
@@ -205,7 +240,7 @@ function popup(contents) {
         globalPopupInner.classList.add('global-popup-inner');
         globalPopup.appendChild(globalPopupInner);
 
-        gradioApp().appendChild(globalPopup);
+        gradioApp().querySelector('.main').appendChild(globalPopup);
     }
 
     globalPopupInner.innerHTML = '';
@@ -263,3 +298,43 @@ function extraNetworksRequestMetadata(event, extraPage, cardName) {
 
     event.stopPropagation();
 }
+
+var extraPageUserMetadataEditors = {};
+
+function extraNetworksEditUserMetadata(event, tabname, extraPage, cardName) {
+    var id = tabname + '_' + extraPage + '_edit_user_metadata';
+
+    var editor = extraPageUserMetadataEditors[id];
+    if (!editor) {
+        editor = {};
+        editor.page = gradioApp().getElementById(id);
+        editor.nameTextarea = gradioApp().querySelector("#" + id + "_name" + ' textarea');
+        editor.button = gradioApp().querySelector("#" + id + "_button");
+        extraPageUserMetadataEditors[id] = editor;
+    }
+
+    editor.nameTextarea.value = cardName;
+    updateInput(editor.nameTextarea);
+
+    editor.button.click();
+
+    popup(editor.page);
+
+    event.stopPropagation();
+}
+
+function extraNetworksRefreshSingleCard(page, tabname, name) {
+    requestGet("./sd_extra_networks/get-single-card", {page: page, tabname: tabname, name: name}, function(data) {
+        if (data && data.html) {
+            var card = gradioApp().querySelector('.card[data-name=' + JSON.stringify(name) + ']'); // likely using the wrong stringify function
+
+            var newDiv = document.createElement('DIV');
+            newDiv.innerHTML = data.html;
+            var newCard = newDiv.firstElementChild;
+
+            newCard.style = '';
+            card.parentElement.insertBefore(newCard, card);
+            card.parentElement.removeChild(card);
+        }
+    });
+}

javascript/hints.js

@@ -15,7 +15,7 @@ var titles = {
     "CFG Scale": "Classifier Free Guidance Scale - how strongly the image should conform to prompt - lower values produce more creative results",
     "Seed": "A value that determines the output of random number generator - if you create an image with same parameters and seed as another image, you'll get the same result",
     "\u{1f3b2}\ufe0f": "Set seed to -1, which will cause a new random number to be used every time",
-    "\u267b\ufe0f": "Reuse seed from last generation, mostly useful if it was randomed",
+    "\u267b\ufe0f": "Reuse seed from last generation, mostly useful if it was randomized",
     "\u2199\ufe0f": "Read generation parameters from prompt or last generation if prompt is empty into user interface.",
     "\u{1f4c2}": "Open images output directory",
     "\u{1f4be}": "Save style",
@@ -84,8 +84,6 @@ var titles = {
     "Checkpoint name": "Loads weights from checkpoint before making images. You can either use hash or a part of filename (as seen in settings) for checkpoint name. Recommended to use with Y axis for less switching.",
     "Inpainting conditioning mask strength": "Only applies to inpainting models. Determines how strongly to mask off the original image for inpainting and img2img. 1.0 means fully masked, which is the default behaviour. 0.0 means a fully unmasked conditioning. Lower values will help preserve the overall composition of the image, but will struggle with large changes.",
 
-    "vram": "Torch active: Peak amount of VRAM used by Torch during generation, excluding cached data.\nTorch reserved: Peak amount of VRAM allocated by Torch, including all active and cached data.\nSys VRAM: Peak amount of VRAM allocation across all applications / total GPU VRAM (peak utilization%).",
-
     "Eta noise seed delta": "If this values is non-zero, it will be added to seed and used to initialize RNG for noises when using samplers with Eta. You can use this to produce even more variation of images, or you can use this to match images of other software if you know what you are doing.",
 
     "Filename word regex": "This regular expression will be used extract words from filename, and they will be joined using the option below into label text used for training. Leave empty to keep filename text as it is.",
@@ -110,9 +108,8 @@ var titles = {
     "Upscale by": "Adjusts the size of the image by multiplying the original width and height by the selected value. Ignored if either Resize width to or Resize height to are non-zero.",
     "Resize width to": "Resizes image to this width. If 0, width is inferred from either of two nearby sliders.",
     "Resize height to": "Resizes image to this height. If 0, height is inferred from either of two nearby sliders.",
-    "Multiplier for extra networks": "When adding extra network such as Hypernetwork or Lora to prompt, use this multiplier for it.",
     "Discard weights with matching name": "Regular expression; if weights's name matches it, the weights is not written to the resulting checkpoint. Use ^model_ema to discard EMA weights.",
-    "Extra networks tab order": "Comma-separated list of tab names; tabs listed here will appear in the extra networks UI first and in order lsited.",
+    "Extra networks tab order": "Comma-separated list of tab names; tabs listed here will appear in the extra networks UI first and in order listed.",
     "Negative Guidance minimum sigma": "Skip negative prompt for steps where image is already mostly denoised; the higher this value, the more skips there will be; provides increased performance in exchange for minor quality reduction."
 };
 
@@ -193,3 +190,14 @@ onUiUpdate(function(mutationRecords) {
         tooltipCheckTimer = setTimeout(processTooltipCheckNodes, 1000);
     }
 });
+
+onUiLoaded(function() {
+    for (var comp of window.gradio_config.components) {
+        if (comp.props.webui_tooltip && comp.props.elem_id) {
+            var elem = gradioApp().getElementById(comp.props.elem_id);
+            if (elem) {
+                elem.title = comp.props.webui_tooltip;
+            }
+        }
+    }
+});

javascript/localStorage.js

@@ -0,0 +1,26 @@
+
+function localSet(k, v) {
+    try {
+        localStorage.setItem(k, v);
+    } catch (e) {
+        console.warn(`Failed to save ${k} to localStorage: ${e}`);
+    }
+}
+
+function localGet(k, def) {
+    try {
+        return localStorage.getItem(k);
+    } catch (e) {
+        console.warn(`Failed to load ${k} from localStorage: ${e}`);
+    }
+
+    return def;
+}
+
+function localRemove(k) {
+    try {
+        return localStorage.removeItem(k);
+    } catch (e) {
+        console.warn(`Failed to remove ${k} from localStorage: ${e}`);
+    }
+}

javascript/localization.js

@@ -11,11 +11,11 @@ var ignore_ids_for_localization = {
     train_hypernetwork: 'OPTION',
     txt2img_styles: 'OPTION',
     img2img_styles: 'OPTION',
-    setting_random_artist_categories: 'SPAN',
+    setting_random_artist_categories: 'OPTION',
-    setting_face_restoration_model: 'SPAN',
+    setting_face_restoration_model: 'OPTION',
-    setting_realesrgan_enabled_models: 'SPAN',
+    setting_realesrgan_enabled_models: 'OPTION',
-    extras_upscaler_1: 'SPAN',
+    extras_upscaler_1: 'OPTION',
-    extras_upscaler_2: 'SPAN',
+    extras_upscaler_2: 'OPTION',
 };
 
 var re_num = /^[.\d]+$/;

javascript/ui.js

@@ -152,11 +152,11 @@ function submit() {
     showSubmitButtons('txt2img', false);
 
     var id = randomId();
-    localStorage.setItem("txt2img_task_id", id);
+    localSet("txt2img_task_id", id);
 
     requestProgress(id, gradioApp().getElementById('txt2img_gallery_container'), gradioApp().getElementById('txt2img_gallery'), function() {
         showSubmitButtons('txt2img', true);
-        localStorage.removeItem("txt2img_task_id");
+        localRemove("txt2img_task_id");
         showRestoreProgressButton('txt2img', false);
     });
 
@@ -171,11 +171,11 @@ function submit_img2img() {
     showSubmitButtons('img2img', false);
 
     var id = randomId();
-    localStorage.setItem("img2img_task_id", id);
+    localSet("img2img_task_id", id);
 
     requestProgress(id, gradioApp().getElementById('img2img_gallery_container'), gradioApp().getElementById('img2img_gallery'), function() {
         showSubmitButtons('img2img', true);
-        localStorage.removeItem("img2img_task_id");
+        localRemove("img2img_task_id");
         showRestoreProgressButton('img2img', false);
     });
 
@@ -189,9 +189,7 @@ function submit_img2img() {
 
 function restoreProgressTxt2img() {
     showRestoreProgressButton("txt2img", false);
-    var id = localStorage.getItem("txt2img_task_id");
+    var id = localGet("txt2img_task_id");
-
-    id = localStorage.getItem("txt2img_task_id");
 
     if (id) {
         requestProgress(id, gradioApp().getElementById('txt2img_gallery_container'), gradioApp().getElementById('txt2img_gallery'), function() {
@@ -205,7 +203,7 @@ function restoreProgressTxt2img() {
 function restoreProgressImg2img() {
     showRestoreProgressButton("img2img", false);
 
-    var id = localStorage.getItem("img2img_task_id");
+    var id = localGet("img2img_task_id");
 
     if (id) {
         requestProgress(id, gradioApp().getElementById('img2img_gallery_container'), gradioApp().getElementById('img2img_gallery'), function() {
@@ -218,8 +216,8 @@ function restoreProgressImg2img() {
 
 
 onUiLoaded(function() {
-    showRestoreProgressButton('txt2img', localStorage.getItem("txt2img_task_id"));
+    showRestoreProgressButton('txt2img', localGet("txt2img_task_id"));
-    showRestoreProgressButton('img2img', localStorage.getItem("img2img_task_id"));
+    showRestoreProgressButton('img2img', localGet("img2img_task_id"));
 });
 
 

launch.py

@@ -1,6 +1,5 @@
 from modules import launch_utils
 
-
 args = launch_utils.args
 python = launch_utils.python
 git = launch_utils.git
@@ -18,6 +17,7 @@ run_pip = launch_utils.run_pip
 check_run_python = launch_utils.check_run_python
 git_clone = launch_utils.git_clone
 git_pull_recursive = launch_utils.git_pull_recursive
+list_extensions = launch_utils.list_extensions
 run_extension_installer = launch_utils.run_extension_installer
 prepare_environment = launch_utils.prepare_environment
 configure_for_tests = launch_utils.configure_for_tests
@@ -25,6 +25,9 @@ start = launch_utils.start
 
 
 def main():
+    launch_utils.startup_timer.record("initial startup")
+
+    with launch_utils.startup_timer.subcategory("prepare environment"):
         if not args.skip_prepare_environment:
             prepare_environment()
 

modules/api/api.py

@@ -1,5 +1,6 @@
 import base64
 import io
+import os
 import time
 import datetime
 import uvicorn
@@ -14,7 +15,7 @@ from fastapi.encoders import jsonable_encoder
 from secrets import compare_digest
 
 import modules.shared as shared
-from modules import sd_samplers, deepbooru, sd_hijack, images, scripts, ui, postprocessing, errors
+from modules import sd_samplers, deepbooru, sd_hijack, images, scripts, ui, postprocessing, errors, restart, shared_items
 from modules.api import models
 from modules.shared import opts
 from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, process_images
@@ -22,7 +23,7 @@ from modules.textual_inversion.textual_inversion import create_embedding, train_
 from modules.textual_inversion.preprocess import preprocess
 from modules.hypernetworks.hypernetwork import create_hypernetwork, train_hypernetwork
 from PIL import PngImagePlugin,Image
-from modules.sd_models import checkpoints_list, unload_model_weights, reload_model_weights
+from modules.sd_models import checkpoints_list, unload_model_weights, reload_model_weights, checkpoint_aliases
 from modules.sd_vae import vae_dict
 from modules.sd_models_config import find_checkpoint_config_near_filename
 from modules.realesrgan_model import get_realesrgan_models
@@ -30,13 +31,7 @@ from modules import devices
 from typing import Dict, List, Any
 import piexif
 import piexif.helper
-
+from contextlib import closing
-
-def upscaler_to_index(name: str):
-    try:
-        return [x.name.lower() for x in shared.sd_upscalers].index(name.lower())
-    except Exception as e:
-        raise HTTPException(status_code=400, detail=f"Invalid upscaler, needs to be one of these: {' , '.join([x.name for x in shared.sd_upscalers])}") from e
 
 
 def script_name_to_index(name, scripts):
@@ -84,6 +79,8 @@ def encode_pil_to_base64(image):
             image.save(output_bytes, format="PNG", pnginfo=(metadata if use_metadata else None), quality=opts.jpeg_quality)
 
         elif opts.samples_format.lower() in ("jpg", "jpeg", "webp"):
+            if image.mode == "RGBA":
+                image = image.convert("RGB")
             parameters = image.info.get('parameters', None)
             exif_bytes = piexif.dump({
                 "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(parameters or "", encoding="unicode") }
@@ -102,14 +99,16 @@ def encode_pil_to_base64(image):
 
 
 def api_middleware(app: FastAPI):
-    rich_available = True
+    rich_available = False
     try:
+        if os.environ.get('WEBUI_RICH_EXCEPTIONS', None) is not None:
             import anyio  # importing just so it can be placed on silent list
             import starlette  # importing just so it can be placed on silent list
             from rich.console import Console
             console = Console()
+            rich_available = True
     except Exception:
-        rich_available = False
+        pass
 
     @app.middleware("http")
     async def log_and_time(req: Request, call_next):
@@ -120,14 +119,14 @@ def api_middleware(app: FastAPI):
         endpoint = req.scope.get('path', 'err')
         if shared.cmd_opts.api_log and endpoint.startswith('/sdapi'):
             print('API {t} {code} {prot}/{ver} {method} {endpoint} {cli} {duration}'.format(
-                t = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"),
+                t=datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"),
-                code = res.status_code,
+                code=res.status_code,
-                ver = req.scope.get('http_version', '0.0'),
+                ver=req.scope.get('http_version', '0.0'),
-                cli = req.scope.get('client', ('0:0.0.0', 0))[0],
+                cli=req.scope.get('client', ('0:0.0.0', 0))[0],
-                prot = req.scope.get('scheme', 'err'),
+                prot=req.scope.get('scheme', 'err'),
-                method = req.scope.get('method', 'err'),
+                method=req.scope.get('method', 'err'),
-                endpoint = endpoint,
+                endpoint=endpoint,
-                duration = duration,
+                duration=duration,
             ))
         return res
 
@@ -198,6 +197,7 @@ class Api:
         self.add_api_route("/sdapi/v1/prompt-styles", self.get_prompt_styles, methods=["GET"], response_model=List[models.PromptStyleItem])
         self.add_api_route("/sdapi/v1/embeddings", self.get_embeddings, methods=["GET"], response_model=models.EmbeddingsResponse)
         self.add_api_route("/sdapi/v1/refresh-checkpoints", self.refresh_checkpoints, methods=["POST"])
+        self.add_api_route("/sdapi/v1/refresh-vae", self.refresh_vae, methods=["POST"])
         self.add_api_route("/sdapi/v1/create/embedding", self.create_embedding, methods=["POST"], response_model=models.CreateResponse)
         self.add_api_route("/sdapi/v1/create/hypernetwork", self.create_hypernetwork, methods=["POST"], response_model=models.CreateResponse)
         self.add_api_route("/sdapi/v1/preprocess", self.preprocess, methods=["POST"], response_model=models.PreprocessResponse)
@@ -209,6 +209,11 @@ class Api:
         self.add_api_route("/sdapi/v1/scripts", self.get_scripts_list, methods=["GET"], response_model=models.ScriptsList)
         self.add_api_route("/sdapi/v1/script-info", self.get_script_info, methods=["GET"], response_model=List[models.ScriptInfo])
 
+        if shared.cmd_opts.api_server_stop:
+            self.add_api_route("/sdapi/v1/server-kill", self.kill_webui, methods=["POST"])
+            self.add_api_route("/sdapi/v1/server-restart", self.restart_webui, methods=["POST"])
+            self.add_api_route("/sdapi/v1/server-stop", self.stop_webui, methods=["POST"])
+
         self.default_script_arg_txt2img = []
         self.default_script_arg_img2img = []
 
@@ -324,19 +329,22 @@ class Api:
         args.pop('save_images', None)
 
         with self.queue_lock:
-            p = StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **args)
+            with closing(StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **args)) as p:
                 p.scripts = script_runner
                 p.outpath_grids = opts.outdir_txt2img_grids
                 p.outpath_samples = opts.outdir_txt2img_samples
 
-            shared.state.begin()
+                try:
+                    shared.state.begin(job="scripts_txt2img")
                     if selectable_scripts is not None:
                         p.script_args = script_args
                         processed = scripts.scripts_txt2img.run(p, *p.script_args) # Need to pass args as list here
                     else:
                         p.script_args = tuple(script_args) # Need to pass args as tuple here
                         processed = process_images(p)
+                finally:
                     shared.state.end()
+                    shared.total_tqdm.clear()
 
         b64images = list(map(encode_pil_to_base64, processed.images)) if send_images else []
 
@@ -380,20 +388,23 @@ class Api:
         args.pop('save_images', None)
 
         with self.queue_lock:
-            p = StableDiffusionProcessingImg2Img(sd_model=shared.sd_model, **args)
+            with closing(StableDiffusionProcessingImg2Img(sd_model=shared.sd_model, **args)) as p:
                 p.init_images = [decode_base64_to_image(x) for x in init_images]
                 p.scripts = script_runner
                 p.outpath_grids = opts.outdir_img2img_grids
                 p.outpath_samples = opts.outdir_img2img_samples
 
-            shared.state.begin()
+                try:
+                    shared.state.begin(job="scripts_img2img")
                     if selectable_scripts is not None:
                         p.script_args = script_args
                         processed = scripts.scripts_img2img.run(p, *p.script_args) # Need to pass args as list here
                     else:
                         p.script_args = tuple(script_args) # Need to pass args as tuple here
                         processed = process_images(p)
+                finally:
                     shared.state.end()
+                    shared.total_tqdm.clear()
 
         b64images = list(map(encode_pil_to_base64, processed.images)) if send_images else []
 
@@ -517,6 +528,10 @@ class Api:
         return options
 
     def set_config(self, req: Dict[str, Any]):
+        checkpoint_name = req.get("sd_model_checkpoint", None)
+        if checkpoint_name is not None and checkpoint_name not in checkpoint_aliases:
+            raise RuntimeError(f"model {checkpoint_name!r} not found")
+
         for k, v in req.items():
             shared.opts.set(k, v)
 
@@ -593,48 +608,51 @@ class Api:
         }
 
     def refresh_checkpoints(self):
+        with self.queue_lock:
             shared.refresh_checkpoints()
 
+    def refresh_vae(self):
+        with self.queue_lock:
+            shared_items.refresh_vae_list()
+
     def create_embedding(self, args: dict):
         try:
-            shared.state.begin()
+            shared.state.begin(job="create_embedding")
             filename = create_embedding(**args) # create empty embedding
             sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings() # reload embeddings so new one can be immediately used
-            shared.state.end()
             return models.CreateResponse(info=f"create embedding filename: {filename}")
         except AssertionError as e:
-            shared.state.end()
             return models.TrainResponse(info=f"create embedding error: {e}")
+        finally:
+            shared.state.end()
+
 
     def create_hypernetwork(self, args: dict):
         try:
-            shared.state.begin()
+            shared.state.begin(job="create_hypernetwork")
             filename = create_hypernetwork(**args) # create empty embedding
-            shared.state.end()
             return models.CreateResponse(info=f"create hypernetwork filename: {filename}")
         except AssertionError as e:
-            shared.state.end()
             return models.TrainResponse(info=f"create hypernetwork error: {e}")
+        finally:
+            shared.state.end()
 
     def preprocess(self, args: dict):
         try:
-            shared.state.begin()
+            shared.state.begin(job="preprocess")
             preprocess(**args) # quick operation unless blip/booru interrogation is enabled
             shared.state.end()
-            return models.PreprocessResponse(info = 'preprocess complete')
+            return models.PreprocessResponse(info='preprocess complete')
         except KeyError as e:
-            shared.state.end()
             return models.PreprocessResponse(info=f"preprocess error: invalid token: {e}")
-        except AssertionError as e:
+        except Exception as e:
-            shared.state.end()
             return models.PreprocessResponse(info=f"preprocess error: {e}")
-        except FileNotFoundError as e:
+        finally:
             shared.state.end()
-            return models.PreprocessResponse(info=f'preprocess error: {e}')
 
     def train_embedding(self, args: dict):
         try:
-            shared.state.begin()
+            shared.state.begin(job="train_embedding")
             apply_optimizations = shared.opts.training_xattention_optimizations
             error = None
             filename = ''
@@ -647,15 +665,15 @@ class Api:
             finally:
                 if not apply_optimizations:
                     sd_hijack.apply_optimizations()
-                shared.state.end()
             return models.TrainResponse(info=f"train embedding complete: filename: {filename} error: {error}")
-        except AssertionError as msg:
+        except Exception as msg:
-            shared.state.end()
             return models.TrainResponse(info=f"train embedding error: {msg}")
+        finally:
+            shared.state.end()
 
     def train_hypernetwork(self, args: dict):
         try:
-            shared.state.begin()
+            shared.state.begin(job="train_hypernetwork")
             shared.loaded_hypernetworks = []
             apply_optimizations = shared.opts.training_xattention_optimizations
             error = None
@@ -673,9 +691,10 @@ class Api:
                     sd_hijack.apply_optimizations()
                 shared.state.end()
             return models.TrainResponse(info=f"train embedding complete: filename: {filename} error: {error}")
-        except AssertionError:
+        except Exception as exc:
+            return models.TrainResponse(info=f"train embedding error: {exc}")
+        finally:
             shared.state.end()
-            return models.TrainResponse(info=f"train embedding error: {error}")
 
     def get_memory(self):
         try:
@@ -712,6 +731,19 @@ class Api:
             cuda = {'error': f'{err}'}
         return models.MemoryResponse(ram=ram, cuda=cuda)
 
-    def launch(self, server_name, port):
+    def launch(self, server_name, port, root_path):
         self.app.include_router(self.router)
-        uvicorn.run(self.app, host=server_name, port=port, timeout_keep_alive=0)
+        uvicorn.run(self.app, host=server_name, port=port, timeout_keep_alive=shared.cmd_opts.timeout_keep_alive, root_path=root_path)
+
+    def kill_webui(self):
+        restart.stop_program()
+
+    def restart_webui(self):
+        if restart.is_restartable():
+            restart.restart_program()
+        return Response(status_code=501)
+
+    def stop_webui(request):
+        shared.state.server_command = "stop"
+        return Response("Stopping.")
+

modules/api/models.py

@@ -1,4 +1,5 @@
 import inspect
+
 from pydantic import BaseModel, Field, create_model
 from typing import Any, Optional
 from typing_extensions import Literal
@@ -207,11 +208,10 @@ class PreprocessResponse(BaseModel):
 fields = {}
 for key, metadata in opts.data_labels.items():
     value = opts.data.get(key)
-    optType = opts.typemap.get(type(metadata.default), type(value))
+    optType = opts.typemap.get(type(metadata.default), type(metadata.default)) if metadata.default else Any
 
-    if (metadata is not None):
+    if metadata is not None:
-        fields.update({key: (Optional[optType], Field(
+        fields.update({key: (Optional[optType], Field(default=metadata.default, description=metadata.label))})
-            default=metadata.default ,description=metadata.label))})
     else:
         fields.update({key: (Optional[optType], Field())})
 
@@ -274,10 +274,6 @@ class PromptStyleItem(BaseModel):
     prompt: Optional[str] = Field(title="Prompt")
     negative_prompt: Optional[str] = Field(title="Negative Prompt")
 
-class ArtistItem(BaseModel):
-    name: str = Field(title="Name")
-    score: float = Field(title="Score")
-    category: str = Field(title="Category")
 
 class EmbeddingItem(BaseModel):
     step: Optional[int] = Field(title="Step", description="The number of steps that were used to train this embedding, if available")

modules/cache.py

@@ -0,0 +1,120 @@
+import json
+import os.path
+import threading
+import time
+
+from modules.paths import data_path, script_path
+
+cache_filename = os.path.join(data_path, "cache.json")
+cache_data = None
+cache_lock = threading.Lock()
+
+dump_cache_after = None
+dump_cache_thread = None
+
+
+def dump_cache():
+    """
+    Marks cache for writing to disk. 5 seconds after no one else flags the cache for writing, it is written.
+    """
+
+    global dump_cache_after
+    global dump_cache_thread
+
+    def thread_func():
+        global dump_cache_after
+        global dump_cache_thread
+
+        while dump_cache_after is not None and time.time() < dump_cache_after:
+            time.sleep(1)
+
+        with cache_lock:
+            with open(cache_filename, "w", encoding="utf8") as file:
+                json.dump(cache_data, file, indent=4)
+
+            dump_cache_after = None
+            dump_cache_thread = None
+
+    with cache_lock:
+        dump_cache_after = time.time() + 5
+        if dump_cache_thread is None:
+            dump_cache_thread = threading.Thread(name='cache-writer', target=thread_func)
+            dump_cache_thread.start()
+
+
+def cache(subsection):
+    """
+    Retrieves or initializes a cache for a specific subsection.
+
+    Parameters:
+        subsection (str): The subsection identifier for the cache.
+
+    Returns:
+        dict: The cache data for the specified subsection.
+    """
+
+    global cache_data
+
+    if cache_data is None:
+        with cache_lock:
+            if cache_data is None:
+                if not os.path.isfile(cache_filename):
+                    cache_data = {}
+                else:
+                    try:
+                        with open(cache_filename, "r", encoding="utf8") as file:
+                            cache_data = json.load(file)
+                    except Exception:
+                        os.replace(cache_filename, os.path.join(script_path, "tmp", "cache.json"))
+                        print('[ERROR] issue occurred while trying to read cache.json, move current cache to tmp/cache.json and create new cache')
+                        cache_data = {}
+
+    s = cache_data.get(subsection, {})
+    cache_data[subsection] = s
+
+    return s
+
+
+def cached_data_for_file(subsection, title, filename, func):
+    """
+    Retrieves or generates data for a specific file, using a caching mechanism.
+
+    Parameters:
+        subsection (str): The subsection of the cache to use.
+        title (str): The title of the data entry in the subsection of the cache.
+        filename (str): The path to the file to be checked for modifications.
+        func (callable): A function that generates the data if it is not available in the cache.
+
+    Returns:
+        dict or None: The cached or generated data, or None if data generation fails.
+
+    The `cached_data_for_file` function implements a caching mechanism for data stored in files.
+    It checks if the data associated with the given `title` is present in the cache and compares the
+    modification time of the file with the cached modification time. If the file has been modified,
+    the cache is considered invalid and the data is regenerated using the provided `func`.
+    Otherwise, the cached data is returned.
+
+    If the data generation fails, None is returned to indicate the failure. Otherwise, the generated
+    or cached data is returned as a dictionary.
+    """
+
+    existing_cache = cache(subsection)
+    ondisk_mtime = os.path.getmtime(filename)
+
+    entry = existing_cache.get(title)
+    if entry:
+        cached_mtime = entry.get("mtime", 0)
+        if ondisk_mtime > cached_mtime:
+            entry = None
+
+    if not entry or 'value' not in entry:
+        value = func()
+        if value is None:
+            return None
+
+        entry = {'mtime': ondisk_mtime, 'value': value}
+        existing_cache[title] = entry
+
+        dump_cache()
+
+    return entry['value']

modules/call_queue.py

@@ -1,8 +1,9 @@
+from functools import wraps
 import html
 import threading
 import time
 
-from modules import shared, progress, errors
+from modules import shared, progress, errors, devices
 
 queue_lock = threading.Lock()
 
@@ -18,6 +19,7 @@ def wrap_queued_call(func):
 
 
 def wrap_gradio_gpu_call(func, extra_outputs=None):
+    @wraps(func)
     def f(*args, **kwargs):
 
         # if the first argument is a string that says "task(...)", it is treated as a job id
@@ -28,7 +30,7 @@ def wrap_gradio_gpu_call(func, extra_outputs=None):
             id_task = None
 
         with queue_lock:
-            shared.state.begin()
+            shared.state.begin(job=id_task)
             progress.start_task(id_task)
 
             try:
@@ -45,6 +47,7 @@ def wrap_gradio_gpu_call(func, extra_outputs=None):
 
 
 def wrap_gradio_call(func, extra_outputs=None, add_stats=False):
+    @wraps(func)
     def f(*args, extra_outputs_array=extra_outputs, **kwargs):
         run_memmon = shared.opts.memmon_poll_rate > 0 and not shared.mem_mon.disabled and add_stats
         if run_memmon:
@@ -72,6 +75,8 @@ def wrap_gradio_call(func, extra_outputs=None, add_stats=False):
             error_message = f'{type(e).__name__}: {e}'
             res = extra_outputs_array + [f"<div class='error'>{html.escape(error_message)}</div>"]
 
+        devices.torch_gc()
+
         shared.state.skipped = False
         shared.state.interrupted = False
         shared.state.job_count = 0
@@ -82,9 +87,9 @@ def wrap_gradio_call(func, extra_outputs=None, add_stats=False):
         elapsed = time.perf_counter() - t
         elapsed_m = int(elapsed // 60)
         elapsed_s = elapsed % 60
-        elapsed_text = f"{elapsed_s:.2f}s"
+        elapsed_text = f"{elapsed_s:.1f} sec."
         if elapsed_m > 0:
-            elapsed_text = f"{elapsed_m}m "+elapsed_text
+            elapsed_text = f"{elapsed_m} min. "+elapsed_text
 
         if run_memmon:
             mem_stats = {k: -(v//-(1024*1024)) for k, v in shared.mem_mon.stop().items()}
@@ -92,14 +97,22 @@ def wrap_gradio_call(func, extra_outputs=None, add_stats=False):
             reserved_peak = mem_stats['reserved_peak']
             sys_peak = mem_stats['system_peak']
             sys_total = mem_stats['total']
-            sys_pct = round(sys_peak/max(sys_total, 1) * 100, 2)
+            sys_pct = sys_peak/max(sys_total, 1) * 100
 
-            vram_html = f"<p class='vram'>Torch active/reserved: {active_peak}/{reserved_peak} MiB, <wbr>Sys VRAM: {sys_peak}/{sys_total} MiB ({sys_pct}%)</p>"
+            toltip_a = "Active: peak amount of video memory used during generation (excluding cached data)"
+            toltip_r = "Reserved: total amout of video memory allocated by the Torch library "
+            toltip_sys = "System: peak amout of video memory allocated by all running programs, out of total capacity"
+
+            text_a = f"<abbr title='{toltip_a}'>A</abbr>: <span class='measurement'>{active_peak/1024:.2f} GB</span>"
+            text_r = f"<abbr title='{toltip_r}'>R</abbr>: <span class='measurement'>{reserved_peak/1024:.2f} GB</span>"
+            text_sys = f"<abbr title='{toltip_sys}'>Sys</abbr>: <span class='measurement'>{sys_peak/1024:.1f}/{sys_total/1024:g} GB</span> ({sys_pct:.1f}%)"
+
+            vram_html = f"<p class='vram'>{text_a}, <wbr>{text_r}, <wbr>{text_sys}</p>"
         else:
             vram_html = ''
 
         # last item is always HTML
-        res[-1] += f"<div class='performance'><p class='time'>Time taken: <wbr>{elapsed_text}</p>{vram_html}</div>"
+        res[-1] += f"<div class='performance'><p class='time'>Time taken: <wbr><span class='measurement'>{elapsed_text}</span></p>{vram_html}</div>"
 
         return tuple(res)
 

modules/cmd_args.py

@@ -13,8 +13,10 @@ parser.add_argument("--reinstall-xformers", action='store_true', help="launch.py
 parser.add_argument("--reinstall-torch", action='store_true', help="launch.py argument: install the appropriate version of torch even if you have some version already installed")
 parser.add_argument("--update-check", action='store_true', help="launch.py argument: check for updates at startup")
 parser.add_argument("--test-server", action='store_true', help="launch.py argument: configure server for testing")
+parser.add_argument("--log-startup", action='store_true', help="launch.py argument: print a detailed log of what's happening at startup")
 parser.add_argument("--skip-prepare-environment", action='store_true', help="launch.py argument: skip all environment preparation")
 parser.add_argument("--skip-install", action='store_true', help="launch.py argument: skip installation of packages")
+parser.add_argument("--do-not-download-clip", action='store_true', help="do not download CLIP model even if it's not included in the checkpoint")
 parser.add_argument("--data-dir", type=str, default=os.path.dirname(os.path.dirname(os.path.realpath(__file__))), help="base path where all user data is stored")
 parser.add_argument("--config", type=str, default=sd_default_config, help="path to config which constructs model",)
 parser.add_argument("--ckpt", type=str, default=sd_model_file, help="path to checkpoint of stable diffusion model; if specified, this checkpoint will be added to the list of checkpoints and loaded",)
@@ -65,6 +67,7 @@ parser.add_argument("--opt-sdp-no-mem-attention", action='store_true', help="pre
 parser.add_argument("--disable-opt-split-attention", action='store_true', help="prefer no cross-attention layer optimization for automatic choice of optimization")
 parser.add_argument("--disable-nan-check", action='store_true', help="do not check if produced images/latent spaces have nans; useful for running without a checkpoint in CI")
 parser.add_argument("--use-cpu", nargs='+', help="use CPU as torch device for specified modules", default=[], type=str.lower)
+parser.add_argument("--disable-model-loading-ram-optimization", action='store_true', help="disable an optimization that reduces RAM use when loading a model")
 parser.add_argument("--listen", action='store_true', help="launch gradio with 0.0.0.0 as server name, allowing to respond to network requests")
 parser.add_argument("--port", type=int, help="launch gradio with given server port, you need root/admin rights for ports < 1024, defaults to 7860 if available", default=None)
 parser.add_argument("--show-negative-prompt", action='store_true', help="does not do anything", default=False)
@@ -107,3 +110,7 @@ parser.add_argument("--no-hashing", action='store_true', help="disable sha256 ha
 parser.add_argument("--no-download-sd-model", action='store_true', help="don't download SD1.5 model even if no model is found in --ckpt-dir", default=False)
 parser.add_argument('--subpath', type=str, help='customize the subpath for gradio, use with reverse proxy')
 parser.add_argument('--add-stop-route', action='store_true', help='add /_stop route to stop server')
+parser.add_argument('--api-server-stop', action='store_true', help='enable server stop/restart/kill via api')
+parser.add_argument('--timeout-keep-alive', type=int, default=30, help='set timeout_keep_alive for uvicorn')
+parser.add_argument("--disable-all-extensions", action='store_true', help="prevent all extensions from running regardless of any other settings", default=False)
+parser.add_argument("--disable-extra-extensions", action='store_true', help=" prevent all extensions except built-in from running regardless of any other settings", default=False)

modules/codeformer_model.py

@@ -15,14 +15,11 @@ model_dir = "Codeformer"
 model_path = os.path.join(models_path, model_dir)
 model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
 
-have_codeformer = False
 codeformer = None
 
 
 def setup_model(dirname):
-    global model_path
+    os.makedirs(model_path, exist_ok=True)
-    if not os.path.exists(model_path):
-        os.makedirs(model_path)
 
     path = modules.paths.paths.get("CodeFormer", None)
     if path is None:
@@ -102,7 +99,7 @@ def setup_model(dirname):
                             output = self.net(cropped_face_t, w=w if w is not None else shared.opts.code_former_weight, adain=True)[0]
                             restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
                         del output
-                        torch.cuda.empty_cache()
+                        devices.torch_gc()
                     except Exception:
                         errors.report('Failed inference for CodeFormer', exc_info=True)
                         restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))
@@ -125,9 +122,6 @@ def setup_model(dirname):
 
                 return restored_img
 
-        global have_codeformer
-        have_codeformer = True
-
         global codeformer
         codeformer = FaceRestorerCodeFormer(dirname)
         shared.face_restorers.append(codeformer)

modules/devices.py

@@ -3,7 +3,7 @@ import contextlib
 from functools import lru_cache
 
 import torch
-from modules import errors
+from modules import errors, rng_philox
 
 if sys.platform == "darwin":
     from modules import mac_specific
@@ -15,13 +15,6 @@ def has_mps() -> bool:
     else:
         return mac_specific.has_mps
 
-def extract_device_id(args, name):
-    for x in range(len(args)):
-        if name in args[x]:
-            return args[x + 1]
-
-    return None
-
 
 def get_cuda_device_string():
     from modules import shared
@@ -56,11 +49,15 @@ def get_device_for(task):
 
 
 def torch_gc():
+
     if torch.cuda.is_available():
         with torch.cuda.device(get_cuda_device_string()):
             torch.cuda.empty_cache()
             torch.cuda.ipc_collect()
 
+    if has_mps():
+        mac_specific.torch_mps_gc()
+
 
 def enable_tf32():
     if torch.cuda.is_available():
@@ -74,14 +71,17 @@ def enable_tf32():
         torch.backends.cudnn.allow_tf32 = True
 
 
-
 errors.run(enable_tf32, "Enabling TF32")
 
-cpu = torch.device("cpu")
+cpu: torch.device = torch.device("cpu")
-device = device_interrogate = device_gfpgan = device_esrgan = device_codeformer = None
+device: torch.device = None
-dtype = torch.float16
+device_interrogate: torch.device = None
-dtype_vae = torch.float16
+device_gfpgan: torch.device = None
-dtype_unet = torch.float16
+device_esrgan: torch.device = None
+device_codeformer: torch.device = None
+dtype: torch.dtype = torch.float16
+dtype_vae: torch.dtype = torch.float16
+dtype_unet: torch.dtype = torch.float16
 unet_needs_upcast = False
 
 
@@ -93,23 +93,87 @@ def cond_cast_float(input):
     return input.float() if unet_needs_upcast else input
 
 
+nv_rng = None
+
+
 def randn(seed, shape):
+    """Generate a tensor with random numbers from a normal distribution using seed.
+
+    Uses the seed parameter to set the global torch seed; to generate more with that seed, use randn_like/randn_without_seed."""
+
     from modules.shared import opts
 
-    torch.manual_seed(seed)
+    manual_seed(seed)
+
+    if opts.randn_source == "NV":
+        return torch.asarray(nv_rng.randn(shape), device=device)
+
     if opts.randn_source == "CPU" or device.type == 'mps':
         return torch.randn(shape, device=cpu).to(device)
+
     return torch.randn(shape, device=device)
 
 
+def randn_local(seed, shape):
+    """Generate a tensor with random numbers from a normal distribution using seed.
+
+    Does not change the global random number generator. You can only generate the seed's first tensor using this function."""
+
+    from modules.shared import opts
+
+    if opts.randn_source == "NV":
+        rng = rng_philox.Generator(seed)
+        return torch.asarray(rng.randn(shape), device=device)
+
+    local_device = cpu if opts.randn_source == "CPU" or device.type == 'mps' else device
+    local_generator = torch.Generator(local_device).manual_seed(int(seed))
+    return torch.randn(shape, device=local_device, generator=local_generator).to(device)
+
+
+def randn_like(x):
+    """Generate a tensor with random numbers from a normal distribution using the previously initialized genrator.
+
+    Use either randn() or manual_seed() to initialize the generator."""
+
+    from modules.shared import opts
+
+    if opts.randn_source == "NV":
+        return torch.asarray(nv_rng.randn(x.shape), device=x.device, dtype=x.dtype)
+
+    if opts.randn_source == "CPU" or x.device.type == 'mps':
+        return torch.randn_like(x, device=cpu).to(x.device)
+
+    return torch.randn_like(x)
+
+
 def randn_without_seed(shape):
+    """Generate a tensor with random numbers from a normal distribution using the previously initialized genrator.
+
+    Use either randn() or manual_seed() to initialize the generator."""
+
     from modules.shared import opts
 
+    if opts.randn_source == "NV":
+        return torch.asarray(nv_rng.randn(shape), device=device)
+
     if opts.randn_source == "CPU" or device.type == 'mps':
         return torch.randn(shape, device=cpu).to(device)
+
     return torch.randn(shape, device=device)
 
 
+def manual_seed(seed):
+    """Set up a global random number generator using the specified seed."""
+    from modules.shared import opts
+
+    if opts.randn_source == "NV":
+        global nv_rng
+        nv_rng = rng_philox.Generator(seed)
+        return
+
+    torch.manual_seed(seed)
+
+
 def autocast(disable=False):
     from modules import shared
 

modules/errors.py

@@ -14,7 +14,8 @@ def record_exception():
     if exception_records and exception_records[-1] == e:
         return
 
-    exception_records.append((e, tb))
+    from modules import sysinfo
+    exception_records.append(sysinfo.format_exception(e, tb))
 
     if len(exception_records) > 5:
         exception_records.pop(0)
@@ -83,3 +84,53 @@ def run(code, task):
         code()
     except Exception as e:
         display(task, e)
+
+
+def check_versions():
+    from packaging import version
+    from modules import shared
+
+    import torch
+    import gradio
+
+    expected_torch_version = "2.0.0"
+    expected_xformers_version = "0.0.20"
+    expected_gradio_version = "3.39.0"
+
+    if version.parse(torch.__version__) < version.parse(expected_torch_version):
+        print_error_explanation(f"""
+You are running torch {torch.__version__}.
+The program is tested to work with torch {expected_torch_version}.
+To reinstall the desired version, run with commandline flag --reinstall-torch.
+Beware that this will cause a lot of large files to be downloaded, as well as
+there are reports of issues with training tab on the latest version.
+
+Use --skip-version-check commandline argument to disable this check.
+        """.strip())
+
+    if shared.xformers_available:
+        import xformers
+
+        if version.parse(xformers.__version__) < version.parse(expected_xformers_version):
+            print_error_explanation(f"""
+You are running xformers {xformers.__version__}.
+The program is tested to work with xformers {expected_xformers_version}.
+To reinstall the desired version, run with commandline flag --reinstall-xformers.
+
+Use --skip-version-check commandline argument to disable this check.
+            """.strip())
+
+    if gradio.__version__ != expected_gradio_version:
+        print_error_explanation(f"""
+You are running gradio {gradio.__version__}.
+The program is designed to work with gradio {expected_gradio_version}.
+Using a different version of gradio is extremely likely to break the program.
+
+Reasons why you have the mismatched gradio version can be:
+  - you use --skip-install flag.
+  - you use webui.py to start the program instead of launch.py.
+  - an extension installs the incompatible gradio version.
+
+Use --skip-version-check commandline argument to disable this check.
+        """.strip())
+

modules/esrgan_model.py

@@ -1,15 +1,13 @@
-import os
+import sys
 
 import numpy as np
 import torch
 from PIL import Image
-from basicsr.utils.download_util import load_file_from_url
 
 import modules.esrgan_model_arch as arch
 from modules import modelloader, images, devices
-from modules.upscaler import Upscaler, UpscalerData
 from modules.shared import opts
-
+from modules.upscaler import Upscaler, UpscalerData
 
 
 def mod2normal(state_dict):
@@ -134,7 +132,7 @@ class UpscalerESRGAN(Upscaler):
             scaler_data = UpscalerData(self.model_name, self.model_url, self, 4)
             scalers.append(scaler_data)
         for file in model_paths:
-            if "http" in file:
+            if file.startswith("http"):
                 name = self.model_name
             else:
                 name = modelloader.friendly_name(file)
@@ -143,26 +141,25 @@ class UpscalerESRGAN(Upscaler):
             self.scalers.append(scaler_data)
 
     def do_upscale(self, img, selected_model):
+        try:
             model = self.load_model(selected_model)
-        if model is None:
+        except Exception as e:
+            print(f"Unable to load ESRGAN model {selected_model}: {e}", file=sys.stderr)
             return img
         model.to(devices.device_esrgan)
         img = esrgan_upscale(model, img)
         return img
 
     def load_model(self, path: str):
-        if "http" in path:
+        if path.startswith("http"):
-            filename = load_file_from_url(
+            # TODO: this doesn't use `path` at all?
+            filename = modelloader.load_file_from_url(
                 url=self.model_url,
                 model_dir=self.model_download_path,
                 file_name=f"{self.model_name}.pth",
-                progress=True,
             )
         else:
             filename = path
-        if not os.path.exists(filename) or filename is None:
-            print(f"Unable to load {self.model_path} from {filename}")
-            return None
 
         state_dict = torch.load(filename, map_location='cpu' if devices.device_esrgan.type == 'mps' else None)
 

modules/extensions.py

@@ -1,20 +1,19 @@
 import os
 import threading
 
-from modules import shared, errors
+from modules import shared, errors, cache
 from modules.gitpython_hack import Repo
 from modules.paths_internal import extensions_dir, extensions_builtin_dir, script_path  # noqa: F401
 
 extensions = []
 
-if not os.path.exists(extensions_dir):
+os.makedirs(extensions_dir, exist_ok=True)
-    os.makedirs(extensions_dir)
 
 
 def active():
-    if shared.opts.disable_all_extensions == "all":
+    if shared.cmd_opts.disable_all_extensions or shared.opts.disable_all_extensions == "all":
         return []
-    elif shared.opts.disable_all_extensions == "extra":
+    elif shared.cmd_opts.disable_extra_extensions or shared.opts.disable_all_extensions == "extra":
         return [x for x in extensions if x.enabled and x.is_builtin]
     else:
         return [x for x in extensions if x.enabled]
@@ -22,6 +21,7 @@ def active():
 
 class Extension:
     lock = threading.Lock()
+    cached_fields = ['remote', 'commit_date', 'branch', 'commit_hash', 'version']
 
     def __init__(self, name, path, enabled=True, is_builtin=False):
         self.name = name
@@ -37,16 +37,32 @@ class Extension:
         self.remote = None
         self.have_info_from_repo = False
 
+    def to_dict(self):
+        return {x: getattr(self, x) for x in self.cached_fields}
+
+    def from_dict(self, d):
+        for field in self.cached_fields:
+            setattr(self, field, d[field])
+
     def read_info_from_repo(self):
         if self.is_builtin or self.have_info_from_repo:
             return
 
+        def read_from_repo():
             with self.lock:
                 if self.have_info_from_repo:
                     return
 
                 self.do_read_info_from_repo()
 
+                return self.to_dict()
+        try:
+            d = cache.cached_data_for_file('extensions-git', self.name, os.path.join(self.path, ".git"), read_from_repo)
+            self.from_dict(d)
+        except FileNotFoundError:
+            pass
+        self.status = 'unknown' if self.status == '' else self.status
+
     def do_read_info_from_repo(self):
         repo = None
         try:
@@ -59,7 +75,6 @@ class Extension:
             self.remote = None
         else:
             try:
-                self.status = 'unknown'
                 self.remote = next(repo.remote().urls, None)
                 commit = repo.head.commit
                 self.commit_date = commit.committed_date
@@ -126,8 +141,12 @@ def list_extensions():
     if not os.path.isdir(extensions_dir):
         return
 
-    if shared.opts.disable_all_extensions == "all":
+    if shared.cmd_opts.disable_all_extensions:
+        print("*** \"--disable-all-extensions\" arg was used, will not load any extensions ***")
+    elif shared.opts.disable_all_extensions == "all":
         print("*** \"Disable all extensions\" option was set, will not load any extensions ***")
+    elif shared.cmd_opts.disable_extra_extensions:
+        print("*** \"--disable-extra-extensions\" arg was used, will only load built-in extensions ***")
     elif shared.opts.disable_all_extensions == "extra":
         print("*** \"Disable all extensions\" option was set, will only load built-in extensions ***")
 

modules/extra_networks.py

@@ -1,19 +1,27 @@
+import json
+import os
 import re
 from collections import defaultdict
 
 from modules import errors
 
 extra_network_registry = {}
+extra_network_aliases = {}
 
 
 def initialize():
     extra_network_registry.clear()
+    extra_network_aliases.clear()
 
 
 def register_extra_network(extra_network):
     extra_network_registry[extra_network.name] = extra_network
 
 
+def register_extra_network_alias(extra_network, alias):
+    extra_network_aliases[alias] = extra_network
+
+
 def register_default_extra_networks():
     from modules.extra_networks_hypernet import ExtraNetworkHypernet
     register_extra_network(ExtraNetworkHypernet())
@@ -82,20 +90,26 @@ def activate(p, extra_network_data):
     """call activate for extra networks in extra_network_data in specified order, then call
     activate for all remaining registered networks with an empty argument list"""
 
+    activated = []
+
     for extra_network_name, extra_network_args in extra_network_data.items():
         extra_network = extra_network_registry.get(extra_network_name, None)
+
+        if extra_network is None:
+            extra_network = extra_network_aliases.get(extra_network_name, None)
+
         if extra_network is None:
             print(f"Skipping unknown extra network: {extra_network_name}")
             continue
 
         try:
             extra_network.activate(p, extra_network_args)
+            activated.append(extra_network)
         except Exception as e:
             errors.display(e, f"activating extra network {extra_network_name} with arguments {extra_network_args}")
 
     for extra_network_name, extra_network in extra_network_registry.items():
-        args = extra_network_data.get(extra_network_name, None)
+        if extra_network in activated:
-        if args is not None:
             continue
 
         try:
@@ -103,6 +117,9 @@ def activate(p, extra_network_data):
         except Exception as e:
             errors.display(e, f"activating extra network {extra_network_name}")
 
+    if p.scripts is not None:
+        p.scripts.after_extra_networks_activate(p, batch_number=p.iteration, prompts=p.prompts, seeds=p.seeds, subseeds=p.subseeds, extra_network_data=extra_network_data)
+
 
 def deactivate(p, extra_network_data):
     """call deactivate for extra networks in extra_network_data in specified order, then call
@@ -162,3 +179,20 @@ def parse_prompts(prompts):
 
     return res, extra_data
 
+
+def get_user_metadata(filename):
+    if filename is None:
+        return {}
+
+    basename, ext = os.path.splitext(filename)
+    metadata_filename = basename + '.json'
+
+    metadata = {}
+    try:
+        if os.path.isfile(metadata_filename):
+            with open(metadata_filename, "r", encoding="utf8") as file:
+                metadata = json.load(file)
+    except Exception as e:
+        errors.display(e, f"reading extra network user metadata from {metadata_filename}")
+
+    return metadata

modules/extras.py

@@ -7,7 +7,7 @@ import json
 import torch
 import tqdm
 
-from modules import shared, images, sd_models, sd_vae, sd_models_config
+from modules import shared, images, sd_models, sd_vae, sd_models_config, errors
 from modules.ui_common import plaintext_to_html
 import gradio as gr
 import safetensors.torch
@@ -72,9 +72,21 @@ def to_half(tensor, enable):
     return tensor
 
 
-def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format, config_source, bake_in_vae, discard_weights, save_metadata):
+def read_metadata(primary_model_name, secondary_model_name, tertiary_model_name):
-    shared.state.begin()
+    metadata = {}
-    shared.state.job = 'model-merge'
+
+    for checkpoint_name in [primary_model_name, secondary_model_name, tertiary_model_name]:
+        checkpoint_info = sd_models.checkpoints_list.get(checkpoint_name, None)
+        if checkpoint_info is None:
+            continue
+
+        metadata.update(checkpoint_info.metadata)
+
+    return json.dumps(metadata, indent=4, ensure_ascii=False)
+
+
+def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format, config_source, bake_in_vae, discard_weights, save_metadata, add_merge_recipe, copy_metadata_fields, metadata_json):
+    shared.state.begin(job="model-merge")
 
     def fail(message):
         shared.state.textinfo = message
@@ -242,11 +254,25 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
     shared.state.textinfo = "Saving"
     print(f"Saving to {output_modelname}...")
 
-    metadata = None
+    metadata = {}
+
+    if save_metadata and copy_metadata_fields:
+        if primary_model_info:
+            metadata.update(primary_model_info.metadata)
+        if secondary_model_info:
+            metadata.update(secondary_model_info.metadata)
+        if tertiary_model_info:
+            metadata.update(tertiary_model_info.metadata)
 
     if save_metadata:
-        metadata = {"format": "pt"}
+        try:
+            metadata.update(json.loads(metadata_json))
+        except Exception as e:
+            errors.display(e, "readin metadata from json")
 
+        metadata["format"] = "pt"
+
+    if save_metadata and add_merge_recipe:
         merge_recipe = {
             "type": "webui", # indicate this model was merged with webui's built-in merger
             "primary_model_hash": primary_model_info.sha256,
@@ -262,7 +288,6 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
             "is_inpainting": result_is_inpainting_model,
             "is_instruct_pix2pix": result_is_instruct_pix2pix_model
         }
-        metadata["sd_merge_recipe"] = json.dumps(merge_recipe)
 
         sd_merge_models = {}
 
@@ -282,11 +307,12 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
         if tertiary_model_info:
             add_model_metadata(tertiary_model_info)
 
+        metadata["sd_merge_recipe"] = json.dumps(merge_recipe)
         metadata["sd_merge_models"] = json.dumps(sd_merge_models)
 
     _, extension = os.path.splitext(output_modelname)
     if extension.lower() == ".safetensors":
-        safetensors.torch.save_file(theta_0, output_modelname, metadata=metadata)
+        safetensors.torch.save_file(theta_0, output_modelname, metadata=metadata if len(metadata)>0 else None)
     else:
         torch.save(theta_0, output_modelname)
 

modules/generation_parameters_copypaste.py

@@ -174,31 +174,6 @@ def send_image_and_dimensions(x):
     return img, w, h
 
 
-
-def find_hypernetwork_key(hypernet_name, hypernet_hash=None):
-    """Determines the config parameter name to use for the hypernet based on the parameters in the infotext.
-
-    Example: an infotext provides "Hypernet: ke-ta" and "Hypernet hash: 1234abcd". For the "Hypernet" config
-    parameter this means there should be an entry that looks like "ke-ta-10000(1234abcd)" to set it to.
-
-    If the infotext has no hash, then a hypernet with the same name will be selected instead.
-    """
-    hypernet_name = hypernet_name.lower()
-    if hypernet_hash is not None:
-        # Try to match the hash in the name
-        for hypernet_key in shared.hypernetworks.keys():
-            result = re_hypernet_hash.search(hypernet_key)
-            if result is not None and result[1] == hypernet_hash:
-                return hypernet_key
-    else:
-        # Fall back to a hypernet with the same name
-        for hypernet_key in shared.hypernetworks.keys():
-            if hypernet_key.lower().startswith(hypernet_name):
-                return hypernet_key
-
-    return None
-
-
 def restore_old_hires_fix_params(res):
     """for infotexts that specify old First pass size parameter, convert it into
     width, height, and hr scale"""
@@ -305,6 +280,9 @@ Steps: 20, Sampler: Euler a, CFG scale: 7, Seed: 965400086, Size: 512x512, Model
     if "Hires sampler" not in res:
         res["Hires sampler"] = "Use same sampler"
 
+    if "Hires checkpoint" not in res:
+        res["Hires checkpoint"] = "Use same checkpoint"
+
     if "Hires prompt" not in res:
         res["Hires prompt"] = ""
 
@@ -329,13 +307,15 @@ Steps: 20, Sampler: Euler a, CFG scale: 7, Seed: 965400086, Size: 512x512, Model
     if "Schedule rho" not in res:
         res["Schedule rho"] = 0
 
+    if "VAE Encoder" not in res:
+        res["VAE Encoder"] = "Full"
+
+    if "VAE Decoder" not in res:
+        res["VAE Decoder"] = "Full"
+
     return res
 
 
-settings_map = {}
-
-
-
 infotext_to_setting_name_mapping = [
     ('Clip skip', 'CLIP_stop_at_last_layers', ),
     ('Conditional mask weight', 'inpainting_mask_weight'),
@@ -348,6 +328,10 @@ infotext_to_setting_name_mapping = [
     ('Noise multiplier', 'initial_noise_multiplier'),
     ('Eta', 'eta_ancestral'),
     ('Eta DDIM', 'eta_ddim'),
+    ('Sigma churn', 's_churn'),
+    ('Sigma tmin', 's_tmin'),
+    ('Sigma tmax', 's_tmax'),
+    ('Sigma noise', 's_noise'),
     ('Discard penultimate sigma', 'always_discard_next_to_last_sigma'),
     ('UniPC variant', 'uni_pc_variant'),
     ('UniPC skip type', 'uni_pc_skip_type'),
@@ -357,6 +341,11 @@ infotext_to_setting_name_mapping = [
     ('Token merging ratio hr', 'token_merging_ratio_hr'),
     ('RNG', 'randn_source'),
     ('NGMS', 's_min_uncond'),
+    ('Pad conds', 'pad_cond_uncond'),
+    ('VAE Encoder', 'sd_vae_encode_method'),
+    ('VAE Decoder', 'sd_vae_decode_method'),
+    ('Refiner', 'sd_refiner_checkpoint'),
+    ('Refiner switch at', 'sd_refiner_switch_at'),
 ]
 
 

modules/gfpgan_model.py

@@ -25,7 +25,7 @@ def gfpgann():
         return None
 
     models = modelloader.load_models(model_path, model_url, user_path, ext_filter="GFPGAN")
-    if len(models) == 1 and "http" in models[0]:
+    if len(models) == 1 and models[0].startswith("http"):
         model_file = models[0]
     elif len(models) != 0:
         latest_file = max(models, key=os.path.getctime)
@@ -70,11 +70,8 @@ gfpgan_constructor = None
 
 
 def setup_model(dirname):
-    global model_path
-    if not os.path.exists(model_path):
-        os.makedirs(model_path)
-
     try:
+        os.makedirs(model_path, exist_ok=True)
         from gfpgan import GFPGANer
         from facexlib import detection, parsing  # noqa: F401
         global user_path

modules/gradio_extensons.py

@@ -0,0 +1,60 @@
+import gradio as gr
+
+from modules import scripts
+
+def add_classes_to_gradio_component(comp):
+    """
+    this adds gradio-* to the component for css styling (ie gradio-button to gr.Button), as well as some others
+    """
+
+    comp.elem_classes = [f"gradio-{comp.get_block_name()}", *(comp.elem_classes or [])]
+
+    if getattr(comp, 'multiselect', False):
+        comp.elem_classes.append('multiselect')
+
+
+def IOComponent_init(self, *args, **kwargs):
+    self.webui_tooltip = kwargs.pop('tooltip', None)
+
+    if scripts.scripts_current is not None:
+        scripts.scripts_current.before_component(self, **kwargs)
+
+    scripts.script_callbacks.before_component_callback(self, **kwargs)
+
+    res = original_IOComponent_init(self, *args, **kwargs)
+
+    add_classes_to_gradio_component(self)
+
+    scripts.script_callbacks.after_component_callback(self, **kwargs)
+
+    if scripts.scripts_current is not None:
+        scripts.scripts_current.after_component(self, **kwargs)
+
+    return res
+
+
+def Block_get_config(self):
+    config = original_Block_get_config(self)
+
+    webui_tooltip = getattr(self, 'webui_tooltip', None)
+    if webui_tooltip:
+        config["webui_tooltip"] = webui_tooltip
+
+    return config
+
+
+def BlockContext_init(self, *args, **kwargs):
+    res = original_BlockContext_init(self, *args, **kwargs)
+
+    add_classes_to_gradio_component(self)
+
+    return res
+
+
+original_IOComponent_init = gr.components.IOComponent.__init__
+original_Block_get_config = gr.blocks.Block.get_config
+original_BlockContext_init = gr.blocks.BlockContext.__init__
+
+gr.components.IOComponent.__init__ = IOComponent_init
+gr.blocks.Block.get_config = Block_get_config
+gr.blocks.BlockContext.__init__ = BlockContext_init

modules/hashes.py

@@ -1,38 +1,11 @@
 import hashlib
-import json
 import os.path
 
-import filelock
-
 from modules import shared
-from modules.paths import data_path
+import modules.cache
 
-
+dump_cache = modules.cache.dump_cache
-cache_filename = os.path.join(data_path, "cache.json")
+cache = modules.cache.cache
-cache_data = None
-
-
-def dump_cache():
-    with filelock.FileLock(f"{cache_filename}.lock"):
-        with open(cache_filename, "w", encoding="utf8") as file:
-            json.dump(cache_data, file, indent=4)
-
-
-def cache(subsection):
-    global cache_data
-
-    if cache_data is None:
-        with filelock.FileLock(f"{cache_filename}.lock"):
-            if not os.path.isfile(cache_filename):
-                cache_data = {}
-            else:
-                with open(cache_filename, "r", encoding="utf8") as file:
-                    cache_data = json.load(file)
-
-    s = cache_data.get(subsection, {})
-    cache_data[subsection] = s
-
-    return s
 
 
 def calculate_sha256(filename):

modules/hypernetworks/hypernetwork.py

@@ -3,13 +3,14 @@ import glob
 import html
 import os
 import inspect
+from contextlib import closing
 
 import modules.textual_inversion.dataset
 import torch
 import tqdm
 from einops import rearrange, repeat
 from ldm.util import default
-from modules import devices, processing, sd_models, shared, sd_samplers, hashes, sd_hijack_checkpoint, errors
+from modules import devices, sd_models, shared, sd_samplers, hashes, sd_hijack_checkpoint, errors
 from modules.textual_inversion import textual_inversion, logging
 from modules.textual_inversion.learn_schedule import LearnRateScheduler
 from torch import einsum
@@ -353,17 +354,6 @@ def load_hypernetworks(names, multipliers=None):
         shared.loaded_hypernetworks.append(hypernetwork)
 
 
-def find_closest_hypernetwork_name(search: str):
-    if not search:
-        return None
-    search = search.lower()
-    applicable = [name for name in shared.hypernetworks if search in name.lower()]
-    if not applicable:
-        return None
-    applicable = sorted(applicable, key=lambda name: len(name))
-    return applicable[0]
-
-
 def apply_single_hypernetwork(hypernetwork, context_k, context_v, layer=None):
     hypernetwork_layers = (hypernetwork.layers if hypernetwork is not None else {}).get(context_k.shape[2], None)
 
@@ -388,7 +378,7 @@ def apply_hypernetworks(hypernetworks, context, layer=None):
     return context_k, context_v
 
 
-def attention_CrossAttention_forward(self, x, context=None, mask=None):
+def attention_CrossAttention_forward(self, x, context=None, mask=None, **kwargs):
     h = self.heads
 
     q = self.to_q(x)
@@ -446,18 +436,6 @@ def statistics(data):
     return total_information, recent_information
 
 
-def report_statistics(loss_info:dict):
-    keys = sorted(loss_info.keys(), key=lambda x: sum(loss_info[x]) / len(loss_info[x]))
-    for key in keys:
-        try:
-            print("Loss statistics for file " + key)
-            info, recent = statistics(list(loss_info[key]))
-            print(info)
-            print(recent)
-        except Exception as e:
-            print(e)
-
-
 def create_hypernetwork(name, enable_sizes, overwrite_old, layer_structure=None, activation_func=None, weight_init=None, add_layer_norm=False, use_dropout=False, dropout_structure=None):
     # Remove illegal characters from name.
     name = "".join( x for x in name if (x.isalnum() or x in "._- "))
@@ -491,8 +469,7 @@ def create_hypernetwork(name, enable_sizes, overwrite_old, layer_structure=None,
 
 
 def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradient_step, data_root, log_directory, training_width, training_height, varsize, steps, clip_grad_mode, clip_grad_value, shuffle_tags, tag_drop_out, latent_sampling_method, use_weight, create_image_every, save_hypernetwork_every, template_filename, preview_from_txt2img, preview_prompt, preview_negative_prompt, preview_steps, preview_sampler_index, preview_cfg_scale, preview_seed, preview_width, preview_height):
-    # images allows training previews to have infotext. Importing it at the top causes a circular import problem.
+    from modules import images, processing
-    from modules import images
 
     save_hypernetwork_every = save_hypernetwork_every or 0
     create_image_every = create_image_every or 0
@@ -734,6 +711,7 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi
 
                     preview_text = p.prompt
 
+                    with closing(p):
                         processed = processing.process_images(p)
                         image = processed.images[0] if len(processed.images) > 0 else None
 
@@ -770,7 +748,6 @@ Last saved image: {html.escape(last_saved_image)}<br/>
         pbar.leave = False
         pbar.close()
         hypernetwork.eval()
-        #report_statistics(loss_dict)
         sd_hijack_checkpoint.remove()
 
 

modules/images.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import datetime
 
 import pytz
@@ -10,7 +12,7 @@ import re
 import numpy as np
 import piexif
 import piexif.helper
-from PIL import Image, ImageFont, ImageDraw, PngImagePlugin
+from PIL import Image, ImageFont, ImageDraw, ImageColor, PngImagePlugin
 import string
 import json
 import hashlib
@@ -139,6 +141,11 @@ class GridAnnotation:
 
 
 def draw_grid_annotations(im, width, height, hor_texts, ver_texts, margin=0):
+
+    color_active = ImageColor.getcolor(opts.grid_text_active_color, 'RGB')
+    color_inactive = ImageColor.getcolor(opts.grid_text_inactive_color, 'RGB')
+    color_background = ImageColor.getcolor(opts.grid_background_color, 'RGB')
+
     def wrap(drawing, text, font, line_length):
         lines = ['']
         for word in text.split():
@@ -168,9 +175,6 @@ def draw_grid_annotations(im, width, height, hor_texts, ver_texts, margin=0):
 
     fnt = get_font(fontsize)
 
-    color_active = (0, 0, 0)
-    color_inactive = (153, 153, 153)
-
     pad_left = 0 if sum([sum([len(line.text) for line in lines]) for lines in ver_texts]) == 0 else width * 3 // 4
 
     cols = im.width // width
@@ -179,7 +183,7 @@ def draw_grid_annotations(im, width, height, hor_texts, ver_texts, margin=0):
     assert cols == len(hor_texts), f'bad number of horizontal texts: {len(hor_texts)}; must be {cols}'
     assert rows == len(ver_texts), f'bad number of vertical texts: {len(ver_texts)}; must be {rows}'
 
-    calc_img = Image.new("RGB", (1, 1), "white")
+    calc_img = Image.new("RGB", (1, 1), color_background)
     calc_d = ImageDraw.Draw(calc_img)
 
     for texts, allowed_width in zip(hor_texts + ver_texts, [width] * len(hor_texts) + [pad_left] * len(ver_texts)):
@@ -200,7 +204,7 @@ def draw_grid_annotations(im, width, height, hor_texts, ver_texts, margin=0):
 
     pad_top = 0 if sum(hor_text_heights) == 0 else max(hor_text_heights) + line_spacing * 2
 
-    result = Image.new("RGB", (im.width + pad_left + margin * (cols-1), im.height + pad_top + margin * (rows-1)), "white")
+    result = Image.new("RGB", (im.width + pad_left + margin * (cols-1), im.height + pad_top + margin * (rows-1)), color_background)
 
     for row in range(rows):
         for col in range(cols):
@@ -302,17 +306,19 @@ def resize_image(resize_mode, im, width, height, upscaler_name=None):
 
         if ratio < src_ratio:
             fill_height = height // 2 - src_h // 2
+            if fill_height > 0:
                 res.paste(resized.resize((width, fill_height), box=(0, 0, width, 0)), box=(0, 0))
                 res.paste(resized.resize((width, fill_height), box=(0, resized.height, width, resized.height)), box=(0, fill_height + src_h))
         elif ratio > src_ratio:
             fill_width = width // 2 - src_w // 2
+            if fill_width > 0:
                 res.paste(resized.resize((fill_width, height), box=(0, 0, 0, height)), box=(0, 0))
                 res.paste(resized.resize((fill_width, height), box=(resized.width, 0, resized.width, height)), box=(fill_width + src_w, 0))
 
     return res
 
 
-invalid_filename_chars = '<>:"/\\|?*\n'
+invalid_filename_chars = '<>:"/\\|?*\n\r\t'
 invalid_filename_prefix = ' '
 invalid_filename_postfix = ' .'
 re_nonletters = re.compile(r'[\s' + string.punctuation + ']+')
@@ -357,7 +363,7 @@ class FilenameGenerator:
         'styles': lambda self: self.p and sanitize_filename_part(", ".join([style for style in self.p.styles if not style == "None"]) or "None", replace_spaces=False),
         'sampler': lambda self: self.p and sanitize_filename_part(self.p.sampler_name, replace_spaces=False),
         'model_hash': lambda self: getattr(self.p, "sd_model_hash", shared.sd_model.sd_model_hash),
-        'model_name': lambda self: sanitize_filename_part(shared.sd_model.sd_checkpoint_info.model_name, replace_spaces=False),
+        'model_name': lambda self: sanitize_filename_part(shared.sd_model.sd_checkpoint_info.name_for_extra, replace_spaces=False),
         'date': lambda self: datetime.datetime.now().strftime('%Y-%m-%d'),
         'datetime': lambda self, *args: self.datetime(*args),  # accepts formats: [datetime], [datetime<Format>], [datetime<Format><Time Zone>]
         'job_timestamp': lambda self: getattr(self.p, "job_timestamp", shared.state.job_timestamp),
@@ -372,8 +378,9 @@ class FilenameGenerator:
         'hasprompt': lambda self, *args: self.hasprompt(*args),  # accepts formats:[hasprompt<prompt1|default><prompt2>..]
         'clip_skip': lambda self: opts.data["CLIP_stop_at_last_layers"],
         'denoising': lambda self: self.p.denoising_strength if self.p and self.p.denoising_strength else NOTHING_AND_SKIP_PREVIOUS_TEXT,
+        'user': lambda self: self.p.user,
         'vae_filename': lambda self: self.get_vae_filename(),
-
+        'none': lambda self: '', # Overrides the default so you can get just the sequence number
     }
     default_time_format = '%Y%m%d%H%M%S'
 
@@ -497,13 +504,23 @@ def get_next_sequence_number(path, basename):
     return result + 1
 
 
-def save_image_with_geninfo(image, geninfo, filename, extension=None, existing_pnginfo=None):
+def save_image_with_geninfo(image, geninfo, filename, extension=None, existing_pnginfo=None, pnginfo_section_name='parameters'):
+    """
+    Saves image to filename, including geninfo as text information for generation info.
+    For PNG images, geninfo is added to existing pnginfo dictionary using the pnginfo_section_name argument as key.
+    For JPG images, there's no dictionary and geninfo just replaces the EXIF description.
+    """
+
     if extension is None:
         extension = os.path.splitext(filename)[1]
 
     image_format = Image.registered_extensions()[extension]
 
     if extension.lower() == '.png':
+        existing_pnginfo = existing_pnginfo or {}
+        if opts.enable_pnginfo:
+            existing_pnginfo[pnginfo_section_name] = geninfo
+
         if opts.enable_pnginfo:
             pnginfo_data = PngImagePlugin.PngInfo()
             for k, v in (existing_pnginfo or {}).items():
@@ -585,13 +602,13 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
         else:
             file_decoration = opts.samples_filename_pattern or "[seed]-[prompt_spaces]"
 
+        file_decoration = namegen.apply(file_decoration) + suffix
+
         add_number = opts.save_images_add_number or file_decoration == ''
 
         if file_decoration != "" and add_number:
             file_decoration = f"-{file_decoration}"
 
-        file_decoration = namegen.apply(file_decoration) + suffix
-
         if add_number:
             basecount = get_next_sequence_number(path, basename)
             fullfn = None
@@ -622,7 +639,7 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
         """
         temp_file_path = f"{filename_without_extension}.tmp"
 
-        save_image_with_geninfo(image_to_save, info, temp_file_path, extension, params.pnginfo)
+        save_image_with_geninfo(image_to_save, info, temp_file_path, extension, existing_pnginfo=params.pnginfo, pnginfo_section_name=pnginfo_section_name)
 
         os.replace(temp_file_path, filename_without_extension + extension)
 
@@ -639,12 +656,18 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
     oversize = image.width > opts.target_side_length or image.height > opts.target_side_length
     if opts.export_for_4chan and (oversize or os.stat(fullfn).st_size > opts.img_downscale_threshold * 1024 * 1024):
         ratio = image.width / image.height
-
+        resize_to = None
         if oversize and ratio > 1:
-            image = image.resize((round(opts.target_side_length), round(image.height * opts.target_side_length / image.width)), LANCZOS)
+            resize_to = round(opts.target_side_length), round(image.height * opts.target_side_length / image.width)
         elif oversize:
-            image = image.resize((round(image.width * opts.target_side_length / image.height), round(opts.target_side_length)), LANCZOS)
+            resize_to = round(image.width * opts.target_side_length / image.height), round(opts.target_side_length)
 
+        if resize_to is not None:
+            try:
+                # Resizing image with LANCZOS could throw an exception if e.g. image mode is I;16
+                image = image.resize(resize_to, LANCZOS)
+            except Exception:
+                image = image.resize(resize_to)
         try:
             _atomically_save_image(image, fullfn_without_extension, ".jpg")
         except Exception as e:
@@ -662,8 +685,15 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
     return fullfn, txt_fullfn
 
 
-def read_info_from_image(image):
+IGNORED_INFO_KEYS = {
-    items = image.info or {}
+    'jfif', 'jfif_version', 'jfif_unit', 'jfif_density', 'dpi', 'exif',
+    'loop', 'background', 'timestamp', 'duration', 'progressive', 'progression',
+    'icc_profile', 'chromaticity', 'photoshop',
+}
+
+
+def read_info_from_image(image: Image.Image) -> tuple[str | None, dict]:
+    items = (image.info or {}).copy()
 
     geninfo = items.pop('parameters', None)
 
@@ -679,9 +709,7 @@ def read_info_from_image(image):
             items['exif comment'] = exif_comment
             geninfo = exif_comment
 
-    for field in ['jfif', 'jfif_version', 'jfif_unit', 'jfif_density', 'dpi', 'exif',
+    for field in IGNORED_INFO_KEYS:
-                    'loop', 'background', 'timestamp', 'duration', 'progressive', 'progression',
-                    'icc_profile', 'chromaticity']:
         items.pop(field, None)
 
     if items.get("Software", None) == "NovelAI":

modules/img2img.py

@@ -1,11 +1,13 @@
 import os
+from contextlib import closing
 from pathlib import Path
 
 import numpy as np
-from PIL import Image, ImageOps, ImageFilter, ImageEnhance, ImageChops, UnidentifiedImageError
+from PIL import Image, ImageOps, ImageFilter, ImageEnhance, UnidentifiedImageError
+import gradio as gr
 
-from modules import sd_samplers
+from modules import sd_samplers, images as imgutil
-from modules.generation_parameters_copypaste import create_override_settings_dict
+from modules.generation_parameters_copypaste import create_override_settings_dict, parse_generation_parameters
 from modules.processing import Processed, StableDiffusionProcessingImg2Img, process_images
 from modules.shared import opts, state
 import modules.shared as shared
@@ -14,10 +16,11 @@ from modules.ui import plaintext_to_html
 import modules.scripts
 
 
-def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args, to_scale=False, scale_by=1.0):
+def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args, to_scale=False, scale_by=1.0, use_png_info=False, png_info_props=None, png_info_dir=None):
+    output_dir = output_dir.strip()
     processing.fix_seed(p)
 
-    images = shared.listfiles(input_dir)
+    images = list(shared.walk_files(input_dir, allowed_extensions=(".png", ".jpg", ".jpeg", ".webp", ".tif", ".tiff")))
 
     is_inpaint_batch = False
     if inpaint_mask_dir:
@@ -29,13 +32,16 @@ def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args, to_scale=Fal
 
     print(f"Will process {len(images)} images, creating {p.n_iter * p.batch_size} new images for each.")
 
-    save_normally = output_dir == ''
-
-    p.do_not_save_grid = True
-    p.do_not_save_samples = not save_normally
-
     state.job_count = len(images) * p.n_iter
 
+    # extract "default" params to use in case getting png info fails
+    prompt = p.prompt
+    negative_prompt = p.negative_prompt
+    seed = p.seed
+    cfg_scale = p.cfg_scale
+    sampler_name = p.sampler_name
+    steps = p.steps
+
     for i, image in enumerate(images):
         state.job = f"{i+1} out of {len(images)}"
         if state.skipped:
@@ -79,25 +85,38 @@ def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args, to_scale=Fal
             mask_image = Image.open(mask_image_path)
             p.image_mask = mask_image
 
+        if use_png_info:
+            try:
+                info_img = img
+                if png_info_dir:
+                    info_img_path = os.path.join(png_info_dir, os.path.basename(image))
+                    info_img = Image.open(info_img_path)
+                geninfo, _ = imgutil.read_info_from_image(info_img)
+                parsed_parameters = parse_generation_parameters(geninfo)
+                parsed_parameters = {k: v for k, v in parsed_parameters.items() if k in (png_info_props or {})}
+            except Exception:
+                parsed_parameters = {}
+
+            p.prompt = prompt + (" " + parsed_parameters["Prompt"] if "Prompt" in parsed_parameters else "")
+            p.negative_prompt = negative_prompt + (" " + parsed_parameters["Negative prompt"] if "Negative prompt" in parsed_parameters else "")
+            p.seed = int(parsed_parameters.get("Seed", seed))
+            p.cfg_scale = float(parsed_parameters.get("CFG scale", cfg_scale))
+            p.sampler_name = parsed_parameters.get("Sampler", sampler_name)
+            p.steps = int(parsed_parameters.get("Steps", steps))
+
         proc = modules.scripts.scripts_img2img.run(p, *args)
         if proc is None:
-            proc = process_images(p)
+            if output_dir:
-
+                p.outpath_samples = output_dir
-        for n, processed_image in enumerate(proc.images):
+                p.override_settings['save_to_dirs'] = False
-            filename = image_path.name
+                if p.n_iter > 1 or p.batch_size > 1:
-
+                    p.override_settings['samples_filename_pattern'] = f'{image_path.stem}-[generation_number]'
-            if n > 0:
+                else:
-                left, right = os.path.splitext(filename)
+                    p.override_settings['samples_filename_pattern'] = f'{image_path.stem}'
-                filename = f"{left}-{n}{right}"
+            process_images(p)
-
-            if not save_normally:
-                os.makedirs(output_dir, exist_ok=True)
-                if processed_image.mode == 'RGBA':
-                    processed_image = processed_image.convert("RGB")
-                processed_image.save(os.path.join(output_dir, filename))
 
 
-def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_styles, init_img, sketch, init_img_with_mask, inpaint_color_sketch, inpaint_color_sketch_orig, init_img_inpaint, init_mask_inpaint, steps: int, sampler_index: int, mask_blur: int, mask_alpha: float, inpainting_fill: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, image_cfg_scale: float, denoising_strength: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, selected_scale_tab: int, height: int, width: int, scale_by: float, resize_mode: int, inpaint_full_res: bool, inpaint_full_res_padding: int, inpainting_mask_invert: int, img2img_batch_input_dir: str, img2img_batch_output_dir: str, img2img_batch_inpaint_mask_dir: str, override_settings_texts, *args):
+def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_styles, init_img, sketch, init_img_with_mask, inpaint_color_sketch, inpaint_color_sketch_orig, init_img_inpaint, init_mask_inpaint, steps: int, sampler_index: int, mask_blur: int, mask_alpha: float, inpainting_fill: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, image_cfg_scale: float, denoising_strength: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, selected_scale_tab: int, height: int, width: int, scale_by: float, resize_mode: int, inpaint_full_res: bool, inpaint_full_res_padding: int, inpainting_mask_invert: int, img2img_batch_input_dir: str, img2img_batch_output_dir: str, img2img_batch_inpaint_mask_dir: str, override_settings_texts, img2img_batch_use_png_info: bool, img2img_batch_png_info_props: list, img2img_batch_png_info_dir: str, request: gr.Request, *args):
     override_settings = create_override_settings_dict(override_settings_texts)
 
     is_batch = mode == 5
@@ -110,9 +129,7 @@ def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_s
         mask = None
     elif mode == 2:  # inpaint
         image, mask = init_img_with_mask["image"], init_img_with_mask["mask"]
-        alpha_mask = ImageOps.invert(image.split()[-1]).convert('L').point(lambda x: 255 if x > 0 else 0, mode='1')
+        mask = mask.split()[-1].convert("L").point(lambda x: 255 if x > 128 else 0)
-        mask = mask.convert('L').point(lambda x: 255 if x > 128 else 0, mode='1')
-        mask = ImageChops.lighter(alpha_mask, mask).convert('L')
         image = image.convert("RGB")
     elif mode == 3:  # inpaint sketch
         image = inpaint_color_sketch
@@ -180,16 +197,19 @@ def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_s
     p.scripts = modules.scripts.scripts_img2img
     p.script_args = args
 
+    p.user = request.username
+
     if shared.cmd_opts.enable_console_prompts:
         print(f"\nimg2img: {prompt}", file=shared.progress_print_out)
 
     if mask:
         p.extra_generation_params["Mask blur"] = mask_blur
 
+    with closing(p):
         if is_batch:
             assert not shared.cmd_opts.hide_ui_dir_config, "Launched with --hide-ui-dir-config, batch img2img disabled"
 
-        process_batch(p, img2img_batch_input_dir, img2img_batch_output_dir, img2img_batch_inpaint_mask_dir, args, to_scale=selected_scale_tab == 1, scale_by=scale_by)
+            process_batch(p, img2img_batch_input_dir, img2img_batch_output_dir, img2img_batch_inpaint_mask_dir, args, to_scale=selected_scale_tab == 1, scale_by=scale_by, use_png_info=img2img_batch_use_png_info, png_info_props=img2img_batch_png_info_props, png_info_dir=img2img_batch_png_info_dir)
 
             processed = Processed(p, [], p.seed, "")
         else:
@@ -197,8 +217,6 @@ def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_s
             if processed is None:
                 processed = process_images(p)
 
-    p.close()
-
     shared.total_tqdm.clear()
 
     generation_info_js = processed.js()
@@ -208,4 +226,4 @@ def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_s
     if opts.do_not_show_images:
         processed.images = []
 
-    return processed.images, generation_info_js, plaintext_to_html(processed.info), plaintext_to_html(processed.comments)
+    return processed.images, generation_info_js, plaintext_to_html(processed.info), plaintext_to_html(processed.comments, classname="comments")

modules/interrogate.py

@@ -184,8 +184,7 @@ class InterrogateModels:
 
     def interrogate(self, pil_image):
         res = ""
-        shared.state.begin()
+        shared.state.begin(job="interrogate")
-        shared.state.job = 'interrogate'
         try:
             if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
                 lowvram.send_everything_to_cpu()

modules/launch_utils.py

@@ -1,4 +1,5 @@
 # this scripts installs necessary requirements and launches main program in webui.py
+import re
 import subprocess
 import os
 import sys
@@ -9,6 +10,7 @@ from functools import lru_cache
 
 from modules import cmd_args, errors
 from modules.paths_internal import script_path, extensions_dir
+from modules.timer import startup_timer
 
 args, _ = cmd_args.parser.parse_known_args()
 
@@ -69,10 +71,12 @@ def git_tag():
         return subprocess.check_output([git, "describe", "--tags"], shell=False, encoding='utf8').strip()
     except Exception:
         try:
-            from pathlib import Path
+
-            changelog_md = Path(__file__).parent.parent / "CHANGELOG.md"
+            changelog_md = os.path.join(os.path.dirname(os.path.dirname(__file__)), "CHANGELOG.md")
-            with changelog_md.open(encoding="utf-8") as file:
+            with open(changelog_md, "r", encoding="utf-8") as file:
-                return next((line.strip() for line in file if line.strip()), "<none>")
+                line = next((line.strip() for line in file if line.strip()), "<none>")
+                line = line.replace("## ", "")
+                return line
         except Exception:
             return "<none>"
 
@@ -142,15 +146,15 @@ def git_clone(url, dir, name, commithash=None):
         if commithash is None:
             return
 
-        current_hash = run(f'"{git}" -C "{dir}" rev-parse HEAD', None, f"Couldn't determine {name}'s hash: {commithash}").strip()
+        current_hash = run(f'"{git}" -C "{dir}" rev-parse HEAD', None, f"Couldn't determine {name}'s hash: {commithash}", live=False).strip()
         if current_hash == commithash:
             return
 
         run(f'"{git}" -C "{dir}" fetch', f"Fetching updates for {name}...", f"Couldn't fetch {name}")
-        run(f'"{git}" -C "{dir}" checkout {commithash}', f"Checking out commit for {name} with hash: {commithash}...", f"Couldn't checkout commit {commithash} for {name}")
+        run(f'"{git}" -C "{dir}" checkout {commithash}', f"Checking out commit for {name} with hash: {commithash}...", f"Couldn't checkout commit {commithash} for {name}", live=True)
         return
 
-    run(f'"{git}" clone "{url}" "{dir}"', f"Cloning {name} into {dir}...", f"Couldn't clone {name}")
+    run(f'"{git}" clone "{url}" "{dir}"', f"Cloning {name} into {dir}...", f"Couldn't clone {name}", live=True)
 
     if commithash is not None:
         run(f'"{git}" -C "{dir}" checkout {commithash}', None, "Couldn't checkout {name}'s hash: {commithash}")
@@ -190,7 +194,7 @@ def run_extension_installer(extension_dir):
 
     try:
         env = os.environ.copy()
-        env['PYTHONPATH'] = os.path.abspath(".")
+        env['PYTHONPATH'] = f"{os.path.abspath('.')}{os.pathsep}{env.get('PYTHONPATH', '')}"
 
         print(run(f'"{python}" "{path_installer}"', errdesc=f"Error running install.py for extension {extension_dir}", custom_env=env))
     except Exception as e:
@@ -220,8 +224,51 @@ def run_extensions_installers(settings_file):
     if not os.path.isdir(extensions_dir):
         return
 
+    with startup_timer.subcategory("run extensions installers"):
         for dirname_extension in list_extensions(settings_file):
-        run_extension_installer(os.path.join(extensions_dir, dirname_extension))
+            path = os.path.join(extensions_dir, dirname_extension)
+
+            if os.path.isdir(path):
+                run_extension_installer(path)
+                startup_timer.record(dirname_extension)
+
+
+re_requirement = re.compile(r"\s*([-_a-zA-Z0-9]+)\s*(?:==\s*([-+_.a-zA-Z0-9]+))?\s*")
+
+
+def requirements_met(requirements_file):
+    """
+    Does a simple parse of a requirements.txt file to determine if all rerqirements in it
+    are already installed. Returns True if so, False if not installed or parsing fails.
+    """
+
+    import importlib.metadata
+    import packaging.version
+
+    with open(requirements_file, "r", encoding="utf8") as file:
+        for line in file:
+            if line.strip() == "":
+                continue
+
+            m = re.match(re_requirement, line)
+            if m is None:
+                return False
+
+            package = m.group(1).strip()
+            version_required = (m.group(2) or "").strip()
+
+            if version_required == "":
+                continue
+
+            try:
+                version_installed = importlib.metadata.version(package)
+            except Exception:
+                return False
+
+            if packaging.version.parse(version_required) != packaging.version.parse(version_installed):
+                return False
+
+    return True
 
 
 def prepare_environment():
@@ -235,11 +282,13 @@ def prepare_environment():
     openclip_package = os.environ.get('OPENCLIP_PACKAGE', "https://github.com/mlfoundations/open_clip/archive/bb6e834e9c70d9c27d0dc3ecedeebeaeb1ffad6b.zip")
 
     stable_diffusion_repo = os.environ.get('STABLE_DIFFUSION_REPO', "https://github.com/Stability-AI/stablediffusion.git")
+    stable_diffusion_xl_repo = os.environ.get('STABLE_DIFFUSION_XL_REPO', "https://github.com/Stability-AI/generative-models.git")
     k_diffusion_repo = os.environ.get('K_DIFFUSION_REPO', 'https://github.com/crowsonkb/k-diffusion.git')
     codeformer_repo = os.environ.get('CODEFORMER_REPO', 'https://github.com/sczhou/CodeFormer.git')
     blip_repo = os.environ.get('BLIP_REPO', 'https://github.com/salesforce/BLIP.git')
 
     stable_diffusion_commit_hash = os.environ.get('STABLE_DIFFUSION_COMMIT_HASH', "cf1d67a6fd5ea1aa600c4df58e5b47da45f6bdbf")
+    stable_diffusion_xl_commit_hash = os.environ.get('STABLE_DIFFUSION_XL_COMMIT_HASH', "5c10deee76adad0032b412294130090932317a87")
     k_diffusion_commit_hash = os.environ.get('K_DIFFUSION_COMMIT_HASH', "c9fe758757e022f05ca5a53fa8fac28889e4f1cf")
     codeformer_commit_hash = os.environ.get('CODEFORMER_COMMIT_HASH', "c5b4593074ba6214284d6acd5f1719b6c5d739af")
     blip_commit_hash = os.environ.get('BLIP_COMMIT_HASH', "48211a1594f1321b00f14c9f7a5b4813144b2fb9")
@@ -247,15 +296,18 @@ def prepare_environment():
     try:
         # the existance of this file is a signal to webui.sh/bat that webui needs to be restarted when it stops execution
         os.remove(os.path.join(script_path, "tmp", "restart"))
-        os.environ.setdefault('SD_WEBUI_RESTARTING ', '1')
+        os.environ.setdefault('SD_WEBUI_RESTARTING', '1')
     except OSError:
         pass
 
     if not args.skip_python_version_check:
         check_python_version()
 
+    startup_timer.record("checks")
+
     commit = commit_hash()
     tag = git_tag()
+    startup_timer.record("git version info")
 
     print(f"Python {sys.version}")
     print(f"Version: {tag}")
@@ -263,21 +315,27 @@ def prepare_environment():
 
     if args.reinstall_torch or not is_installed("torch") or not is_installed("torchvision"):
         run(f'"{python}" -m {torch_command}', "Installing torch and torchvision", "Couldn't install torch", live=True)
+        startup_timer.record("install torch")
 
     if not args.skip_torch_cuda_test and not check_run_python("import torch; assert torch.cuda.is_available()"):
         raise RuntimeError(
             'Torch is not able to use GPU; '
             'add --skip-torch-cuda-test to COMMANDLINE_ARGS variable to disable this check'
         )
+    startup_timer.record("torch GPU test")
+
 
     if not is_installed("gfpgan"):
         run_pip(f"install {gfpgan_package}", "gfpgan")
+        startup_timer.record("install gfpgan")
 
     if not is_installed("clip"):
         run_pip(f"install {clip_package}", "clip")
+        startup_timer.record("install clip")
 
     if not is_installed("open_clip"):
         run_pip(f"install {openclip_package}", "open_clip")
+        startup_timer.record("install open_clip")
 
     if (not is_installed("xformers") or args.reinstall_xformers) and args.xformers:
         if platform.system() == "Windows":
@@ -291,36 +349,49 @@ def prepare_environment():
         elif platform.system() == "Linux":
             run_pip(f"install -U -I --no-deps {xformers_package}", "xformers")
 
+        startup_timer.record("install xformers")
+
     if not is_installed("ngrok") and args.ngrok:
         run_pip("install ngrok", "ngrok")
+        startup_timer.record("install ngrok")
 
     os.makedirs(os.path.join(script_path, dir_repos), exist_ok=True)
 
     git_clone(stable_diffusion_repo, repo_dir('stable-diffusion-stability-ai'), "Stable Diffusion", stable_diffusion_commit_hash)
+    git_clone(stable_diffusion_xl_repo, repo_dir('generative-models'), "Stable Diffusion XL", stable_diffusion_xl_commit_hash)
     git_clone(k_diffusion_repo, repo_dir('k-diffusion'), "K-diffusion", k_diffusion_commit_hash)
     git_clone(codeformer_repo, repo_dir('CodeFormer'), "CodeFormer", codeformer_commit_hash)
     git_clone(blip_repo, repo_dir('BLIP'), "BLIP", blip_commit_hash)
 
+    startup_timer.record("clone repositores")
+
     if not is_installed("lpips"):
         run_pip(f"install -r \"{os.path.join(repo_dir('CodeFormer'), 'requirements.txt')}\"", "requirements for CodeFormer")
+        startup_timer.record("install CodeFormer requirements")
 
     if not os.path.isfile(requirements_file):
         requirements_file = os.path.join(script_path, requirements_file)
+
+    if not requirements_met(requirements_file):
         run_pip(f"install -r \"{requirements_file}\"", "requirements")
+        startup_timer.record("install requirements")
 
     run_extensions_installers(settings_file=args.ui_settings_file)
 
     if args.update_check:
         version_check(commit)
+        startup_timer.record("check version")
 
     if args.update_all_extensions:
         git_pull_recursive(extensions_dir)
+        startup_timer.record("update extensions")
 
     if "--exit" in sys.argv:
         print("Exiting because of --exit argument")
         exit(0)
 
 
+
 def configure_for_tests():
     if "--api" not in sys.argv:
         sys.argv.append("--api")

modules/lowvram.py

@@ -15,6 +15,9 @@ def send_everything_to_cpu():
 
 
 def setup_for_low_vram(sd_model, use_medvram):
+    if getattr(sd_model, 'lowvram', False):
+        return
+
     sd_model.lowvram = True
 
     parents = {}
@@ -53,19 +56,50 @@ def setup_for_low_vram(sd_model, use_medvram):
         send_me_to_gpu(first_stage_model, None)
         return first_stage_model_decode(z)
 
-    # for SD1, cond_stage_model is CLIP and its NN is in the tranformer frield, but for SD2, it's open clip, and it's in model field
+    to_remain_in_cpu = [
-    if hasattr(sd_model.cond_stage_model, 'model'):
+        (sd_model, 'first_stage_model'),
-        sd_model.cond_stage_model.transformer = sd_model.cond_stage_model.model
+        (sd_model, 'depth_model'),
+        (sd_model, 'embedder'),
+        (sd_model, 'model'),
+        (sd_model, 'embedder'),
+    ]
 
-    # remove several big modules: cond, first_stage, depth/embedder (if applicable), and unet from the model and then
+    is_sdxl = hasattr(sd_model, 'conditioner')
-    # send the model to GPU. Then put modules back. the modules will be in CPU.
+    is_sd2 = not is_sdxl and hasattr(sd_model.cond_stage_model, 'model')
-    stored = sd_model.cond_stage_model.transformer, sd_model.first_stage_model, getattr(sd_model, 'depth_model', None), getattr(sd_model, 'embedder', None), sd_model.model
+
-    sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.depth_model, sd_model.embedder, sd_model.model = None, None, None, None, None
+    if is_sdxl:
+        to_remain_in_cpu.append((sd_model, 'conditioner'))
+    elif is_sd2:
+        to_remain_in_cpu.append((sd_model.cond_stage_model, 'model'))
+    else:
+        to_remain_in_cpu.append((sd_model.cond_stage_model, 'transformer'))
+
+    # remove several big modules: cond, first_stage, depth/embedder (if applicable), and unet from the model
+    stored = []
+    for obj, field in to_remain_in_cpu:
+        module = getattr(obj, field, None)
+        stored.append(module)
+        setattr(obj, field, None)
+
+    # send the model to GPU.
     sd_model.to(devices.device)
-    sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.depth_model, sd_model.embedder, sd_model.model = stored
+
+    # put modules back. the modules will be in CPU.
+    for (obj, field), module in zip(to_remain_in_cpu, stored):
+        setattr(obj, field, module)
 
     # register hooks for those the first three models
+    if is_sdxl:
+        sd_model.conditioner.register_forward_pre_hook(send_me_to_gpu)
+    elif is_sd2:
+        sd_model.cond_stage_model.model.register_forward_pre_hook(send_me_to_gpu)
+        sd_model.cond_stage_model.model.token_embedding.register_forward_pre_hook(send_me_to_gpu)
+        parents[sd_model.cond_stage_model.model] = sd_model.cond_stage_model
+        parents[sd_model.cond_stage_model.model.token_embedding] = sd_model.cond_stage_model
+    else:
         sd_model.cond_stage_model.transformer.register_forward_pre_hook(send_me_to_gpu)
+        parents[sd_model.cond_stage_model.transformer] = sd_model.cond_stage_model
+
     sd_model.first_stage_model.register_forward_pre_hook(send_me_to_gpu)
     sd_model.first_stage_model.encode = first_stage_model_encode_wrap
     sd_model.first_stage_model.decode = first_stage_model_decode_wrap
@@ -73,11 +107,6 @@ def setup_for_low_vram(sd_model, use_medvram):
         sd_model.depth_model.register_forward_pre_hook(send_me_to_gpu)
     if sd_model.embedder:
         sd_model.embedder.register_forward_pre_hook(send_me_to_gpu)
-    parents[sd_model.cond_stage_model.transformer] = sd_model.cond_stage_model
-
-    if hasattr(sd_model.cond_stage_model, 'model'):
-        sd_model.cond_stage_model.model = sd_model.cond_stage_model.transformer
-        del sd_model.cond_stage_model.transformer
 
     if use_medvram:
         sd_model.model.register_forward_pre_hook(send_me_to_gpu)

modules/mac_specific.py

@@ -1,12 +1,19 @@
+import logging
+
 import torch
 import platform
 from modules.sd_hijack_utils import CondFunc
 from packaging import version
 
+log = logging.getLogger(__name__)
 
-# has_mps is only available in nightly pytorch (for now) and macOS 12.3+.
+
-# check `getattr` and try it for compatibility
+# before torch version 1.13, has_mps is only available in nightly pytorch and macOS 12.3+,
+# use check `getattr` and try it for compatibility.
+# in torch version 1.13, backends.mps.is_available() and backends.mps.is_built() are introduced in to check mps availabilty,
+# since torch 2.0.1+ nightly build, getattr(torch, 'has_mps', False) was deprecated, see https://github.com/pytorch/pytorch/pull/103279
 def check_for_mps() -> bool:
+    if version.parse(torch.__version__) <= version.parse("2.0.1"):
         if not getattr(torch, 'has_mps', False):
             return False
         try:
@@ -14,9 +21,25 @@ def check_for_mps() -> bool:
             return True
         except Exception:
             return False
+    else:
+        return torch.backends.mps.is_available() and torch.backends.mps.is_built()
+
+
 has_mps = check_for_mps()
 
 
+def torch_mps_gc() -> None:
+    try:
+        from modules.shared import state
+        if state.current_latent is not None:
+            log.debug("`current_latent` is set, skipping MPS garbage collection")
+            return
+        from torch.mps import empty_cache
+        empty_cache()
+    except Exception:
+        log.warning("MPS garbage collection failed", exc_info=True)
+
+
 # MPS workaround for https://github.com/pytorch/pytorch/issues/89784
 def cumsum_fix(input, cumsum_func, *args, **kwargs):
     if input.device.type == 'mps':

modules/modelloader.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 import shutil
 import importlib
@@ -8,6 +10,29 @@ from modules.upscaler import Upscaler, UpscalerLanczos, UpscalerNearest, Upscale
 from modules.paths import script_path, models_path
 
 
+def load_file_from_url(
+    url: str,
+    *,
+    model_dir: str,
+    progress: bool = True,
+    file_name: str | None = None,
+) -> str:
+    """Download a file from `url` into `model_dir`, using the file present if possible.
+
+    Returns the path to the downloaded file.
+    """
+    os.makedirs(model_dir, exist_ok=True)
+    if not file_name:
+        parts = urlparse(url)
+        file_name = os.path.basename(parts.path)
+    cached_file = os.path.abspath(os.path.join(model_dir, file_name))
+    if not os.path.exists(cached_file):
+        print(f'Downloading: "{url}" to {cached_file}\n')
+        from torch.hub import download_url_to_file
+        download_url_to_file(url, cached_file, progress=progress)
+    return cached_file
+
+
 def load_models(model_path: str, model_url: str = None, command_path: str = None, ext_filter=None, download_name=None, ext_blacklist=None) -> list:
     """
     A one-and done loader to try finding the desired models in specified directories.
@@ -46,9 +71,7 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None
 
         if model_url is not None and len(output) == 0:
             if download_name is not None:
-                from basicsr.utils.download_util import load_file_from_url
+                output.append(load_file_from_url(model_url, model_dir=places[0], file_name=download_name))
-                dl = load_file_from_url(model_url, places[0], True, download_name)
-                output.append(dl)
             else:
                 output.append(model_url)
 
@@ -59,7 +82,7 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None
 
 
 def friendly_name(file: str):
-    if "http" in file:
+    if file.startswith("http"):
         file = urlparse(file).path
 
     file = os.path.basename(file)
@@ -95,8 +118,7 @@ def cleanup_models():
 
 def move_files(src_path: str, dest_path: str, ext_filter: str = None):
     try:
-        if not os.path.exists(dest_path):
+        os.makedirs(dest_path, exist_ok=True)
-            os.makedirs(dest_path)
         if os.path.exists(src_path):
             for file in os.listdir(src_path):
                 fullpath = os.path.join(src_path, file)

modules/paths.py

@@ -5,6 +5,21 @@ from modules.paths_internal import models_path, script_path, data_path, extensio
 import modules.safe  # noqa: F401
 
 
+def mute_sdxl_imports():
+    """create fake modules that SDXL wants to import but doesn't actually use for our purposes"""
+
+    class Dummy:
+        pass
+
+    module = Dummy()
+    module.LPIPS = None
+    sys.modules['taming.modules.losses.lpips'] = module
+
+    module = Dummy()
+    module.StableDataModuleFromConfig = None
+    sys.modules['sgm.data'] = module
+
+
 # data_path = cmd_opts_pre.data
 sys.path.insert(0, script_path)
 
@@ -18,8 +33,11 @@ for possible_sd_path in possible_sd_paths:
 
 assert sd_path is not None, f"Couldn't find Stable Diffusion in any of: {possible_sd_paths}"
 
+mute_sdxl_imports()
+
 path_dirs = [
     (sd_path, 'ldm', 'Stable Diffusion', []),
+    (os.path.join(sd_path, '../generative-models'), 'sgm', 'Stable Diffusion XL', ["sgm"]),
     (os.path.join(sd_path, '../CodeFormer'), 'inference_codeformer.py', 'CodeFormer', []),
     (os.path.join(sd_path, '../BLIP'), 'models/blip.py', 'BLIP', []),
     (os.path.join(sd_path, '../k-diffusion'), 'k_diffusion/sampling.py', 'k_diffusion', ["atstart"]),
@@ -35,20 +53,13 @@ for d, must_exist, what, options in path_dirs:
         d = os.path.abspath(d)
         if "atstart" in options:
             sys.path.insert(0, d)
+        elif "sgm" in options:
+            # Stable Diffusion XL repo has scripts dir with __init__.py in it which ruins every extension's scripts dir, so we
+            # import sgm and remove it from sys.path so that when a script imports scripts.something, it doesbn't use sgm's scripts dir.
+
+            sys.path.insert(0, d)
+            import sgm  # noqa: F401
+            sys.path.pop(0)
         else:
             sys.path.append(d)
         paths[what] = d
-
-
-class Prioritize:
-    def __init__(self, name):
-        self.name = name
-        self.path = None
-
-    def __enter__(self):
-        self.path = sys.path.copy()
-        sys.path = [paths[self.name]] + sys.path
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        sys.path = self.path
-        self.path = None

modules/postprocessing.py

@@ -9,8 +9,7 @@ from modules.shared import opts
 def run_postprocessing(extras_mode, image, image_folder, input_dir, output_dir, show_extras_results, *args, save_output: bool = True):
     devices.torch_gc()
 
-    shared.state.begin()
+    shared.state.begin(job="extras")
-    shared.state.job = 'extras'
 
     image_data = []
     image_names = []
@@ -54,7 +53,9 @@ def run_postprocessing(extras_mode, image, image_folder, input_dir, output_dir,
     for image, name in zip(image_data, image_names):
         shared.state.textinfo = name
 
-        existing_pnginfo = image.info or {}
+        parameters, existing_pnginfo = images.read_info_from_image(image)
+        if parameters:
+            existing_pnginfo["parameters"] = parameters
 
         pp = scripts_postprocessing.PostprocessedImage(image.convert("RGB"))
 

modules/processing.py

@@ -14,8 +14,9 @@ from skimage import exposure
 from typing import Any, Dict, List
 
 import modules.sd_hijack
-from modules import devices, prompt_parser, masking, sd_samplers, lowvram, generation_parameters_copypaste, extra_networks, sd_vae_approx, scripts, sd_samplers_common, sd_unet
+from modules import devices, prompt_parser, masking, sd_samplers, lowvram, generation_parameters_copypaste, extra_networks, sd_vae_approx, scripts, sd_samplers_common, sd_unet, errors
 from modules.sd_hijack import model_hijack
+from modules.sd_samplers_common import images_tensor_to_samples, decode_first_stage, approximation_indexes
 from modules.shared import opts, cmd_opts, state
 import modules.shared as shared
 import modules.paths as paths
@@ -83,7 +84,7 @@ def txt2img_image_conditioning(sd_model, x, width, height):
 
         # The "masked-image" in this case will just be all zeros since the entire image is masked.
         image_conditioning = torch.zeros(x.shape[0], 3, height, width, device=x.device)
-        image_conditioning = sd_model.get_first_stage_encoding(sd_model.encode_first_stage(image_conditioning))
+        image_conditioning = images_tensor_to_samples(image_conditioning, approximation_indexes.get(opts.sd_vae_encode_method))
 
         # Add the fake full 1s mask to the first dimension.
         image_conditioning = torch.nn.functional.pad(image_conditioning, (0, 0, 0, 0, 1, 0), value=1.0)
@@ -109,7 +110,7 @@ class StableDiffusionProcessing:
     cached_uc = [None, None]
     cached_c = [None, None]
 
-    def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prompt: str = "", styles: List[str] = None, seed: int = -1, subseed: int = -1, subseed_strength: float = 0, seed_resize_from_h: int = -1, seed_resize_from_w: int = -1, seed_enable_extras: bool = True, sampler_name: str = None, batch_size: int = 1, n_iter: int = 1, steps: int = 50, cfg_scale: float = 7.0, width: int = 512, height: int = 512, restore_faces: bool = False, tiling: bool = False, do_not_save_samples: bool = False, do_not_save_grid: bool = False, extra_generation_params: Dict[Any, Any] = None, overlay_images: Any = None, negative_prompt: str = None, eta: float = None, do_not_reload_embeddings: bool = False, denoising_strength: float = 0, ddim_discretize: str = None, s_min_uncond: float = 0.0, s_churn: float = 0.0, s_tmax: float = None, s_tmin: float = 0.0, s_noise: float = 1.0, override_settings: Dict[str, Any] = None, override_settings_restore_afterwards: bool = True, sampler_index: int = None, script_args: list = None):
+    def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prompt: str = "", styles: List[str] = None, seed: int = -1, subseed: int = -1, subseed_strength: float = 0, seed_resize_from_h: int = -1, seed_resize_from_w: int = -1, seed_enable_extras: bool = True, sampler_name: str = None, batch_size: int = 1, n_iter: int = 1, steps: int = 50, cfg_scale: float = 7.0, width: int = 512, height: int = 512, restore_faces: bool = False, tiling: bool = False, do_not_save_samples: bool = False, do_not_save_grid: bool = False, extra_generation_params: Dict[Any, Any] = None, overlay_images: Any = None, negative_prompt: str = None, eta: float = None, do_not_reload_embeddings: bool = False, denoising_strength: float = 0, ddim_discretize: str = None, s_min_uncond: float = 0.0, s_churn: float = 0.0, s_tmax: float = None, s_tmin: float = 0.0, s_noise: float = None, override_settings: Dict[str, Any] = None, override_settings_restore_afterwards: bool = True, sampler_index: int = None, script_args: list = None):
         if sampler_index is not None:
             print("sampler_index argument for StableDiffusionProcessing does not do anything; use sampler_name", file=sys.stderr)
 
@@ -147,8 +148,8 @@ class StableDiffusionProcessing:
         self.s_min_uncond = s_min_uncond or opts.s_min_uncond
         self.s_churn = s_churn or opts.s_churn
         self.s_tmin = s_tmin or opts.s_tmin
-        self.s_tmax = s_tmax or float('inf')  # not representable as a standard ui option
+        self.s_tmax = (s_tmax if s_tmax is not None else opts.s_tmax) or float('inf')
-        self.s_noise = s_noise or opts.s_noise
+        self.s_noise = s_noise if s_noise is not None else opts.s_noise
         self.override_settings = {k: v for k, v in (override_settings or {}).items() if k not in shared.restricted_opts}
         self.override_settings_restore_afterwards = override_settings_restore_afterwards
         self.is_using_inpainting_conditioning = False
@@ -177,6 +178,8 @@ class StableDiffusionProcessing:
         self.extra_network_data = None
         self.seeds = None
         self.subseeds = None
+        self.recorded_checkpoint = None
+        self.recorded_checkpoint_hash = None
 
         self.step_multiplier = 1
         self.cached_uc = StableDiffusionProcessing.cached_uc
@@ -184,6 +187,9 @@ class StableDiffusionProcessing:
         self.uc = None
         self.c = None
 
+        self.user = None
+        self.image_conditioning = None
+
     @property
     def sd_model(self):
         return shared.sd_model
@@ -200,7 +206,7 @@ class StableDiffusionProcessing:
         midas_in = torch.from_numpy(transformed["midas_in"][None, ...]).to(device=shared.device)
         midas_in = repeat(midas_in, "1 ... -> n ...", n=self.batch_size)
 
-        conditioning_image = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(source_image))
+        conditioning_image = images_tensor_to_samples(source_image*0.5+0.5, approximation_indexes.get(opts.sd_vae_encode_method))
         conditioning = torch.nn.functional.interpolate(
             self.sd_model.depth_model(midas_in),
             size=conditioning_image.shape[2:],
@@ -213,7 +219,7 @@ class StableDiffusionProcessing:
         return conditioning
 
     def edit_image_conditioning(self, source_image):
-        conditioning_image = self.sd_model.encode_first_stage(source_image).mode()
+        conditioning_image = images_tensor_to_samples(source_image*0.5+0.5, approximation_indexes.get(opts.sd_vae_encode_method))
 
         return conditioning_image
 
@@ -273,10 +279,10 @@ class StableDiffusionProcessing:
         if self.sd_model.cond_stage_key == "edit":
             return self.edit_image_conditioning(source_image)
 
-        if self.sampler.conditioning_key in {'hybrid', 'concat'}:
+        if self.sd_model.model.conditioning_key in {'hybrid', 'concat'}:
             return self.inpainting_image_conditioning(source_image, latent_image, image_mask=image_mask)
 
-        if self.sampler.conditioning_key == "crossattn-adm":
+        if self.sd_model.model.conditioning_key == "crossattn-adm":
             return self.unclip_image_conditioning(source_image)
 
         # Dummy zero conditioning if we're not using inpainting or depth model.
@@ -292,7 +298,7 @@ class StableDiffusionProcessing:
         self.sampler = None
         self.c = None
         self.uc = None
-        if not opts.experimental_persistent_cond_cache:
+        if not opts.persistent_cond_cache:
             StableDiffusionProcessing.cached_c = [None, None]
             StableDiffusionProcessing.cached_uc = [None, None]
 
@@ -316,6 +322,21 @@ class StableDiffusionProcessing:
         self.all_prompts = [shared.prompt_styles.apply_styles_to_prompt(x, self.styles) for x in self.all_prompts]
         self.all_negative_prompts = [shared.prompt_styles.apply_negative_styles_to_prompt(x, self.styles) for x in self.all_negative_prompts]
 
+    def cached_params(self, required_prompts, steps, extra_network_data):
+        """Returns parameters that invalidate the cond cache if changed"""
+
+        return (
+            required_prompts,
+            steps,
+            opts.CLIP_stop_at_last_layers,
+            shared.sd_model.sd_checkpoint_info,
+            extra_network_data,
+            opts.sdxl_crop_left,
+            opts.sdxl_crop_top,
+            self.width,
+            self.height,
+        )
+
     def get_conds_with_caching(self, function, required_prompts, steps, caches, extra_network_data):
         """
         Returns the result of calling function(shared.sd_model, required_prompts, steps)
@@ -328,8 +349,11 @@ class StableDiffusionProcessing:
 
         caches is a list with items described above.
         """
+
+        cached_params = self.cached_params(required_prompts, steps, extra_network_data)
+
         for cache in caches:
-            if cache[0] is not None and (required_prompts, steps, opts.CLIP_stop_at_last_layers, shared.sd_model.sd_checkpoint_info, extra_network_data) == cache[0]:
+            if cache[0] is not None and cached_params == cache[0]:
                 return cache[1]
 
         cache = caches[0]
@@ -337,18 +361,73 @@ class StableDiffusionProcessing:
         with devices.autocast():
             cache[1] = function(shared.sd_model, required_prompts, steps)
 
-        cache[0] = (required_prompts, steps, opts.CLIP_stop_at_last_layers, shared.sd_model.sd_checkpoint_info, extra_network_data)
+        cache[0] = cached_params
         return cache[1]
 
     def setup_conds(self):
+        prompts = prompt_parser.SdConditioning(self.prompts, width=self.width, height=self.height)
+        negative_prompts = prompt_parser.SdConditioning(self.negative_prompts, width=self.width, height=self.height, is_negative_prompt=True)
+
         sampler_config = sd_samplers.find_sampler_config(self.sampler_name)
         self.step_multiplier = 2 if sampler_config and sampler_config.options.get("second_order", False) else 1
-        self.uc = self.get_conds_with_caching(prompt_parser.get_learned_conditioning, self.negative_prompts, self.steps * self.step_multiplier, [self.cached_uc], self.extra_network_data)
+        self.uc = self.get_conds_with_caching(prompt_parser.get_learned_conditioning, negative_prompts, self.steps * self.step_multiplier, [self.cached_uc], self.extra_network_data)
-        self.c = self.get_conds_with_caching(prompt_parser.get_multicond_learned_conditioning, self.prompts, self.steps * self.step_multiplier, [self.cached_c], self.extra_network_data)
+        self.c = self.get_conds_with_caching(prompt_parser.get_multicond_learned_conditioning, prompts, self.steps * self.step_multiplier, [self.cached_c], self.extra_network_data)
 
     def parse_extra_network_prompts(self):
         self.prompts, self.extra_network_data = extra_networks.parse_prompts(self.prompts)
 
+    def save_samples(self) -> bool:
+        """Returns whether generated images need to be written to disk"""
+        return opts.samples_save and not self.do_not_save_samples and (opts.save_incomplete_images or not state.interrupted and not state.skipped)
+
+    def run_refiner(self, samples):
+        shared.state.nextjob()
+
+        stopped_at = self.sampler.stop_at
+        noisy_output = self.sampler.noisy_output
+        self.sampler = None
+
+        a_is_sdxl = shared.sd_model.is_sdxl
+        decoded_noisy = decode_latent_batch(shared.sd_model, noisy_output, target_device=devices.cpu, check_for_nans=True)
+
+        refiner_checkpoint_info = sd_models.get_closet_checkpoint_match(shared.opts.sd_refiner_checkpoint)
+        if refiner_checkpoint_info is None:
+            raise Exception(f'Could not find checkpoint with name {shared.opts.sd_refiner_checkpoint}')
+
+        self.recorded_checkpoint = shared.sd_model.sd_checkpoint_info.name_for_extra
+        self.recorded_checkpoint_hash = shared.sd_model.sd_model_hash
+        self.extra_generation_params['Refiner'] = refiner_checkpoint_info.short_title
+        self.extra_generation_params['Refiner switch at'] = shared.opts.sd_refiner_switch_at
+
+        with sd_models.SkipWritingToConfig():
+            sd_models.reload_model_weights(info=refiner_checkpoint_info)
+
+        devices.torch_gc()
+        self.setup_conds()
+
+        b_is_sdxl = shared.sd_model.is_sdxl
+
+        if a_is_sdxl != b_is_sdxl:
+            decoded_noisy = torch.stack(decoded_noisy).float()
+            decoded_noisy = torch.clamp((decoded_noisy + 1.0) / 2.0, min=0.0, max=1.0)
+            noisy_latent = images_tensor_to_samples(decoded_noisy, approximation_indexes.get(opts.sd_vae_encode_method), shared.sd_model)
+        else:
+            noisy_latent = noisy_output
+
+        x = torch.zeros_like(noisy_latent)
+
+        with devices.without_autocast() if devices.unet_needs_upcast else devices.autocast():
+            denoising_strength = self.denoising_strength
+
+            self.denoising_strength = 1.0 - (stopped_at + 1) / self.steps
+            self.image_conditioning = txt2img_image_conditioning(shared.sd_model, noisy_latent, self.width, self.height)
+            self.sampler = sd_samplers.create_sampler(self.sampler_name, shared.sd_model)
+            samples = self.sampler.sample_img2img(self, noisy_latent, x, self.c, self.uc, image_conditioning=self.image_conditioning, steps=max(1, self.steps - stopped_at - 1))
+
+            self.denoising_strength = denoising_strength
+
+        return samples
+
 
 class Processed:
     def __init__(self, p: StableDiffusionProcessing, images_list, seed=-1, info="", subseed=None, all_prompts=None, all_negative_prompts=None, all_seeds=None, all_subseeds=None, index_of_first_image=0, infotexts=None, comments=""):
@@ -474,7 +553,7 @@ def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, see
         noise_shape = shape if seed_resize_from_h <= 0 or seed_resize_from_w <= 0 else (shape[0], seed_resize_from_h//8, seed_resize_from_w//8)
 
         subnoise = None
-        if subseeds is not None:
+        if subseeds is not None and subseed_strength != 0:
             subseed = 0 if i >= len(subseeds) else subseeds[i]
 
             subnoise = devices.randn(subseed, noise_shape)
@@ -506,7 +585,7 @@ def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, see
             cnt = p.sampler.number_of_needed_noises(p)
 
             if eta_noise_seed_delta > 0:
-                torch.manual_seed(seed + eta_noise_seed_delta)
+                devices.manual_seed(seed + eta_noise_seed_delta)
 
             for j in range(cnt):
                 sampler_noises[j].append(devices.randn_without_seed(tuple(noise_shape)))
@@ -520,11 +599,45 @@ def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, see
     return x
 
 
-def decode_first_stage(model, x):
+class DecodedSamples(list):
-    with devices.autocast(disable=x.dtype == devices.dtype_vae):
+    already_decoded = True
-        x = model.decode_first_stage(x)
 
-    return x
+
+def decode_latent_batch(model, batch, target_device=None, check_for_nans=False):
+    if getattr(batch, 'already_decoded', False):
+        return batch
+
+    samples = DecodedSamples()
+
+    for i in range(batch.shape[0]):
+        sample = decode_first_stage(model, batch[i:i + 1])[0]
+
+        if check_for_nans:
+            try:
+                devices.test_for_nans(sample, "vae")
+            except devices.NansException as e:
+                if devices.dtype_vae == torch.float32 or not shared.opts.auto_vae_precision:
+                    raise e
+
+                errors.print_error_explanation(
+                    "A tensor with all NaNs was produced in VAE.\n"
+                    "Web UI will now convert VAE into 32-bit float and retry.\n"
+                    "To disable this behavior, disable the 'Automaticlly revert VAE to 32-bit floats' setting.\n"
+                    "To always start with 32-bit VAE, use --no-half-vae commandline flag."
+                )
+
+                devices.dtype_vae = torch.float32
+                model.first_stage_model.to(devices.dtype_vae)
+                batch = batch.to(devices.dtype_vae)
+
+                sample = decode_first_stage(model, batch[i:i + 1])[0]
+
+        if target_device is not None:
+            sample = sample.to(target_device)
+
+        samples.append(sample)
+
+    return samples
 
 
 def get_fixed_seed(seed):
@@ -549,9 +662,13 @@ def program_version():
     return res
 
 
-def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iteration=0, position_in_batch=0):
+def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iteration=0, position_in_batch=0, use_main_prompt=False, index=None, all_negative_prompts=None):
+    if index is None:
         index = position_in_batch + iteration * p.batch_size
 
+    if all_negative_prompts is None:
+        all_negative_prompts = p.all_negative_prompts
+
     clip_skip = getattr(p, 'clip_skip', opts.CLIP_stop_at_last_layers)
     enable_hr = getattr(p, 'enable_hr', False)
     token_merging_ratio = p.get_token_merging_ratio()
@@ -566,14 +683,14 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iter
         "Sampler": p.sampler_name,
         "CFG scale": p.cfg_scale,
         "Image CFG scale": getattr(p, 'image_cfg_scale', None),
-        "Seed": all_seeds[index],
+        "Seed": p.all_seeds[0] if use_main_prompt else all_seeds[index],
         "Face restoration": (opts.face_restoration_model if p.restore_faces else None),
         "Size": f"{p.width}x{p.height}",
-        "Model hash": getattr(p, 'sd_model_hash', None if not opts.add_model_hash_to_info or not shared.sd_model.sd_model_hash else shared.sd_model.sd_model_hash),
+        "Model hash": getattr(p, 'sd_model_hash', None if not opts.add_model_hash_to_info or not shared.sd_model.sd_model_hash else p.recorded_checkpoint_hash or shared.sd_model.sd_model_hash),
-        "Model": (None if not 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": (None if not opts.add_model_name_to_info else p.recorded_checkpoint or shared.sd_model.sd_checkpoint_info.name_for_extra),
-        "Variation seed": (None if p.subseed_strength == 0 else all_subseeds[index]),
+        "Variation seed": (None if p.subseed_strength == 0 else (p.all_subseeds[0] if use_main_prompt else all_subseeds[index])),
         "Variation seed strength": (None if p.subseed_strength == 0 else p.subseed_strength),
-        "Seed resize from": (None if p.seed_resize_from_w == 0 or p.seed_resize_from_h == 0 else f"{p.seed_resize_from_w}x{p.seed_resize_from_h}"),
+        "Seed resize from": (None if p.seed_resize_from_w <= 0 or p.seed_resize_from_h <= 0 else f"{p.seed_resize_from_w}x{p.seed_resize_from_h}"),
         "Denoising strength": getattr(p, 'denoising_strength', None),
         "Conditional mask weight": getattr(p, "inpainting_mask_weight", shared.opts.inpainting_mask_weight) if p.is_using_inpainting_conditioning else None,
         "Clip skip": None if clip_skip <= 1 else clip_skip,
@@ -581,17 +698,19 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iter
         "Token merging ratio": None if token_merging_ratio == 0 else token_merging_ratio,
         "Token merging ratio hr": None if not enable_hr or token_merging_ratio_hr == 0 else token_merging_ratio_hr,
         "Init image hash": getattr(p, 'init_img_hash', None),
-        "RNG": opts.randn_source if opts.randn_source != "GPU" else None,
+        "RNG": opts.randn_source if opts.randn_source != "GPU" and opts.randn_source != "NV" else None,
         "NGMS": None if p.s_min_uncond == 0 else p.s_min_uncond,
         **p.extra_generation_params,
         "Version": program_version() if opts.add_version_to_infotext else None,
+        "User": p.user if opts.add_user_name_to_info else None,
     }
 
     generation_params_text = ", ".join([k if k == v else f'{k}: {generation_parameters_copypaste.quote(v)}' for k, v in generation_params.items() if v is not None])
 
-    negative_prompt_text = f"\nNegative prompt: {p.all_negative_prompts[index]}" if p.all_negative_prompts[index] else ""
+    prompt_text = p.prompt if use_main_prompt else all_prompts[index]
+    negative_prompt_text = f"\nNegative prompt: {all_negative_prompts[index]}" if all_negative_prompts[index] else ""
 
-    return f"{all_prompts[index]}{negative_prompt_text}\n{generation_params_text}".strip()
+    return f"{prompt_text}{negative_prompt_text}\n{generation_params_text}".strip()
 
 
 def process_images(p: StableDiffusionProcessing) -> Processed:
@@ -601,8 +720,12 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
     stored_opts = {k: opts.data[k] for k in p.override_settings.keys()}
 
     try:
+        # after running refiner, the refiner model is not unloaded - webui swaps back to main model here
+        if shared.sd_model.sd_checkpoint_info.title != opts.sd_model_checkpoint:
+            sd_models.reload_model_weights()
+
         # if no checkpoint override or the override checkpoint can't be found, remove override entry and load opts checkpoint
-        if sd_models.checkpoint_alisases.get(p.override_settings.get('sd_model_checkpoint')) is None:
+        if sd_models.checkpoint_aliases.get(p.override_settings.get('sd_model_checkpoint')) is None:
             p.override_settings.pop('sd_model_checkpoint', None)
             sd_models.reload_model_weights()
 
@@ -663,9 +786,6 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
     else:
         p.all_subseeds = [int(subseed) + x for x in range(len(p.all_prompts))]
 
-    def infotext(iteration=0, position_in_batch=0):
-        return create_infotext(p, p.all_prompts, p.all_seeds, p.all_subseeds, comments, iteration, position_in_batch)
-
     if os.path.exists(cmd_opts.embeddings_dir) and not p.do_not_reload_embeddings:
         model_hijack.embedding_db.load_textual_inversion_embeddings()
 
@@ -675,6 +795,8 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
     infotexts = []
     output_images = []
 
+    have_refiner = shared.opts.sd_refiner_switch_at < 1.0 and shared.sd_model.sd_checkpoint_info.title != shared.opts.sd_refiner_checkpoint
+
     with torch.no_grad(), p.sd_model.ema_scope():
         with devices.autocast():
             p.init(p.all_prompts, p.all_seeds, p.all_subseeds)
@@ -688,6 +810,10 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
         if state.job_count == -1:
             state.job_count = p.n_iter
 
+        if have_refiner:
+            state.job_count *= 2
+            shared.total_tqdm.updateTotal(p.steps * state.job_count // 2)
+
         for n in range(p.n_iter):
             p.iteration = n
 
@@ -697,6 +823,8 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             if state.interrupted:
                 break
 
+            sd_models.reload_model_weights()  # model can be changed for example by refiner
+
             p.prompts = p.all_prompts[n * p.batch_size:(n + 1) * p.batch_size]
             p.negative_prompts = p.all_negative_prompts[n * p.batch_size:(n + 1) * p.batch_size]
             p.seeds = p.all_seeds[n * p.batch_size:(n + 1) * p.batch_size]
@@ -728,19 +856,29 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
 
             p.setup_conds()
 
-            if len(model_hijack.comments) > 0:
             for comment in model_hijack.comments:
                 comments[comment] = 1
 
+            p.extra_generation_params.update(model_hijack.extra_generation_params)
+
             if p.n_iter > 1:
                 shared.state.job = f"Batch {n+1} out of {p.n_iter}"
 
             with devices.without_autocast() if devices.unet_needs_upcast else devices.autocast():
+                p.sampler = sd_samplers.create_sampler(p.sampler_name, p.sd_model)
+
+                if have_refiner:
+                    p.sampler.stop_at = max(1, int(shared.opts.sd_refiner_switch_at * p.steps - 1))
+
                 samples_ddim = p.sample(conditioning=p.c, unconditional_conditioning=p.uc, seeds=p.seeds, subseeds=p.subseeds, subseed_strength=p.subseed_strength, prompts=p.prompts)
 
-            x_samples_ddim = [decode_first_stage(p.sd_model, samples_ddim[i:i+1].to(dtype=devices.dtype_vae))[0].cpu() for i in range(samples_ddim.size(0))]
+            if opts.sd_vae_decode_method != 'Full':
-            for x in x_samples_ddim:
+                p.extra_generation_params['VAE Decoder'] = opts.sd_vae_decode_method
-                devices.test_for_nans(x, "vae")
+
+            if have_refiner:
+                samples_ddim = p.run_refiner(samples_ddim)
+
+            x_samples_ddim = decode_latent_batch(p.sd_model, samples_ddim, target_device=devices.cpu, check_for_nans=True)
 
             x_samples_ddim = torch.stack(x_samples_ddim).float()
             x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
@@ -755,6 +893,18 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             if p.scripts is not None:
                 p.scripts.postprocess_batch(p, x_samples_ddim, batch_number=n)
 
+                p.prompts = p.all_prompts[n * p.batch_size:(n + 1) * p.batch_size]
+                p.negative_prompts = p.all_negative_prompts[n * p.batch_size:(n + 1) * p.batch_size]
+
+                batch_params = scripts.PostprocessBatchListArgs(list(x_samples_ddim))
+                p.scripts.postprocess_batch_list(p, batch_params, batch_number=n)
+                x_samples_ddim = batch_params.images
+
+            def infotext(index=0, use_main_prompt=False):
+                return create_infotext(p, p.prompts, p.seeds, p.subseeds, use_main_prompt=use_main_prompt, index=index, all_negative_prompts=p.negative_prompts)
+
+            save_samples = p.save_samples()
+
             for i, x_sample in enumerate(x_samples_ddim):
                 p.batch_index = i
 
@@ -762,8 +912,8 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                 x_sample = x_sample.astype(np.uint8)
 
                 if p.restore_faces:
-                    if opts.save and not p.do_not_save_samples and opts.save_images_before_face_restoration:
+                    if save_samples and opts.save_images_before_face_restoration:
-                        images.save_image(Image.fromarray(x_sample), p.outpath_samples, "", p.seeds[i], p.prompts[i], opts.samples_format, info=infotext(n, i), p=p, suffix="-before-face-restoration")
+                        images.save_image(Image.fromarray(x_sample), p.outpath_samples, "", p.seeds[i], p.prompts[i], opts.samples_format, info=infotext(i), p=p, suffix="-before-face-restoration")
 
                     devices.torch_gc()
 
@@ -776,33 +926,31 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                     pp = scripts.PostprocessImageArgs(image)
                     p.scripts.postprocess_image(p, pp)
                     image = pp.image
-
                 if p.color_corrections is not None and i < len(p.color_corrections):
-                    if opts.save and not p.do_not_save_samples and opts.save_images_before_color_correction:
+                    if save_samples and opts.save_images_before_color_correction:
                         image_without_cc = apply_overlay(image, p.paste_to, i, p.overlay_images)
-                        images.save_image(image_without_cc, p.outpath_samples, "", p.seeds[i], p.prompts[i], opts.samples_format, info=infotext(n, i), p=p, suffix="-before-color-correction")
+                        images.save_image(image_without_cc, p.outpath_samples, "", p.seeds[i], p.prompts[i], opts.samples_format, info=infotext(i), p=p, suffix="-before-color-correction")
                     image = apply_color_correction(p.color_corrections[i], image)
 
                 image = apply_overlay(image, p.paste_to, i, p.overlay_images)
 
-                if opts.samples_save and not p.do_not_save_samples:
+                if save_samples:
-                    images.save_image(image, p.outpath_samples, "", p.seeds[i], p.prompts[i], opts.samples_format, info=infotext(n, i), p=p)
+                    images.save_image(image, p.outpath_samples, "", p.seeds[i], p.prompts[i], opts.samples_format, info=infotext(i), p=p)
 
-                text = infotext(n, i)
+                text = infotext(i)
                 infotexts.append(text)
                 if opts.enable_pnginfo:
                     image.info["parameters"] = text
                 output_images.append(image)
-
+                if save_samples and hasattr(p, 'mask_for_overlay') and p.mask_for_overlay and any([opts.save_mask, opts.save_mask_composite, opts.return_mask, opts.return_mask_composite]):
-                if hasattr(p, 'mask_for_overlay') and p.mask_for_overlay and any([opts.save_mask, opts.save_mask_composite, opts.return_mask, opts.return_mask_composite]):
                     image_mask = p.mask_for_overlay.convert('RGB')
                     image_mask_composite = Image.composite(image.convert('RGBA').convert('RGBa'), Image.new('RGBa', image.size), images.resize_image(2, p.mask_for_overlay, image.width, image.height).convert('L')).convert('RGBA')
 
                     if opts.save_mask:
-                        images.save_image(image_mask, p.outpath_samples, "", p.seeds[i], p.prompts[i], opts.samples_format, info=infotext(n, i), p=p, suffix="-mask")
+                        images.save_image(image_mask, p.outpath_samples, "", p.seeds[i], p.prompts[i], opts.samples_format, info=infotext(i), p=p, suffix="-mask")
 
                     if opts.save_mask_composite:
-                        images.save_image(image_mask_composite, p.outpath_samples, "", p.seeds[i], p.prompts[i], opts.samples_format, info=infotext(n, i), p=p, suffix="-mask-composite")
+                        images.save_image(image_mask_composite, p.outpath_samples, "", p.seeds[i], p.prompts[i], opts.samples_format, info=infotext(i), p=p, suffix="-mask-composite")
 
                     if opts.return_mask:
                         output_images.append(image_mask)
@@ -824,15 +972,14 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             grid = images.image_grid(output_images, p.batch_size)
 
             if opts.return_grid:
-                text = infotext()
+                text = infotext(use_main_prompt=True)
                 infotexts.insert(0, text)
                 if opts.enable_pnginfo:
                     grid.info["parameters"] = text
                 output_images.insert(0, grid)
                 index_of_first_image = 1
-
             if opts.grid_save:
-                images.save_image(grid, p.outpath_grids, "grid", p.all_seeds[0], p.all_prompts[0], opts.grid_format, info=infotext(), short_filename=not opts.grid_extended_filename, p=p, grid=True)
+                images.save_image(grid, p.outpath_grids, "grid", p.all_seeds[0], p.all_prompts[0], opts.grid_format, info=infotext(use_main_prompt=True), short_filename=not opts.grid_extended_filename, p=p, grid=True)
 
     if not p.disable_extra_networks and p.extra_network_data:
         extra_networks.deactivate(p, p.extra_network_data)
@@ -843,7 +990,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
         p,
         images_list=output_images,
         seed=p.all_seeds[0],
-        info=infotext(),
+        info=infotexts[0],
         comments="".join(f"{comment}\n" for comment in comments),
         subseed=p.all_subseeds[0],
         index_of_first_image=index_of_first_image,
@@ -873,7 +1020,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
     cached_hr_uc = [None, None]
     cached_hr_c = [None, None]
 
-    def __init__(self, enable_hr: bool = False, denoising_strength: float = 0.75, firstphase_width: int = 0, firstphase_height: int = 0, hr_scale: float = 2.0, hr_upscaler: str = None, hr_second_pass_steps: int = 0, hr_resize_x: int = 0, hr_resize_y: int = 0, hr_sampler_name: str = None, hr_prompt: str = '', hr_negative_prompt: str = '', **kwargs):
+    def __init__(self, enable_hr: bool = False, denoising_strength: float = 0.75, firstphase_width: int = 0, firstphase_height: int = 0, hr_scale: float = 2.0, hr_upscaler: str = None, hr_second_pass_steps: int = 0, hr_resize_x: int = 0, hr_resize_y: int = 0, hr_checkpoint_name: str = None, hr_sampler_name: str = None, hr_prompt: str = '', hr_negative_prompt: str = '', **kwargs):
         super().__init__(**kwargs)
         self.enable_hr = enable_hr
         self.denoising_strength = denoising_strength
@@ -884,11 +1031,14 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
         self.hr_resize_y = hr_resize_y
         self.hr_upscale_to_x = hr_resize_x
         self.hr_upscale_to_y = hr_resize_y
+        self.hr_checkpoint_name = hr_checkpoint_name
+        self.hr_checkpoint_info = None
         self.hr_sampler_name = hr_sampler_name
         self.hr_prompt = hr_prompt
         self.hr_negative_prompt = hr_negative_prompt
         self.all_hr_prompts = None
         self.all_hr_negative_prompts = None
+        self.latent_scale_mode = None
 
         if firstphase_width != 0 or firstphase_height != 0:
             self.hr_upscale_to_x = self.width
@@ -911,6 +1061,14 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 
     def init(self, all_prompts, all_seeds, all_subseeds):
         if self.enable_hr:
+            if self.hr_checkpoint_name:
+                self.hr_checkpoint_info = sd_models.get_closet_checkpoint_match(self.hr_checkpoint_name)
+
+                if self.hr_checkpoint_info is None:
+                    raise Exception(f'Could not find checkpoint with name {self.hr_checkpoint_name}')
+
+                self.extra_generation_params["Hires checkpoint"] = self.hr_checkpoint_info.short_title
+
             if self.hr_sampler_name is not None and self.hr_sampler_name != self.sampler_name:
                 self.extra_generation_params["Hires sampler"] = self.hr_sampler_name
 
@@ -920,6 +1078,11 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
             if tuple(self.hr_negative_prompt) != tuple(self.negative_prompt):
                 self.extra_generation_params["Hires negative prompt"] = self.hr_negative_prompt
 
+            self.latent_scale_mode = shared.latent_upscale_modes.get(self.hr_upscaler, None) if self.hr_upscaler is not None else shared.latent_upscale_modes.get(shared.latent_upscale_default_mode, "nearest")
+            if self.enable_hr and self.latent_scale_mode is None:
+                if not any(x.name == self.hr_upscaler for x in shared.sd_upscalers):
+                    raise Exception(f"could not find upscaler named {self.hr_upscaler}")
+
             if opts.use_old_hires_fix_width_height and self.applied_old_hires_behavior_to != (self.width, self.height):
                 self.hr_resize_x = self.width
                 self.hr_resize_y = self.height
@@ -958,14 +1121,6 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
                     self.truncate_x = (self.hr_upscale_to_x - target_w) // opt_f
                     self.truncate_y = (self.hr_upscale_to_y - target_h) // opt_f
 
-            # special case: the user has chosen to do nothing
-            if self.hr_upscale_to_x == self.width and self.hr_upscale_to_y == self.height:
-                self.enable_hr = False
-                self.denoising_strength = None
-                self.extra_generation_params.pop("Hires upscale", None)
-                self.extra_generation_params.pop("Hires resize", None)
-                return
-
             if not state.processing_has_refined_job_count:
                 if state.job_count == -1:
                     state.job_count = self.n_iter
@@ -981,19 +1136,32 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
                 self.extra_generation_params["Hires upscaler"] = self.hr_upscaler
 
     def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts):
-        self.sampler = sd_samplers.create_sampler(self.sampler_name, self.sd_model)
-
-        latent_scale_mode = shared.latent_upscale_modes.get(self.hr_upscaler, None) if self.hr_upscaler is not None else shared.latent_upscale_modes.get(shared.latent_upscale_default_mode, "nearest")
-        if self.enable_hr and latent_scale_mode is None:
-            if not any(x.name == self.hr_upscaler for x in shared.sd_upscalers):
-                raise Exception(f"could not find upscaler named {self.hr_upscaler}")
-
         x = create_random_tensors([opt_C, self.height // opt_f, self.width // opt_f], seeds=seeds, subseeds=subseeds, subseed_strength=self.subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w, p=self)
         samples = self.sampler.sample(self, x, conditioning, unconditional_conditioning, image_conditioning=self.txt2img_image_conditioning(x))
+        del x
 
         if not self.enable_hr:
             return samples
 
+        if self.latent_scale_mode is None:
+            decoded_samples = torch.stack(decode_latent_batch(self.sd_model, samples, target_device=devices.cpu, check_for_nans=True)).to(dtype=torch.float32)
+        else:
+            decoded_samples = None
+
+        current = shared.sd_model.sd_checkpoint_info
+        try:
+            if self.hr_checkpoint_info is not None:
+                self.sampler = None
+                sd_models.reload_model_weights(info=self.hr_checkpoint_info)
+                devices.torch_gc()
+
+            return self.sample_hr_pass(samples, decoded_samples, seeds, subseeds, subseed_strength, prompts)
+        finally:
+            self.sampler = None
+            sd_models.reload_model_weights(info=current)
+            devices.torch_gc()
+
+    def sample_hr_pass(self, samples, decoded_samples, seeds, subseeds, subseed_strength, prompts):
         self.is_hr_pass = True
 
         target_width = self.hr_upscale_to_x
@@ -1002,20 +1170,27 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
         def save_intermediate(image, index):
             """saves image before applying hires fix, if enabled in options; takes as an argument either an image or batch with latent space images"""
 
-            if not opts.save or self.do_not_save_samples or not opts.save_images_before_highres_fix:
+            if not self.save_samples() or not opts.save_images_before_highres_fix:
                 return
 
             if not isinstance(image, Image.Image):
                 image = sd_samplers.sample_to_image(image, index, approximation=0)
 
             info = create_infotext(self, self.all_prompts, self.all_seeds, self.all_subseeds, [], iteration=self.iteration, position_in_batch=index)
-            images.save_image(image, self.outpath_samples, "", seeds[index], prompts[index], opts.samples_format, info=info, suffix="-before-highres-fix")
+            images.save_image(image, self.outpath_samples, "", seeds[index], prompts[index], opts.samples_format, info=info, p=self, suffix="-before-highres-fix")
 
-        if latent_scale_mode is not None:
+        img2img_sampler_name = self.hr_sampler_name or self.sampler_name
+
+        if self.sampler_name in ['PLMS', 'UniPC']:  # PLMS/UniPC do not support img2img so we just silently switch to DDIM
+            img2img_sampler_name = 'DDIM'
+
+        self.sampler = sd_samplers.create_sampler(img2img_sampler_name, self.sd_model)
+
+        if self.latent_scale_mode is not None:
             for i in range(samples.shape[0]):
                 save_intermediate(samples, i)
 
-            samples = torch.nn.functional.interpolate(samples, size=(target_height // opt_f, target_width // opt_f), mode=latent_scale_mode["mode"], antialias=latent_scale_mode["antialias"])
+            samples = torch.nn.functional.interpolate(samples, size=(target_height // opt_f, target_width // opt_f), mode=self.latent_scale_mode["mode"], antialias=self.latent_scale_mode["antialias"])
 
             # Avoid making the inpainting conditioning unless necessary as
             # this does need some extra compute to decode / encode the image again.
@@ -1024,7 +1199,6 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
             else:
                 image_conditioning = self.txt2img_image_conditioning(samples)
         else:
-            decoded_samples = decode_first_stage(self.sd_model, samples)
             lowres_samples = torch.clamp((decoded_samples + 1.0) / 2.0, min=0.0, max=1.0)
 
             batch_images = []
@@ -1041,28 +1215,21 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
                 batch_images.append(image)
 
             decoded_samples = torch.from_numpy(np.array(batch_images))
-            decoded_samples = decoded_samples.to(shared.device)
+            decoded_samples = decoded_samples.to(shared.device, dtype=devices.dtype_vae)
-            decoded_samples = 2. * decoded_samples - 1.
 
-            samples = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(decoded_samples))
+            if opts.sd_vae_encode_method != 'Full':
+                self.extra_generation_params['VAE Encoder'] = opts.sd_vae_encode_method
+            samples = images_tensor_to_samples(decoded_samples, approximation_indexes.get(opts.sd_vae_encode_method))
 
             image_conditioning = self.img2img_image_conditioning(decoded_samples, samples)
 
         shared.state.nextjob()
 
-        img2img_sampler_name = self.hr_sampler_name or self.sampler_name
-
-        if self.sampler_name in ['PLMS', 'UniPC']:  # PLMS/UniPC do not support img2img so we just silently switch to DDIM
-            img2img_sampler_name = 'DDIM'
-
-        self.sampler = sd_samplers.create_sampler(img2img_sampler_name, self.sd_model)
-
         samples = samples[:, :, self.truncate_y//2:samples.shape[2]-(self.truncate_y+1)//2, self.truncate_x//2:samples.shape[3]-(self.truncate_x+1)//2]
 
         noise = create_random_tensors(samples.shape[1:], seeds=seeds, subseeds=subseeds, subseed_strength=subseed_strength, p=self)
 
         # GC now before running the next img2img to prevent running out of memory
-        x = None
         devices.torch_gc()
 
         if not self.disable_extra_networks:
@@ -1074,19 +1241,24 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 
         sd_models.apply_token_merging(self.sd_model, self.get_token_merging_ratio(for_hr=True))
 
+        if self.scripts is not None:
+            self.scripts.before_hr(self)
+
         samples = self.sampler.sample_img2img(self, samples, noise, self.hr_c, self.hr_uc, steps=self.hr_second_pass_steps or self.steps, image_conditioning=image_conditioning)
 
         sd_models.apply_token_merging(self.sd_model, self.get_token_merging_ratio())
 
+        decoded_samples = decode_latent_batch(self.sd_model, samples, target_device=devices.cpu, check_for_nans=True)
+
         self.is_hr_pass = False
 
-        return samples
+        return decoded_samples
 
     def close(self):
         super().close()
         self.hr_c = None
         self.hr_uc = None
-        if not opts.experimental_persistent_cond_cache:
+        if not opts.persistent_cond_cache:
             StableDiffusionProcessingTxt2Img.cached_hr_uc = [None, None]
             StableDiffusionProcessingTxt2Img.cached_hr_c = [None, None]
 
@@ -1119,8 +1291,11 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
         if self.hr_c is not None:
             return
 
-        self.hr_uc = self.get_conds_with_caching(prompt_parser.get_learned_conditioning, self.hr_negative_prompts, self.steps * self.step_multiplier, [self.cached_hr_uc, self.cached_uc], self.hr_extra_network_data)
+        hr_prompts = prompt_parser.SdConditioning(self.hr_prompts, width=self.hr_upscale_to_x, height=self.hr_upscale_to_y)
-        self.hr_c = self.get_conds_with_caching(prompt_parser.get_multicond_learned_conditioning, self.hr_prompts, self.steps * self.step_multiplier, [self.cached_hr_c, self.cached_c], self.hr_extra_network_data)
+        hr_negative_prompts = prompt_parser.SdConditioning(self.hr_negative_prompts, width=self.hr_upscale_to_x, height=self.hr_upscale_to_y, is_negative_prompt=True)
+
+        self.hr_uc = self.get_conds_with_caching(prompt_parser.get_learned_conditioning, hr_negative_prompts, self.steps * self.step_multiplier, [self.cached_hr_uc, self.cached_uc], self.hr_extra_network_data)
+        self.hr_c = self.get_conds_with_caching(prompt_parser.get_multicond_learned_conditioning, hr_prompts, self.steps * self.step_multiplier, [self.cached_hr_c, self.cached_c], self.hr_extra_network_data)
 
     def setup_conds(self):
         super().setup_conds()
@@ -1128,7 +1303,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
         self.hr_uc = None
         self.hr_c = None
 
-        if self.enable_hr:
+        if self.enable_hr and self.hr_checkpoint_info is None:
             if shared.opts.hires_fix_use_firstpass_conds:
                 self.calculate_hr_conds()
 
@@ -1182,7 +1357,6 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
         self.image_conditioning = None
 
     def init(self, all_prompts, all_seeds, all_subseeds):
-        self.sampler = sd_samplers.create_sampler(self.sampler_name, self.sd_model)
         crop_region = None
 
         image_mask = self.image_mask
@@ -1279,10 +1453,13 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
             raise RuntimeError(f"bad number of images passed: {len(imgs)}; expecting {self.batch_size} or less")
 
         image = torch.from_numpy(batch_images)
-        image = 2. * image - 1.
+        image = image.to(shared.device, dtype=devices.dtype_vae)
-        image = image.to(shared.device)
 
-        self.init_latent = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(image))
+        if opts.sd_vae_encode_method != 'Full':
+            self.extra_generation_params['VAE Encoder'] = opts.sd_vae_encode_method
+
+        self.init_latent = images_tensor_to_samples(image, approximation_indexes.get(opts.sd_vae_encode_method), self.sd_model)
+        devices.torch_gc()
 
         if self.resize_mode == 3:
             self.init_latent = torch.nn.functional.interpolate(self.init_latent, size=(self.height // opt_f, self.width // opt_f), mode="bilinear")

modules/prompt_parser.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import re
 from collections import namedtuple
 from typing import List
@@ -17,8 +19,8 @@ prompt: (emphasized | scheduled | alternate | plain | WHITESPACE)*
 !emphasized: "(" prompt ")"
         | "(" prompt ":" prompt ")"
         | "[" prompt "]"
-scheduled: "[" [prompt ":"] prompt ":" [WHITESPACE] NUMBER "]"
+scheduled: "[" [prompt ":"] prompt ":" [WHITESPACE] NUMBER [WHITESPACE] "]"
-alternate: "[" prompt ("|" prompt)+ "]"
+alternate: "[" prompt ("|" [prompt])+ "]"
 WHITESPACE: /\s+/
 plain: /([^\\\[\]():|]|\\.)+/
 %import common.SIGNED_NUMBER -> NUMBER
@@ -51,6 +53,10 @@ def get_learned_conditioning_prompt_schedules(prompts, steps):
     [[3, '((a][:b:c '], [10, '((a][:b:c d']]
     >>> g("[a|(b:1.1)]")
     [[1, 'a'], [2, '(b:1.1)'], [3, 'a'], [4, '(b:1.1)'], [5, 'a'], [6, '(b:1.1)'], [7, 'a'], [8, '(b:1.1)'], [9, 'a'], [10, '(b:1.1)']]
+    >>> g("[fe|]male")
+    [[1, 'female'], [2, 'male'], [3, 'female'], [4, 'male'], [5, 'female'], [6, 'male'], [7, 'female'], [8, 'male'], [9, 'female'], [10, 'male']]
+    >>> g("[fe|||]male")
+    [[1, 'female'], [2, 'male'], [3, 'male'], [4, 'male'], [5, 'female'], [6, 'male'], [7, 'male'], [8, 'male'], [9, 'female'], [10, 'male']]
     """
 
     def collect_steps(steps, tree):
@@ -58,11 +64,11 @@ def get_learned_conditioning_prompt_schedules(prompts, steps):
 
         class CollectSteps(lark.Visitor):
             def scheduled(self, tree):
-                tree.children[-1] = float(tree.children[-1])
+                tree.children[-2] = float(tree.children[-2])
-                if tree.children[-1] < 1:
+                if tree.children[-2] < 1:
-                    tree.children[-1] *= steps
+                    tree.children[-2] *= steps
-                tree.children[-1] = min(steps, int(tree.children[-1]))
+                tree.children[-2] = min(steps, int(tree.children[-2]))
-                res.append(tree.children[-1])
+                res.append(tree.children[-2])
 
             def alternate(self, tree):
                 res.extend(range(1, steps+1))
@@ -73,10 +79,11 @@ def get_learned_conditioning_prompt_schedules(prompts, steps):
     def at_step(step, tree):
         class AtStep(lark.Transformer):
             def scheduled(self, args):
-                before, after, _, when = args
+                before, after, _, when, _ = args
                 yield before or () if step <= when else after
             def alternate(self, args):
-                yield next(args[(step - 1)%len(args)])
+                args = ["" if not arg else arg for arg in args]
+                yield args[(step - 1) % len(args)]
             def start(self, args):
                 def flatten(x):
                     if type(x) == str:
@@ -109,7 +116,25 @@ def get_learned_conditioning_prompt_schedules(prompts, steps):
 ScheduledPromptConditioning = namedtuple("ScheduledPromptConditioning", ["end_at_step", "cond"])
 
 
-def get_learned_conditioning(model, prompts, steps):
+class SdConditioning(list):
+    """
+    A list with prompts for stable diffusion's conditioner model.
+    Can also specify width and height of created image - SDXL needs it.
+    """
+    def __init__(self, prompts, is_negative_prompt=False, width=None, height=None, copy_from=None):
+        super().__init__()
+        self.extend(prompts)
+
+        if copy_from is None:
+            copy_from = prompts
+
+        self.is_negative_prompt = is_negative_prompt or getattr(copy_from, 'is_negative_prompt', False)
+        self.width = width or getattr(copy_from, 'width', None)
+        self.height = height or getattr(copy_from, 'height', None)
+
+
+
+def get_learned_conditioning(model, prompts: SdConditioning | list[str], steps):
     """converts a list of prompts into a list of prompt schedules - each schedule is a list of ScheduledPromptConditioning, specifying the comdition (cond),
     and the sampling step at which this condition is to be replaced by the next one.
 
@@ -139,12 +164,17 @@ def get_learned_conditioning(model, prompts, steps):
             res.append(cached)
             continue
 
-        texts = [x[1] for x in prompt_schedule]
+        texts = SdConditioning([x[1] for x in prompt_schedule], copy_from=prompts)
         conds = model.get_learned_conditioning(texts)
 
         cond_schedule = []
         for i, (end_at_step, _) in enumerate(prompt_schedule):
-            cond_schedule.append(ScheduledPromptConditioning(end_at_step, conds[i]))
+            if isinstance(conds, dict):
+                cond = {k: v[i] for k, v in conds.items()}
+            else:
+                cond = conds[i]
+
+            cond_schedule.append(ScheduledPromptConditioning(end_at_step, cond))
 
         cache[prompt] = cond_schedule
         res.append(cond_schedule)
@@ -153,13 +183,15 @@ def get_learned_conditioning(model, prompts, steps):
 
 
 re_AND = re.compile(r"\bAND\b")
-re_weight = re.compile(r"^(.*?)(?:\s*:\s*([-+]?(?:\d+\.?|\d*\.\d+)))?\s*$")
+re_weight = re.compile(r"^((?:\s|.)*?)(?:\s*:\s*([-+]?(?:\d+\.?|\d*\.\d+)))?\s*$")
 
-def get_multicond_prompt_list(prompts):
+
+def get_multicond_prompt_list(prompts: SdConditioning | list[str]):
     res_indexes = []
 
-    prompt_flat_list = []
     prompt_indexes = {}
+    prompt_flat_list = SdConditioning(prompts)
+    prompt_flat_list.clear()
 
     for prompt in prompts:
         subprompts = re_AND.split(prompt)
@@ -196,6 +228,7 @@ class MulticondLearnedConditioning:
         self.shape: tuple = shape  # the shape field is needed to send this object to DDIM/PLMS
         self.batch: List[List[ComposableScheduledPromptConditioning]] = batch
 
+
 def get_multicond_learned_conditioning(model, prompts, steps) -> MulticondLearnedConditioning:
     """same as get_learned_conditioning, but returns a list of ScheduledPromptConditioning along with the weight objects for each prompt.
     For each prompt, the list is obtained by splitting the prompt using the AND separator.
@@ -214,20 +247,57 @@ def get_multicond_learned_conditioning(model, prompts, steps) -> MulticondLearne
     return MulticondLearnedConditioning(shape=(len(prompts),), batch=res)
 
 
+class DictWithShape(dict):
+    def __init__(self, x, shape):
+        super().__init__()
+        self.update(x)
+
+    @property
+    def shape(self):
+        return self["crossattn"].shape
+
+
 def reconstruct_cond_batch(c: List[List[ScheduledPromptConditioning]], current_step):
     param = c[0][0].cond
+    is_dict = isinstance(param, dict)
+
+    if is_dict:
+        dict_cond = param
+        res = {k: torch.zeros((len(c),) + param.shape, device=param.device, dtype=param.dtype) for k, param in dict_cond.items()}
+        res = DictWithShape(res, (len(c),) + dict_cond['crossattn'].shape)
+    else:
         res = torch.zeros((len(c),) + param.shape, device=param.device, dtype=param.dtype)
+
     for i, cond_schedule in enumerate(c):
         target_index = 0
         for current, entry in enumerate(cond_schedule):
             if current_step <= entry.end_at_step:
                 target_index = current
                 break
+
+        if is_dict:
+            for k, param in cond_schedule[target_index].cond.items():
+                res[k][i] = param
+        else:
             res[i] = cond_schedule[target_index].cond
 
     return res
 
 
+def stack_conds(tensors):
+    # if prompts have wildly different lengths above the limit we'll get tensors of different shapes
+    # and won't be able to torch.stack them. So this fixes that.
+    token_count = max([x.shape[0] for x in tensors])
+    for i in range(len(tensors)):
+        if tensors[i].shape[0] != token_count:
+            last_vector = tensors[i][-1:]
+            last_vector_repeated = last_vector.repeat([token_count - tensors[i].shape[0], 1])
+            tensors[i] = torch.vstack([tensors[i], last_vector_repeated])
+
+    return torch.stack(tensors)
+
+
+
 def reconstruct_multicond_batch(c: MulticondLearnedConditioning, current_step):
     param = c.batch[0][0].schedules[0].cond
 
@@ -249,16 +319,14 @@ def reconstruct_multicond_batch(c: MulticondLearnedConditioning, current_step):
 
         conds_list.append(conds_for_batch)
 
-    # if prompts have wildly different lengths above the limit we'll get tensors fo different shapes
+    if isinstance(tensors[0], dict):
-    # and won't be able to torch.stack them. So this fixes that.
+        keys = list(tensors[0].keys())
-    token_count = max([x.shape[0] for x in tensors])
+        stacked = {k: stack_conds([x[k] for x in tensors]) for k in keys}
-    for i in range(len(tensors)):
+        stacked = DictWithShape(stacked, stacked['crossattn'].shape)
-        if tensors[i].shape[0] != token_count:
+    else:
-            last_vector = tensors[i][-1:]
+        stacked = stack_conds(tensors).to(device=param.device, dtype=param.dtype)
-            last_vector_repeated = last_vector.repeat([token_count - tensors[i].shape[0], 1])
-            tensors[i] = torch.vstack([tensors[i], last_vector_repeated])
 
-    return conds_list, torch.stack(tensors).to(device=param.device, dtype=param.dtype)
+    return conds_list, stacked
 
 
 re_attention = re.compile(r"""
@@ -270,7 +338,7 @@ re_attention = re.compile(r"""
 \\|
 \(|
 \[|
-:([+-]?[.\d]+)\)|
+:\s*([+-]?[.\d]+)\s*\)|
 \)|
 ]|
 [^\\()\[\]:]+|

modules/realesrgan_model.py

@@ -2,7 +2,6 @@ import os
 
 import numpy as np
 from PIL import Image
-from basicsr.utils.download_util import load_file_from_url
 from realesrgan import RealESRGANer
 
 from modules.upscaler import Upscaler, UpscalerData
@@ -43,9 +42,10 @@ class UpscalerRealESRGAN(Upscaler):
         if not self.enable:
             return img
 
+        try:
             info = self.load_model(path)
-        if not os.path.exists(info.local_data_path):
+        except Exception:
-            print(f"Unable to load RealESRGAN model: {info.name}")
+            errors.report(f"Unable to load RealESRGAN model {path}", exc_info=True)
             return img
 
         upsampler = RealESRGANer(
@@ -63,20 +63,17 @@ class UpscalerRealESRGAN(Upscaler):
         return image
 
     def load_model(self, path):
-        try:
+        for scaler in self.scalers:
-            info = next(iter([scaler for scaler in self.scalers if scaler.data_path == path]), None)
+            if scaler.data_path == path:
-
+                if scaler.local_data_path.startswith("http"):
-            if info is None:
+                    scaler.local_data_path = modelloader.load_file_from_url(
-                print(f"Unable to find model info: {path}")
+                        scaler.data_path,
-                return None
+                        model_dir=self.model_download_path,
-
+                    )
-            if info.local_data_path.startswith("http"):
+                if not os.path.exists(scaler.local_data_path):
-                info.local_data_path = load_file_from_url(url=info.data_path, model_dir=self.model_download_path, progress=True)
+                    raise FileNotFoundError(f"RealESRGAN data missing: {scaler.local_data_path}")
-
+                return scaler
-            return info
+        raise ValueError(f"Unable to find model info: {path}")
-        except Exception:
-            errors.report("Error making Real-ESRGAN models list", exc_info=True)
-        return None
 
     def load_models(self, _):
         return get_realesrgan_models(self)

modules/rng_philox.py

@@ -0,0 +1,102 @@
+"""RNG imitiating torch cuda randn on CPU. You are welcome.
+
+Usage:
+
+```
+g = Generator(seed=0)
+print(g.randn(shape=(3, 4)))
+```
+
+Expected output:
+```
+[[-0.92466259 -0.42534415 -2.6438457   0.14518388]
+ [-0.12086647 -0.57972564 -0.62285122 -0.32838709]
+ [-1.07454231 -0.36314407 -1.67105067  2.26550497]]
+```
+"""
+
+import numpy as np
+
+philox_m = [0xD2511F53, 0xCD9E8D57]
+philox_w = [0x9E3779B9, 0xBB67AE85]
+
+two_pow32_inv = np.array([2.3283064e-10], dtype=np.float32)
+two_pow32_inv_2pi = np.array([2.3283064e-10 * 6.2831855], dtype=np.float32)
+
+
+def uint32(x):
+    """Converts (N,) np.uint64 array into (2, N) np.unit32 array."""
+    return x.view(np.uint32).reshape(-1, 2).transpose(1, 0)
+
+
+def philox4_round(counter, key):
+    """A single round of the Philox 4x32 random number generator."""
+
+    v1 = uint32(counter[0].astype(np.uint64) * philox_m[0])
+    v2 = uint32(counter[2].astype(np.uint64) * philox_m[1])
+
+    counter[0] = v2[1] ^ counter[1] ^ key[0]
+    counter[1] = v2[0]
+    counter[2] = v1[1] ^ counter[3] ^ key[1]
+    counter[3] = v1[0]
+
+
+def philox4_32(counter, key, rounds=10):
+    """Generates 32-bit random numbers using the Philox 4x32 random number generator.
+
+    Parameters:
+        counter (numpy.ndarray): A 4xN array of 32-bit integers representing the counter values (offset into generation).
+        key (numpy.ndarray): A 2xN array of 32-bit integers representing the key values (seed).
+        rounds (int): The number of rounds to perform.
+
+    Returns:
+        numpy.ndarray: A 4xN array of 32-bit integers containing the generated random numbers.
+    """
+
+    for _ in range(rounds - 1):
+        philox4_round(counter, key)
+
+        key[0] = key[0] + philox_w[0]
+        key[1] = key[1] + philox_w[1]
+
+    philox4_round(counter, key)
+    return counter
+
+
+def box_muller(x, y):
+    """Returns just the first out of two numbers generated by Box–Muller transform algorithm."""
+    u = x * two_pow32_inv + two_pow32_inv / 2
+    v = y * two_pow32_inv_2pi + two_pow32_inv_2pi / 2
+
+    s = np.sqrt(-2.0 * np.log(u))
+
+    r1 = s * np.sin(v)
+    return r1.astype(np.float32)
+
+
+class Generator:
+    """RNG that produces same outputs as torch.randn(..., device='cuda') on CPU"""
+
+    def __init__(self, seed):
+        self.seed = seed
+        self.offset = 0
+
+    def randn(self, shape):
+        """Generate a sequence of n standard normal random variables using the Philox 4x32 random number generator and the Box-Muller transform."""
+
+        n = 1
+        for x in shape:
+            n *= x
+
+        counter = np.zeros((4, n), dtype=np.uint32)
+        counter[0] = self.offset
+        counter[2] = np.arange(n, dtype=np.uint32)  # up to 2^32 numbers can be generated - if you want more you'd need to spill into counter[3]
+        self.offset += 1
+
+        key = np.empty(n, dtype=np.uint64)
+        key.fill(self.seed)
+        key = uint32(key)
+
+        g = philox4_32(counter, key)
+
+        return box_muller(g[0], g[1]).reshape(shape)  # discard g[2] and g[3]

modules/script_loading.py

@@ -12,11 +12,12 @@ def load_module(path):
     return module
 
 
-def preload_extensions(extensions_dir, parser):
+def preload_extensions(extensions_dir, parser, extension_list=None):
     if not os.path.isdir(extensions_dir):
         return
 
-    for dirname in sorted(os.listdir(extensions_dir)):
+    extensions = extension_list if extension_list is not None else os.listdir(extensions_dir)
+    for dirname in sorted(extensions):
         preload_script = os.path.join(extensions_dir, dirname, "preload.py")
         if not os.path.isfile(preload_script):
             continue

modules/scripts.py

@@ -1,6 +1,7 @@
 import os
 import re
 import sys
+import inspect
 from collections import namedtuple
 
 import gradio as gr
@@ -15,6 +16,11 @@ class PostprocessImageArgs:
         self.image = image
 
 
+class PostprocessBatchListArgs:
+    def __init__(self, images):
+        self.images = images
+
+
 class Script:
     name = None
     """script's internal name derived from title"""
@@ -116,6 +122,21 @@ class Script:
 
         pass
 
+    def after_extra_networks_activate(self, p, *args, **kwargs):
+        """
+        Called after extra networks activation, before conds calculation
+        allow modification of the network after extra networks activation been applied
+        won't be call if p.disable_extra_networks
+
+        **kwargs will have those items:
+          - batch_number - index of current batch, from 0 to number of batches-1
+          - prompts - list of prompts for current batch; you can change contents of this list but changing the number of entries will likely break things
+          - seeds - list of seeds for current batch
+          - subseeds - list of subseeds for current batch
+          - extra_network_data - list of ExtraNetworkParams for current stage
+        """
+        pass
+
     def process_batch(self, p, *args, **kwargs):
         """
         Same as process(), but called for every batch.
@@ -140,6 +161,25 @@ class Script:
 
         pass
 
+    def postprocess_batch_list(self, p, pp: PostprocessBatchListArgs, *args, **kwargs):
+        """
+        Same as postprocess_batch(), but receives batch images as a list of 3D tensors instead of a 4D tensor.
+        This is useful when you want to update the entire batch instead of individual images.
+
+        You can modify the postprocessing object (pp) to update the images in the batch, remove images, add images, etc.
+        If the number of images is different from the batch size when returning,
+        then the script has the responsibility to also update the following attributes in the processing object (p):
+          - p.prompts
+          - p.negative_prompts
+          - p.seeds
+          - p.subseeds
+
+        **kwargs will have same items as process_batch, and also:
+          - batch_number - index of current batch, from 0 to number of batches-1
+        """
+
+        pass
+
     def postprocess_image(self, p, pp: PostprocessImageArgs, *args):
         """
         Called for every image after it has been generated.
@@ -186,6 +226,11 @@ class Script:
 
         return f'script_{tabname}{title}_{item_id}'
 
+    def before_hr(self, p, *args):
+        """
+        This function is called before hires fix start.
+        """
+        pass
 
 current_basedir = paths.script_path
 
@@ -249,7 +294,7 @@ def load_scripts():
 
     def register_scripts_from_module(module):
         for script_class in module.__dict__.values():
-            if type(script_class) != type:
+            if not inspect.isclass(script_class):
                 continue
 
             if issubclass(script_class, Script):
@@ -483,6 +528,14 @@ class ScriptRunner:
             except Exception:
                 errors.report(f"Error running before_process_batch: {script.filename}", exc_info=True)
 
+    def after_extra_networks_activate(self, p, **kwargs):
+        for script in self.alwayson_scripts:
+            try:
+                script_args = p.script_args[script.args_from:script.args_to]
+                script.after_extra_networks_activate(p, *script_args, **kwargs)
+            except Exception:
+                errors.report(f"Error running after_extra_networks_activate: {script.filename}", exc_info=True)
+
     def process_batch(self, p, **kwargs):
         for script in self.alwayson_scripts:
             try:
@@ -507,6 +560,14 @@ class ScriptRunner:
             except Exception:
                 errors.report(f"Error running postprocess_batch: {script.filename}", exc_info=True)
 
+    def postprocess_batch_list(self, p, pp: PostprocessBatchListArgs, **kwargs):
+        for script in self.alwayson_scripts:
+            try:
+                script_args = p.script_args[script.args_from:script.args_to]
+                script.postprocess_batch_list(p, pp, *script_args, **kwargs)
+            except Exception:
+                errors.report(f"Error running postprocess_batch_list: {script.filename}", exc_info=True)
+
     def postprocess_image(self, p, pp: PostprocessImageArgs):
         for script in self.alwayson_scripts:
             try:
@@ -548,6 +609,15 @@ class ScriptRunner:
                     self.scripts[si].args_to = args_to
 
 
+    def before_hr(self, p):
+        for script in self.alwayson_scripts:
+            try:
+                script_args = p.script_args[script.args_from:script.args_to]
+                script.before_hr(p, *script_args)
+            except Exception:
+                errors.report(f"Error running before_hr: {script.filename}", exc_info=True)
+
+
 scripts_txt2img: ScriptRunner = None
 scripts_img2img: ScriptRunner = None
 scripts_postproc: scripts_postprocessing.ScriptPostprocessingRunner = None
@@ -561,49 +631,3 @@ def reload_script_body_only():
 
 
 reload_scripts = load_scripts  # compatibility alias
-
-
-def add_classes_to_gradio_component(comp):
-    """
-    this adds gradio-* to the component for css styling (ie gradio-button to gr.Button), as well as some others
-    """
-
-    comp.elem_classes = [f"gradio-{comp.get_block_name()}", *(comp.elem_classes or [])]
-
-    if getattr(comp, 'multiselect', False):
-        comp.elem_classes.append('multiselect')
-
-
-
-def IOComponent_init(self, *args, **kwargs):
-    if scripts_current is not None:
-        scripts_current.before_component(self, **kwargs)
-
-    script_callbacks.before_component_callback(self, **kwargs)
-
-    res = original_IOComponent_init(self, *args, **kwargs)
-
-    add_classes_to_gradio_component(self)
-
-    script_callbacks.after_component_callback(self, **kwargs)
-
-    if scripts_current is not None:
-        scripts_current.after_component(self, **kwargs)
-
-    return res
-
-
-original_IOComponent_init = gr.components.IOComponent.__init__
-gr.components.IOComponent.__init__ = IOComponent_init
-
-
-def BlockContext_init(self, *args, **kwargs):
-    res = original_BlockContext_init(self, *args, **kwargs)
-
-    add_classes_to_gradio_component(self)
-
-    return res
-
-
-original_BlockContext_init = gr.blocks.BlockContext.__init__
-gr.blocks.BlockContext.__init__ = BlockContext_init

modules/sd_disable_initialization.py

@@ -3,8 +3,31 @@ import open_clip
 import torch
 import transformers.utils.hub
 
+from modules import shared
 
-class DisableInitialization:
+
+class ReplaceHelper:
+    def __init__(self):
+        self.replaced = []
+
+    def replace(self, obj, field, func):
+        original = getattr(obj, field, None)
+        if original is None:
+            return None
+
+        self.replaced.append((obj, field, original))
+        setattr(obj, field, func)
+
+        return original
+
+    def restore(self):
+        for obj, field, original in self.replaced:
+            setattr(obj, field, original)
+
+        self.replaced.clear()
+
+
+class DisableInitialization(ReplaceHelper):
     """
     When an object of this class enters a `with` block, it starts:
     - preventing torch's layer initialization functions from working
@@ -21,7 +44,7 @@ class DisableInitialization:
     """
 
     def __init__(self, disable_clip=True):
-        self.replaced = []
+        super().__init__()
         self.disable_clip = disable_clip
 
     def replace(self, obj, field, func):
@@ -86,8 +109,81 @@ class DisableInitialization:
             self.transformers_utils_hub_get_from_cache = self.replace(transformers.utils.hub, 'get_from_cache', transformers_utils_hub_get_from_cache)
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        for obj, field, original in self.replaced:
+        self.restore()
-            setattr(obj, field, original)
 
-        self.replaced.clear()
 
+class InitializeOnMeta(ReplaceHelper):
+    """
+    Context manager that causes all parameters for linear/conv2d/mha layers to be allocated on meta device,
+    which results in those parameters having no values and taking no memory. model.to() will be broken and
+    will need to be repaired by using LoadStateDictOnMeta below when loading params from state dict.
+
+    Usage:
+    ```
+    with sd_disable_initialization.InitializeOnMeta():
+        sd_model = instantiate_from_config(sd_config.model)
+    ```
+    """
+
+    def __enter__(self):
+        if shared.cmd_opts.disable_model_loading_ram_optimization:
+            return
+
+        def set_device(x):
+            x["device"] = "meta"
+            return x
+
+        linear_init = self.replace(torch.nn.Linear, '__init__', lambda *args, **kwargs: linear_init(*args, **set_device(kwargs)))
+        conv2d_init = self.replace(torch.nn.Conv2d, '__init__', lambda *args, **kwargs: conv2d_init(*args, **set_device(kwargs)))
+        mha_init = self.replace(torch.nn.MultiheadAttention, '__init__', lambda *args, **kwargs: mha_init(*args, **set_device(kwargs)))
+        self.replace(torch.nn.Module, 'to', lambda *args, **kwargs: None)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.restore()
+
+
+class LoadStateDictOnMeta(ReplaceHelper):
+    """
+    Context manager that allows to read parameters from state_dict into a model that has some of its parameters in the meta device.
+    As those parameters are read from state_dict, they will be deleted from it, so by the end state_dict will be mostly empty, to save memory.
+    Meant to be used together with InitializeOnMeta above.
+
+    Usage:
+    ```
+    with sd_disable_initialization.LoadStateDictOnMeta(state_dict):
+        model.load_state_dict(state_dict, strict=False)
+    ```
+    """
+
+    def __init__(self, state_dict, device):
+        super().__init__()
+        self.state_dict = state_dict
+        self.device = device
+
+    def __enter__(self):
+        if shared.cmd_opts.disable_model_loading_ram_optimization:
+            return
+
+        sd = self.state_dict
+        device = self.device
+
+        def load_from_state_dict(original, self, state_dict, prefix, *args, **kwargs):
+            params = [(name, param) for name, param in self._parameters.items() if param is not None and param.is_meta]
+
+            for name, param in params:
+                if param.is_meta:
+                    self._parameters[name] = torch.nn.parameter.Parameter(torch.zeros_like(param, device=device), requires_grad=param.requires_grad)
+
+            original(self, state_dict, prefix, *args, **kwargs)
+
+            for name, _ in params:
+                key = prefix + name
+                if key in sd:
+                    del sd[key]
+
+        linear_load_from_state_dict = self.replace(torch.nn.Linear, '_load_from_state_dict', lambda *args, **kwargs: load_from_state_dict(linear_load_from_state_dict, *args, **kwargs))
+        conv2d_load_from_state_dict = self.replace(torch.nn.Conv2d, '_load_from_state_dict', lambda *args, **kwargs: load_from_state_dict(conv2d_load_from_state_dict, *args, **kwargs))
+        mha_load_from_state_dict = self.replace(torch.nn.MultiheadAttention, '_load_from_state_dict', lambda *args, **kwargs: load_from_state_dict(mha_load_from_state_dict, *args, **kwargs))
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.restore()

modules/sd_hijack.py

@@ -2,11 +2,10 @@ import torch
 from torch.nn.functional import silu
 from types import MethodType
 
-import modules.textual_inversion.textual_inversion
 from modules import devices, sd_hijack_optimizations, shared, script_callbacks, errors, sd_unet
 from modules.hypernetworks import hypernetwork
 from modules.shared import cmd_opts
-from modules import sd_hijack_clip, sd_hijack_open_clip, sd_hijack_unet, sd_hijack_xlmr, xlmr
+from modules import sd_hijack_clip, sd_hijack_open_clip, sd_hijack_unet, sd_hijack_xlmr, xlmr, sd_hijack_inpainting
 
 import ldm.modules.attention
 import ldm.modules.diffusionmodules.model
@@ -15,6 +14,11 @@ import ldm.models.diffusion.ddim
 import ldm.models.diffusion.plms
 import ldm.modules.encoders.modules
 
+import sgm.modules.attention
+import sgm.modules.diffusionmodules.model
+import sgm.modules.diffusionmodules.openaimodel
+import sgm.modules.encoders.modules
+
 attention_CrossAttention_forward = ldm.modules.attention.CrossAttention.forward
 diffusionmodules_model_nonlinearity = ldm.modules.diffusionmodules.model.nonlinearity
 diffusionmodules_model_AttnBlock_forward = ldm.modules.diffusionmodules.model.AttnBlock.forward
@@ -25,8 +29,12 @@ ldm.modules.attention.MemoryEfficientCrossAttention = ldm.modules.attention.Cros
 ldm.modules.attention.BasicTransformerBlock.ATTENTION_MODES["softmax-xformers"] = ldm.modules.attention.CrossAttention
 
 # silence new console spam from SD2
-ldm.modules.attention.print = lambda *args: None
+ldm.modules.attention.print = shared.ldm_print
-ldm.modules.diffusionmodules.model.print = lambda *args: None
+ldm.modules.diffusionmodules.model.print = shared.ldm_print
+ldm.util.print = shared.ldm_print
+ldm.models.diffusion.ddpm.print = shared.ldm_print
+
+sd_hijack_inpainting.do_inpainting_hijack()
 
 optimizers = []
 current_optimizer: sd_hijack_optimizations.SdOptimization = None
@@ -56,6 +64,9 @@ def apply_optimizations(option=None):
     ldm.modules.diffusionmodules.model.nonlinearity = silu
     ldm.modules.diffusionmodules.openaimodel.th = sd_hijack_unet.th
 
+    sgm.modules.diffusionmodules.model.nonlinearity = silu
+    sgm.modules.diffusionmodules.openaimodel.th = sd_hijack_unet.th
+
     if current_optimizer is not None:
         current_optimizer.undo()
         current_optimizer = None
@@ -89,6 +100,10 @@ def undo_optimizations():
     ldm.modules.attention.CrossAttention.forward = hypernetwork.attention_CrossAttention_forward
     ldm.modules.diffusionmodules.model.AttnBlock.forward = diffusionmodules_model_AttnBlock_forward
 
+    sgm.modules.diffusionmodules.model.nonlinearity = diffusionmodules_model_nonlinearity
+    sgm.modules.attention.CrossAttention.forward = hypernetwork.attention_CrossAttention_forward
+    sgm.modules.diffusionmodules.model.AttnBlock.forward = diffusionmodules_model_AttnBlock_forward
+
 
 def fix_checkpoint():
     """checkpoints are now added and removed in embedding/hypernet code, since torch doesn't want
@@ -147,15 +162,18 @@ def undo_weighted_forward(sd_model):
 
 class StableDiffusionModelHijack:
     fixes = None
-    comments = []
     layers = None
     circular_enabled = False
     clip = None
     optimization_method = None
 
-    embedding_db = modules.textual_inversion.textual_inversion.EmbeddingDatabase()
-
     def __init__(self):
+        import modules.textual_inversion.textual_inversion
+
+        self.extra_generation_params = {}
+        self.comments = []
+
+        self.embedding_db = modules.textual_inversion.textual_inversion.EmbeddingDatabase()
         self.embedding_db.add_embedding_dir(cmd_opts.embeddings_dir)
 
     def apply_optimizations(self, option=None):
@@ -166,6 +184,32 @@ class StableDiffusionModelHijack:
             undo_optimizations()
 
     def hijack(self, m):
+        conditioner = getattr(m, 'conditioner', None)
+        if conditioner:
+            text_cond_models = []
+
+            for i in range(len(conditioner.embedders)):
+                embedder = conditioner.embedders[i]
+                typename = type(embedder).__name__
+                if typename == 'FrozenOpenCLIPEmbedder':
+                    embedder.model.token_embedding = EmbeddingsWithFixes(embedder.model.token_embedding, self)
+                    conditioner.embedders[i] = sd_hijack_open_clip.FrozenOpenCLIPEmbedderWithCustomWords(embedder, self)
+                    text_cond_models.append(conditioner.embedders[i])
+                if typename == 'FrozenCLIPEmbedder':
+                    model_embeddings = embedder.transformer.text_model.embeddings
+                    model_embeddings.token_embedding = EmbeddingsWithFixes(model_embeddings.token_embedding, self)
+                    conditioner.embedders[i] = sd_hijack_clip.FrozenCLIPEmbedderForSDXLWithCustomWords(embedder, self)
+                    text_cond_models.append(conditioner.embedders[i])
+                if typename == 'FrozenOpenCLIPEmbedder2':
+                    embedder.model.token_embedding = EmbeddingsWithFixes(embedder.model.token_embedding, self, textual_inversion_key='clip_g')
+                    conditioner.embedders[i] = sd_hijack_open_clip.FrozenOpenCLIPEmbedder2WithCustomWords(embedder, self)
+                    text_cond_models.append(conditioner.embedders[i])
+
+            if len(text_cond_models) == 1:
+                m.cond_stage_model = text_cond_models[0]
+            else:
+                m.cond_stage_model = conditioner
+
         if type(m.cond_stage_model) == xlmr.BertSeriesModelWithTransformation:
             model_embeddings = m.cond_stage_model.roberta.embeddings
             model_embeddings.token_embedding = EmbeddingsWithFixes(model_embeddings.word_embeddings, self)
@@ -203,7 +247,7 @@ class StableDiffusionModelHijack:
         ldm.modules.diffusionmodules.openaimodel.UNetModel.forward = sd_unet.UNetModel_forward
 
     def undo_hijack(self, m):
-        if type(m.cond_stage_model) == xlmr.BertSeriesModelWithTransformation:
+        if type(m.cond_stage_model) == sd_hijack_xlmr.FrozenXLMREmbedderWithCustomWords:
             m.cond_stage_model = m.cond_stage_model.wrapped
 
         elif type(m.cond_stage_model) == sd_hijack_clip.FrozenCLIPEmbedderWithCustomWords:
@@ -236,6 +280,7 @@ class StableDiffusionModelHijack:
 
     def clear_comments(self):
         self.comments = []
+        self.extra_generation_params = {}
 
     def get_prompt_lengths(self, text):
         if self.clip is None:
@@ -251,10 +296,11 @@ class StableDiffusionModelHijack:
 
 
 class EmbeddingsWithFixes(torch.nn.Module):
-    def __init__(self, wrapped, embeddings):
+    def __init__(self, wrapped, embeddings, textual_inversion_key='clip_l'):
         super().__init__()
         self.wrapped = wrapped
         self.embeddings = embeddings
+        self.textual_inversion_key = textual_inversion_key
 
     def forward(self, input_ids):
         batch_fixes = self.embeddings.fixes
@@ -268,7 +314,8 @@ class EmbeddingsWithFixes(torch.nn.Module):
         vecs = []
         for fixes, tensor in zip(batch_fixes, inputs_embeds):
             for offset, embedding in fixes:
-                emb = devices.cond_cast_unet(embedding.vec)
+                vec = embedding.vec[self.textual_inversion_key] if isinstance(embedding.vec, dict) else embedding.vec
+                emb = devices.cond_cast_unet(vec)
                 emb_len = min(tensor.shape[0] - offset - 1, emb.shape[0])
                 tensor = torch.cat([tensor[0:offset + 1], emb[0:emb_len], tensor[offset + 1 + emb_len:]])
 

modules/sd_hijack_clip.py

@@ -42,6 +42,10 @@ class FrozenCLIPEmbedderWithCustomWordsBase(torch.nn.Module):
         self.hijack: sd_hijack.StableDiffusionModelHijack = hijack
         self.chunk_length = 75
 
+        self.is_trainable = getattr(wrapped, 'is_trainable', False)
+        self.input_key = getattr(wrapped, 'input_key', 'txt')
+        self.legacy_ucg_val = None
+
     def empty_chunk(self):
         """creates an empty PromptChunk and returns it"""
 
@@ -157,7 +161,7 @@ class FrozenCLIPEmbedderWithCustomWordsBase(torch.nn.Module):
                     position += 1
                     continue
 
-                emb_len = int(embedding.vec.shape[0])
+                emb_len = int(embedding.vectors)
                 if len(chunk.tokens) + emb_len > self.chunk_length:
                     next_chunk()
 
@@ -199,8 +203,9 @@ class FrozenCLIPEmbedderWithCustomWordsBase(torch.nn.Module):
         """
         Accepts an array of texts; Passes texts through transformers network to create a tensor with numerical representation of those texts.
         Returns a tensor with shape of (B, T, C), where B is length of the array; T is length, in tokens, of texts (including padding) - T will
-        be a multiple of 77; and C is dimensionality of each token - for SD1 it's 768, and for SD2 it's 1024.
+        be a multiple of 77; and C is dimensionality of each token - for SD1 it's 768, for SD2 it's 1024, and for SDXL it's 1280.
         An example shape returned by this function can be: (2, 77, 768).
+        For SDXL, instead of returning one tensor avobe, it returns a tuple with two: the other one with shape (B, 1280) with pooled values.
         Webui usually sends just one text at a time through this function - the only time when texts is an array with more than one elemenet
         is when you do prompt editing: "a picture of a [cat:dog:0.4] eating ice cream"
         """
@@ -229,10 +234,24 @@ class FrozenCLIPEmbedderWithCustomWordsBase(torch.nn.Module):
             z = self.process_tokens(tokens, multipliers)
             zs.append(z)
 
-        if len(used_embeddings) > 0:
+        if opts.textual_inversion_add_hashes_to_infotext and used_embeddings:
-            embeddings_list = ", ".join([f'{name} [{embedding.checksum()}]' for name, embedding in used_embeddings.items()])
+            hashes = []
-            self.hijack.comments.append(f"Used embeddings: {embeddings_list}")
+            for name, embedding in used_embeddings.items():
+                shorthash = embedding.shorthash
+                if not shorthash:
+                    continue
 
+                name = name.replace(":", "").replace(",", "")
+                hashes.append(f"{name}: {shorthash}")
+
+            if hashes:
+                if self.hijack.extra_generation_params.get("TI hashes"):
+                    hashes.append(self.hijack.extra_generation_params.get("TI hashes"))
+                self.hijack.extra_generation_params["TI hashes"] = ", ".join(hashes)
+
+        if getattr(self.wrapped, 'return_pooled', False):
+            return torch.hstack(zs), zs[0].pooled
+        else:
             return torch.hstack(zs)
 
     def process_tokens(self, remade_batch_tokens, batch_multipliers):
@@ -253,6 +272,8 @@ class FrozenCLIPEmbedderWithCustomWordsBase(torch.nn.Module):
 
         z = self.encode_with_transformers(tokens)
 
+        pooled = getattr(z, 'pooled', None)
+
         # restoring original mean is likely not correct, but it seems to work well to prevent artifacts that happen otherwise
         batch_multipliers = torch.asarray(batch_multipliers).to(devices.device)
         original_mean = z.mean()
@@ -260,6 +281,9 @@ class FrozenCLIPEmbedderWithCustomWordsBase(torch.nn.Module):
         new_mean = z.mean()
         z = z * (original_mean / new_mean)
 
+        if pooled is not None:
+            z.pooled = pooled
+
         return z
 
 
@@ -315,3 +339,18 @@ class FrozenCLIPEmbedderWithCustomWords(FrozenCLIPEmbedderWithCustomWordsBase):
         embedded = embedding_layer.token_embedding.wrapped(ids.to(embedding_layer.token_embedding.wrapped.weight.device)).squeeze(0)
 
         return embedded
+
+
+class FrozenCLIPEmbedderForSDXLWithCustomWords(FrozenCLIPEmbedderWithCustomWords):
+    def __init__(self, wrapped, hijack):
+        super().__init__(wrapped, hijack)
+
+    def encode_with_transformers(self, tokens):
+        outputs = self.wrapped.transformer(input_ids=tokens, output_hidden_states=self.wrapped.layer == "hidden")
+
+        if self.wrapped.layer == "last":
+            z = outputs.last_hidden_state
+        else:
+            z = outputs.hidden_states[self.wrapped.layer_idx]
+
+        return z

modules/sd_hijack_inpainting.py

@@ -92,6 +92,4 @@ def p_sample_plms(self, x, c, t, index, repeat_noise=False, use_original_steps=F
 
 
 def do_inpainting_hijack():
-    # p_sample_plms is needed because PLMS can't work with dicts as conditionings
-
     ldm.models.diffusion.plms.PLMSSampler.p_sample_plms = p_sample_plms

modules/sd_hijack_open_clip.py

@@ -35,3 +35,37 @@ class FrozenOpenCLIPEmbedderWithCustomWords(sd_hijack_clip.FrozenCLIPEmbedderWit
         embedded = self.wrapped.model.token_embedding.wrapped(ids).squeeze(0)
 
         return embedded
+
+
+class FrozenOpenCLIPEmbedder2WithCustomWords(sd_hijack_clip.FrozenCLIPEmbedderWithCustomWordsBase):
+    def __init__(self, wrapped, hijack):
+        super().__init__(wrapped, hijack)
+
+        self.comma_token = [v for k, v in tokenizer.encoder.items() if k == ',</w>'][0]
+        self.id_start = tokenizer.encoder["<start_of_text>"]
+        self.id_end = tokenizer.encoder["<end_of_text>"]
+        self.id_pad = 0
+
+    def tokenize(self, texts):
+        assert not opts.use_old_emphasis_implementation, 'Old emphasis implementation not supported for Open Clip'
+
+        tokenized = [tokenizer.encode(text) for text in texts]
+
+        return tokenized
+
+    def encode_with_transformers(self, tokens):
+        d = self.wrapped.encode_with_transformer(tokens)
+        z = d[self.wrapped.layer]
+
+        pooled = d.get("pooled")
+        if pooled is not None:
+            z.pooled = pooled
+
+        return z
+
+    def encode_embedding_init_text(self, init_text, nvpt):
+        ids = tokenizer.encode(init_text)
+        ids = torch.asarray([ids], device=devices.device, dtype=torch.int)
+        embedded = self.wrapped.model.token_embedding.wrapped(ids.to(self.wrapped.model.token_embedding.wrapped.weight.device)).squeeze(0)
+
+        return embedded

modules/sd_hijack_optimizations.py

@@ -14,7 +14,11 @@ from modules.hypernetworks import hypernetwork
 import ldm.modules.attention
 import ldm.modules.diffusionmodules.model
 
+import sgm.modules.attention
+import sgm.modules.diffusionmodules.model
+
 diffusionmodules_model_AttnBlock_forward = ldm.modules.diffusionmodules.model.AttnBlock.forward
+sgm_diffusionmodules_model_AttnBlock_forward = sgm.modules.diffusionmodules.model.AttnBlock.forward
 
 
 class SdOptimization:
@@ -39,6 +43,9 @@ class SdOptimization:
         ldm.modules.attention.CrossAttention.forward = hypernetwork.attention_CrossAttention_forward
         ldm.modules.diffusionmodules.model.AttnBlock.forward = diffusionmodules_model_AttnBlock_forward
 
+        sgm.modules.attention.CrossAttention.forward = hypernetwork.attention_CrossAttention_forward
+        sgm.modules.diffusionmodules.model.AttnBlock.forward = sgm_diffusionmodules_model_AttnBlock_forward
+
 
 class SdOptimizationXformers(SdOptimization):
     name = "xformers"
@@ -51,6 +58,8 @@ class SdOptimizationXformers(SdOptimization):
     def apply(self):
         ldm.modules.attention.CrossAttention.forward = xformers_attention_forward
         ldm.modules.diffusionmodules.model.AttnBlock.forward = xformers_attnblock_forward
+        sgm.modules.attention.CrossAttention.forward = xformers_attention_forward
+        sgm.modules.diffusionmodules.model.AttnBlock.forward = xformers_attnblock_forward
 
 
 class SdOptimizationSdpNoMem(SdOptimization):
@@ -65,6 +74,8 @@ class SdOptimizationSdpNoMem(SdOptimization):
     def apply(self):
         ldm.modules.attention.CrossAttention.forward = scaled_dot_product_no_mem_attention_forward
         ldm.modules.diffusionmodules.model.AttnBlock.forward = sdp_no_mem_attnblock_forward
+        sgm.modules.attention.CrossAttention.forward = scaled_dot_product_no_mem_attention_forward
+        sgm.modules.diffusionmodules.model.AttnBlock.forward = sdp_no_mem_attnblock_forward
 
 
 class SdOptimizationSdp(SdOptimizationSdpNoMem):
@@ -76,6 +87,8 @@ class SdOptimizationSdp(SdOptimizationSdpNoMem):
     def apply(self):
         ldm.modules.attention.CrossAttention.forward = scaled_dot_product_attention_forward
         ldm.modules.diffusionmodules.model.AttnBlock.forward = sdp_attnblock_forward
+        sgm.modules.attention.CrossAttention.forward = scaled_dot_product_attention_forward
+        sgm.modules.diffusionmodules.model.AttnBlock.forward = sdp_attnblock_forward
 
 
 class SdOptimizationSubQuad(SdOptimization):
@@ -86,6 +99,8 @@ class SdOptimizationSubQuad(SdOptimization):
     def apply(self):
         ldm.modules.attention.CrossAttention.forward = sub_quad_attention_forward
         ldm.modules.diffusionmodules.model.AttnBlock.forward = sub_quad_attnblock_forward
+        sgm.modules.attention.CrossAttention.forward = sub_quad_attention_forward
+        sgm.modules.diffusionmodules.model.AttnBlock.forward = sub_quad_attnblock_forward
 
 
 class SdOptimizationV1(SdOptimization):
@@ -94,9 +109,9 @@ class SdOptimizationV1(SdOptimization):
     cmd_opt = "opt_split_attention_v1"
     priority = 10
 
-
     def apply(self):
         ldm.modules.attention.CrossAttention.forward = split_cross_attention_forward_v1
+        sgm.modules.attention.CrossAttention.forward = split_cross_attention_forward_v1
 
 
 class SdOptimizationInvokeAI(SdOptimization):
@@ -109,6 +124,7 @@ class SdOptimizationInvokeAI(SdOptimization):
 
     def apply(self):
         ldm.modules.attention.CrossAttention.forward = split_cross_attention_forward_invokeAI
+        sgm.modules.attention.CrossAttention.forward = split_cross_attention_forward_invokeAI
 
 
 class SdOptimizationDoggettx(SdOptimization):
@@ -119,6 +135,8 @@ class SdOptimizationDoggettx(SdOptimization):
     def apply(self):
         ldm.modules.attention.CrossAttention.forward = split_cross_attention_forward
         ldm.modules.diffusionmodules.model.AttnBlock.forward = cross_attention_attnblock_forward
+        sgm.modules.attention.CrossAttention.forward = split_cross_attention_forward
+        sgm.modules.diffusionmodules.model.AttnBlock.forward = cross_attention_attnblock_forward
 
 
 def list_optimizers(res):
@@ -155,7 +173,7 @@ def get_available_vram():
 
 
 # see https://github.com/basujindal/stable-diffusion/pull/117 for discussion
-def split_cross_attention_forward_v1(self, x, context=None, mask=None):
+def split_cross_attention_forward_v1(self, x, context=None, mask=None, **kwargs):
     h = self.heads
 
     q_in = self.to_q(x)
@@ -196,7 +214,7 @@ def split_cross_attention_forward_v1(self, x, context=None, mask=None):
 
 
 # taken from https://github.com/Doggettx/stable-diffusion and modified
-def split_cross_attention_forward(self, x, context=None, mask=None):
+def split_cross_attention_forward(self, x, context=None, mask=None, **kwargs):
     h = self.heads
 
     q_in = self.to_q(x)
@@ -238,9 +256,9 @@ def split_cross_attention_forward(self, x, context=None, mask=None):
             raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). '
                                f'Need: {mem_required / 64 / gb:0.1f}GB free, Have:{mem_free_total / gb:0.1f}GB free')
 
-        slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]
+        slice_size = q.shape[1] // steps
         for i in range(0, q.shape[1], slice_size):
-            end = i + slice_size
+            end = min(i + slice_size, q.shape[1])
             s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k)
 
             s2 = s1.softmax(dim=-1, dtype=q.dtype)
@@ -262,11 +280,13 @@ def split_cross_attention_forward(self, x, context=None, mask=None):
 # -- Taken from https://github.com/invoke-ai/InvokeAI and modified --
 mem_total_gb = psutil.virtual_memory().total // (1 << 30)
 
+
 def einsum_op_compvis(q, k, v):
     s = einsum('b i d, b j d -> b i j', q, k)
     s = s.softmax(dim=-1, dtype=s.dtype)
     return einsum('b i j, b j d -> b i d', s, v)
 
+
 def einsum_op_slice_0(q, k, v, slice_size):
     r = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)
     for i in range(0, q.shape[0], slice_size):
@@ -274,6 +294,7 @@ def einsum_op_slice_0(q, k, v, slice_size):
         r[i:end] = einsum_op_compvis(q[i:end], k[i:end], v[i:end])
     return r
 
+
 def einsum_op_slice_1(q, k, v, slice_size):
     r = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)
     for i in range(0, q.shape[1], slice_size):
@@ -281,6 +302,7 @@ def einsum_op_slice_1(q, k, v, slice_size):
         r[:, i:end] = einsum_op_compvis(q[:, i:end], k, v)
     return r
 
+
 def einsum_op_mps_v1(q, k, v):
     if q.shape[0] * q.shape[1] <= 2**16: # (512x512) max q.shape[1]: 4096
         return einsum_op_compvis(q, k, v)
@@ -290,12 +312,14 @@ def einsum_op_mps_v1(q, k, v):
             slice_size -= 1
         return einsum_op_slice_1(q, k, v, slice_size)
 
+
 def einsum_op_mps_v2(q, k, v):
     if mem_total_gb > 8 and q.shape[0] * q.shape[1] <= 2**16:
         return einsum_op_compvis(q, k, v)
     else:
         return einsum_op_slice_0(q, k, v, 1)
 
+
 def einsum_op_tensor_mem(q, k, v, max_tensor_mb):
     size_mb = q.shape[0] * q.shape[1] * k.shape[1] * q.element_size() // (1 << 20)
     if size_mb <= max_tensor_mb:
@@ -305,6 +329,7 @@ def einsum_op_tensor_mem(q, k, v, max_tensor_mb):
         return einsum_op_slice_0(q, k, v, q.shape[0] // div)
     return einsum_op_slice_1(q, k, v, max(q.shape[1] // div, 1))
 
+
 def einsum_op_cuda(q, k, v):
     stats = torch.cuda.memory_stats(q.device)
     mem_active = stats['active_bytes.all.current']
@@ -315,6 +340,7 @@ def einsum_op_cuda(q, k, v):
     # Divide factor of safety as there's copying and fragmentation
     return einsum_op_tensor_mem(q, k, v, mem_free_total / 3.3 / (1 << 20))
 
+
 def einsum_op(q, k, v):
     if q.device.type == 'cuda':
         return einsum_op_cuda(q, k, v)
@@ -328,7 +354,8 @@ def einsum_op(q, k, v):
     # Tested on i7 with 8MB L3 cache.
     return einsum_op_tensor_mem(q, k, v, 32)
 
-def split_cross_attention_forward_invokeAI(self, x, context=None, mask=None):
+
+def split_cross_attention_forward_invokeAI(self, x, context=None, mask=None, **kwargs):
     h = self.heads
 
     q = self.to_q(x)
@@ -356,7 +383,7 @@ def split_cross_attention_forward_invokeAI(self, x, context=None, mask=None):
 
 # Based on Birch-san's modified implementation of sub-quadratic attention from https://github.com/Birch-san/diffusers/pull/1
 # The sub_quad_attention_forward function is under the MIT License listed under Memory Efficient Attention in the Licenses section of the web UI interface
-def sub_quad_attention_forward(self, x, context=None, mask=None):
+def sub_quad_attention_forward(self, x, context=None, mask=None, **kwargs):
     assert mask is None, "attention-mask not currently implemented for SubQuadraticCrossAttnProcessor."
 
     h = self.heads
@@ -392,6 +419,7 @@ def sub_quad_attention_forward(self, x, context=None, mask=None):
 
     return x
 
+
 def sub_quad_attention(q, k, v, q_chunk_size=1024, kv_chunk_size=None, kv_chunk_size_min=None, chunk_threshold=None, use_checkpoint=True):
     bytes_per_token = torch.finfo(q.dtype).bits//8
     batch_x_heads, q_tokens, _ = q.shape
@@ -442,7 +470,7 @@ def get_xformers_flash_attention_op(q, k, v):
     return None
 
 
-def xformers_attention_forward(self, x, context=None, mask=None):
+def xformers_attention_forward(self, x, context=None, mask=None, **kwargs):
     h = self.heads
     q_in = self.to_q(x)
     context = default(context, x)
@@ -465,9 +493,10 @@ def xformers_attention_forward(self, x, context=None, mask=None):
     out = rearrange(out, 'b n h d -> b n (h d)', h=h)
     return self.to_out(out)
 
+
 # Based on Diffusers usage of scaled dot product attention from https://github.com/huggingface/diffusers/blob/c7da8fd23359a22d0df2741688b5b4f33c26df21/src/diffusers/models/cross_attention.py
 # The scaled_dot_product_attention_forward function contains parts of code under Apache-2.0 license listed under Scaled Dot Product Attention in the Licenses section of the web UI interface
-def scaled_dot_product_attention_forward(self, x, context=None, mask=None):
+def scaled_dot_product_attention_forward(self, x, context=None, mask=None, **kwargs):
     batch_size, sequence_length, inner_dim = x.shape
 
     if mask is not None:
@@ -507,10 +536,12 @@ def scaled_dot_product_attention_forward(self, x, context=None, mask=None):
     hidden_states = self.to_out[1](hidden_states)
     return hidden_states
 
-def scaled_dot_product_no_mem_attention_forward(self, x, context=None, mask=None):
+
+def scaled_dot_product_no_mem_attention_forward(self, x, context=None, mask=None, **kwargs):
     with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False):
         return scaled_dot_product_attention_forward(self, x, context, mask)
 
+
 def cross_attention_attnblock_forward(self, x):
         h_ = x
         h_ = self.norm(h_)
@@ -569,6 +600,7 @@ def cross_attention_attnblock_forward(self, x):
 
         return h3
 
+
 def xformers_attnblock_forward(self, x):
     try:
         h_ = x
@@ -592,6 +624,7 @@ def xformers_attnblock_forward(self, x):
     except NotImplementedError:
         return cross_attention_attnblock_forward(self, x)
 
+
 def sdp_attnblock_forward(self, x):
     h_ = x
     h_ = self.norm(h_)
@@ -612,10 +645,12 @@ def sdp_attnblock_forward(self, x):
     out = self.proj_out(out)
     return x + out
 
+
 def sdp_no_mem_attnblock_forward(self, x):
     with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False):
         return sdp_attnblock_forward(self, x)
 
+
 def sub_quad_attnblock_forward(self, x):
     h_ = x
     h_ = self.norm(h_)

modules/sd_hijack_unet.py

@@ -39,7 +39,10 @@ def apply_model(orig_func, self, x_noisy, t, cond, **kwargs):
 
     if isinstance(cond, dict):
         for y in cond.keys():
+            if isinstance(cond[y], list):
                 cond[y] = [x.to(devices.dtype_unet) if isinstance(x, torch.Tensor) else x for x in cond[y]]
+            else:
+                cond[y] = cond[y].to(devices.dtype_unet) if isinstance(cond[y], torch.Tensor) else cond[y]
 
     with devices.autocast():
         return orig_func(self, x_noisy.to(devices.dtype_unet), t.to(devices.dtype_unet), cond, **kwargs).float()
@@ -77,3 +80,6 @@ first_stage_sub = lambda orig_func, self, x, **kwargs: orig_func(self, x.to(devi
 CondFunc('ldm.models.diffusion.ddpm.LatentDiffusion.decode_first_stage', first_stage_sub, first_stage_cond)
 CondFunc('ldm.models.diffusion.ddpm.LatentDiffusion.encode_first_stage', first_stage_sub, first_stage_cond)
 CondFunc('ldm.models.diffusion.ddpm.LatentDiffusion.get_first_stage_encoding', lambda orig_func, *args, **kwargs: orig_func(*args, **kwargs).float(), first_stage_cond)
+
+CondFunc('sgm.modules.diffusionmodules.wrappers.OpenAIWrapper.forward', apply_model, unet_needs_upcast)
+CondFunc('sgm.modules.diffusionmodules.openaimodel.timestep_embedding', lambda orig_func, timesteps, *args, **kwargs: orig_func(timesteps, *args, **kwargs).to(torch.float32 if timesteps.dtype == torch.int64 else devices.dtype_unet), unet_needs_upcast)

modules/sd_models.py

@@ -14,8 +14,7 @@ import ldm.modules.midas as midas
 
 from ldm.util import instantiate_from_config
 
-from modules import paths, shared, modelloader, devices, script_callbacks, sd_vae, sd_disable_initialization, errors, hashes, sd_models_config, sd_unet
+from modules import paths, shared, modelloader, devices, script_callbacks, sd_vae, sd_disable_initialization, errors, hashes, sd_models_config, sd_unet, sd_models_xl, cache
-from modules.sd_hijack_inpainting import do_inpainting_hijack
 from modules.timer import Timer
 import tomesd
 
@@ -23,7 +22,8 @@ model_dir = "Stable-diffusion"
 model_path = os.path.abspath(os.path.join(paths.models_path, model_dir))
 
 checkpoints_list = {}
-checkpoint_alisases = {}
+checkpoint_aliases = {}
+checkpoint_alisases = checkpoint_aliases  # for compatibility with old name
 checkpoints_loaded = collections.OrderedDict()
 
 
@@ -32,6 +32,8 @@ class CheckpointInfo:
         self.filename = filename
         abspath = os.path.abspath(filename)
 
+        self.is_safetensors = os.path.splitext(filename)[1].lower() == ".safetensors"
+
         if shared.cmd_opts.ckpt_dir is not None and abspath.startswith(shared.cmd_opts.ckpt_dir):
             name = abspath.replace(shared.cmd_opts.ckpt_dir, '')
         elif abspath.startswith(model_path):
@@ -42,6 +44,19 @@ class CheckpointInfo:
         if name.startswith("\\") or name.startswith("/"):
             name = name[1:]
 
+        def read_metadata():
+            metadata = read_metadata_from_safetensors(filename)
+            self.modelspec_thumbnail = metadata.pop('modelspec.thumbnail', None)
+
+            return metadata
+
+        self.metadata = {}
+        if self.is_safetensors:
+            try:
+                self.metadata = cache.cached_data_for_file('safetensors-metadata', "checkpoint/" + name, filename, read_metadata)
+            except Exception as e:
+                errors.display(e, f"reading metadata for {filename}")
+
         self.name = name
         self.name_for_extra = os.path.splitext(os.path.basename(filename))[0]
         self.model_name = os.path.splitext(name.replace("/", "_").replace("\\", "_"))[0]
@@ -51,22 +66,14 @@ class CheckpointInfo:
         self.shorthash = self.sha256[0:10] if self.sha256 else None
 
         self.title = name if self.shorthash is None else f'{name} [{self.shorthash}]'
+        self.short_title = self.name_for_extra if self.shorthash is None else f'{self.name_for_extra} [{self.shorthash}]'
 
-        self.ids = [self.hash, self.model_name, self.title, name, f'{name} [{self.hash}]'] + ([self.shorthash, self.sha256, f'{self.name} [{self.shorthash}]'] if self.shorthash else [])
+        self.ids = [self.hash, self.model_name, self.title, name, self.name_for_extra, f'{name} [{self.hash}]'] + ([self.shorthash, self.sha256, f'{self.name} [{self.shorthash}]'] if self.shorthash else [])
-
-        self.metadata = {}
-
-        _, ext = os.path.splitext(self.filename)
-        if ext.lower() == ".safetensors":
-            try:
-                self.metadata = read_metadata_from_safetensors(filename)
-            except Exception as e:
-                errors.display(e, f"reading checkpoint metadata: {filename}")
 
     def register(self):
         checkpoints_list[self.title] = self
         for id in self.ids:
-            checkpoint_alisases[id] = self
+            checkpoint_aliases[id] = self
 
     def calculate_shorthash(self):
         self.sha256 = hashes.sha256(self.filename, f"checkpoint/{self.name}")
@@ -78,8 +85,9 @@ class CheckpointInfo:
         if self.shorthash not in self.ids:
             self.ids += [self.shorthash, self.sha256, f'{self.name} [{self.shorthash}]']
 
-        checkpoints_list.pop(self.title)
+        checkpoints_list.pop(self.title, None)
         self.title = f'{self.name} [{self.shorthash}]'
+        self.short_title = f'{self.name_for_extra} [{self.shorthash}]'
         self.register()
 
         return self.shorthash
@@ -95,25 +103,18 @@ except Exception:
 
 
 def setup_model():
-    if not os.path.exists(model_path):
+    os.makedirs(model_path, exist_ok=True)
-        os.makedirs(model_path)
 
     enable_midas_autodownload()
 
 
-def checkpoint_tiles():
+def checkpoint_tiles(use_short=False):
-    def convert(name):
+    return [x.short_title if use_short else x.title for x in checkpoints_list.values()]
-        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():
     checkpoints_list.clear()
-    checkpoint_alisases.clear()
+    checkpoint_aliases.clear()
 
     cmd_ckpt = shared.cmd_opts.ckpt
     if shared.cmd_opts.no_download_sd_model or cmd_ckpt != shared.sd_model_file or os.path.exists(cmd_ckpt):
@@ -131,13 +132,16 @@ def list_models():
     elif cmd_ckpt is not None and cmd_ckpt != shared.default_sd_model_file:
         print(f"Checkpoint in --ckpt argument not found (Possible it was moved to {model_path}: {cmd_ckpt}", file=sys.stderr)
 
-    for filename in sorted(model_list, key=str.lower):
+    for filename in model_list:
         checkpoint_info = CheckpointInfo(filename)
         checkpoint_info.register()
 
 
+re_strip_checksum = re.compile(r"\s*\[[^]]+]\s*$")
+
+
 def get_closet_checkpoint_match(search_string):
-    checkpoint_info = checkpoint_alisases.get(search_string, None)
+    checkpoint_info = checkpoint_aliases.get(search_string, None)
     if checkpoint_info is not None:
         return checkpoint_info
 
@@ -145,6 +149,11 @@ def get_closet_checkpoint_match(search_string):
     if found:
         return found[0]
 
+    search_string_without_checksum = re.sub(re_strip_checksum, '', search_string)
+    found = sorted([info for info in checkpoints_list.values() if search_string_without_checksum in info.title], key=lambda x: len(x.title))
+    if found:
+        return found[0]
+
     return None
 
 
@@ -167,7 +176,7 @@ def select_checkpoint():
     """Raises `FileNotFoundError` if no checkpoints are found."""
     model_checkpoint = shared.opts.sd_model_checkpoint
 
-    checkpoint_info = checkpoint_alisases.get(model_checkpoint, None)
+    checkpoint_info = checkpoint_aliases.get(model_checkpoint, None)
     if checkpoint_info is not None:
         return checkpoint_info
 
@@ -248,7 +257,12 @@ def read_state_dict(checkpoint_file, print_global_state=False, map_location=None
     _, extension = os.path.splitext(checkpoint_file)
     if extension.lower() == ".safetensors":
         device = map_location or shared.weight_load_location or devices.get_optimal_device_name()
+
+        if not shared.opts.disable_mmap_load_safetensors:
             pl_sd = safetensors.torch.load_file(checkpoint_file, device=device)
+        else:
+            pl_sd = safetensors.torch.load(open(checkpoint_file, 'rb').read())
+            pl_sd = {k: v.to(device) for k, v in pl_sd.items()}
     else:
         pl_sd = torch.load(checkpoint_file, map_location=map_location or shared.weight_load_location)
 
@@ -275,22 +289,46 @@ def get_checkpoint_state_dict(checkpoint_info: CheckpointInfo, timer):
     return res
 
 
+class SkipWritingToConfig:
+    """This context manager prevents load_model_weights from writing checkpoint name to the config when it loads weight."""
+
+    skip = False
+    previous = None
+
+    def __enter__(self):
+        self.previous = SkipWritingToConfig.skip
+        SkipWritingToConfig.skip = True
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        SkipWritingToConfig.skip = self.previous
+
+
 def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer):
     sd_model_hash = checkpoint_info.calculate_shorthash()
     timer.record("calculate hash")
 
+    if not SkipWritingToConfig.skip:
         shared.opts.data["sd_model_checkpoint"] = checkpoint_info.title
 
     if state_dict is None:
         state_dict = get_checkpoint_state_dict(checkpoint_info, timer)
 
+    model.is_sdxl = hasattr(model, 'conditioner')
+    model.is_sd2 = not model.is_sdxl and hasattr(model.cond_stage_model, 'model')
+    model.is_sd1 = not model.is_sdxl and not model.is_sd2
+
+    if model.is_sdxl:
+        sd_models_xl.extend_sdxl(model)
+
     model.load_state_dict(state_dict, strict=False)
-    del state_dict
     timer.record("apply weights to model")
 
     if shared.opts.sd_checkpoint_cache > 0:
         # cache newly loaded model
-        checkpoints_loaded[checkpoint_info] = model.state_dict().copy()
+        checkpoints_loaded[checkpoint_info] = state_dict
+
+    del state_dict
 
     if shared.cmd_opts.opt_channelslast:
         model.to(memory_format=torch.channels_last)
@@ -314,7 +352,7 @@ def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer
 
         timer.record("apply half()")
 
-    devices.dtype_unet = model.model.diffusion_model.dtype
+    devices.dtype_unet = torch.float16 if model.is_sdxl and not shared.cmd_opts.no_half else model.model.diffusion_model.dtype
     devices.unet_needs_upcast = shared.cmd_opts.upcast_sampling and devices.dtype == torch.float16 and devices.dtype_unet == torch.float16
 
     model.first_stage_model.to(devices.dtype_vae)
@@ -329,6 +367,7 @@ def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer
     model.sd_checkpoint_info = checkpoint_info
     shared.opts.data["sd_checkpoint_hash"] = checkpoint_info.sha256
 
+    if hasattr(model, 'logvar'):
         model.logvar = model.logvar.to(devices.device)  # fix for training
 
     sd_vae.delete_base_vae()
@@ -386,6 +425,7 @@ def repair_config(sd_config):
     if not hasattr(sd_config.model.params, "use_ema"):
         sd_config.model.params.use_ema = False
 
+    if hasattr(sd_config.model.params, 'unet_config'):
         if shared.cmd_opts.no_half:
             sd_config.model.params.unet_config.params.use_fp16 = False
         elif shared.cmd_opts.upcast_sampling:
@@ -402,11 +442,14 @@ def repair_config(sd_config):
 
 sd1_clip_weight = 'cond_stage_model.transformer.text_model.embeddings.token_embedding.weight'
 sd2_clip_weight = 'cond_stage_model.model.transformer.resblocks.0.attn.in_proj_weight'
+sdxl_clip_weight = 'conditioner.embedders.1.model.ln_final.weight'
+sdxl_refiner_clip_weight = 'conditioner.embedders.0.model.ln_final.weight'
 
 
 class SdModelData:
     def __init__(self):
         self.sd_model = None
+        self.loaded_sd_models = []
         self.was_loaded_at_least_once = False
         self.lock = threading.Lock()
 
@@ -421,6 +464,7 @@ class SdModelData:
 
                 try:
                     load_model()
+
                 except Exception as e:
                     errors.display(e, "loading stable diffusion model", full_traceback=True)
                     print("", file=sys.stderr)
@@ -432,31 +476,76 @@ class SdModelData:
     def set_sd_model(self, v):
         self.sd_model = v
 
+        try:
+            self.loaded_sd_models.remove(v)
+        except ValueError:
+            pass
+
+        if v is not None:
+            self.loaded_sd_models.insert(0, v)
+
 
 model_data = SdModelData()
 
 
-def load_model(checkpoint_info=None, already_loaded_state_dict=None):
+def get_empty_cond(sd_model):
-    from modules import lowvram, sd_hijack
+    from modules import extra_networks, processing
-    checkpoint_info = checkpoint_info or select_checkpoint()
+
+    p = processing.StableDiffusionProcessingTxt2Img()
+    extra_networks.activate(p, {})
+
+    if hasattr(sd_model, 'conditioner'):
+        d = sd_model.get_learned_conditioning([""])
+        return d['crossattn']
+    else:
+        return sd_model.cond_stage_model([""])
+
+
+def send_model_to_cpu(m):
+    from modules import lowvram
+
+    if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
+        lowvram.send_everything_to_cpu()
+    else:
+        m.to(devices.cpu)
 
-    if model_data.sd_model:
-        sd_hijack.model_hijack.undo_hijack(model_data.sd_model)
-        model_data.sd_model = None
-        gc.collect()
     devices.torch_gc()
 
-    do_inpainting_hijack()
+
+def send_model_to_device(m):
+    from modules import lowvram
+
+    if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
+        lowvram.setup_for_low_vram(m, shared.cmd_opts.medvram)
+    else:
+        m.to(shared.device)
+
+
+def send_model_to_trash(m):
+    m.to(device="meta")
+    devices.torch_gc()
+
+
+def load_model(checkpoint_info=None, already_loaded_state_dict=None):
+    from modules import sd_hijack
+    checkpoint_info = checkpoint_info or select_checkpoint()
 
     timer = Timer()
 
+    if model_data.sd_model:
+        send_model_to_trash(model_data.sd_model)
+        model_data.sd_model = None
+        devices.torch_gc()
+
+    timer.record("unload existing model")
+
     if already_loaded_state_dict is not None:
         state_dict = already_loaded_state_dict
     else:
         state_dict = get_checkpoint_state_dict(checkpoint_info, timer)
 
     checkpoint_config = sd_models_config.find_checkpoint_config(state_dict, checkpoint_info)
-    clip_is_included_into_sd = sd1_clip_weight in state_dict or sd2_clip_weight in state_dict
+    clip_is_included_into_sd = any(x for x in [sd1_clip_weight, sd2_clip_weight, sdxl_clip_weight, sdxl_refiner_clip_weight] if x in state_dict)
 
     timer.record("find config")
 
@@ -469,26 +558,28 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None):
 
     sd_model = None
     try:
-        with sd_disable_initialization.DisableInitialization(disable_clip=clip_is_included_into_sd):
+        with sd_disable_initialization.DisableInitialization(disable_clip=clip_is_included_into_sd or shared.cmd_opts.do_not_download_clip):
+            with sd_disable_initialization.InitializeOnMeta():
                 sd_model = instantiate_from_config(sd_config.model)
-    except Exception:
+
-        pass
+    except Exception as e:
+        errors.display(e, "creating model quickly", full_traceback=True)
 
     if sd_model is None:
         print('Failed to create model quickly; will retry using slow method.', file=sys.stderr)
+
+        with sd_disable_initialization.InitializeOnMeta():
             sd_model = instantiate_from_config(sd_config.model)
 
     sd_model.used_config = checkpoint_config
 
     timer.record("create model")
 
+    with sd_disable_initialization.LoadStateDictOnMeta(state_dict, devices.cpu):
         load_model_weights(sd_model, checkpoint_info, state_dict, timer)
+    timer.record("load weights from state dict")
 
-    if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
+    send_model_to_device(sd_model)
-        lowvram.setup_for_low_vram(sd_model, shared.cmd_opts.medvram)
-    else:
-        sd_model.to(shared.device)
-
     timer.record("move model to device")
 
     sd_hijack.model_hijack.hijack(sd_model)
@@ -496,7 +587,7 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None):
     timer.record("hijack")
 
     sd_model.eval()
-    model_data.sd_model = sd_model
+    model_data.set_sd_model(sd_model)
     model_data.was_loaded_at_least_once = True
 
     sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings(force_reload=True)  # Reload embeddings after model load as they may or may not fit the model
@@ -508,7 +599,7 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None):
     timer.record("scripts callbacks")
 
     with devices.autocast(), torch.no_grad():
-        sd_model.cond_stage_model_empty_prompt = sd_model.cond_stage_model([""])
+        sd_model.cond_stage_model_empty_prompt = get_empty_cond(sd_model)
 
     timer.record("calculate empty prompt")
 
@@ -517,10 +608,61 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None):
     return sd_model
 
 
+def reuse_model_from_already_loaded(sd_model, checkpoint_info, timer):
+    """
+    Checks if the desired checkpoint from checkpoint_info is not already loaded in model_data.loaded_sd_models.
+    If it is loaded, returns that (moving it to GPU if necessary, and moving the currently loadded model to CPU if necessary).
+    If not, returns the model that can be used to load weights from checkpoint_info's file.
+    If no such model exists, returns None.
+    Additionaly deletes loaded models that are over the limit set in settings (sd_checkpoints_limit).
+    """
+
+    already_loaded = None
+    for i in reversed(range(len(model_data.loaded_sd_models))):
+        loaded_model = model_data.loaded_sd_models[i]
+        if loaded_model.sd_checkpoint_info.filename == checkpoint_info.filename:
+            already_loaded = loaded_model
+            continue
+
+        if len(model_data.loaded_sd_models) > shared.opts.sd_checkpoints_limit > 0:
+            print(f"Unloading model {len(model_data.loaded_sd_models)} over the limit of {shared.opts.sd_checkpoints_limit}: {loaded_model.sd_checkpoint_info.title}")
+            model_data.loaded_sd_models.pop()
+            send_model_to_trash(loaded_model)
+            timer.record("send model to trash")
+
+        if shared.opts.sd_checkpoints_keep_in_cpu:
+            send_model_to_cpu(sd_model)
+            timer.record("send model to cpu")
+
+    if already_loaded is not None:
+        send_model_to_device(already_loaded)
+        timer.record("send model to device")
+
+        model_data.set_sd_model(already_loaded)
+        print(f"Using already loaded model {already_loaded.sd_checkpoint_info.title}: done in {timer.summary()}")
+        return model_data.sd_model
+    elif shared.opts.sd_checkpoints_limit > 1 and len(model_data.loaded_sd_models) < shared.opts.sd_checkpoints_limit:
+        print(f"Loading model {checkpoint_info.title} ({len(model_data.loaded_sd_models) + 1} out of {shared.opts.sd_checkpoints_limit})")
+
+        model_data.sd_model = None
+        load_model(checkpoint_info)
+        return model_data.sd_model
+    elif len(model_data.loaded_sd_models) > 0:
+        sd_model = model_data.loaded_sd_models.pop()
+        model_data.sd_model = sd_model
+
+        print(f"Reusing loaded model {sd_model.sd_checkpoint_info.title} to load {checkpoint_info.title}")
+        return sd_model
+    else:
+        return None
+
+
 def reload_model_weights(sd_model=None, info=None):
-    from modules import lowvram, devices, sd_hijack
+    from modules import devices, sd_hijack
     checkpoint_info = info or select_checkpoint()
 
+    timer = Timer()
+
     if not sd_model:
         sd_model = model_data.sd_model
 
@@ -529,19 +671,17 @@ def reload_model_weights(sd_model=None, info=None):
     else:
         current_checkpoint_info = sd_model.sd_checkpoint_info
         if sd_model.sd_model_checkpoint == checkpoint_info.filename:
-            return
+            return sd_model
 
+    sd_model = reuse_model_from_already_loaded(sd_model, checkpoint_info, timer)
+    if sd_model is not None and sd_model.sd_checkpoint_info.filename == checkpoint_info.filename:
+        return sd_model
+
+    if sd_model is not None:
         sd_unet.apply_unet("None")
-
+        send_model_to_cpu(sd_model)
-        if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
-            lowvram.send_everything_to_cpu()
-        else:
-            sd_model.to(devices.cpu)
-
         sd_hijack.model_hijack.undo_hijack(sd_model)
 
-    timer = Timer()
-
     state_dict = get_checkpoint_state_dict(checkpoint_info, timer)
 
     checkpoint_config = sd_models_config.find_checkpoint_config(state_dict, checkpoint_info)
@@ -549,7 +689,9 @@ def reload_model_weights(sd_model=None, info=None):
     timer.record("find config")
 
     if sd_model is None or checkpoint_config != sd_model.used_config:
-        del sd_model
+        if sd_model is not None:
+            send_model_to_trash(sd_model)
+
         load_model(checkpoint_info, already_loaded_state_dict=state_dict)
         return model_data.sd_model
 
@@ -572,6 +714,9 @@ def reload_model_weights(sd_model=None, info=None):
 
     print(f"Weights loaded in {timer.summary()}.")
 
+    model_data.set_sd_model(sd_model)
+    sd_unet.apply_unet()
+
     return sd_model
 
 
@@ -586,7 +731,6 @@ def unload_model_weights(sd_model=None, info=None):
         sd_model = None
         gc.collect()
         devices.torch_gc()
-        torch.cuda.empty_cache()
 
     print(f"Unloaded weights {timer.summary()}.")
 

modules/sd_models_config.py

@@ -6,12 +6,15 @@ from modules import shared, paths, sd_disable_initialization
 
 sd_configs_path = shared.sd_configs_path
 sd_repo_configs_path = os.path.join(paths.paths['Stable Diffusion'], "configs", "stable-diffusion")
+sd_xl_repo_configs_path = os.path.join(paths.paths['Stable Diffusion XL'], "configs", "inference")
 
 
 config_default = shared.sd_default_config
 config_sd2 = os.path.join(sd_repo_configs_path, "v2-inference.yaml")
 config_sd2v = os.path.join(sd_repo_configs_path, "v2-inference-v.yaml")
 config_sd2_inpainting = os.path.join(sd_repo_configs_path, "v2-inpainting-inference.yaml")
+config_sdxl = os.path.join(sd_xl_repo_configs_path, "sd_xl_base.yaml")
+config_sdxl_refiner = os.path.join(sd_xl_repo_configs_path, "sd_xl_refiner.yaml")
 config_depth_model = os.path.join(sd_repo_configs_path, "v2-midas-inference.yaml")
 config_unclip = os.path.join(sd_repo_configs_path, "v2-1-stable-unclip-l-inference.yaml")
 config_unopenclip = os.path.join(sd_repo_configs_path, "v2-1-stable-unclip-h-inference.yaml")
@@ -68,7 +71,11 @@ def guess_model_config_from_state_dict(sd, filename):
     diffusion_model_input = sd.get('model.diffusion_model.input_blocks.0.0.weight', None)
     sd2_variations_weight = sd.get('embedder.model.ln_final.weight', None)
 
-    if sd.get('depth_model.model.pretrained.act_postprocess3.0.project.0.bias', None) is not None:
+    if sd.get('conditioner.embedders.1.model.ln_final.weight', None) is not None:
+        return config_sdxl
+    if sd.get('conditioner.embedders.0.model.ln_final.weight', None) is not None:
+        return config_sdxl_refiner
+    elif sd.get('depth_model.model.pretrained.act_postprocess3.0.project.0.bias', None) is not None:
         return config_depth_model
     elif sd2_variations_weight is not None and sd2_variations_weight.shape[0] == 768:
         return config_unclip

modules/sd_models_xl.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+import torch
+
+import sgm.models.diffusion
+import sgm.modules.diffusionmodules.denoiser_scaling
+import sgm.modules.diffusionmodules.discretizer
+from modules import devices, shared, prompt_parser
+
+
+def get_learned_conditioning(self: sgm.models.diffusion.DiffusionEngine, batch: prompt_parser.SdConditioning | list[str]):
+    for embedder in self.conditioner.embedders:
+        embedder.ucg_rate = 0.0
+
+    width = getattr(batch, 'width', 1024)
+    height = getattr(batch, 'height', 1024)
+    is_negative_prompt = getattr(batch, 'is_negative_prompt', False)
+    aesthetic_score = shared.opts.sdxl_refiner_low_aesthetic_score if is_negative_prompt else shared.opts.sdxl_refiner_high_aesthetic_score
+
+    devices_args = dict(device=devices.device, dtype=devices.dtype)
+
+    sdxl_conds = {
+        "txt": batch,
+        "original_size_as_tuple": torch.tensor([height, width], **devices_args).repeat(len(batch), 1),
+        "crop_coords_top_left": torch.tensor([shared.opts.sdxl_crop_top, shared.opts.sdxl_crop_left], **devices_args).repeat(len(batch), 1),
+        "target_size_as_tuple": torch.tensor([height, width], **devices_args).repeat(len(batch), 1),
+        "aesthetic_score": torch.tensor([aesthetic_score], **devices_args).repeat(len(batch), 1),
+    }
+
+    force_zero_negative_prompt = is_negative_prompt and all(x == '' for x in batch)
+    c = self.conditioner(sdxl_conds, force_zero_embeddings=['txt'] if force_zero_negative_prompt else [])
+
+    return c
+
+
+def apply_model(self: sgm.models.diffusion.DiffusionEngine, x, t, cond):
+    return self.model(x, t, cond)
+
+
+def get_first_stage_encoding(self, x):  # SDXL's encode_first_stage does everything so get_first_stage_encoding is just there for compatibility
+    return x
+
+
+sgm.models.diffusion.DiffusionEngine.get_learned_conditioning = get_learned_conditioning
+sgm.models.diffusion.DiffusionEngine.apply_model = apply_model
+sgm.models.diffusion.DiffusionEngine.get_first_stage_encoding = get_first_stage_encoding
+
+
+def encode_embedding_init_text(self: sgm.modules.GeneralConditioner, init_text, nvpt):
+    res = []
+
+    for embedder in [embedder for embedder in self.embedders if hasattr(embedder, 'encode_embedding_init_text')]:
+        encoded = embedder.encode_embedding_init_text(init_text, nvpt)
+        res.append(encoded)
+
+    return torch.cat(res, dim=1)
+
+
+def tokenize(self: sgm.modules.GeneralConditioner, texts):
+    for embedder in [embedder for embedder in self.embedders if hasattr(embedder, 'tokenize')]:
+        return embedder.tokenize(texts)
+
+    raise AssertionError('no tokenizer available')
+
+
+
+def process_texts(self, texts):
+    for embedder in [embedder for embedder in self.embedders if hasattr(embedder, 'process_texts')]:
+        return embedder.process_texts(texts)
+
+
+def get_target_prompt_token_count(self, token_count):
+    for embedder in [embedder for embedder in self.embedders if hasattr(embedder, 'get_target_prompt_token_count')]:
+        return embedder.get_target_prompt_token_count(token_count)
+
+
+# those additions to GeneralConditioner make it possible to use it as model.cond_stage_model from SD1.5 in exist
+sgm.modules.GeneralConditioner.encode_embedding_init_text = encode_embedding_init_text
+sgm.modules.GeneralConditioner.tokenize = tokenize
+sgm.modules.GeneralConditioner.process_texts = process_texts
+sgm.modules.GeneralConditioner.get_target_prompt_token_count = get_target_prompt_token_count
+
+
+def extend_sdxl(model):
+    """this adds a bunch of parameters to make SDXL model look a bit more like SD1.5 to the rest of the codebase."""
+
+    dtype = next(model.model.diffusion_model.parameters()).dtype
+    model.model.diffusion_model.dtype = dtype
+    model.model.conditioning_key = 'crossattn'
+    model.cond_stage_key = 'txt'
+    # model.cond_stage_model will be set in sd_hijack
+
+    model.parameterization = "v" if isinstance(model.denoiser.scaling, sgm.modules.diffusionmodules.denoiser_scaling.VScaling) else "eps"
+
+    discretization = sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization()
+    model.alphas_cumprod = torch.asarray(discretization.alphas_cumprod, device=devices.device, dtype=dtype)
+
+    model.conditioner.wrapped = torch.nn.Module()
+
+
+sgm.modules.attention.print = shared.ldm_print
+sgm.modules.diffusionmodules.model.print = shared.ldm_print
+sgm.modules.diffusionmodules.openaimodel.print = shared.ldm_print
+sgm.modules.encoders.modules.print = shared.ldm_print
+
+# this gets the code to load the vanilla attention that we override
+sgm.modules.attention.SDP_IS_AVAILABLE = True
+sgm.modules.attention.XFORMERS_IS_AVAILABLE = False

modules/sd_samplers.py

@@ -28,6 +28,9 @@ def create_sampler(name, model):
 
     assert config is not None, f'bad sampler name: {name}'
 
+    if model.is_sdxl and config.options.get("no_sdxl", False):
+        raise Exception(f"Sampler {config.name} is not supported for SDXL")
+
     sampler = config.constructor(model)
     sampler.config = config
 

modules/sd_samplers_common.py

@@ -2,10 +2,8 @@ from collections import namedtuple
 import numpy as np
 import torch
 from PIL import Image
-from modules import devices, processing, images, sd_vae_approx, sd_samplers, sd_vae_taesd
+from modules import devices, images, sd_vae_approx, sd_samplers, sd_vae_taesd, shared
-
 from modules.shared import opts, state
-import modules.shared as shared
 
 SamplerData = namedtuple('SamplerData', ['name', 'constructor', 'aliases', 'options'])
 
@@ -25,19 +23,29 @@ def setup_img2img_steps(p, steps=None):
 approximation_indexes = {"Full": 0, "Approx NN": 1, "Approx cheap": 2, "TAESD": 3}
 
 
-def single_sample_to_image(sample, approximation=None):
+def samples_to_images_tensor(sample, approximation=None, model=None):
+    '''latents -> images [-1, 1]'''
     if approximation is None:
         approximation = approximation_indexes.get(opts.show_progress_type, 0)
 
     if approximation == 2:
-        x_sample = sd_vae_approx.cheap_approximation(sample) * 0.5 + 0.5
+        x_sample = sd_vae_approx.cheap_approximation(sample)
     elif approximation == 1:
-        x_sample = sd_vae_approx.model()(sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach() * 0.5 + 0.5
+        x_sample = sd_vae_approx.model()(sample.to(devices.device, devices.dtype)).detach()
     elif approximation == 3:
         x_sample = sample * 1.5
-        x_sample = sd_vae_taesd.model()(x_sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach()
+        x_sample = sd_vae_taesd.decoder_model()(x_sample.to(devices.device, devices.dtype)).detach()
+        x_sample = x_sample * 2 - 1
     else:
-        x_sample = processing.decode_first_stage(shared.sd_model, sample.unsqueeze(0))[0] * 0.5 + 0.5
+        if model is None:
+            model = shared.sd_model
+        x_sample = model.decode_first_stage(sample.to(model.first_stage_model.dtype))
+
+    return x_sample
+
+
+def single_sample_to_image(sample, approximation=None):
+    x_sample = samples_to_images_tensor(sample.unsqueeze(0), approximation)[0] * 0.5 + 0.5
 
     x_sample = torch.clamp(x_sample, min=0.0, max=1.0)
     x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)
@@ -46,6 +54,12 @@ def single_sample_to_image(sample, approximation=None):
     return Image.fromarray(x_sample)
 
 
+def decode_first_stage(model, x):
+    x = x.to(devices.dtype_vae)
+    approx_index = approximation_indexes.get(opts.sd_vae_decode_method, 0)
+    return samples_to_images_tensor(x, approx_index, model)
+
+
 def sample_to_image(samples, index=0, approximation=None):
     return single_sample_to_image(samples[index], approximation)
 
@@ -54,6 +68,24 @@ def samples_to_image_grid(samples, approximation=None):
     return images.image_grid([single_sample_to_image(sample, approximation) for sample in samples])
 
 
+def images_tensor_to_samples(image, approximation=None, model=None):
+    '''image[0, 1] -> latent'''
+    if approximation is None:
+        approximation = approximation_indexes.get(opts.sd_vae_encode_method, 0)
+
+    if approximation == 3:
+        image = image.to(devices.device, devices.dtype)
+        x_latent = sd_vae_taesd.encoder_model()(image)
+    else:
+        if model is None:
+            model = shared.sd_model
+        image = image.to(shared.device, dtype=devices.dtype_vae)
+        image = image * 2 - 1
+        x_latent = model.get_first_stage_encoding(model.encode_first_stage(image))
+
+    return x_latent
+
+
 def store_latent(decoded):
     state.current_latent = decoded
 
@@ -85,11 +117,13 @@ class InterruptedException(BaseException):
     pass
 
 
-if opts.randn_source == "CPU":
+def replace_torchsde_browinan():
     import torchsde._brownian.brownian_interval
 
     def torchsde_randn(size, dtype, device, seed):
-        generator = torch.Generator(devices.cpu).manual_seed(int(seed))
+        return devices.randn_local(seed, size).to(device=device, dtype=dtype)
-        return torch.randn(size, dtype=dtype, device=devices.cpu, generator=generator).to(device)
 
     torchsde._brownian.brownian_interval._randn = torchsde_randn
+
+
+replace_torchsde_browinan()

modules/sd_samplers_compvis.py

@@ -11,9 +11,9 @@ import modules.models.diffusion.uni_pc
 
 
 samplers_data_compvis = [
-    sd_samplers_common.SamplerData('DDIM', lambda model: VanillaStableDiffusionSampler(ldm.models.diffusion.ddim.DDIMSampler, model), [], {"default_eta_is_0": True, "uses_ensd": True}),
+    sd_samplers_common.SamplerData('DDIM', lambda model: VanillaStableDiffusionSampler(ldm.models.diffusion.ddim.DDIMSampler, model), [], {"default_eta_is_0": True, "uses_ensd": True, "no_sdxl": True}),
-    sd_samplers_common.SamplerData('PLMS', lambda model: VanillaStableDiffusionSampler(ldm.models.diffusion.plms.PLMSSampler, model), [], {}),
+    sd_samplers_common.SamplerData('PLMS', lambda model: VanillaStableDiffusionSampler(ldm.models.diffusion.plms.PLMSSampler, model), [], {"no_sdxl": True}),
-    sd_samplers_common.SamplerData('UniPC', lambda model: VanillaStableDiffusionSampler(modules.models.diffusion.uni_pc.UniPCSampler, model), [], {}),
+    sd_samplers_common.SamplerData('UniPC', lambda model: VanillaStableDiffusionSampler(modules.models.diffusion.uni_pc.UniPCSampler, model), [], {"no_sdxl": True}),
 ]
 
 
@@ -44,7 +44,7 @@ class VanillaStableDiffusionSampler:
         return 0
 
     def launch_sampling(self, steps, func):
-        state.sampling_steps = steps
+        state.sampling_steps = self.stop_at if self.stop_at is not None else steps
         state.sampling_step = 0
 
         try:

modules/sd_samplers_extra.py

@@ -0,0 +1,74 @@
+import torch
+import tqdm
+import k_diffusion.sampling
+
+
+@torch.no_grad()
+def restart_sampler(model, x, sigmas, extra_args=None, callback=None, disable=None, s_noise=1., restart_list=None):
+    """Implements restart sampling in Restart Sampling for Improving Generative Processes (2023)
+    Restart_list format: {min_sigma: [ restart_steps, restart_times, max_sigma]}
+    If restart_list is None: will choose restart_list automatically, otherwise will use the given restart_list
+    """
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+    step_id = 0
+    from k_diffusion.sampling import to_d, get_sigmas_karras
+
+    def heun_step(x, old_sigma, new_sigma, second_order=True):
+        nonlocal step_id
+        denoised = model(x, old_sigma * s_in, **extra_args)
+        d = to_d(x, old_sigma, denoised)
+        if callback is not None:
+            callback({'x': x, 'i': step_id, 'sigma': new_sigma, 'sigma_hat': old_sigma, 'denoised': denoised})
+        dt = new_sigma - old_sigma
+        if new_sigma == 0 or not second_order:
+            # Euler method
+            x = x + d * dt
+        else:
+            # Heun's method
+            x_2 = x + d * dt
+            denoised_2 = model(x_2, new_sigma * s_in, **extra_args)
+            d_2 = to_d(x_2, new_sigma, denoised_2)
+            d_prime = (d + d_2) / 2
+            x = x + d_prime * dt
+        step_id += 1
+        return x
+
+    steps = sigmas.shape[0] - 1
+    if restart_list is None:
+        if steps >= 20:
+            restart_steps = 9
+            restart_times = 1
+            if steps >= 36:
+                restart_steps = steps // 4
+                restart_times = 2
+            sigmas = get_sigmas_karras(steps - restart_steps * restart_times, sigmas[-2].item(), sigmas[0].item(), device=sigmas.device)
+            restart_list = {0.1: [restart_steps + 1, restart_times, 2]}
+        else:
+            restart_list = {}
+
+    restart_list = {int(torch.argmin(abs(sigmas - key), dim=0)): value for key, value in restart_list.items()}
+
+    step_list = []
+    for i in range(len(sigmas) - 1):
+        step_list.append((sigmas[i], sigmas[i + 1]))
+        if i + 1 in restart_list:
+            restart_steps, restart_times, restart_max = restart_list[i + 1]
+            min_idx = i + 1
+            max_idx = int(torch.argmin(abs(sigmas - restart_max), dim=0))
+            if max_idx < min_idx:
+                sigma_restart = get_sigmas_karras(restart_steps, sigmas[min_idx].item(), sigmas[max_idx].item(), device=sigmas.device)[:-1]
+                while restart_times > 0:
+                    restart_times -= 1
+                    step_list.extend([(old_sigma, new_sigma) for (old_sigma, new_sigma) in zip(sigma_restart[:-1], sigma_restart[1:])])
+
+    last_sigma = None
+    for old_sigma, new_sigma in tqdm.tqdm(step_list, disable=disable):
+        if last_sigma is None:
+            last_sigma = old_sigma
+        elif last_sigma < old_sigma:
+            x = x + k_diffusion.sampling.torch.randn_like(x) * s_noise * (old_sigma ** 2 - last_sigma ** 2) ** 0.5
+        x = heun_step(x, old_sigma, new_sigma)
+        last_sigma = new_sigma
+
+    return x

modules/sd_samplers_kdiffusion.py

@@ -2,8 +2,9 @@ from collections import deque
 import torch
 import inspect
 import k_diffusion.sampling
-from modules import prompt_parser, devices, sd_samplers_common
+from modules import prompt_parser, devices, sd_samplers_common, sd_samplers_extra
 
+from modules.processing import StableDiffusionProcessing
 from modules.shared import opts, state
 import modules.shared as shared
 from modules.script_callbacks import CFGDenoiserParams, cfg_denoiser_callback
@@ -30,12 +31,15 @@ samplers_k_diffusion = [
     ('DPM++ 2M Karras', 'sample_dpmpp_2m', ['k_dpmpp_2m_ka'], {'scheduler': 'karras'}),
     ('DPM++ SDE Karras', 'sample_dpmpp_sde', ['k_dpmpp_sde_ka'], {'scheduler': 'karras', "second_order": True, "brownian_noise": True}),
     ('DPM++ 2M SDE Karras', 'sample_dpmpp_2m_sde', ['k_dpmpp_2m_sde_ka'], {'scheduler': 'karras', "brownian_noise": True}),
+    ('DPM++ 2M SDE Exponential', 'sample_dpmpp_2m_sde', ['k_dpmpp_2m_sde_exp'], {'scheduler': 'exponential', "brownian_noise": True}),
+    ('Restart', sd_samplers_extra.restart_sampler, ['restart'], {'scheduler': 'karras'}),
 ]
 
+
 samplers_data_k_diffusion = [
     sd_samplers_common.SamplerData(label, lambda model, funcname=funcname: KDiffusionSampler(funcname, model), aliases, options)
     for label, funcname, aliases, options in samplers_k_diffusion
-    if hasattr(k_diffusion.sampling, funcname)
+    if callable(funcname) or hasattr(k_diffusion.sampling, funcname)
 ]
 
 sampler_extra_params = {
@@ -53,6 +57,28 @@ k_diffusion_scheduler = {
 }
 
 
+def catenate_conds(conds):
+    if not isinstance(conds[0], dict):
+        return torch.cat(conds)
+
+    return {key: torch.cat([x[key] for x in conds]) for key in conds[0].keys()}
+
+
+def subscript_cond(cond, a, b):
+    if not isinstance(cond, dict):
+        return cond[a:b]
+
+    return {key: vec[a:b] for key, vec in cond.items()}
+
+
+def pad_cond(tensor, repeats, empty):
+    if not isinstance(tensor, dict):
+        return torch.cat([tensor, empty.repeat((tensor.shape[0], repeats, 1))], axis=1)
+
+    tensor['crossattn'] = pad_cond(tensor['crossattn'], repeats, empty)
+    return tensor
+
+
 class CFGDenoiser(torch.nn.Module):
     """
     Classifier free guidance denoiser. A wrapper for stable diffusion model (specifically for unet)
@@ -69,6 +95,7 @@ class CFGDenoiser(torch.nn.Module):
         self.init_latent = None
         self.step = 0
         self.image_cfg_scale = None
+        self.padded_cond_uncond = False
 
     def combine_denoised(self, x_out, conds_list, uncond, cond_scale):
         denoised_uncond = x_out[-uncond.shape[0]:]
@@ -104,10 +131,13 @@ class CFGDenoiser(torch.nn.Module):
 
         if shared.sd_model.model.conditioning_key == "crossattn-adm":
             image_uncond = torch.zeros_like(image_cond)
-            make_condition_dict = lambda c_crossattn, c_adm: {"c_crossattn": c_crossattn, "c_adm": c_adm}
+            make_condition_dict = lambda c_crossattn, c_adm: {"c_crossattn": [c_crossattn], "c_adm": c_adm}
         else:
             image_uncond = image_cond
-            make_condition_dict = lambda c_crossattn, c_concat: {"c_crossattn": c_crossattn, "c_concat": [c_concat]}
+            if isinstance(uncond, dict):
+                make_condition_dict = lambda c_crossattn, c_concat: {**c_crossattn, "c_concat": [c_concat]}
+            else:
+                make_condition_dict = lambda c_crossattn, c_concat: {"c_crossattn": [c_crossattn], "c_concat": [c_concat]}
 
         if not is_edit_model:
             x_in = torch.cat([torch.stack([x[i] for _ in range(n)]) for i, n in enumerate(repeats)] + [x])
@@ -133,32 +163,34 @@ class CFGDenoiser(torch.nn.Module):
             x_in = x_in[:-batch_size]
             sigma_in = sigma_in[:-batch_size]
 
-        # TODO add infotext entry
+        self.padded_cond_uncond = False
         if shared.opts.pad_cond_uncond and tensor.shape[1] != uncond.shape[1]:
             empty = shared.sd_model.cond_stage_model_empty_prompt
             num_repeats = (tensor.shape[1] - uncond.shape[1]) // empty.shape[1]
 
             if num_repeats < 0:
-                tensor = torch.cat([tensor, empty.repeat((tensor.shape[0], -num_repeats, 1))], axis=1)
+                tensor = pad_cond(tensor, -num_repeats, empty)
+                self.padded_cond_uncond = True
             elif num_repeats > 0:
-                uncond = torch.cat([uncond, empty.repeat((uncond.shape[0], num_repeats, 1))], axis=1)
+                uncond = pad_cond(uncond, num_repeats, empty)
+                self.padded_cond_uncond = True
 
         if tensor.shape[1] == uncond.shape[1] or skip_uncond:
             if is_edit_model:
-                cond_in = torch.cat([tensor, uncond, uncond])
+                cond_in = catenate_conds([tensor, uncond, uncond])
             elif skip_uncond:
                 cond_in = tensor
             else:
-                cond_in = torch.cat([tensor, uncond])
+                cond_in = catenate_conds([tensor, uncond])
 
             if shared.batch_cond_uncond:
-                x_out = self.inner_model(x_in, sigma_in, cond=make_condition_dict([cond_in], image_cond_in))
+                x_out = self.inner_model(x_in, sigma_in, cond=make_condition_dict(cond_in, image_cond_in))
             else:
                 x_out = torch.zeros_like(x_in)
                 for batch_offset in range(0, x_out.shape[0], batch_size):
                     a = batch_offset
                     b = a + batch_size
-                    x_out[a:b] = self.inner_model(x_in[a:b], sigma_in[a:b], cond=make_condition_dict([cond_in[a:b]], image_cond_in[a:b]))
+                    x_out[a:b] = self.inner_model(x_in[a:b], sigma_in[a:b], cond=make_condition_dict(subscript_cond(cond_in, a, b), image_cond_in[a:b]))
         else:
             x_out = torch.zeros_like(x_in)
             batch_size = batch_size*2 if shared.batch_cond_uncond else batch_size
@@ -167,14 +199,14 @@ class CFGDenoiser(torch.nn.Module):
                 b = min(a + batch_size, tensor.shape[0])
 
                 if not is_edit_model:
-                    c_crossattn = [tensor[a:b]]
+                    c_crossattn = subscript_cond(tensor, a, b)
                 else:
                     c_crossattn = torch.cat([tensor[a:b]], uncond)
 
                 x_out[a:b] = self.inner_model(x_in[a:b], sigma_in[a:b], cond=make_condition_dict(c_crossattn, image_cond_in[a:b]))
 
             if not skip_uncond:
-                x_out[-uncond.shape[0]:] = self.inner_model(x_in[-uncond.shape[0]:], sigma_in[-uncond.shape[0]:], cond=make_condition_dict([uncond], image_cond_in[-uncond.shape[0]:]))
+                x_out[-uncond.shape[0]:] = self.inner_model(x_in[-uncond.shape[0]:], sigma_in[-uncond.shape[0]:], cond=make_condition_dict(uncond, image_cond_in[-uncond.shape[0]:]))
 
         denoised_image_indexes = [x[0][0] for x in conds_list]
         if skip_uncond:
@@ -230,10 +262,7 @@ class TorchHijack:
             if noise.shape == x.shape:
                 return noise
 
-        if opts.randn_source == "CPU" or x.device.type == 'mps':
+        return devices.randn_like(x)
-            return torch.randn_like(x, device=devices.cpu).to(x.device)
-        else:
-            return torch.randn_like(x)
 
 
 class KDiffusionSampler:
@@ -242,16 +271,25 @@ class KDiffusionSampler:
 
         self.model_wrap = denoiser(sd_model, quantize=shared.opts.enable_quantization)
         self.funcname = funcname
-        self.func = getattr(k_diffusion.sampling, self.funcname)
+        self.func = funcname if callable(funcname) else getattr(k_diffusion.sampling, self.funcname)
         self.extra_params = sampler_extra_params.get(funcname, [])
         self.model_wrap_cfg = CFGDenoiser(self.model_wrap)
         self.sampler_noises = None
         self.stop_at = None
+        self.noisy_output = None
         self.eta = None
         self.config = None  # set by the function calling the constructor
         self.last_latent = None
         self.s_min_uncond = None
 
+        # NOTE: These are also defined in the StableDiffusionProcessing class.
+        # They should have been here to begin with but we're going to
+        # leave that class __init__ signature alone.
+        self.s_churn = 0.0
+        self.s_tmin = 0.0
+        self.s_tmax = float('inf')
+        self.s_noise = 1.0
+
         self.conditioning_key = sd_model.model.conditioning_key
 
     def callback_state(self, d):
@@ -260,6 +298,7 @@ class KDiffusionSampler:
         if opts.live_preview_content == "Combined":
             sd_samplers_common.store_latent(latent)
         self.last_latent = latent
+        self.noisy_output = d['x']
 
         if self.stop_at is not None and step > self.stop_at:
             raise sd_samplers_common.InterruptedException
@@ -268,7 +307,7 @@ class KDiffusionSampler:
         shared.total_tqdm.update()
 
     def launch_sampling(self, steps, func):
-        state.sampling_steps = steps
+        state.sampling_steps = self.stop_at if self.stop_at is not None else steps
         state.sampling_step = 0
 
         try:
@@ -286,7 +325,7 @@ class KDiffusionSampler:
     def number_of_needed_noises(self, p):
         return p.steps
 
-    def initialize(self, p):
+    def initialize(self, p: StableDiffusionProcessing):
         self.model_wrap_cfg.mask = p.mask if hasattr(p, 'mask') else None
         self.model_wrap_cfg.nmask = p.nmask if hasattr(p, 'nmask') else None
         self.model_wrap_cfg.step = 0
@@ -307,6 +346,29 @@ class KDiffusionSampler:
 
             extra_params_kwargs['eta'] = self.eta
 
+        if len(self.extra_params) > 0:
+            s_churn = getattr(opts, 's_churn', p.s_churn)
+            s_tmin = getattr(opts, 's_tmin', p.s_tmin)
+            s_tmax = getattr(opts, 's_tmax', p.s_tmax) or self.s_tmax # 0 = inf
+            s_noise = getattr(opts, 's_noise', p.s_noise)
+
+            if s_churn != self.s_churn:
+                extra_params_kwargs['s_churn'] = s_churn
+                p.s_churn = s_churn
+                p.extra_generation_params['Sigma churn'] = s_churn
+            if s_tmin != self.s_tmin:
+                extra_params_kwargs['s_tmin'] = s_tmin
+                p.s_tmin = s_tmin
+                p.extra_generation_params['Sigma tmin'] = s_tmin
+            if s_tmax != self.s_tmax:
+                extra_params_kwargs['s_tmax'] = s_tmax
+                p.s_tmax = s_tmax
+                p.extra_generation_params['Sigma tmax'] = s_tmax
+            if s_noise != self.s_noise:
+                extra_params_kwargs['s_noise'] = s_noise
+                p.s_noise = s_noise
+                p.extra_generation_params['Sigma noise'] = s_noise
+
         return extra_params_kwargs
 
     def get_sigmas(self, p, steps):
@@ -348,6 +410,9 @@ class KDiffusionSampler:
             sigma_min, sigma_max = (0.1, 10) if opts.use_old_karras_scheduler_sigmas else (self.model_wrap.sigmas[0].item(), self.model_wrap.sigmas[-1].item())
 
             sigmas = k_diffusion.sampling.get_sigmas_karras(n=steps, sigma_min=sigma_min, sigma_max=sigma_max, device=shared.device)
+        elif self.config is not None and self.config.options.get('scheduler', None) == 'exponential':
+            m_sigma_min, m_sigma_max = (self.model_wrap.sigmas[0].item(), self.model_wrap.sigmas[-1].item())
+            sigmas = k_diffusion.sampling.get_sigmas_exponential(n=steps, sigma_min=m_sigma_min, sigma_max=m_sigma_max, device=shared.device)
         else:
             sigmas = self.model_wrap.get_sigmas(steps)
 
@@ -405,6 +470,9 @@ class KDiffusionSampler:
 
         samples = self.launch_sampling(t_enc + 1, lambda: self.func(self.model_wrap_cfg, xi, extra_args=extra_args, disable=False, callback=self.callback_state, **extra_params_kwargs))
 
+        if self.model_wrap_cfg.padded_cond_uncond:
+            p.extra_generation_params["Pad conds"] = True
+
         return samples
 
     def sample(self, p, x, conditioning, unconditional_conditioning, steps=None, image_conditioning=None):
@@ -438,5 +506,8 @@ class KDiffusionSampler:
             's_min_uncond': self.s_min_uncond
         }, disable=False, callback=self.callback_state, **extra_params_kwargs))
 
+        if self.model_wrap_cfg.padded_cond_uncond:
+            p.extra_generation_params["Pad conds"] = True
+
         return samples
 

modules/sd_vae.py

@@ -1,6 +1,6 @@
 import os
 import collections
-from modules import paths, shared, devices, script_callbacks, sd_models
+from modules import paths, shared, devices, script_callbacks, sd_models, extra_networks
 import glob
 from copy import deepcopy
 
@@ -16,6 +16,7 @@ checkpoint_info = None
 
 checkpoints_loaded = collections.OrderedDict()
 
+
 def get_base_vae(model):
     if base_vae is not None and checkpoint_info == model.sd_checkpoint_info and model:
         return base_vae
@@ -50,6 +51,7 @@ def get_filename(filepath):
 
 
 def refresh_vae_list():
+    global vae_dict
     vae_dict.clear()
 
     paths = [
@@ -83,6 +85,8 @@ def refresh_vae_list():
         name = get_filename(filepath)
         vae_dict[name] = filepath
 
+    vae_dict = dict(sorted(vae_dict.items(), key=lambda item: shared.natural_sort_key(item[0])))
+
 
 def find_vae_near_checkpoint(checkpoint_file):
     checkpoint_path = os.path.basename(checkpoint_file).rsplit('.', 1)[0]
@@ -97,6 +101,16 @@ def resolve_vae(checkpoint_file):
     if shared.cmd_opts.vae_path is not None:
         return shared.cmd_opts.vae_path, 'from commandline argument'
 
+    metadata = extra_networks.get_user_metadata(checkpoint_file)
+    vae_metadata = metadata.get("vae", None)
+    if vae_metadata is not None and vae_metadata != "Automatic":
+        if vae_metadata == "None":
+            return None, None
+
+        vae_from_metadata = vae_dict.get(vae_metadata, None)
+        if vae_from_metadata is not None:
+            return vae_from_metadata, "from user metadata"
+
     is_automatic = shared.opts.sd_vae in {"Automatic", "auto"}  # "auto" for people with old config
 
     vae_near_checkpoint = find_vae_near_checkpoint(checkpoint_file)

modules/sd_vae_approx.py

@@ -2,9 +2,9 @@ import os
 
 import torch
 from torch import nn
-from modules import devices, paths
+from modules import devices, paths, shared
 
-sd_vae_approx_model = None
+sd_vae_approx_models = {}
 
 
 class VAEApprox(nn.Module):
@@ -31,31 +31,56 @@ class VAEApprox(nn.Module):
         return x
 
 
-def model():
+def download_model(model_path, model_url):
-    global sd_vae_approx_model
-
-    if sd_vae_approx_model is None:
-        model_path = os.path.join(paths.models_path, "VAE-approx", "model.pt")
-        sd_vae_approx_model = VAEApprox()
     if not os.path.exists(model_path):
-            model_path = os.path.join(paths.script_path, "models", "VAE-approx", "model.pt")
+        os.makedirs(os.path.dirname(model_path), exist_ok=True)
-        sd_vae_approx_model.load_state_dict(torch.load(model_path, map_location='cpu' if devices.device.type != 'cuda' else None))
-        sd_vae_approx_model.eval()
-        sd_vae_approx_model.to(devices.device, devices.dtype)
 
-    return sd_vae_approx_model
+        print(f'Downloading VAEApprox model to: {model_path}')
+        torch.hub.download_url_to_file(model_url, model_path)
+
+
+def model():
+    model_name = "vaeapprox-sdxl.pt" if getattr(shared.sd_model, 'is_sdxl', False) else "model.pt"
+    loaded_model = sd_vae_approx_models.get(model_name)
+
+    if loaded_model is None:
+        model_path = os.path.join(paths.models_path, "VAE-approx", model_name)
+        if not os.path.exists(model_path):
+            model_path = os.path.join(paths.script_path, "models", "VAE-approx", model_name)
+
+        if not os.path.exists(model_path):
+            model_path = os.path.join(paths.models_path, "VAE-approx", model_name)
+            download_model(model_path, 'https://github.com/AUTOMATIC1111/stable-diffusion-webui/releases/download/v1.0.0-pre/' + model_name)
+
+        loaded_model = VAEApprox()
+        loaded_model.load_state_dict(torch.load(model_path, map_location='cpu' if devices.device.type != 'cuda' else None))
+        loaded_model.eval()
+        loaded_model.to(devices.device, devices.dtype)
+        sd_vae_approx_models[model_name] = loaded_model
+
+    return loaded_model
 
 
 def cheap_approximation(sample):
     # https://discuss.huggingface.co/t/decoding-latents-to-rgb-without-upscaling/23204/2
 
-    coefs = torch.tensor([
+    if shared.sd_model.is_sdxl:
-        [0.298, 0.207, 0.208],
+        coeffs = [
-        [0.187, 0.286, 0.173],
+            [ 0.3448,  0.4168,  0.4395],
+            [-0.1953, -0.0290,  0.0250],
+            [ 0.1074,  0.0886, -0.0163],
+            [-0.3730, -0.2499, -0.2088],
+        ]
+    else:
+        coeffs = [
+            [ 0.298,  0.207,  0.208],
+            [ 0.187,  0.286,  0.173],
             [-0.158,  0.189,  0.264],
             [-0.184, -0.271, -0.473],
-    ]).to(sample.device)
+        ]
 
-    x_sample = torch.einsum("lxy,lr -> rxy", sample, coefs)
+    coefs = torch.tensor(coeffs).to(sample.device)
+
+    x_sample = torch.einsum("...lxy,lr -> ...rxy", sample, coefs)
 
     return x_sample

modules/sd_vae_taesd.py

@@ -8,9 +8,9 @@ import os
 import torch
 import torch.nn as nn
 
-from modules import devices, paths_internal
+from modules import devices, paths_internal, shared
 
-sd_vae_taesd = None
+sd_vae_taesd_models = {}
 
 
 def conv(n_in, n_out, **kwargs):
@@ -44,7 +44,17 @@ def decoder():
     )
 
 
-class TAESD(nn.Module):
+def encoder():
+    return nn.Sequential(
+        conv(3, 64), Block(64, 64),
+        conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64),
+        conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64),
+        conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64),
+        conv(64, 4),
+    )
+
+
+class TAESDDecoder(nn.Module):
     latent_magnitude = 3
     latent_shift = 0.5
 
@@ -55,34 +65,60 @@ class TAESD(nn.Module):
         self.decoder.load_state_dict(
             torch.load(decoder_path, map_location='cpu' if devices.device.type != 'cuda' else None))
 
-    @staticmethod
+
-    def unscale_latents(x):
+class TAESDEncoder(nn.Module):
-        """[0, 1] -> raw latents"""
+    latent_magnitude = 3
-        return x.sub(TAESD.latent_shift).mul(2 * TAESD.latent_magnitude)
+    latent_shift = 0.5
+
+    def __init__(self, encoder_path="taesd_encoder.pth"):
+        """Initialize pretrained TAESD on the given device from the given checkpoints."""
+        super().__init__()
+        self.encoder = encoder()
+        self.encoder.load_state_dict(
+            torch.load(encoder_path, map_location='cpu' if devices.device.type != 'cuda' else None))
 
 
-def download_model(model_path):
+def download_model(model_path, model_url):
-    model_url = 'https://github.com/madebyollin/taesd/raw/main/taesd_decoder.pth'
-
     if not os.path.exists(model_path):
         os.makedirs(os.path.dirname(model_path), exist_ok=True)
 
-        print(f'Downloading TAESD decoder to: {model_path}')
+        print(f'Downloading TAESD model to: {model_path}')
         torch.hub.download_url_to_file(model_url, model_path)
 
 
-def model():
+def decoder_model():
-    global sd_vae_taesd
+    model_name = "taesdxl_decoder.pth" if getattr(shared.sd_model, 'is_sdxl', False) else "taesd_decoder.pth"
+    loaded_model = sd_vae_taesd_models.get(model_name)
 
-    if sd_vae_taesd is None:
+    if loaded_model is None:
-        model_path = os.path.join(paths_internal.models_path, "VAE-taesd", "taesd_decoder.pth")
+        model_path = os.path.join(paths_internal.models_path, "VAE-taesd", model_name)
-        download_model(model_path)
+        download_model(model_path, 'https://github.com/madebyollin/taesd/raw/main/' + model_name)
 
         if os.path.exists(model_path):
-            sd_vae_taesd = TAESD(model_path)
+            loaded_model = TAESDDecoder(model_path)
-            sd_vae_taesd.eval()
+            loaded_model.eval()
-            sd_vae_taesd.to(devices.device, devices.dtype)
+            loaded_model.to(devices.device, devices.dtype)
+            sd_vae_taesd_models[model_name] = loaded_model
         else:
             raise FileNotFoundError('TAESD model not found')
 
-    return sd_vae_taesd.decoder
+    return loaded_model.decoder
+
+
+def encoder_model():
+    model_name = "taesdxl_encoder.pth" if getattr(shared.sd_model, 'is_sdxl', False) else "taesd_encoder.pth"
+    loaded_model = sd_vae_taesd_models.get(model_name)
+
+    if loaded_model is None:
+        model_path = os.path.join(paths_internal.models_path, "VAE-taesd", model_name)
+        download_model(model_path, 'https://github.com/madebyollin/taesd/raw/main/' + model_name)
+
+        if os.path.exists(model_path):
+            loaded_model = TAESDEncoder(model_path)
+            loaded_model.eval()
+            loaded_model.to(devices.device, devices.dtype)
+            sd_vae_taesd_models[model_name] = loaded_model
+        else:
+            raise FileNotFoundError('TAESD model not found')
+
+    return loaded_model.encoder

modules/shared.py

@@ -1,14 +1,17 @@
 import datetime
 import json
 import os
+import re
 import sys
 import threading
 import time
+import logging
 
 import gradio as gr
 import torch
 import tqdm
 
+import launch
 import modules.interrogate
 import modules.memmon
 import modules.styles
@@ -18,11 +21,13 @@ from modules.paths_internal import models_path, script_path, data_path, sd_confi
 from ldm.models.diffusion.ddpm import LatentDiffusion
 from typing import Optional
 
+log = logging.getLogger(__name__)
+
 demo = None
 
 parser = cmd_args.parser
 
-script_loading.preload_extensions(extensions_dir, parser)
+script_loading.preload_extensions(extensions_dir, parser, extension_list=launch.list_extensions(launch.args.ui_settings_file))
 script_loading.preload_extensions(extensions_builtin_dir, parser)
 
 if os.environ.get('IGNORE_CMD_ARGS_ERRORS', None) is None:
@@ -46,15 +51,34 @@ restricted_opts = {
 
 # https://huggingface.co/datasets/freddyaboulton/gradio-theme-subdomains/resolve/main/subdomains.json
 gradio_hf_hub_themes = [
+    "gradio/base",
     "gradio/glass",
     "gradio/monochrome",
     "gradio/seafoam",
     "gradio/soft",
-    "freddyaboulton/dracula_revamped",
     "gradio/dracula_test",
     "abidlabs/dracula_test",
+    "abidlabs/Lime",
     "abidlabs/pakistan",
+    "Ama434/neutral-barlow",
     "dawood/microsoft_windows",
+    "finlaymacklon/smooth_slate",
+    "Franklisi/darkmode",
+    "freddyaboulton/dracula_revamped",
+    "freddyaboulton/test-blue",
+    "gstaff/xkcd",
+    "Insuz/Mocha",
+    "Insuz/SimpleIndigo",
+    "JohnSmith9982/small_and_pretty",
+    "nota-ai/theme",
+    "nuttea/Softblue",
+    "ParityError/Anime",
+    "reilnuud/polite",
+    "remilia/Ghostly",
+    "rottenlittlecreature/Moon_Goblin",
+    "step-3-profit/Midnight-Deep",
+    "Taithrah/Minimal",
+    "ysharma/huggingface",
     "ysharma/steampunk"
 ]
 
@@ -144,12 +168,15 @@ class State:
     def request_restart(self) -> None:
         self.interrupt()
         self.server_command = "restart"
+        log.info("Received restart request")
 
     def skip(self):
         self.skipped = True
+        log.info("Received skip request")
 
     def interrupt(self):
         self.interrupted = True
+        log.info("Received interrupt request")
 
     def nextjob(self):
         if opts.live_previews_enable and opts.show_progress_every_n_steps == -1:
@@ -173,7 +200,7 @@ class State:
 
         return obj
 
-    def begin(self):
+    def begin(self, job: str = "(unknown)"):
         self.sampling_step = 0
         self.job_count = -1
         self.processing_has_refined_job_count = False
@@ -187,10 +214,13 @@ class State:
         self.interrupted = False
         self.textinfo = None
         self.time_start = time.time()
-
+        self.job = job
         devices.torch_gc()
+        log.info("Starting job %s", job)
 
     def end(self):
+        duration = time.time() - self.time_start
+        log.info("Ending job %s (%.2f seconds)", self.job, duration)
         self.job = ""
         self.job_count = 0
 
@@ -209,6 +239,8 @@ class State:
             return
 
         import modules.sd_samplers
+
+        try:
             if opts.show_progress_grid:
                 self.assign_current_image(modules.sd_samplers.samples_to_image_grid(self.current_latent))
             else:
@@ -216,6 +248,11 @@ class State:
 
             self.current_image_sampling_step = self.sampling_step
 
+        except Exception:
+            # when switching models during genration, VAE would be on CPU, so creating an image will fail.
+            # we silently ignore this error
+            errors.record_exception()
+
     def assign_current_image(self, image):
         self.current_image = image
         self.id_live_preview += 1
@@ -241,6 +278,7 @@ class OptionInfo:
         self.onchange = onchange
         self.section = section
         self.refresh = refresh
+        self.do_not_save = False
 
         self.comment_before = comment_before
         """HTML text that will be added after label in UI"""
@@ -268,8 +306,17 @@ class OptionInfo:
         self.comment_after += " <span class='info'>(requires restart)</span>"
         return self
 
+    def needs_reload_ui(self):
+        self.comment_after += " <span class='info'>(requires Reload UI)</span>"
+        return self
 
 
+class OptionHTML(OptionInfo):
+    def __init__(self, text):
+        super().__init__(str(text).strip(), label='', component=lambda **kwargs: gr.HTML(elem_classes="settings-info", **kwargs))
+
+        self.do_not_save = True
+
 
 def options_section(section_identifier, options_dict):
     for v in options_dict.values():
@@ -311,6 +358,10 @@ options_templates.update(options_section(('saving-images', "Saving images/grids"
     "grid_prevent_empty_spots": OptionInfo(False, "Prevent empty spots in grid (when set to autodetect)"),
     "grid_zip_filename_pattern": OptionInfo("", "Archive filename pattern", component_args=hide_dirs).link("wiki", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Custom-Images-Filename-Name-and-Subdirectory"),
     "n_rows": OptionInfo(-1, "Grid row count; use -1 for autodetect and 0 for it to be same as batch size", gr.Slider, {"minimum": -1, "maximum": 16, "step": 1}),
+    "font": OptionInfo("", "Font for image grids that have text"),
+    "grid_text_active_color": OptionInfo("#000000", "Text color for image grids", ui_components.FormColorPicker, {}),
+    "grid_text_inactive_color": OptionInfo("#999999", "Inactive text color for image grids", ui_components.FormColorPicker, {}),
+    "grid_background_color": OptionInfo("#ffffff", "Background color for image grids", ui_components.FormColorPicker, {}),
 
     "enable_pnginfo": OptionInfo(True, "Save text information about generation parameters as chunks to png files"),
     "save_txt": OptionInfo(False, "Create a text file next to every image with generation parameters."),
@@ -334,6 +385,7 @@ options_templates.update(options_section(('saving-images', "Saving images/grids"
     "temp_dir":  OptionInfo("", "Directory for temporary images; leave empty for default"),
     "clean_temp_dir_at_start": OptionInfo(False, "Cleanup non-default temporary directory when starting webui"),
 
+    "save_incomplete_images": OptionInfo(False, "Save incomplete images").info("save images that has been interrupted in mid-generation; even if not saved, they will still show up in webui output."),
 }))
 
 options_templates.update(options_section(('saving-paths', "Paths for saving"), {
@@ -370,12 +422,15 @@ options_templates.update(options_section(('face-restoration', "Face restoration"
 }))
 
 options_templates.update(options_section(('system', "System"), {
-    "show_warnings": OptionInfo(False, "Show warnings in console."),
+    "show_warnings": OptionInfo(False, "Show warnings in console.").needs_reload_ui(),
+    "show_gradio_deprecation_warnings": OptionInfo(True, "Show gradio deprecation warnings in console.").needs_reload_ui(),
     "memmon_poll_rate": OptionInfo(8, "VRAM usage polls per second during generation.", gr.Slider, {"minimum": 0, "maximum": 40, "step": 1}).info("0 = disable"),
     "samples_log_stdout": OptionInfo(False, "Always print all generation info to standard output"),
     "multiple_tqdm": OptionInfo(True, "Add a second progress bar to the console that shows progress for an entire job."),
     "print_hypernet_extra": OptionInfo(False, "Print extra hypernetwork information to console."),
     "list_hidden_files": OptionInfo(True, "Load models/files in hidden directories").info("directory is hidden if its name starts with \".\""),
+    "disable_mmap_load_safetensors": OptionInfo(False, "Disable memmapping for loading .safetensors files.").info("fixes very slow loading speed in some cases"),
+    "hide_ldm_prints": OptionInfo(True, "Prevent Stability-AI's ldm/sgm modules from printing noise to console."),
 }))
 
 options_templates.update(options_section(('training', "Training"), {
@@ -395,33 +450,65 @@ options_templates.update(options_section(('training', "Training"), {
 
 options_templates.update(options_section(('sd', "Stable Diffusion"), {
     "sd_model_checkpoint": OptionInfo(None, "Stable Diffusion checkpoint", gr.Dropdown, lambda: {"choices": list_checkpoint_tiles()}, refresh=refresh_checkpoints),
-    "sd_checkpoint_cache": OptionInfo(0, "Checkpoints to cache in RAM", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}),
+    "sd_checkpoints_limit": OptionInfo(1, "Maximum number of checkpoints loaded at the same time", gr.Slider, {"minimum": 1, "maximum": 10, "step": 1}),
+    "sd_checkpoints_keep_in_cpu": OptionInfo(True, "Only keep one model on device").info("will keep models other than the currently used one in RAM rather than VRAM"),
+    "sd_checkpoint_cache": OptionInfo(0, "Checkpoints to cache in RAM", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}).info("obsolete; set to 0 and use the two settings above instead"),
+    "sd_unet": OptionInfo("Automatic", "SD Unet", gr.Dropdown, lambda: {"choices": shared_items.sd_unet_items()}, refresh=shared_items.refresh_unet_list).info("choose Unet model: Automatic = use one with same filename as checkpoint; None = use Unet from checkpoint"),
+    "enable_quantization": OptionInfo(False, "Enable quantization in K samplers for sharper and cleaner results. This may change existing seeds").needs_reload_ui(),
+    "enable_emphasis": OptionInfo(True, "Enable emphasis").info("use (text) to make model pay more attention to text and [text] to make it pay less attention"),
+    "enable_batch_seeds": OptionInfo(True, "Make K-diffusion samplers produce same images in a batch as when making a single image"),
+    "comma_padding_backtrack": OptionInfo(20, "Prompt word wrap length limit", gr.Slider, {"minimum": 0, "maximum": 74, "step": 1}).info("in tokens - for texts shorter than specified, if they don't fit into 75 token limit, move them to the next 75 token chunk"),
+    "CLIP_stop_at_last_layers": OptionInfo(1, "Clip skip", gr.Slider, {"minimum": 1, "maximum": 12, "step": 1}).link("wiki", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#clip-skip").info("ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer"),
+    "upcast_attn": OptionInfo(False, "Upcast cross attention layer to float32"),
+    "randn_source": OptionInfo("GPU", "Random number generator source.", gr.Radio, {"choices": ["GPU", "CPU", "NV"]}).info("changes seeds drastically; use CPU to produce the same picture across different videocard vendors; use NV to produce same picture as on NVidia videocards"),
+    "sd_refiner_checkpoint": OptionInfo(None, "Refiner checkpoint", gr.Dropdown, lambda: {"choices": list_checkpoint_tiles()}, refresh=refresh_checkpoints).info("switch to another model in the middle of generation"),
+    "sd_refiner_switch_at": OptionInfo(1.0, "Refiner switch at", gr.Slider, {"minimum": 0.01, "maximum": 1.0, "step": 0.01}).info("fraction of sampling steps when the swtch to refiner model should happen; 1=never, 0.5=switch in the middle of generation"),
+}))
+
+options_templates.update(options_section(('sdxl', "Stable Diffusion XL"), {
+    "sdxl_crop_top": OptionInfo(0, "crop top coordinate"),
+    "sdxl_crop_left": OptionInfo(0, "crop left coordinate"),
+    "sdxl_refiner_low_aesthetic_score": OptionInfo(2.5, "SDXL low aesthetic score", gr.Number).info("used for refiner model negative prompt"),
+    "sdxl_refiner_high_aesthetic_score": OptionInfo(6.0, "SDXL high aesthetic score", gr.Number).info("used for refiner model prompt"),
+}))
+
+options_templates.update(options_section(('vae', "VAE"), {
+    "sd_vae_explanation": OptionHTML("""
+<abbr title='Variational autoencoder'>VAE</abbr> is a neural network that transforms a standard <abbr title='red/green/blue'>RGB</abbr>
+image into latent space representation and back. Latent space representation is what stable diffusion is working on during sampling
+(i.e. when the progress bar is between empty and full). For txt2img, VAE is used to create a resulting image after the sampling is finished.
+For img2img, VAE is used to process user's input image before the sampling, and to create an image after sampling.
+"""),
     "sd_vae_checkpoint_cache": OptionInfo(0, "VAE Checkpoints to cache in RAM", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}),
     "sd_vae": OptionInfo("Automatic", "SD VAE", gr.Dropdown, lambda: {"choices": shared_items.sd_vae_items()}, refresh=shared_items.refresh_vae_list).info("choose VAE model: Automatic = use one with same filename as checkpoint; None = use VAE from checkpoint"),
     "sd_vae_as_default": OptionInfo(True, "Ignore selected VAE for stable diffusion checkpoints that have their own .vae.pt next to them"),
-    "sd_unet": OptionInfo("Automatic", "SD Unet", gr.Dropdown, lambda: {"choices": shared_items.sd_unet_items()}, refresh=shared_items.refresh_unet_list).info("choose Unet model: Automatic = use one with same filename as checkpoint; None = use Unet from checkpoint"),
+    "auto_vae_precision": OptionInfo(True, "Automaticlly revert VAE to 32-bit floats").info("triggers when a tensor with NaNs is produced in VAE; disabling the option in this case will result in a black square image"),
+    "sd_vae_encode_method": OptionInfo("Full", "VAE type for encode", gr.Radio, {"choices": ["Full", "TAESD"]}).info("method to encode image to latent (use in img2img, hires-fix or inpaint mask)"),
+    "sd_vae_decode_method": OptionInfo("Full", "VAE type for decode", gr.Radio, {"choices": ["Full", "TAESD"]}).info("method to decode latent to image"),
+}))
+
+options_templates.update(options_section(('img2img', "img2img"), {
     "inpainting_mask_weight": OptionInfo(1.0, "Inpainting conditioning mask strength", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
     "initial_noise_multiplier": OptionInfo(1.0, "Noise multiplier for img2img", gr.Slider, {"minimum": 0.5, "maximum": 1.5, "step": 0.01}),
     "img2img_color_correction": OptionInfo(False, "Apply color correction to img2img results to match original colors."),
     "img2img_fix_steps": OptionInfo(False, "With img2img, do exactly the amount of steps the slider specifies.").info("normally you'd do less with less denoising"),
-    "img2img_background_color": OptionInfo("#ffffff", "With img2img, fill image's transparent parts with this color.", ui_components.FormColorPicker, {}),
+    "img2img_background_color": OptionInfo("#ffffff", "With img2img, fill transparent parts of the input image with this color.", ui_components.FormColorPicker, {}),
-    "enable_quantization": OptionInfo(False, "Enable quantization in K samplers for sharper and cleaner results. This may change existing seeds. Requires restart to apply."),
+    "img2img_editor_height": OptionInfo(720, "Height of the image editor", gr.Slider, {"minimum": 80, "maximum": 1600, "step": 1}).info("in pixels").needs_reload_ui(),
-    "enable_emphasis": OptionInfo(True, "Enable emphasis").info("use (text) to make model pay more attention to text and [text] to make it pay less attention"),
+    "img2img_sketch_default_brush_color": OptionInfo("#ffffff", "Sketch initial brush color", ui_components.FormColorPicker, {}).info("default brush color of img2img sketch").needs_reload_ui(),
-    "enable_batch_seeds": OptionInfo(True, "Make K-diffusion samplers produce same images in a batch as when making a single image"),
+    "img2img_inpaint_mask_brush_color": OptionInfo("#ffffff", "Inpaint mask brush color", ui_components.FormColorPicker,  {}).info("brush color of inpaint mask").needs_reload_ui(),
-    "comma_padding_backtrack": OptionInfo(20, "Prompt word wrap length limit", gr.Slider, {"minimum": 0, "maximum": 74, "step": 1}).info("in tokens - for texts shorter than specified, if they don't fit into 75 token limit, move them to the next 75 token chunk"),
+    "img2img_inpaint_sketch_default_brush_color": OptionInfo("#ffffff", "Inpaint sketch initial brush color", ui_components.FormColorPicker, {}).info("default brush color of img2img inpaint sketch").needs_reload_ui(),
-    "CLIP_stop_at_last_layers": OptionInfo(1, "Clip skip", gr.Slider, {"minimum": 1, "maximum": 12, "step": 1}).link("wiki", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#clip-skip").info("ignore last layers of CLIP nrtwork; 1 ignores none, 2 ignores one layer"),
+    "return_mask": OptionInfo(False, "For inpainting, include the greyscale mask in results for web"),
-    "upcast_attn": OptionInfo(False, "Upcast cross attention layer to float32"),
+    "return_mask_composite": OptionInfo(False, "For inpainting, include masked composite in results for web"),
-    "randn_source": OptionInfo("GPU", "Random number generator source.", gr.Radio, {"choices": ["GPU", "CPU"]}).info("changes seeds drastically; use CPU to produce the same picture across different videocard vendors"),
 }))
 
 options_templates.update(options_section(('optimizations', "Optimizations"), {
     "cross_attention_optimization": OptionInfo("Automatic", "Cross attention optimization", gr.Dropdown, lambda: {"choices": shared_items.cross_attention_optimizations()}),
-    "s_min_uncond": OptionInfo(0.0, "Negative Guidance minimum sigma", gr.Slider, {"minimum": 0.0, "maximum": 4.0, "step": 0.01}).link("PR", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/9177").info("skip negative prompt for some steps when the image is almost ready; 0=disable, higher=faster"),
+    "s_min_uncond": OptionInfo(0.0, "Negative Guidance minimum sigma", gr.Slider, {"minimum": 0.0, "maximum": 15.0, "step": 0.01}).link("PR", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/9177").info("skip negative prompt for some steps when the image is almost ready; 0=disable, higher=faster"),
     "token_merging_ratio": OptionInfo(0.0, "Token merging ratio", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}).link("PR", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/9256").info("0=disable, higher=faster"),
     "token_merging_ratio_img2img": OptionInfo(0.0, "Token merging ratio for img2img", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}).info("only applies if non-zero and overrides above"),
     "token_merging_ratio_hr": OptionInfo(0.0, "Token merging ratio for high-res pass", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}).info("only applies if non-zero and overrides above"),
     "pad_cond_uncond": OptionInfo(False, "Pad prompt/negative prompt to be same length").info("improves performance when prompt and negative prompt have different lengths; changes seeds"),
-    "experimental_persistent_cond_cache": OptionInfo(False, "persistent cond cache").info("Experimental, keep cond caches across jobs, reduce overhead."),
+    "persistent_cond_cache": OptionInfo(True, "Persistent cond cache").info("Do not recalculate conds from prompts if prompts have not changed since previous calculation"),
 }))
 
 options_templates.update(options_section(('compatibility', "Compatibility"), {
@@ -433,7 +520,7 @@ options_templates.update(options_section(('compatibility', "Compatibility"), {
     "hires_fix_use_firstpass_conds": OptionInfo(False, "For hires fix, calculate conds of second pass using extra networks of first pass."),
 }))
 
-options_templates.update(options_section(('interrogate', "Interrogate Options"), {
+options_templates.update(options_section(('interrogate', "Interrogate"), {
     "interrogate_keep_models_in_memory": OptionInfo(False, "Keep models in VRAM"),
     "interrogate_return_ranks": OptionInfo(False, "Include ranks of model tags matches in results.").info("booru only"),
     "interrogate_clip_num_beams": OptionInfo(1, "BLIP: num_beams", gr.Slider, {"minimum": 1, "maximum": 16, "step": 1}),
@@ -451,48 +538,51 @@ options_templates.update(options_section(('interrogate', "Interrogate Options"),
 options_templates.update(options_section(('extra_networks', "Extra Networks"), {
     "extra_networks_show_hidden_directories": OptionInfo(True, "Show hidden directories").info("directory is hidden if its name starts with \".\"."),
     "extra_networks_hidden_models": OptionInfo("When searched", "Show cards for models in hidden directories", gr.Radio, {"choices": ["Always", "When searched", "Never"]}).info('"When searched" option will only show the item when the search string has 4 characters or more'),
-    "extra_networks_default_view": OptionInfo("cards", "Default view for Extra Networks", gr.Dropdown, {"choices": ["cards", "thumbs"]}),
+    "extra_networks_default_multiplier": OptionInfo(1.0, "Default multiplier for extra networks", gr.Slider, {"minimum": 0.0, "maximum": 2.0, "step": 0.01}),
-    "extra_networks_default_multiplier": OptionInfo(1.0, "Multiplier for extra networks", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
     "extra_networks_card_width": OptionInfo(0, "Card width for Extra Networks").info("in pixels"),
     "extra_networks_card_height": OptionInfo(0, "Card height for Extra Networks").info("in pixels"),
+    "extra_networks_card_text_scale": OptionInfo(1.0, "Card text scale", gr.Slider, {"minimum": 0.0, "maximum": 2.0, "step": 0.01}).info("1 = original size"),
+    "extra_networks_card_show_desc": OptionInfo(True, "Show description on card"),
     "extra_networks_add_text_separator": OptionInfo(" ", "Extra networks separator").info("extra text to add before <...> when adding extra network to prompt"),
-    "ui_extra_networks_tab_reorder": OptionInfo("", "Extra networks tab order").needs_restart(),
+    "ui_extra_networks_tab_reorder": OptionInfo("", "Extra networks tab order").needs_reload_ui(),
+    "textual_inversion_print_at_load": OptionInfo(False, "Print a list of Textual Inversion embeddings when loading model"),
+    "textual_inversion_add_hashes_to_infotext": OptionInfo(True, "Add Textual Inversion hashes to infotext"),
     "sd_hypernetwork": OptionInfo("None", "Add hypernetwork to prompt", gr.Dropdown, lambda: {"choices": ["None", *hypernetworks]}, refresh=reload_hypernetworks),
 }))
 
 options_templates.update(options_section(('ui', "User interface"), {
-    "localization": OptionInfo("None", "Localization", gr.Dropdown, lambda: {"choices": ["None"] + list(localization.localizations.keys())}, refresh=lambda: localization.list_localizations(cmd_opts.localizations_dir)).needs_restart(),
+    "localization": OptionInfo("None", "Localization", gr.Dropdown, lambda: {"choices": ["None"] + list(localization.localizations.keys())}, refresh=lambda: localization.list_localizations(cmd_opts.localizations_dir)).needs_reload_ui(),
-    "gradio_theme": OptionInfo("Default", "Gradio theme", ui_components.DropdownEditable, lambda: {"choices": ["Default"] + gradio_hf_hub_themes}).needs_restart(),
+    "gradio_theme": OptionInfo("Default", "Gradio theme", ui_components.DropdownEditable, lambda: {"choices": ["Default"] + gradio_hf_hub_themes}).info("you can also manually enter any of themes from the <a href='https://huggingface.co/spaces/gradio/theme-gallery'>gallery</a>.").needs_reload_ui(),
-    "img2img_editor_height": OptionInfo(720, "img2img: height of image editor", gr.Slider, {"minimum": 80, "maximum": 1600, "step": 1}).info("in pixels").needs_restart(),
+    "gradio_themes_cache": OptionInfo(True, "Cache gradio themes locally").info("disable to update the selected Gradio theme"),
     "return_grid": OptionInfo(True, "Show grid in results for web"),
-    "return_mask": OptionInfo(False, "For inpainting, include the greyscale mask in results for web"),
-    "return_mask_composite": OptionInfo(False, "For inpainting, include masked composite in results for web"),
     "do_not_show_images": OptionInfo(False, "Do not show any images in results for web"),
     "send_seed": OptionInfo(True, "Send seed when sending prompt or image to other interface"),
     "send_size": OptionInfo(True, "Send size when sending prompt or image to another interface"),
-    "font": OptionInfo("", "Font for image grids that have text"),
     "js_modal_lightbox": OptionInfo(True, "Enable full page image viewer"),
     "js_modal_lightbox_initially_zoomed": OptionInfo(True, "Show images zoomed in by default in full page image viewer"),
     "js_modal_lightbox_gamepad": OptionInfo(False, "Navigate image viewer with gamepad"),
     "js_modal_lightbox_gamepad_repeat": OptionInfo(250, "Gamepad repeat period, in milliseconds"),
     "show_progress_in_title": OptionInfo(True, "Show generation progress in window title."),
-    "samplers_in_dropdown": OptionInfo(True, "Use dropdown for sampler selection instead of radio group").needs_restart(),
+    "samplers_in_dropdown": OptionInfo(True, "Use dropdown for sampler selection instead of radio group").needs_reload_ui(),
-    "dimensions_and_batch_together": OptionInfo(True, "Show Width/Height and Batch sliders in same row").needs_restart(),
+    "dimensions_and_batch_together": OptionInfo(True, "Show Width/Height and Batch sliders in same row").needs_reload_ui(),
     "keyedit_precision_attention": OptionInfo(0.1, "Ctrl+up/down precision when editing (attention:1.1)", gr.Slider, {"minimum": 0.01, "maximum": 0.2, "step": 0.001}),
     "keyedit_precision_extra": OptionInfo(0.05, "Ctrl+up/down precision when editing <extra networks:0.9>", gr.Slider, {"minimum": 0.01, "maximum": 0.2, "step": 0.001}),
     "keyedit_delimiters": OptionInfo(".,\\/!?%^*;:{}=`~()", "Ctrl+up/down word delimiters"),
-    "quicksettings_list": OptionInfo(["sd_model_checkpoint"], "Quicksettings list", ui_components.DropdownMulti, lambda: {"choices": list(opts.data_labels.keys())}).js("info", "settingsHintsShowQuicksettings").info("setting entries that appear at the top of page rather than in settings tab").needs_restart(),
+    "keyedit_move": OptionInfo(True, "Alt+left/right moves prompt elements"),
-    "ui_tab_order": OptionInfo([], "UI tab order", ui_components.DropdownMulti, lambda: {"choices": list(tab_names)}).needs_restart(),
+    "quicksettings_list": OptionInfo(["sd_model_checkpoint"], "Quicksettings list", ui_components.DropdownMulti, lambda: {"choices": list(opts.data_labels.keys())}).js("info", "settingsHintsShowQuicksettings").info("setting entries that appear at the top of page rather than in settings tab").needs_reload_ui(),
-    "hidden_tabs": OptionInfo([], "Hidden UI tabs", ui_components.DropdownMulti, lambda: {"choices": list(tab_names)}).needs_restart(),
+    "ui_tab_order": OptionInfo([], "UI tab order", ui_components.DropdownMulti, lambda: {"choices": list(tab_names)}).needs_reload_ui(),
-    "ui_reorder_list": OptionInfo([], "txt2img/img2img UI item order", ui_components.DropdownMulti, lambda: {"choices": list(shared_items.ui_reorder_categories())}).info("selected items appear first").needs_restart(),
+    "hidden_tabs": OptionInfo([], "Hidden UI tabs", ui_components.DropdownMulti, lambda: {"choices": list(tab_names)}).needs_reload_ui(),
-    "hires_fix_show_sampler": OptionInfo(False, "Hires fix: show hires sampler selection").needs_restart(),
+    "ui_reorder_list": OptionInfo([], "txt2img/img2img UI item order", ui_components.DropdownMulti, lambda: {"choices": list(shared_items.ui_reorder_categories())}).info("selected items appear first").needs_reload_ui(),
-    "hires_fix_show_prompts": OptionInfo(False, "Hires fix: show hires prompt and negative prompt").needs_restart(),
+    "hires_fix_show_sampler": OptionInfo(False, "Hires fix: show hires checkpoint and sampler selection").needs_reload_ui(),
-    "disable_token_counters": OptionInfo(False, "Disable prompt token counters").needs_restart(),
+    "hires_fix_show_prompts": OptionInfo(False, "Hires fix: show hires prompt and negative prompt").needs_reload_ui(),
+    "disable_token_counters": OptionInfo(False, "Disable prompt token counters").needs_reload_ui(),
 }))
 
+
 options_templates.update(options_section(('infotext', "Infotext"), {
     "add_model_hash_to_info": OptionInfo(True, "Add model hash to generation information"),
     "add_model_name_to_info": OptionInfo(True, "Add model name to generation information"),
+    "add_user_name_to_info": OptionInfo(False, "Add user name to generation information when authenticated"),
     "add_version_to_infotext": OptionInfo(True, "Add program version to generation information"),
     "disable_weights_auto_swap": OptionInfo(True, "Disregard checkpoint information from pasted infotext").info("when reading generation parameters from text into UI"),
     "infotext_styles": OptionInfo("Apply if any", "Infer styles from prompts of pasted infotext", gr.Radio, {"choices": ["Ignore", "Apply", "Discard", "Apply if any"]}).info("when reading generation parameters from text into UI)").html("""<ul style='margin-left: 1.5em'>
@@ -516,12 +606,13 @@ options_templates.update(options_section(('ui', "Live previews"), {
 }))
 
 options_templates.update(options_section(('sampler-params', "Sampler parameters"), {
-    "hide_samplers": OptionInfo([], "Hide samplers in user interface", gr.CheckboxGroup, lambda: {"choices": [x.name for x in list_samplers()]}).needs_restart(),
+    "hide_samplers": OptionInfo([], "Hide samplers in user interface", gr.CheckboxGroup, lambda: {"choices": [x.name for x in list_samplers()]}).needs_reload_ui(),
     "eta_ddim": OptionInfo(0.0, "Eta for DDIM", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}).info("noise multiplier; higher = more unperdictable results"),
     "eta_ancestral": OptionInfo(1.0, "Eta for ancestral samplers", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}).info("noise multiplier; applies to Euler a and other samplers that have a in them"),
     "ddim_discretize": OptionInfo('uniform', "img2img DDIM discretize", gr.Radio, {"choices": ['uniform', 'quad']}),
-    's_churn': OptionInfo(0.0, "sigma churn", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
+    's_churn': OptionInfo(0.0, "sigma churn", gr.Slider, {"minimum": 0.0, "maximum": 100.0, "step": 0.01}),
     's_tmin':  OptionInfo(0.0, "sigma tmin",  gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
+    's_tmax':  OptionInfo(0.0, "sigma tmax",  gr.Slider, {"minimum": 0.0, "maximum": 999.0, "step": 0.01}).info("0 = inf"),
     's_noise': OptionInfo(1.0, "sigma noise", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
     'k_sched_type':  OptionInfo("Automatic", "scheduler type", gr.Dropdown, {"choices": ["Automatic", "karras", "exponential", "polyexponential"]}).info("lets you override the noise schedule for k-diffusion samplers; choosing Automatic disables the three parameters below"),
     'sigma_min': OptionInfo(0.0, "sigma min", gr.Number).info("0 = default (~0.03); minimum noise strength for k-diffusion noise scheduler"),
@@ -566,6 +657,9 @@ class Options:
                 assert not cmd_opts.freeze_settings, "changing settings is disabled"
 
                 info = opts.data_labels.get(key, None)
+                if info.do_not_save:
+                    return
+
                 comp_args = info.component_args if info else None
                 if isinstance(comp_args, dict) and comp_args.get('visible', True) is False:
                     raise RuntimeError(f"not possible to set {key} because it is restricted")
@@ -595,6 +689,9 @@ class Options:
         if oldval == value:
             return False
 
+        if self.data_labels[key].do_not_save:
+            return False
+
         try:
             setattr(self, key, value)
         except RuntimeError:
@@ -771,13 +868,19 @@ def reload_gradio_theme(theme_name=None):
         gradio_theme = gr.themes.Default(**default_theme_args)
     else:
         try:
+            theme_cache_dir = os.path.join(script_path, 'tmp', 'gradio_themes')
+            theme_cache_path = os.path.join(theme_cache_dir, f'{theme_name.replace("/", "_")}.json')
+            if opts.gradio_themes_cache and os.path.exists(theme_cache_path):
+                gradio_theme = gr.themes.ThemeClass.load(theme_cache_path)
+            else:
+                os.makedirs(theme_cache_dir, exist_ok=True)
                 gradio_theme = gr.themes.ThemeClass.from_hub(theme_name)
+                gradio_theme.dump(theme_cache_path)
         except Exception as e:
             errors.display(e, "changing gradio theme")
             gradio_theme = gr.themes.Default(**default_theme_args)
 
 
-
 class TotalTQDM:
     def __init__(self):
         self._tqdm = None
@@ -817,8 +920,12 @@ mem_mon = modules.memmon.MemUsageMonitor("MemMon", device, opts)
 mem_mon.start()
 
 
+def natural_sort_key(s, regex=re.compile('([0-9]+)')):
+    return [int(text) if text.isdigit() else text.lower() for text in regex.split(s)]
+
+
 def listfiles(dirname):
-    filenames = [os.path.join(dirname, x) for x in sorted(os.listdir(dirname), key=str.lower) if not x.startswith(".")]
+    filenames = [os.path.join(dirname, x) for x in sorted(os.listdir(dirname), key=natural_sort_key) if not x.startswith(".")]
     return [file for file in filenames if os.path.isfile(file)]
 
 
@@ -843,8 +950,11 @@ def walk_files(path, allowed_extensions=None):
     if allowed_extensions is not None:
         allowed_extensions = set(allowed_extensions)
 
-    for root, _, files in os.walk(path, followlinks=True):
+    items = list(os.walk(path, followlinks=True))
-        for filename in files:
+    items = sorted(items, key=lambda x: natural_sort_key(x[0]))
+
+    for root, _, files in items:
+        for filename in sorted(files, key=natural_sort_key):
             if allowed_extensions is not None:
                 _, ext = os.path.splitext(filename)
                 if ext not in allowed_extensions:
@@ -854,3 +964,10 @@ def walk_files(path, allowed_extensions=None):
                 continue
 
             yield os.path.join(root, filename)
+
+
+def ldm_print(*args, **kwargs):
+    if opts.hide_ldm_prints:
+        return
+
+    print(*args, **kwargs)

modules/styles.py

@@ -106,10 +106,7 @@ class StyleDatabase:
         if os.path.exists(path):
             shutil.copy(path, f"{path}.bak")
 
-        fd = os.open(path, os.O_RDWR | os.O_CREAT)
+        with open(path, "w", encoding="utf-8-sig", newline='') as file:
-        with os.fdopen(fd, "w", encoding="utf-8-sig", newline='') as file:
-            # _fields is actually part of the public API: typing.NamedTuple is a replacement for collections.NamedTuple,
-            # and collections.NamedTuple has explicit documentation for accessing _fields. Same goes for _asdict()
             writer = csv.DictWriter(file, fieldnames=PromptStyle._fields)
             writer.writeheader()
             writer.writerows(style._asdict() for k, style in self.styles.items())

modules/sysinfo.py

@@ -109,11 +109,15 @@ def format_traceback(tb):
     return [[f"{x.filename}, line {x.lineno}, {x.name}", x.line] for x in traceback.extract_tb(tb)]
 
 
+def format_exception(e, tb):
+    return {"exception": str(e), "traceback": format_traceback(tb)}
+
+
 def get_exceptions():
     try:
         from modules import errors
 
-        return [{"exception": str(e), "traceback": format_traceback(tb)} for e, tb in reversed(errors.exception_records)]
+        return list(reversed(errors.exception_records))
     except Exception as e:
         return str(e)
 

modules/textual_inversion/autocrop.py

@@ -298,8 +298,7 @@ def download_and_cache_models(dirname):
     download_url = 'https://github.com/opencv/opencv_zoo/blob/91fb0290f50896f38a0ab1e558b74b16bc009428/models/face_detection_yunet/face_detection_yunet_2022mar.onnx?raw=true'
     model_file_name = 'face_detection_yunet.onnx'
 
-    if not os.path.exists(dirname):
+    os.makedirs(dirname, exist_ok=True)
-        os.makedirs(dirname)
 
     cache_file = os.path.join(dirname, model_file_name)
     if not os.path.exists(cache_file):

modules/textual_inversion/logging.py

@@ -2,11 +2,51 @@ import datetime
 import json
 import os
 
-saved_params_shared = {"model_name", "model_hash", "initial_step", "num_of_dataset_images", "learn_rate", "batch_size", "clip_grad_mode", "clip_grad_value", "gradient_step", "data_root", "log_directory", "training_width", "training_height", "steps", "create_image_every", "template_file", "gradient_step", "latent_sampling_method"}
+saved_params_shared = {
-saved_params_ti = {"embedding_name", "num_vectors_per_token", "save_embedding_every", "save_image_with_stored_embedding"}
+    "batch_size",
-saved_params_hypernet = {"hypernetwork_name", "layer_structure", "activation_func", "weight_init", "add_layer_norm", "use_dropout", "save_hypernetwork_every"}
+    "clip_grad_mode",
+    "clip_grad_value",
+    "create_image_every",
+    "data_root",
+    "gradient_step",
+    "initial_step",
+    "latent_sampling_method",
+    "learn_rate",
+    "log_directory",
+    "model_hash",
+    "model_name",
+    "num_of_dataset_images",
+    "steps",
+    "template_file",
+    "training_height",
+    "training_width",
+}
+saved_params_ti = {
+    "embedding_name",
+    "num_vectors_per_token",
+    "save_embedding_every",
+    "save_image_with_stored_embedding",
+}
+saved_params_hypernet = {
+    "activation_func",
+    "add_layer_norm",
+    "hypernetwork_name",
+    "layer_structure",
+    "save_hypernetwork_every",
+    "use_dropout",
+    "weight_init",
+}
 saved_params_all = saved_params_shared | saved_params_ti | saved_params_hypernet
-saved_params_previews = {"preview_prompt", "preview_negative_prompt", "preview_steps", "preview_sampler_index", "preview_cfg_scale", "preview_seed", "preview_width", "preview_height"}
+saved_params_previews = {
+    "preview_cfg_scale",
+    "preview_height",
+    "preview_negative_prompt",
+    "preview_prompt",
+    "preview_sampler_index",
+    "preview_seed",
+    "preview_steps",
+    "preview_width",
+}
 
 
 def save_settings_to_file(log_directory, all_params):

modules/textual_inversion/preprocess.py

@@ -7,7 +7,7 @@ from modules import paths, shared, images, deepbooru
 from modules.textual_inversion import autocrop
 
 
-def preprocess(id_task, process_src, process_dst, process_width, process_height, preprocess_txt_action, process_keep_original_size, process_flip, process_split, process_caption, process_caption_deepbooru=False, split_threshold=0.5, overlap_ratio=0.2, process_focal_crop=False, process_focal_crop_face_weight=0.9, process_focal_crop_entropy_weight=0.3, process_focal_crop_edges_weight=0.5, process_focal_crop_debug=False, process_multicrop=None, process_multicrop_mindim=None, process_multicrop_maxdim=None, process_multicrop_minarea=None, process_multicrop_maxarea=None, process_multicrop_objective=None, process_multicrop_threshold=None):
+def preprocess(id_task, process_src, process_dst, process_width, process_height, preprocess_txt_action, process_keep_original_size, process_flip, process_split, process_caption, process_caption_deepbooru=False, split_threshold=0.5, overlap_ratio=0.2, process_focal_crop=False, process_focal_crop_face_weight=0.9, process_focal_crop_entropy_weight=0.15, process_focal_crop_edges_weight=0.5, process_focal_crop_debug=False, process_multicrop=None, process_multicrop_mindim=None, process_multicrop_maxdim=None, process_multicrop_minarea=None, process_multicrop_maxarea=None, process_multicrop_objective=None, process_multicrop_threshold=None):
     try:
         if process_caption:
             shared.interrogator.load()

modules/textual_inversion/textual_inversion.py

@@ -1,5 +1,6 @@
 import os
 from collections import namedtuple
+from contextlib import closing
 
 import torch
 import tqdm
@@ -12,7 +13,7 @@ import numpy as np
 from PIL import Image, PngImagePlugin
 from torch.utils.tensorboard import SummaryWriter
 
-from modules import shared, devices, sd_hijack, processing, sd_models, images, sd_samplers, sd_hijack_checkpoint, errors
+from modules import shared, devices, sd_hijack, sd_models, images, sd_samplers, sd_hijack_checkpoint, errors, hashes
 import modules.textual_inversion.dataset
 from modules.textual_inversion.learn_schedule import LearnRateScheduler
 
@@ -48,6 +49,8 @@ class Embedding:
         self.sd_checkpoint_name = None
         self.optimizer_state_dict = None
         self.filename = None
+        self.hash = None
+        self.shorthash = None
 
     def save(self, filename):
         embedding_data = {
@@ -81,6 +84,10 @@ class Embedding:
         self.cached_checksum = f'{const_hash(self.vec.reshape(-1) * 100) & 0xffff:04x}'
         return self.cached_checksum
 
+    def set_hash(self, v):
+        self.hash = v
+        self.shorthash = self.hash[0:12]
+
 
 class DirWithTextualInversionEmbeddings:
     def __init__(self, path):
@@ -174,30 +181,40 @@ class EmbeddingDatabase:
         else:
             return
 
+
         # textual inversion embeddings
         if 'string_to_param' in data:
             param_dict = data['string_to_param']
             param_dict = getattr(param_dict, '_parameters', param_dict)  # fix for torch 1.12.1 loading saved file from torch 1.11
             assert len(param_dict) == 1, 'embedding file has multiple terms in it'
             emb = next(iter(param_dict.items()))[1]
-        # diffuser concepts
+            vec = emb.detach().to(devices.device, dtype=torch.float32)
-        elif type(data) == dict and type(next(iter(data.values()))) == torch.Tensor:
+            shape = vec.shape[-1]
+            vectors = vec.shape[0]
+        elif type(data) == dict and 'clip_g' in data and 'clip_l' in data:  # SDXL embedding
+            vec = {k: v.detach().to(devices.device, dtype=torch.float32) for k, v in data.items()}
+            shape = data['clip_g'].shape[-1] + data['clip_l'].shape[-1]
+            vectors = data['clip_g'].shape[0]
+        elif type(data) == dict and type(next(iter(data.values()))) == torch.Tensor: # diffuser concepts
             assert len(data.keys()) == 1, 'embedding file has multiple terms in it'
 
             emb = next(iter(data.values()))
             if len(emb.shape) == 1:
                 emb = emb.unsqueeze(0)
+            vec = emb.detach().to(devices.device, dtype=torch.float32)
+            shape = vec.shape[-1]
+            vectors = vec.shape[0]
         else:
             raise Exception(f"Couldn't identify {filename} as neither textual inversion embedding nor diffuser concept.")
 
-        vec = emb.detach().to(devices.device, dtype=torch.float32)
         embedding = Embedding(vec, name)
         embedding.step = data.get('step', None)
         embedding.sd_checkpoint = data.get('sd_checkpoint', None)
         embedding.sd_checkpoint_name = data.get('sd_checkpoint_name', None)
-        embedding.vectors = vec.shape[0]
+        embedding.vectors = vectors
-        embedding.shape = vec.shape[-1]
+        embedding.shape = shape
         embedding.filename = path
+        embedding.set_hash(hashes.sha256(embedding.filename, "textual_inversion/" + name) or '')
 
         if self.expected_shape == -1 or self.expected_shape == embedding.shape:
             self.register_embedding(embedding, shared.sd_model)
@@ -248,7 +265,7 @@ class EmbeddingDatabase:
         self.word_embeddings.update(sorted_word_embeddings)
 
         displayed_embeddings = (tuple(self.word_embeddings.keys()), tuple(self.skipped_embeddings.keys()))
-        if self.previously_displayed_embeddings != displayed_embeddings:
+        if shared.opts.textual_inversion_print_at_load and self.previously_displayed_embeddings != displayed_embeddings:
             self.previously_displayed_embeddings = displayed_embeddings
             print(f"Textual inversion embeddings loaded({len(self.word_embeddings)}): {', '.join(self.word_embeddings.keys())}")
             if self.skipped_embeddings:
@@ -370,6 +387,8 @@ def validate_train_inputs(model_name, learn_rate, batch_size, gradient_step, dat
 
 
 def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_step, data_root, log_directory, training_width, training_height, varsize, steps, clip_grad_mode, clip_grad_value, shuffle_tags, tag_drop_out, latent_sampling_method, use_weight, create_image_every, save_embedding_every, template_filename, save_image_with_stored_embedding, preview_from_txt2img, preview_prompt, preview_negative_prompt, preview_steps, preview_sampler_index, preview_cfg_scale, preview_seed, preview_width, preview_height):
+    from modules import processing
+
     save_embedding_every = save_embedding_every or 0
     create_image_every = create_image_every or 0
     template_file = textual_inversion_templates.get(template_filename, None)
@@ -584,6 +603,7 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st
 
                     preview_text = p.prompt
 
+                    with closing(p):
                         processed = processing.process_images(p)
                         image = processed.images[0] if len(processed.images) > 0 else None
 

modules/timer.py

@@ -1,4 +1,5 @@
 import time
+import argparse
 
 
 class TimerSubcategory:
@@ -11,20 +12,27 @@ class TimerSubcategory:
     def __enter__(self):
         self.start = time.time()
         self.timer.base_category = self.original_base_category + self.category + "/"
+        self.timer.subcategory_level += 1
+
+        if self.timer.print_log:
+            print(f"{'  ' * self.timer.subcategory_level}{self.category}:")
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         elapsed_for_subcategroy = time.time() - self.start
         self.timer.base_category = self.original_base_category
         self.timer.add_time_to_record(self.original_base_category + self.category, elapsed_for_subcategroy)
-        self.timer.record(self.category)
+        self.timer.subcategory_level -= 1
+        self.timer.record(self.category, disable_log=True)
 
 
 class Timer:
-    def __init__(self):
+    def __init__(self, print_log=False):
         self.start = time.time()
         self.records = {}
         self.total = 0
         self.base_category = ''
+        self.print_log = print_log
+        self.subcategory_level = 0
 
     def elapsed(self):
         end = time.time()
@@ -38,13 +46,16 @@ class Timer:
 
         self.records[category] += amount
 
-    def record(self, category, extra_time=0):
+    def record(self, category, extra_time=0, disable_log=False):
         e = self.elapsed()
 
         self.add_time_to_record(self.base_category + category, e + extra_time)
 
         self.total += e + extra_time
 
+        if self.print_log and not disable_log:
+            print(f"{'  ' * self.subcategory_level}{category}: done in {e + extra_time:.3f}s")
+
     def subcategory(self, name):
         self.elapsed()
 
@@ -71,6 +82,10 @@ class Timer:
         self.__init__()
 
 
-startup_timer = Timer()
+parser = argparse.ArgumentParser(add_help=False)
+parser.add_argument("--log-startup", action='store_true', help="print a detailed log of what's happening at startup")
+args = parser.parse_known_args()[0]
+
+startup_timer = Timer(print_log=args.log_startup)
 
 startup_record = None

modules/txt2img.py

@@ -1,13 +1,15 @@
+from contextlib import closing
+
 import modules.scripts
 from modules import sd_samplers, processing
 from modules.generation_parameters_copypaste import create_override_settings_dict
 from modules.shared import opts, cmd_opts
 import modules.shared as shared
 from modules.ui import plaintext_to_html
+import gradio as gr
 
 
-
+def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, steps: int, sampler_index: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, enable_hr: bool, denoising_strength: float, hr_scale: float, hr_upscaler: str, hr_second_pass_steps: int, hr_resize_x: int, hr_resize_y: int, hr_checkpoint_name: str, hr_sampler_index: int, hr_prompt: str, hr_negative_prompt, override_settings_texts, request: gr.Request, *args):
-def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, steps: int, sampler_index: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, enable_hr: bool, denoising_strength: float, hr_scale: float, hr_upscaler: str, hr_second_pass_steps: int, hr_resize_x: int, hr_resize_y: int, hr_sampler_index: int, hr_prompt: str, hr_negative_prompt, override_settings_texts, *args):
     override_settings = create_override_settings_dict(override_settings_texts)
 
     p = processing.StableDiffusionProcessingTxt2Img(
@@ -39,6 +41,7 @@ def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, step
         hr_second_pass_steps=hr_second_pass_steps,
         hr_resize_x=hr_resize_x,
         hr_resize_y=hr_resize_y,
+        hr_checkpoint_name=None if hr_checkpoint_name == 'Use same checkpoint' else hr_checkpoint_name,
         hr_sampler_name=sd_samplers.samplers_for_img2img[hr_sampler_index - 1].name if hr_sampler_index != 0 else None,
         hr_prompt=hr_prompt,
         hr_negative_prompt=hr_negative_prompt,
@@ -48,16 +51,17 @@ def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, step
     p.scripts = modules.scripts.scripts_txt2img
     p.script_args = args
 
+    p.user = request.username
+
     if cmd_opts.enable_console_prompts:
         print(f"\ntxt2img: {prompt}", file=shared.progress_print_out)
 
+    with closing(p):
         processed = modules.scripts.scripts_txt2img.run(p, *args)
 
         if processed is None:
             processed = processing.process_images(p)
 
-    p.close()
-
     shared.total_tqdm.clear()
 
     generation_info_js = processed.js()
@@ -67,4 +71,4 @@ def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, step
     if opts.do_not_show_images:
         processed.images = []
 
-    return processed.images, generation_info_js, plaintext_to_html(processed.info), plaintext_to_html(processed.comments)
+    return processed.images, generation_info_js, plaintext_to_html(processed.info), plaintext_to_html(processed.comments, classname="comments")

modules/ui.py

@@ -12,34 +12,30 @@ import numpy as np
 from PIL import Image, PngImagePlugin  # noqa: F401
 from modules.call_queue import wrap_gradio_gpu_call, wrap_queued_call, wrap_gradio_call
 
-from modules import sd_hijack, sd_models, script_callbacks, ui_extensions, deepbooru, sd_vae, extra_networks, ui_common, ui_postprocessing, progress, ui_loadsave, errors, shared_items, ui_settings, timer, sysinfo
+from modules import gradio_extensons  # noqa: F401
+from modules import sd_hijack, sd_models, script_callbacks, ui_extensions, deepbooru, extra_networks, ui_common, ui_postprocessing, progress, ui_loadsave, errors, shared_items, ui_settings, timer, sysinfo, ui_checkpoint_merger, ui_prompt_styles, scripts
 from modules.ui_components import FormRow, FormGroup, ToolButton, FormHTML
 from modules.paths import script_path
 from modules.ui_common import create_refresh_button
 from modules.ui_gradio_extensions import reload_javascript
 
-
 from modules.shared import opts, cmd_opts
 
-import modules.codeformer_model
 import modules.generation_parameters_copypaste as parameters_copypaste
-import modules.gfpgan_model
+import modules.hypernetworks.ui as hypernetworks_ui
-import modules.hypernetworks.ui
+import modules.textual_inversion.ui as textual_inversion_ui
-import modules.scripts
+import modules.textual_inversion.textual_inversion as textual_inversion
 import modules.shared as shared
-import modules.styles
+import modules.images
-import modules.textual_inversion.ui
 from modules import prompt_parser
 from modules.sd_hijack import model_hijack
 from modules.sd_samplers import samplers, samplers_for_img2img
-from modules.textual_inversion import textual_inversion
-import modules.hypernetworks.ui
 from modules.generation_parameters_copypaste import image_from_url_text
-import modules.extras
 
 create_setting_component = ui_settings.create_setting_component
 
 warnings.filterwarnings("default" if opts.show_warnings else "ignore", category=UserWarning)
+warnings.filterwarnings("default" if opts.show_gradio_deprecation_warnings else "ignore", category=gr.deprecation.GradioDeprecationWarning)
 
 # this is a fix for Windows users. Without it, javascript files will be served with text/html content-type and the browser will not show any UI
 mimetypes.init()
@@ -83,8 +79,7 @@ detect_image_size_symbol = '\U0001F4D0'  # 📐
 up_down_symbol = '\u2195\ufe0f' # ↕️
 
 
-def plaintext_to_html(text):
+plaintext_to_html = ui_common.plaintext_to_html
-    return ui_common.plaintext_to_html(text)
 
 
 def send_gradio_gallery_to_image(x):
@@ -93,19 +88,6 @@ def send_gradio_gallery_to_image(x):
     return image_from_url_text(x[0])
 
 
-def add_style(name: str, prompt: str, negative_prompt: str):
-    if name is None:
-        return [gr_show() for x in range(4)]
-
-    style = modules.styles.PromptStyle(name, prompt, negative_prompt)
-    shared.prompt_styles.styles[style.name] = style
-    # Save all loaded prompt styles: this allows us to update the storage format in the future more easily, because we
-    # reserialize all styles every time we save them
-    shared.prompt_styles.save_styles(shared.styles_filename)
-
-    return [gr.Dropdown.update(visible=True, choices=list(shared.prompt_styles.styles)) for _ in range(2)]
-
-
 def calc_resolution_hires(enable, width, height, hr_scale, hr_resize_x, hr_resize_y):
     from modules import processing, devices
 
@@ -130,13 +112,6 @@ def resize_from_to_html(width, height, scale_by):
     return f"resize: from <span class='resolution'>{width}x{height}</span> to <span class='resolution'>{target_width}x{target_height}</span>"
 
 
-def apply_styles(prompt, prompt_neg, styles):
-    prompt = shared.prompt_styles.apply_styles_to_prompt(prompt, styles)
-    prompt_neg = shared.prompt_styles.apply_negative_styles_to_prompt(prompt_neg, styles)
-
-    return [gr.Textbox.update(value=prompt), gr.Textbox.update(value=prompt_neg), gr.Dropdown.update(value=[])]
-
-
 def process_interrogate(interrogation_function, mode, ii_input_dir, ii_output_dir, *ii_singles):
     if mode in {0, 1, 3, 4}:
         return [interrogation_function(ii_singles[mode]), None]
@@ -155,7 +130,7 @@ def process_interrogate(interrogation_function, mode, ii_input_dir, ii_output_di
             img = Image.open(image)
             filename = os.path.basename(image)
             left, _ = os.path.splitext(filename)
-            print(interrogation_function(img), file=open(os.path.join(ii_output_dir, f"{left}.txt"), 'a'))
+            print(interrogation_function(img), file=open(os.path.join(ii_output_dir, f"{left}.txt"), 'a', encoding='utf-8'))
 
         return [gr.update(), None]
 
@@ -173,7 +148,6 @@ def interrogate_deepbooru(image):
 def create_seed_inputs(target_interface):
     with FormRow(elem_id=f"{target_interface}_seed_row", variant="compact"):
         seed = (gr.Textbox if cmd_opts.use_textbox_seed else gr.Number)(label='Seed', value=-1, elem_id=f"{target_interface}_seed")
-        seed.style(container=False)
         random_seed = ToolButton(random_symbol, elem_id=f"{target_interface}_random_seed", label='Random seed')
         reuse_seed = ToolButton(reuse_symbol, elem_id=f"{target_interface}_reuse_seed", label='Reuse seed')
 
@@ -185,7 +159,6 @@ def create_seed_inputs(target_interface):
     with FormRow(visible=False, elem_id=f"{target_interface}_subseed_row") as seed_extra_row_1:
         seed_extras.append(seed_extra_row_1)
         subseed = gr.Number(label='Variation seed', value=-1, elem_id=f"{target_interface}_subseed")
-        subseed.style(container=False)
         random_subseed = ToolButton(random_symbol, elem_id=f"{target_interface}_random_subseed")
         reuse_subseed = ToolButton(reuse_symbol, elem_id=f"{target_interface}_reuse_subseed")
         subseed_strength = gr.Slider(label='Variation strength', value=0.0, minimum=0, maximum=1, step=0.01, elem_id=f"{target_interface}_subseed_strength")
@@ -268,71 +241,76 @@ def update_token_counter(text, steps):
     return f"<span class='gr-box gr-text-input'>{token_count}/{max_length}</span>"
 
 
-def create_toprow(is_img2img):
+class Toprow:
+    """Creates a top row UI with prompts, generate button, styles, extra little buttons for things, and enables some functionality related to their operation"""
+
+    def __init__(self, is_img2img):
         id_part = "img2img" if is_img2img else "txt2img"
+        self.id_part = id_part
 
         with gr.Row(elem_id=f"{id_part}_toprow", variant="compact"):
             with gr.Column(elem_id=f"{id_part}_prompt_container", scale=6):
                 with gr.Row():
                     with gr.Column(scale=80):
                         with gr.Row():
-                        prompt = gr.Textbox(label="Prompt", elem_id=f"{id_part}_prompt", show_label=False, lines=3, placeholder="Prompt (press Ctrl+Enter or Alt+Enter to generate)", elem_classes=["prompt"])
+                            self.prompt = gr.Textbox(label="Prompt", elem_id=f"{id_part}_prompt", show_label=False, lines=3, placeholder="Prompt (press Ctrl+Enter or Alt+Enter to generate)", elem_classes=["prompt"])
+                            self.prompt_img = gr.File(label="", elem_id=f"{id_part}_prompt_image", file_count="single", type="binary", visible=False)
 
                 with gr.Row():
                     with gr.Column(scale=80):
                         with gr.Row():
-                        negative_prompt = gr.Textbox(label="Negative prompt", elem_id=f"{id_part}_neg_prompt", show_label=False, lines=3, placeholder="Negative prompt (press Ctrl+Enter or Alt+Enter to generate)", elem_classes=["prompt"])
+                            self.negative_prompt = gr.Textbox(label="Negative prompt", elem_id=f"{id_part}_neg_prompt", show_label=False, lines=3, placeholder="Negative prompt (press Ctrl+Enter or Alt+Enter to generate)", elem_classes=["prompt"])
 
-        button_interrogate = None
+            self.button_interrogate = None
-        button_deepbooru = None
+            self.button_deepbooru = None
             if is_img2img:
                 with gr.Column(scale=1, elem_classes="interrogate-col"):
-                button_interrogate = gr.Button('Interrogate\nCLIP', elem_id="interrogate")
+                    self.button_interrogate = gr.Button('Interrogate\nCLIP', elem_id="interrogate")
-                button_deepbooru = gr.Button('Interrogate\nDeepBooru', elem_id="deepbooru")
+                    self.button_deepbooru = gr.Button('Interrogate\nDeepBooru', elem_id="deepbooru")
 
             with gr.Column(scale=1, elem_id=f"{id_part}_actions_column"):
                 with gr.Row(elem_id=f"{id_part}_generate_box", elem_classes="generate-box"):
-                interrupt = gr.Button('Interrupt', elem_id=f"{id_part}_interrupt", elem_classes="generate-box-interrupt")
+                    self.interrupt = gr.Button('Interrupt', elem_id=f"{id_part}_interrupt", elem_classes="generate-box-interrupt")
-                skip = gr.Button('Skip', elem_id=f"{id_part}_skip", elem_classes="generate-box-skip")
+                    self.skip = gr.Button('Skip', elem_id=f"{id_part}_skip", elem_classes="generate-box-skip")
-                submit = gr.Button('Generate', elem_id=f"{id_part}_generate", variant='primary')
+                    self.submit = gr.Button('Generate', elem_id=f"{id_part}_generate", variant='primary')
 
-                skip.click(
+                    self.skip.click(
                         fn=lambda: shared.state.skip(),
                         inputs=[],
                         outputs=[],
                     )
 
-                interrupt.click(
+                    self.interrupt.click(
                         fn=lambda: shared.state.interrupt(),
                         inputs=[],
                         outputs=[],
                     )
 
                 with gr.Row(elem_id=f"{id_part}_tools"):
-                paste = ToolButton(value=paste_symbol, elem_id="paste")
+                    self.paste = ToolButton(value=paste_symbol, elem_id="paste")
-                clear_prompt_button = ToolButton(value=clear_prompt_symbol, elem_id=f"{id_part}_clear_prompt")
+                    self.clear_prompt_button = ToolButton(value=clear_prompt_symbol, elem_id=f"{id_part}_clear_prompt")
-                extra_networks_button = ToolButton(value=extra_networks_symbol, elem_id=f"{id_part}_extra_networks")
+                    self.restore_progress_button = ToolButton(value=restore_progress_symbol, elem_id=f"{id_part}_restore_progress", visible=False)
-                prompt_style_apply = ToolButton(value=apply_style_symbol, elem_id=f"{id_part}_style_apply")
-                save_style = ToolButton(value=save_style_symbol, elem_id=f"{id_part}_style_create")
-                restore_progress_button = ToolButton(value=restore_progress_symbol, elem_id=f"{id_part}_restore_progress", visible=False)
 
-                token_counter = gr.HTML(value="<span>0/75</span>", elem_id=f"{id_part}_token_counter", elem_classes=["token-counter"])
+                    self.token_counter = gr.HTML(value="<span>0/75</span>", elem_id=f"{id_part}_token_counter", elem_classes=["token-counter"])
-                token_button = gr.Button(visible=False, elem_id=f"{id_part}_token_button")
+                    self.token_button = gr.Button(visible=False, elem_id=f"{id_part}_token_button")
-                negative_token_counter = gr.HTML(value="<span>0/75</span>", elem_id=f"{id_part}_negative_token_counter", elem_classes=["token-counter"])
+                    self.negative_token_counter = gr.HTML(value="<span>0/75</span>", elem_id=f"{id_part}_negative_token_counter", elem_classes=["token-counter"])
-                negative_token_button = gr.Button(visible=False, elem_id=f"{id_part}_negative_token_button")
+                    self.negative_token_button = gr.Button(visible=False, elem_id=f"{id_part}_negative_token_button")
 
-                clear_prompt_button.click(
+                    self.clear_prompt_button.click(
                         fn=lambda *x: x,
                         _js="confirm_clear_prompt",
-                    inputs=[prompt, negative_prompt],
+                        inputs=[self.prompt, self.negative_prompt],
-                    outputs=[prompt, negative_prompt],
+                        outputs=[self.prompt, self.negative_prompt],
                     )
 
-            with gr.Row(elem_id=f"{id_part}_styles_row"):
+                self.ui_styles = ui_prompt_styles.UiPromptStyles(id_part, self.prompt, self.negative_prompt)
-                prompt_styles = gr.Dropdown(label="Styles", elem_id=f"{id_part}_styles", choices=[k for k, v in shared.prompt_styles.styles.items()], value=[], multiselect=True)
-                create_refresh_button(prompt_styles, shared.prompt_styles.reload, lambda: {"choices": [k for k, v in shared.prompt_styles.styles.items()]}, f"refresh_{id_part}_styles")
 
-    return prompt, prompt_styles, negative_prompt, submit, button_interrogate, button_deepbooru, prompt_style_apply, save_style, paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button, restore_progress_button
+        self.prompt_img.change(
+            fn=modules.images.image_data,
+            inputs=[self.prompt_img],
+            outputs=[self.prompt, self.prompt_img],
+            show_progress=False,
+        )
 
 
 def setup_progressbar(*args, **kwargs):
@@ -416,22 +394,20 @@ def create_ui():
 
     parameters_copypaste.reset()
 
-    modules.scripts.scripts_current = modules.scripts.scripts_txt2img
+    scripts.scripts_current = scripts.scripts_txt2img
-    modules.scripts.scripts_txt2img.initialize_scripts(is_img2img=False)
+    scripts.scripts_txt2img.initialize_scripts(is_img2img=False)
 
     with gr.Blocks(analytics_enabled=False) as txt2img_interface:
-        txt2img_prompt, txt2img_prompt_styles, txt2img_negative_prompt, submit, _, _, txt2img_prompt_style_apply, txt2img_save_style, txt2img_paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button, restore_progress_button = create_toprow(is_img2img=False)
+        toprow = Toprow(is_img2img=False)
 
         dummy_component = gr.Label(visible=False)
-        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:
+        extra_tabs = gr.Tabs(elem_id="txt2img_extra_tabs")
-            from modules import ui_extra_networks
+        extra_tabs.__enter__()
-            extra_networks_ui = ui_extra_networks.create_ui(extra_networks, extra_networks_button, 'txt2img')
 
-        with gr.Row().style(equal_height=False):
+        with gr.Tab("Generation", id="txt2img_generation") as txt2img_generation_tab, gr.Row().style(equal_height=False):
             with gr.Column(variant='compact', elem_id="txt2img_settings"):
-                modules.scripts.scripts_txt2img.prepare_ui()
+                scripts.scripts_txt2img.prepare_ui()
 
                 for category in ordered_ui_categories():
                     if category == "sampler":
@@ -477,6 +453,10 @@ def create_ui():
                                 hr_resize_y = gr.Slider(minimum=0, maximum=2048, step=8, label="Resize height to", value=0, elem_id="txt2img_hr_resize_y")
 
                             with FormRow(elem_id="txt2img_hires_fix_row3", variant="compact", visible=opts.hires_fix_show_sampler) as hr_sampler_container:
+
+                                hr_checkpoint_name = gr.Dropdown(label='Hires checkpoint', elem_id="hr_checkpoint", choices=["Use same checkpoint"] + modules.sd_models.checkpoint_tiles(use_short=True), value="Use same checkpoint")
+                                create_refresh_button(hr_checkpoint_name, modules.sd_models.list_models, lambda: {"choices": ["Use same checkpoint"] + modules.sd_models.checkpoint_tiles(use_short=True)}, "hr_checkpoint_refresh")
+
                                 hr_sampler_index = gr.Dropdown(label='Hires sampling method', elem_id="hr_sampler", choices=["Use same sampler"] + [x.name for x in samplers_for_img2img], value="Use same sampler", type="index")
 
                             with FormRow(elem_id="txt2img_hires_fix_row4", variant="compact", visible=opts.hires_fix_show_prompts) as hr_prompts_container:
@@ -499,10 +479,10 @@ def create_ui():
 
                     elif category == "scripts":
                         with FormGroup(elem_id="txt2img_script_container"):
-                            custom_inputs = modules.scripts.scripts_txt2img.setup_ui()
+                            custom_inputs = scripts.scripts_txt2img.setup_ui()
 
                     else:
-                        modules.scripts.scripts_txt2img.setup_ui_for_section(category)
+                        scripts.scripts_txt2img.setup_ui_for_section(category)
 
             hr_resolution_preview_inputs = [enable_hr, width, height, hr_scale, hr_resize_x, hr_resize_y]
 
@@ -533,9 +513,9 @@ def create_ui():
                 _js="submit",
                 inputs=[
                     dummy_component,
-                    txt2img_prompt,
+                    toprow.prompt,
-                    txt2img_negative_prompt,
+                    toprow.negative_prompt,
-                    txt2img_prompt_styles,
+                    toprow.ui_styles.dropdown,
                     steps,
                     sampler_index,
                     restore_faces,
@@ -554,6 +534,7 @@ def create_ui():
                     hr_second_pass_steps,
                     hr_resize_x,
                     hr_resize_y,
+                    hr_checkpoint_name,
                     hr_sampler_index,
                     hr_prompt,
                     hr_negative_prompt,
@@ -570,12 +551,12 @@ def create_ui():
                 show_progress=False,
             )
 
-            txt2img_prompt.submit(**txt2img_args)
+            toprow.prompt.submit(**txt2img_args)
-            submit.click(**txt2img_args)
+            toprow.submit.click(**txt2img_args)
 
             res_switch_btn.click(fn=None, _js="function(){switchWidthHeight('txt2img')}", inputs=None, outputs=None, show_progress=False)
 
-            restore_progress_button.click(
+            toprow.restore_progress_button.click(
                 fn=progress.restore_progress,
                 _js="restoreProgressTxt2img",
                 inputs=[dummy_component],
@@ -588,18 +569,6 @@ def create_ui():
                 show_progress=False,
             )
 
-            txt_prompt_img.change(
-                fn=modules.images.image_data,
-                inputs=[
-                    txt_prompt_img
-                ],
-                outputs=[
-                    txt2img_prompt,
-                    txt_prompt_img
-                ],
-                show_progress=False,
-            )
-
             enable_hr.change(
                 fn=lambda x: gr_show(x),
                 inputs=[enable_hr],
@@ -608,8 +577,8 @@ def create_ui():
             )
 
             txt2img_paste_fields = [
-                (txt2img_prompt, "Prompt"),
+                (toprow.prompt, "Prompt"),
-                (txt2img_negative_prompt, "Negative prompt"),
+                (toprow.negative_prompt, "Negative prompt"),
                 (steps, "Steps"),
                 (sampler_index, "Sampler"),
                 (restore_faces, "Face restoration"),
@@ -618,34 +587,36 @@ def create_ui():
                 (width, "Size-1"),
                 (height, "Size-2"),
                 (batch_size, "Batch size"),
+                (seed_checkbox, lambda d: "Variation seed" in d or "Seed resize from-1" in d),
                 (subseed, "Variation seed"),
                 (subseed_strength, "Variation seed strength"),
                 (seed_resize_from_w, "Seed resize from-1"),
                 (seed_resize_from_h, "Seed resize from-2"),
-                (txt2img_prompt_styles, lambda d: d["Styles array"] if isinstance(d.get("Styles array"), list) else gr.update()),
+                (toprow.ui_styles.dropdown, lambda d: d["Styles array"] if isinstance(d.get("Styles array"), list) else gr.update()),
                 (denoising_strength, "Denoising strength"),
-                (enable_hr, lambda d: "Denoising strength" in d),
+                (enable_hr, lambda d: "Denoising strength" in d and ("Hires upscale" in d or "Hires upscaler" in d or "Hires resize-1" in d)),
-                (hr_options, lambda d: gr.Row.update(visible="Denoising strength" in d)),
+                (hr_options, lambda d: gr.Row.update(visible="Denoising strength" in d and ("Hires upscale" in d or "Hires upscaler" in d or "Hires resize-1" in d))),
                 (hr_scale, "Hires upscale"),
                 (hr_upscaler, "Hires upscaler"),
                 (hr_second_pass_steps, "Hires steps"),
                 (hr_resize_x, "Hires resize-1"),
                 (hr_resize_y, "Hires resize-2"),
+                (hr_checkpoint_name, "Hires checkpoint"),
                 (hr_sampler_index, "Hires sampler"),
-                (hr_sampler_container, lambda d: gr.update(visible=True) if d.get("Hires sampler", "Use same sampler") != "Use same sampler" else gr.update()),
+                (hr_sampler_container, lambda d: gr.update(visible=True) if d.get("Hires sampler", "Use same sampler") != "Use same sampler" or d.get("Hires checkpoint", "Use same checkpoint") != "Use same checkpoint" else gr.update()),
                 (hr_prompt, "Hires prompt"),
                 (hr_negative_prompt, "Hires negative prompt"),
                 (hr_prompts_container, lambda d: gr.update(visible=True) if d.get("Hires prompt", "") != "" or d.get("Hires negative prompt", "") != "" else gr.update()),
-                *modules.scripts.scripts_txt2img.infotext_fields
+                *scripts.scripts_txt2img.infotext_fields
             ]
             parameters_copypaste.add_paste_fields("txt2img", None, txt2img_paste_fields, override_settings)
             parameters_copypaste.register_paste_params_button(parameters_copypaste.ParamBinding(
-                paste_button=txt2img_paste, tabname="txt2img", source_text_component=txt2img_prompt, source_image_component=None,
+                paste_button=toprow.paste, tabname="txt2img", source_text_component=toprow.prompt, source_image_component=None,
             ))
 
             txt2img_preview_params = [
-                txt2img_prompt,
+                toprow.prompt,
-                txt2img_negative_prompt,
+                toprow.negative_prompt,
                 steps,
                 sampler_index,
                 cfg_scale,
@@ -654,24 +625,25 @@ def create_ui():
                 height,
             ]
 
-            token_button.click(fn=wrap_queued_call(update_token_counter), inputs=[txt2img_prompt, steps], outputs=[token_counter])
+            toprow.token_button.click(fn=wrap_queued_call(update_token_counter), inputs=[toprow.prompt, steps], outputs=[toprow.token_counter])
-            negative_token_button.click(fn=wrap_queued_call(update_token_counter), inputs=[txt2img_negative_prompt, steps], outputs=[negative_token_counter])
+            toprow.negative_token_button.click(fn=wrap_queued_call(update_token_counter), inputs=[toprow.negative_prompt, steps], outputs=[toprow.negative_token_counter])
 
+        from modules import ui_extra_networks
+        extra_networks_ui = ui_extra_networks.create_ui(txt2img_interface, [txt2img_generation_tab], 'txt2img')
         ui_extra_networks.setup_ui(extra_networks_ui, txt2img_gallery)
 
-    modules.scripts.scripts_current = modules.scripts.scripts_img2img
+        extra_tabs.__exit__()
-    modules.scripts.scripts_img2img.initialize_scripts(is_img2img=True)
+
+    scripts.scripts_current = scripts.scripts_img2img
+    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, restore_progress_button = create_toprow(is_img2img=True)
+        toprow = Toprow(is_img2img=True)
 
-        img2img_prompt_img = gr.File(label="", elem_id="img2img_prompt_image", file_count="single", type="binary", visible=False)
+        extra_tabs = gr.Tabs(elem_id="img2img_extra_tabs")
+        extra_tabs.__enter__()
 
-        with FormRow(variant='compact', elem_id="img2img_extra_networks", visible=False) as extra_networks:
+        with gr.Tab("Generation", id="img2img_generation") as img2img_generation_tab, FormRow().style(equal_height=False):
-            from modules import ui_extra_networks
-            extra_networks_ui_img2img = ui_extra_networks.create_ui(extra_networks, extra_networks_button, 'img2img')
-
-        with FormRow().style(equal_height=False):
             with gr.Column(variant='compact', elem_id="img2img_settings"):
                 copy_image_buttons = []
                 copy_image_destinations = {}
@@ -693,19 +665,19 @@ def create_ui():
                     img2img_selected_tab = gr.State(0)
 
                     with gr.TabItem('img2img', id='img2img', elem_id="img2img_img2img_tab") as tab_img2img:
-                        init_img = gr.Image(label="Image for img2img", elem_id="img2img_image", show_label=False, source="upload", interactive=True, type="pil", tool="editor", image_mode="RGBA").style(height=opts.img2img_editor_height)
+                        init_img = gr.Image(label="Image for img2img", elem_id="img2img_image", show_label=False, source="upload", interactive=True, type="pil", tool="editor", image_mode="RGBA", height=opts.img2img_editor_height)
                         add_copy_image_controls('img2img', init_img)
 
                     with gr.TabItem('Sketch', id='img2img_sketch', elem_id="img2img_img2img_sketch_tab") as tab_sketch:
-                        sketch = gr.Image(label="Image for img2img", elem_id="img2img_sketch", show_label=False, source="upload", interactive=True, type="pil", tool="color-sketch", image_mode="RGBA").style(height=opts.img2img_editor_height)
+                        sketch = gr.Image(label="Image for img2img", elem_id="img2img_sketch", show_label=False, source="upload", interactive=True, type="pil", tool="color-sketch", image_mode="RGBA", height=opts.img2img_editor_height, brush_color=opts.img2img_sketch_default_brush_color)
                         add_copy_image_controls('sketch', sketch)
 
                     with gr.TabItem('Inpaint', id='inpaint', elem_id="img2img_inpaint_tab") as tab_inpaint:
-                        init_img_with_mask = gr.Image(label="Image for inpainting with mask", show_label=False, elem_id="img2maskimg", source="upload", interactive=True, type="pil", tool="sketch", image_mode="RGBA").style(height=opts.img2img_editor_height)
+                        init_img_with_mask = gr.Image(label="Image for inpainting with mask", show_label=False, elem_id="img2maskimg", source="upload", interactive=True, type="pil", tool="sketch", image_mode="RGBA", height=opts.img2img_editor_height, brush_color=opts.img2img_inpaint_mask_brush_color)
                         add_copy_image_controls('inpaint', init_img_with_mask)
 
                     with gr.TabItem('Inpaint sketch', id='inpaint_sketch', elem_id="img2img_inpaint_sketch_tab") as tab_inpaint_color:
-                        inpaint_color_sketch = gr.Image(label="Color sketch inpainting", show_label=False, elem_id="inpaint_sketch", source="upload", interactive=True, type="pil", tool="color-sketch", image_mode="RGBA").style(height=opts.img2img_editor_height)
+                        inpaint_color_sketch = gr.Image(label="Color sketch inpainting", show_label=False, elem_id="inpaint_sketch", source="upload", interactive=True, type="pil", tool="color-sketch", image_mode="RGBA", height=opts.img2img_editor_height, brush_color=opts.img2img_inpaint_sketch_default_brush_color)
                         inpaint_color_sketch_orig = gr.State(None)
                         add_copy_image_controls('inpaint_sketch', inpaint_color_sketch)
 
@@ -733,6 +705,10 @@ def create_ui():
                         img2img_batch_input_dir = gr.Textbox(label="Input directory", **shared.hide_dirs, elem_id="img2img_batch_input_dir")
                         img2img_batch_output_dir = gr.Textbox(label="Output directory", **shared.hide_dirs, elem_id="img2img_batch_output_dir")
                         img2img_batch_inpaint_mask_dir = gr.Textbox(label="Inpaint batch mask directory (required for inpaint batch processing only)", **shared.hide_dirs, elem_id="img2img_batch_inpaint_mask_dir")
+                        with gr.Accordion("PNG info", open=False):
+                            img2img_batch_use_png_info = gr.Checkbox(label="Append png info to prompts", **shared.hide_dirs, elem_id="img2img_batch_use_png_info")
+                            img2img_batch_png_info_dir = gr.Textbox(label="PNG info directory", **shared.hide_dirs, placeholder="Leave empty to use input directory", elem_id="img2img_batch_png_info_dir")
+                            img2img_batch_png_info_props = gr.CheckboxGroup(["Prompt", "Negative prompt", "Seed", "CFG scale", "Sampler", "Steps"], label="Parameters to take from png info", info="Prompts from png info will be appended to prompts set in ui.")
 
                     img2img_tabs = [tab_img2img, tab_sketch, tab_inpaint, tab_inpaint_color, tab_inpaint_upload, tab_batch]
 
@@ -761,7 +737,7 @@ def create_ui():
                 with FormRow():
                     resize_mode = gr.Radio(label="Resize mode", elem_id="resize_mode", choices=["Just resize", "Crop and resize", "Resize and fill", "Just resize (latent upscale)"], type="index", value="Just resize")
 
-                modules.scripts.scripts_img2img.prepare_ui()
+                scripts.scripts_img2img.prepare_ui()
 
                 for category in ordered_ui_categories():
                     if category == "sampler":
@@ -773,7 +749,7 @@ def create_ui():
                                 selected_scale_tab = gr.State(value=0)
 
                                 with gr.Tabs():
-                                    with gr.Tab(label="Resize to") as tab_scale_to:
+                                    with gr.Tab(label="Resize to", elem_id="img2img_tab_resize_to") as tab_scale_to:
                                         with FormRow():
                                             with gr.Column(elem_id="img2img_column_size", scale=4):
                                                 width = gr.Slider(minimum=64, maximum=2048, step=8, label="Width", value=512, elem_id="img2img_width")
@@ -782,7 +758,7 @@ def create_ui():
                                                 res_switch_btn = ToolButton(value=switch_values_symbol, elem_id="img2img_res_switch_btn")
                                                 detect_image_size_btn = ToolButton(value=detect_image_size_symbol, elem_id="img2img_detect_image_size_btn")
 
-                                    with gr.Tab(label="Resize by") as tab_scale_by:
+                                    with gr.Tab(label="Resize by", elem_id="img2img_tab_resize_by") as tab_scale_by:
                                         scale_by = gr.Slider(minimum=0.05, maximum=4.0, step=0.05, label="Scale", value=1.0, elem_id="img2img_scale")
 
                                         with FormRow():
@@ -842,7 +818,7 @@ def create_ui():
 
                     elif category == "scripts":
                         with FormGroup(elem_id="img2img_script_container"):
-                            custom_inputs = modules.scripts.scripts_img2img.setup_ui()
+                            custom_inputs = scripts.scripts_img2img.setup_ui()
 
                     elif category == "inpaint":
                         with FormGroup(elem_id="inpaint_controls", visible=False) as inpaint_controls:
@@ -873,34 +849,22 @@ def create_ui():
                                     outputs=[inpaint_controls, mask_alpha],
                                 )
                     else:
-                        modules.scripts.scripts_img2img.setup_ui_for_section(category)
+                        scripts.scripts_img2img.setup_ui_for_section(category)
 
             img2img_gallery, generation_info, html_info, html_log = create_output_panel("img2img", opts.outdir_img2img_samples)
 
             connect_reuse_seed(seed, reuse_seed, generation_info, dummy_component, is_subseed=False)
             connect_reuse_seed(subseed, reuse_subseed, generation_info, dummy_component, is_subseed=True)
 
-            img2img_prompt_img.change(
-                fn=modules.images.image_data,
-                inputs=[
-                    img2img_prompt_img
-                ],
-                outputs=[
-                    img2img_prompt,
-                    img2img_prompt_img
-                ],
-                show_progress=False,
-            )
-
             img2img_args = dict(
                 fn=wrap_gradio_gpu_call(modules.img2img.img2img, extra_outputs=[None, '', '']),
                 _js="submit_img2img",
                 inputs=[
                     dummy_component,
                     dummy_component,
-                    img2img_prompt,
+                    toprow.prompt,
-                    img2img_negative_prompt,
+                    toprow.negative_prompt,
-                    img2img_prompt_styles,
+                    toprow.ui_styles.dropdown,
                     init_img,
                     sketch,
                     init_img_with_mask,
@@ -934,6 +898,9 @@ def create_ui():
                     img2img_batch_output_dir,
                     img2img_batch_inpaint_mask_dir,
                     override_settings,
+                    img2img_batch_use_png_info,
+                    img2img_batch_png_info_props,
+                    img2img_batch_png_info_dir,
                 ] + custom_inputs,
                 outputs=[
                     img2img_gallery,
@@ -956,11 +923,11 @@ def create_ui():
                     inpaint_color_sketch,
                     init_img_inpaint,
                 ],
-                outputs=[img2img_prompt, dummy_component],
+                outputs=[toprow.prompt, dummy_component],
             )
 
-            img2img_prompt.submit(**img2img_args)
+            toprow.prompt.submit(**img2img_args)
-            submit.click(**img2img_args)
+            toprow.submit.click(**img2img_args)
 
             res_switch_btn.click(fn=None, _js="function(){switchWidthHeight('img2img')}", inputs=None, outputs=None, show_progress=False)
 
@@ -972,7 +939,7 @@ def create_ui():
                 show_progress=False,
             )
 
-            restore_progress_button.click(
+            toprow.restore_progress_button.click(
                 fn=progress.restore_progress,
                 _js="restoreProgressImg2img",
                 inputs=[dummy_component],
@@ -985,46 +952,22 @@ def create_ui():
                 show_progress=False,
             )
 
-            img2img_interrogate.click(
+            toprow.button_interrogate.click(
                 fn=lambda *args: process_interrogate(interrogate, *args),
                 **interrogate_args,
             )
 
-            img2img_deepbooru.click(
+            toprow.button_deepbooru.click(
                 fn=lambda *args: process_interrogate(interrogate_deepbooru, *args),
                 **interrogate_args,
             )
 
-            prompts = [(txt2img_prompt, txt2img_negative_prompt), (img2img_prompt, img2img_negative_prompt)]
+            toprow.token_button.click(fn=update_token_counter, inputs=[toprow.prompt, steps], outputs=[toprow.token_counter])
-            style_dropdowns = [txt2img_prompt_styles, img2img_prompt_styles]
+            toprow.negative_token_button.click(fn=wrap_queued_call(update_token_counter), inputs=[toprow.negative_prompt, steps], outputs=[toprow.negative_token_counter])
-            style_js_funcs = ["update_txt2img_tokens", "update_img2img_tokens"]
-
-            for button, (prompt, negative_prompt) in zip([txt2img_save_style, img2img_save_style], prompts):
-                button.click(
-                    fn=add_style,
-                    _js="ask_for_style_name",
-                    # Have to pass empty dummy component here, because the JavaScript and Python function have to accept
-                    # the same number of parameters, but we only know the style-name after the JavaScript prompt
-                    inputs=[dummy_component, prompt, negative_prompt],
-                    outputs=[txt2img_prompt_styles, img2img_prompt_styles],
-                )
-
-            for button, (prompt, negative_prompt), styles, js_func in zip([txt2img_prompt_style_apply, img2img_prompt_style_apply], prompts, style_dropdowns, style_js_funcs):
-                button.click(
-                    fn=apply_styles,
-                    _js=js_func,
-                    inputs=[prompt, negative_prompt, styles],
-                    outputs=[prompt, negative_prompt, styles],
-                )
-
-            token_button.click(fn=update_token_counter, inputs=[img2img_prompt, steps], outputs=[token_counter])
-            negative_token_button.click(fn=wrap_queued_call(update_token_counter), inputs=[img2img_negative_prompt, steps], outputs=[negative_token_counter])
-
-            ui_extra_networks.setup_ui(extra_networks_ui_img2img, img2img_gallery)
 
             img2img_paste_fields = [
-                (img2img_prompt, "Prompt"),
+                (toprow.prompt, "Prompt"),
-                (img2img_negative_prompt, "Negative prompt"),
+                (toprow.negative_prompt, "Negative prompt"),
                 (steps, "Steps"),
                 (sampler_index, "Sampler"),
                 (restore_faces, "Face restoration"),
@@ -1034,28 +977,35 @@ def create_ui():
                 (width, "Size-1"),
                 (height, "Size-2"),
                 (batch_size, "Batch size"),
+                (seed_checkbox, lambda d: "Variation seed" in d or "Seed resize from-1" in d),
                 (subseed, "Variation seed"),
                 (subseed_strength, "Variation seed strength"),
                 (seed_resize_from_w, "Seed resize from-1"),
                 (seed_resize_from_h, "Seed resize from-2"),
-                (img2img_prompt_styles, lambda d: d["Styles array"] if isinstance(d.get("Styles array"), list) else gr.update()),
+                (toprow.ui_styles.dropdown, lambda d: d["Styles array"] if isinstance(d.get("Styles array"), list) else gr.update()),
                 (denoising_strength, "Denoising strength"),
                 (mask_blur, "Mask blur"),
-                *modules.scripts.scripts_img2img.infotext_fields
+                *scripts.scripts_img2img.infotext_fields
             ]
             parameters_copypaste.add_paste_fields("img2img", init_img, img2img_paste_fields, override_settings)
             parameters_copypaste.add_paste_fields("inpaint", init_img_with_mask, img2img_paste_fields, override_settings)
             parameters_copypaste.register_paste_params_button(parameters_copypaste.ParamBinding(
-                paste_button=img2img_paste, tabname="img2img", source_text_component=img2img_prompt, source_image_component=None,
+                paste_button=toprow.paste, tabname="img2img", source_text_component=toprow.prompt, source_image_component=None,
             ))
 
-    modules.scripts.scripts_current = None
+        from modules import ui_extra_networks
+        extra_networks_ui_img2img = ui_extra_networks.create_ui(img2img_interface, [img2img_generation_tab], 'img2img')
+        ui_extra_networks.setup_ui(extra_networks_ui_img2img, img2img_gallery)
+
+        extra_tabs.__exit__()
+
+    scripts.scripts_current = None
 
     with gr.Blocks(analytics_enabled=False) as extras_interface:
         ui_postprocessing.create_ui()
 
     with gr.Blocks(analytics_enabled=False) as pnginfo_interface:
-        with gr.Row().style(equal_height=False):
+        with gr.Row(equal_height=False):
             with gr.Column(variant='panel'):
                 image = gr.Image(elem_id="pnginfo_image", label="Source", source="upload", interactive=True, type="pil")
 
@@ -1077,64 +1027,13 @@ def create_ui():
             outputs=[html, generation_info, html2],
         )
 
-    def update_interp_description(value):
+    modelmerger_ui = ui_checkpoint_merger.UiCheckpointMerger()
-        interp_description_css = "<p style='margin-bottom: 2.5em'>{}</p>"
-        interp_descriptions = {
-            "No interpolation": interp_description_css.format("No interpolation will be used. Requires one model; A. Allows for format conversion and VAE baking."),
-            "Weighted sum": interp_description_css.format("A weighted sum will be used for interpolation. Requires two models; A and B. The result is calculated as A * (1 - M) + B * M"),
-            "Add difference": interp_description_css.format("The difference between the last two models will be added to the first. Requires three models; A, B and C. The result is calculated as A + (B - C) * M")
-        }
-        return interp_descriptions[value]
-
-    with gr.Blocks(analytics_enabled=False) as modelmerger_interface:
-        with gr.Row().style(equal_height=False):
-            with gr.Column(variant='compact'):
-                interp_description = gr.HTML(value=update_interp_description("Weighted sum"), elem_id="modelmerger_interp_description")
-
-                with FormRow(elem_id="modelmerger_models"):
-                    primary_model_name = gr.Dropdown(modules.sd_models.checkpoint_tiles(), elem_id="modelmerger_primary_model_name", label="Primary model (A)")
-                    create_refresh_button(primary_model_name, modules.sd_models.list_models, lambda: {"choices": modules.sd_models.checkpoint_tiles()}, "refresh_checkpoint_A")
-
-                    secondary_model_name = gr.Dropdown(modules.sd_models.checkpoint_tiles(), elem_id="modelmerger_secondary_model_name", label="Secondary model (B)")
-                    create_refresh_button(secondary_model_name, modules.sd_models.list_models, lambda: {"choices": modules.sd_models.checkpoint_tiles()}, "refresh_checkpoint_B")
-
-                    tertiary_model_name = gr.Dropdown(modules.sd_models.checkpoint_tiles(), elem_id="modelmerger_tertiary_model_name", label="Tertiary model (C)")
-                    create_refresh_button(tertiary_model_name, modules.sd_models.list_models, lambda: {"choices": modules.sd_models.checkpoint_tiles()}, "refresh_checkpoint_C")
-
-                custom_name = gr.Textbox(label="Custom Name (Optional)", elem_id="modelmerger_custom_name")
-                interp_amount = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Multiplier (M) - set to 0 to get model A', value=0.3, elem_id="modelmerger_interp_amount")
-                interp_method = gr.Radio(choices=["No interpolation", "Weighted sum", "Add difference"], value="Weighted sum", label="Interpolation Method", elem_id="modelmerger_interp_method")
-                interp_method.change(fn=update_interp_description, inputs=[interp_method], outputs=[interp_description])
-
-                with FormRow():
-                    checkpoint_format = gr.Radio(choices=["ckpt", "safetensors"], value="safetensors", label="Checkpoint format", elem_id="modelmerger_checkpoint_format")
-                    save_as_half = gr.Checkbox(value=False, label="Save as float16", elem_id="modelmerger_save_as_half")
-                    save_metadata = gr.Checkbox(value=True, label="Save metadata (.safetensors only)", elem_id="modelmerger_save_metadata")
-
-                with FormRow():
-                    with gr.Column():
-                        config_source = gr.Radio(choices=["A, B or C", "B", "C", "Don't"], value="A, B or C", label="Copy config from", type="index", elem_id="modelmerger_config_method")
-
-                    with gr.Column():
-                        with FormRow():
-                            bake_in_vae = gr.Dropdown(choices=["None"] + list(sd_vae.vae_dict), value="None", label="Bake in VAE", elem_id="modelmerger_bake_in_vae")
-                            create_refresh_button(bake_in_vae, sd_vae.refresh_vae_list, lambda: {"choices": ["None"] + list(sd_vae.vae_dict)}, "modelmerger_refresh_bake_in_vae")
-
-                with FormRow():
-                    discard_weights = gr.Textbox(value="", label="Discard weights with matching name", elem_id="modelmerger_discard_weights")
-
-                with gr.Row():
-                    modelmerger_merge = gr.Button(elem_id="modelmerger_merge", value="Merge", variant='primary')
-
-            with gr.Column(variant='compact', elem_id="modelmerger_results_container"):
-                with gr.Group(elem_id="modelmerger_results_panel"):
-                    modelmerger_result = gr.HTML(elem_id="modelmerger_result", show_label=False)
 
     with gr.Blocks(analytics_enabled=False) as train_interface:
-        with gr.Row().style(equal_height=False):
+        with gr.Row(equal_height=False):
             gr.HTML(value="<p style='margin-bottom: 0.7em'>See <b><a href=\"https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Textual-Inversion\">wiki</a></b> for detailed explanation.</p>")
 
-        with gr.Row(variant="compact").style(equal_height=False):
+        with gr.Row(variant="compact", equal_height=False):
             with gr.Tabs(elem_id="train_tabs"):
 
                 with gr.Tab(label="Create embedding", id="create_embedding"):
@@ -1154,7 +1053,7 @@ def create_ui():
                     new_hypernetwork_name = gr.Textbox(label="Name", elem_id="train_new_hypernetwork_name")
                     new_hypernetwork_sizes = gr.CheckboxGroup(label="Modules", value=["768", "320", "640", "1280"], choices=["768", "1024", "320", "640", "1280"], elem_id="train_new_hypernetwork_sizes")
                     new_hypernetwork_layer_structure = gr.Textbox("1, 2, 1", label="Enter hypernetwork layer structure", placeholder="1st and last digit must be 1. ex:'1, 2, 1'", elem_id="train_new_hypernetwork_layer_structure")
-                    new_hypernetwork_activation_func = gr.Dropdown(value="linear", label="Select activation function of hypernetwork. Recommended : Swish / Linear(none)", choices=modules.hypernetworks.ui.keys, elem_id="train_new_hypernetwork_activation_func")
+                    new_hypernetwork_activation_func = gr.Dropdown(value="linear", label="Select activation function of hypernetwork. Recommended : Swish / Linear(none)", choices=hypernetworks_ui.keys, elem_id="train_new_hypernetwork_activation_func")
                     new_hypernetwork_initialization_option = gr.Dropdown(value = "Normal", label="Select Layer weights initialization. Recommended: Kaiming for relu-like, Xavier for sigmoid-like, Normal otherwise", choices=["Normal", "KaimingUniform", "KaimingNormal", "XavierUniform", "XavierNormal"], elem_id="train_new_hypernetwork_initialization_option")
                     new_hypernetwork_add_layer_norm = gr.Checkbox(label="Add layer normalization", elem_id="train_new_hypernetwork_add_layer_norm")
                     new_hypernetwork_use_dropout = gr.Checkbox(label="Use dropout", elem_id="train_new_hypernetwork_use_dropout")
@@ -1294,12 +1193,12 @@ def create_ui():
 
             with gr.Column(elem_id='ti_gallery_container'):
                 ti_output = gr.Text(elem_id="ti_output", value="", show_label=False)
-                gr.Gallery(label='Output', show_label=False, elem_id='ti_gallery').style(columns=4)
+                gr.Gallery(label='Output', show_label=False, elem_id='ti_gallery', columns=4)
                 gr.HTML(elem_id="ti_progress", value="")
                 ti_outcome = gr.HTML(elem_id="ti_error", value="")
 
         create_embedding.click(
-            fn=modules.textual_inversion.ui.create_embedding,
+            fn=textual_inversion_ui.create_embedding,
             inputs=[
                 new_embedding_name,
                 initialization_text,
@@ -1314,7 +1213,7 @@ def create_ui():
         )
 
         create_hypernetwork.click(
-            fn=modules.hypernetworks.ui.create_hypernetwork,
+            fn=hypernetworks_ui.create_hypernetwork,
             inputs=[
                 new_hypernetwork_name,
                 new_hypernetwork_sizes,
@@ -1334,7 +1233,7 @@ def create_ui():
         )
 
         run_preprocess.click(
-            fn=wrap_gradio_gpu_call(modules.textual_inversion.ui.preprocess, extra_outputs=[gr.update()]),
+            fn=wrap_gradio_gpu_call(textual_inversion_ui.preprocess, extra_outputs=[gr.update()]),
             _js="start_training_textual_inversion",
             inputs=[
                 dummy_component,
@@ -1370,7 +1269,7 @@ def create_ui():
         )
 
         train_embedding.click(
-            fn=wrap_gradio_gpu_call(modules.textual_inversion.ui.train_embedding, extra_outputs=[gr.update()]),
+            fn=wrap_gradio_gpu_call(textual_inversion_ui.train_embedding, extra_outputs=[gr.update()]),
             _js="start_training_textual_inversion",
             inputs=[
                 dummy_component,
@@ -1404,7 +1303,7 @@ def create_ui():
         )
 
         train_hypernetwork.click(
-            fn=wrap_gradio_gpu_call(modules.hypernetworks.ui.train_hypernetwork, extra_outputs=[gr.update()]),
+            fn=wrap_gradio_gpu_call(hypernetworks_ui.train_hypernetwork, extra_outputs=[gr.update()]),
             _js="start_training_textual_inversion",
             inputs=[
                 dummy_component,
@@ -1458,7 +1357,7 @@ def create_ui():
         (img2img_interface, "img2img", "img2img"),
         (extras_interface, "Extras", "extras"),
         (pnginfo_interface, "PNG Info", "pnginfo"),
-        (modelmerger_interface, "Checkpoint Merger", "modelmerger"),
+        (modelmerger_ui.blocks, "Checkpoint Merger", "modelmerger"),
         (train_interface, "Train", "train"),
     ]
 
@@ -1510,49 +1409,11 @@ def create_ui():
         settings.text_settings.change(fn=update_image_cfg_scale_visibility, inputs=[], outputs=[image_cfg_scale])
         demo.load(fn=update_image_cfg_scale_visibility, inputs=[], outputs=[image_cfg_scale])
 
-        def modelmerger(*args):
+        modelmerger_ui.setup_ui(dummy_component=dummy_component, sd_model_checkpoint_component=settings.component_dict['sd_model_checkpoint'])
-            try:
-                results = modules.extras.run_modelmerger(*args)
-            except Exception as e:
-                errors.report("Error loading/saving model file", exc_info=True)
-                modules.sd_models.list_models()  # to remove the potentially missing models from the list
-                return [*[gr.Dropdown.update(choices=modules.sd_models.checkpoint_tiles()) for _ in range(4)], f"Error merging checkpoints: {e}"]
-            return results
-
-        modelmerger_merge.click(fn=lambda: '', inputs=[], outputs=[modelmerger_result])
-        modelmerger_merge.click(
-            fn=wrap_gradio_gpu_call(modelmerger, extra_outputs=lambda: [gr.update() for _ in range(4)]),
-            _js='modelmerger',
-            inputs=[
-                dummy_component,
-                primary_model_name,
-                secondary_model_name,
-                tertiary_model_name,
-                interp_method,
-                interp_amount,
-                save_as_half,
-                custom_name,
-                checkpoint_format,
-                config_source,
-                bake_in_vae,
-                discard_weights,
-                save_metadata,
-            ],
-            outputs=[
-                primary_model_name,
-                secondary_model_name,
-                tertiary_model_name,
-                settings.component_dict['sd_model_checkpoint'],
-                modelmerger_result,
-            ]
-        )
 
     loadsave.dump_defaults()
     demo.ui_loadsave = loadsave
 
-    # Required as a workaround for change() event not triggering when loading values from ui-config.json
-    interp_description.value = update_interp_description(interp_method.value)
-
     return demo