61 files changed, 3331 insertions, 615 deletions

diff --git a/modules/api/api.py b/modules/api/api.py
index eb7b1da5..9ffcbd5f 100644
--- a/modules/api/api.py
+++ b/modules/api/api.py
@@ -3,11 +3,14 @@ import io
 import time
 import datetime
 import uvicorn
+import gradio as gr
 from threading import Lock
 from io import BytesIO
-from gradio.processing_utils import decode_base64_to_file
-from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request, Response
+from fastapi import APIRouter, Depends, FastAPI, Request, Response
 from fastapi.security import HTTPBasic, HTTPBasicCredentials
+from fastapi.exceptions import HTTPException
+from fastapi.responses import JSONResponse
+from fastapi.encoders import jsonable_encoder
 from secrets import compare_digest
 
 import modules.shared as shared
@@ -18,7 +21,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
+from modules.sd_models import checkpoints_list, unload_model_weights, reload_model_weights
 from modules.sd_models_config import find_checkpoint_config_near_filename
 from modules.realesrgan_model import get_realesrgan_models
 from modules import devices
@@ -90,6 +93,16 @@ def encode_pil_to_base64(image):
     return base64.b64encode(bytes_data)
 
 def api_middleware(app: FastAPI):
+    rich_available = True
+    try:
+        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()
+    except:
+        import traceback
+        rich_available = False
+
     @app.middleware("http")
     async def log_and_time(req: Request, call_next):
         ts = time.time()
@@ -110,6 +123,36 @@ def api_middleware(app: FastAPI):
             ))
         return res
 
+    def handle_exception(request: Request, e: Exception):
+        err = {
+            "error": type(e).__name__,
+            "detail": vars(e).get('detail', ''),
+            "body": vars(e).get('body', ''),
+            "errors": str(e),
+        }
+        print(f"API error: {request.method}: {request.url} {err}")
+        if not isinstance(e, HTTPException): # do not print backtrace on known httpexceptions
+            if rich_available:
+                console.print_exception(show_locals=True, max_frames=2, extra_lines=1, suppress=[anyio, starlette], word_wrap=False, width=min([console.width, 200]))
+            else:
+                traceback.print_exc()
+        return JSONResponse(status_code=vars(e).get('status_code', 500), content=jsonable_encoder(err))
+
+    @app.middleware("http")
+    async def exception_handling(request: Request, call_next):
+        try:
+            return await call_next(request)
+        except Exception as e:
+            return handle_exception(request, e)
+
+    @app.exception_handler(Exception)
+    async def fastapi_exception_handler(request: Request, e: Exception):
+        return handle_exception(request, e)
+
+    @app.exception_handler(HTTPException)
+    async def http_exception_handler(request: Request, e: HTTPException):
+        return handle_exception(request, e)
+
 
 class Api:
     def __init__(self, app: FastAPI, queue_lock: Lock):
@@ -150,6 +193,12 @@ class Api:
         self.add_api_route("/sdapi/v1/train/embedding", self.train_embedding, methods=["POST"], response_model=TrainResponse)
         self.add_api_route("/sdapi/v1/train/hypernetwork", self.train_hypernetwork, methods=["POST"], response_model=TrainResponse)
         self.add_api_route("/sdapi/v1/memory", self.get_memory, methods=["GET"], response_model=MemoryResponse)
+        self.add_api_route("/sdapi/v1/unload-checkpoint", self.unloadapi, methods=["POST"])
+        self.add_api_route("/sdapi/v1/reload-checkpoint", self.reloadapi, methods=["POST"])
+        self.add_api_route("/sdapi/v1/scripts", self.get_scripts_list, methods=["GET"], response_model=ScriptsList)
+
+        self.default_script_arg_txt2img = []
+        self.default_script_arg_img2img = []
 
     def add_api_route(self, path: str, endpoint, **kwargs):
         if shared.cmd_opts.api_auth:
@@ -163,47 +212,113 @@ class Api:
 
         raise HTTPException(status_code=401, detail="Incorrect username or password", headers={"WWW-Authenticate": "Basic"})
 
-    def get_script(self, script_name, script_runner):
-        if script_name is None:
+    def get_selectable_script(self, script_name, script_runner):
+        if script_name is None or script_name == "":
             return None, None
 
-        if not script_runner.scripts:
-            script_runner.initialize_scripts(False)
-            ui.create_ui()
-
         script_idx = script_name_to_index(script_name, script_runner.selectable_scripts)
         script = script_runner.selectable_scripts[script_idx]
         return script, script_idx
+    
+    def get_scripts_list(self):
+        t2ilist = [str(title.lower()) for title in scripts.scripts_txt2img.titles]
+        i2ilist = [str(title.lower()) for title in scripts.scripts_img2img.titles]
+
+        return ScriptsList(txt2img = t2ilist, img2img = i2ilist)  
+
+    def get_script(self, script_name, script_runner):
+        if script_name is None or script_name == "":
+            return None, None
+        
+        script_idx = script_name_to_index(script_name, script_runner.scripts)
+        return script_runner.scripts[script_idx]
+
+    def init_default_script_args(self, script_runner):
+        #find max idx from the scripts in runner and generate a none array to init script_args
+        last_arg_index = 1
+        for script in script_runner.scripts:
+            if last_arg_index < script.args_to:
+                last_arg_index = script.args_to
+        # None everywhere except position 0 to initialize script args
+        script_args = [None]*last_arg_index
+        script_args[0] = 0
+
+        # get default values
+        with gr.Blocks(): # will throw errors calling ui function without this
+            for script in script_runner.scripts:
+                if script.ui(script.is_img2img):
+                    ui_default_values = []
+                    for elem in script.ui(script.is_img2img):
+                        ui_default_values.append(elem.value)
+                    script_args[script.args_from:script.args_to] = ui_default_values
+        return script_args
+
+    def init_script_args(self, request, default_script_args, selectable_scripts, selectable_idx, script_runner):
+        script_args = default_script_args.copy()
+        # position 0 in script_arg is the idx+1 of the selectable script that is going to be run when using scripts.scripts_*2img.run()
+        if selectable_scripts:
+            script_args[selectable_scripts.args_from:selectable_scripts.args_to] = request.script_args
+            script_args[0] = selectable_idx + 1
+
+        # Now check for always on scripts
+        if request.alwayson_scripts and (len(request.alwayson_scripts) > 0):
+            for alwayson_script_name in request.alwayson_scripts.keys():
+                alwayson_script = self.get_script(alwayson_script_name, script_runner)
+                if alwayson_script == None:
+                    raise HTTPException(status_code=422, detail=f"always on script {alwayson_script_name} not found")
+                # Selectable script in always on script param check
+                if alwayson_script.alwayson == False:
+                    raise HTTPException(status_code=422, detail=f"Cannot have a selectable script in the always on scripts params")
+                # always on script with no arg should always run so you don't really need to add them to the requests
+                if "args" in request.alwayson_scripts[alwayson_script_name]:
+                    # min between arg length in scriptrunner and arg length in the request
+                    for idx in range(0, min((alwayson_script.args_to - alwayson_script.args_from), len(request.alwayson_scripts[alwayson_script_name]["args"]))):
+                        script_args[alwayson_script.args_from + idx] = request.alwayson_scripts[alwayson_script_name]["args"][idx]
+        return script_args
 
     def text2imgapi(self, txt2imgreq: StableDiffusionTxt2ImgProcessingAPI):
-        script, script_idx = self.get_script(txt2imgreq.script_name, scripts.scripts_txt2img)
+        script_runner = scripts.scripts_txt2img
+        if not script_runner.scripts:
+            script_runner.initialize_scripts(False)
+            ui.create_ui()
+        if not self.default_script_arg_txt2img:
+            self.default_script_arg_txt2img = self.init_default_script_args(script_runner)
+        selectable_scripts, selectable_script_idx = self.get_selectable_script(txt2imgreq.script_name, script_runner)
 
-        populate = txt2imgreq.copy(update={ # Override __init__ params
+        populate = txt2imgreq.copy(update={  # Override __init__ params
             "sampler_name": validate_sampler_name(txt2imgreq.sampler_name or txt2imgreq.sampler_index),
-            "do_not_save_samples": True,
-            "do_not_save_grid": True
-            }
-        )
+            "do_not_save_samples": not txt2imgreq.save_images,
+            "do_not_save_grid": not txt2imgreq.save_images,
+        })
         if populate.sampler_name:
             populate.sampler_index = None  # prevent a warning later on
 
         args = vars(populate)
         args.pop('script_name', None)
+        args.pop('script_args', None) # will refeed them to the pipeline directly after initializing them
+        args.pop('alwayson_scripts', None)
+
+        script_args = self.init_script_args(txt2imgreq, self.default_script_arg_txt2img, selectable_scripts, selectable_script_idx, script_runner)
+
+        send_images = args.pop('send_images', True)
+        args.pop('save_images', None)
 
         with self.queue_lock:
             p = StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **args)
+            p.scripts = script_runner
+            p.outpath_grids = opts.outdir_txt2img_grids
+            p.outpath_samples = opts.outdir_txt2img_samples
 
             shared.state.begin()
-            if script is not None:
-                p.outpath_grids = opts.outdir_txt2img_grids
-                p.outpath_samples = opts.outdir_txt2img_samples
-                p.script_args = [script_idx + 1] + [None] * (script.args_from - 1) + p.script_args
-                processed = scripts.scripts_txt2img.run(p, *p.script_args)
+            if selectable_scripts != 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)
             shared.state.end()
 
-        b64images = list(map(encode_pil_to_base64, processed.images))
+        b64images = list(map(encode_pil_to_base64, processed.images)) if send_images else []
 
         return TextToImageResponse(images=b64images, parameters=vars(txt2imgreq), info=processed.js())
 
@@ -212,41 +327,55 @@ class Api:
         if init_images is None:
             raise HTTPException(status_code=404, detail="Init image not found")
 
-        script, script_idx = self.get_script(img2imgreq.script_name, scripts.scripts_img2img)
-
         mask = img2imgreq.mask
         if mask:
             mask = decode_base64_to_image(mask)
 
-        populate = img2imgreq.copy(update={ # Override __init__ params
+        script_runner = scripts.scripts_img2img
+        if not script_runner.scripts:
+            script_runner.initialize_scripts(True)
+            ui.create_ui()
+        if not self.default_script_arg_img2img:
+            self.default_script_arg_img2img = self.init_default_script_args(script_runner)
+        selectable_scripts, selectable_script_idx = self.get_selectable_script(img2imgreq.script_name, script_runner)
+
+        populate = img2imgreq.copy(update={  # Override __init__ params
             "sampler_name": validate_sampler_name(img2imgreq.sampler_name or img2imgreq.sampler_index),
-            "do_not_save_samples": True,
-            "do_not_save_grid": True,
-            "mask": mask
-            }
-        )
+            "do_not_save_samples": not img2imgreq.save_images,
+            "do_not_save_grid": not img2imgreq.save_images,
+            "mask": mask,
+        })
         if populate.sampler_name:
             populate.sampler_index = None  # prevent a warning later on
 
         args = vars(populate)
         args.pop('include_init_images', None)  # this is meant to be done by "exclude": True in model, but it's for a reason that I cannot determine.
         args.pop('script_name', None)
+        args.pop('script_args', None)  # will refeed them to the pipeline directly after initializing them
+        args.pop('alwayson_scripts', None)
+
+        script_args = self.init_script_args(img2imgreq, self.default_script_arg_img2img, selectable_scripts, selectable_script_idx, script_runner)
+
+        send_images = args.pop('send_images', True)
+        args.pop('save_images', None)
 
         with self.queue_lock:
             p = StableDiffusionProcessingImg2Img(sd_model=shared.sd_model, **args)
             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()
-            if script is not None:
-                p.outpath_grids = opts.outdir_img2img_grids
-                p.outpath_samples = opts.outdir_img2img_samples
-                p.script_args = [script_idx + 1] + [None] * (script.args_from - 1) + p.script_args
-                processed = scripts.scripts_img2img.run(p, *p.script_args)
+            if selectable_scripts != 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)
             shared.state.end()
 
-        b64images = list(map(encode_pil_to_base64, processed.images))
+        b64images = list(map(encode_pil_to_base64, processed.images)) if send_images else []
 
         if not img2imgreq.include_init_images:
             img2imgreq.init_images = None
@@ -267,16 +396,11 @@ class Api:
     def extras_batch_images_api(self, req: ExtrasBatchImagesRequest):
         reqDict = setUpscalers(req)
 
-        def prepareFiles(file):
-            file = decode_base64_to_file(file.data, file_path=file.name)
-            file.orig_name = file.name
-            return file
-
-        reqDict['image_folder'] = list(map(prepareFiles, reqDict['imageList']))
-        reqDict.pop('imageList')
+        image_list = reqDict.pop('imageList', [])
+        image_folder = [decode_base64_to_image(x.data) for x in image_list]
 
         with self.queue_lock:
-            result = postprocessing.run_extras(extras_mode=1, image="", input_dir="", output_dir="", save_output=False, **reqDict)
+            result = postprocessing.run_extras(extras_mode=1, image_folder=image_folder, image="", input_dir="", output_dir="", save_output=False, **reqDict)
 
         return ExtrasBatchImagesResponse(images=list(map(encode_pil_to_base64, result[0])), html_info=result[1])
 
@@ -348,6 +472,16 @@ class Api:
 
         return {}
 
+    def unloadapi(self):
+        unload_model_weights()
+
+        return {}
+
+    def reloadapi(self):
+        reload_model_weights()
+
+        return {}
+
     def skip(self):
         shared.state.skip()
 
@@ -498,7 +632,7 @@ class Api:
             if not apply_optimizations:
                 sd_hijack.undo_optimizations()
             try:
-                hypernetwork, filename = train_hypernetwork(*args)
+                hypernetwork, filename = train_hypernetwork(**args)
             except Exception as e:
                 error = e
             finally:
diff --git a/modules/api/models.py b/modules/api/models.py
index cba43d3b..4a70f440 100644
--- a/modules/api/models.py
+++ b/modules/api/models.py
@@ -14,8 +14,8 @@ API_NOT_ALLOWED = [
     "outpath_samples",
     "outpath_grids",
     "sampler_index",
-    "do_not_save_samples",
-    "do_not_save_grid",
+    # "do_not_save_samples",
+    # "do_not_save_grid",
     "extra_generation_params",
     "overlay_images",
     "do_not_reload_embeddings",
@@ -100,13 +100,31 @@ class PydanticModelGenerator:
 StableDiffusionTxt2ImgProcessingAPI = PydanticModelGenerator(
     "StableDiffusionProcessingTxt2Img",
     StableDiffusionProcessingTxt2Img,
-    [{"key": "sampler_index", "type": str, "default": "Euler"}, {"key": "script_name", "type": str, "default": None}, {"key": "script_args", "type": list, "default": []}]
+    [
+        {"key": "sampler_index", "type": str, "default": "Euler"},
+        {"key": "script_name", "type": str, "default": None},
+        {"key": "script_args", "type": list, "default": []},
+        {"key": "send_images", "type": bool, "default": True},
+        {"key": "save_images", "type": bool, "default": False},
+        {"key": "alwayson_scripts", "type": dict, "default": {}},
+    ]
 ).generate_model()
 
 StableDiffusionImg2ImgProcessingAPI = PydanticModelGenerator(
     "StableDiffusionProcessingImg2Img",
     StableDiffusionProcessingImg2Img,
-    [{"key": "sampler_index", "type": str, "default": "Euler"}, {"key": "init_images", "type": list, "default": None}, {"key": "denoising_strength", "type": float, "default": 0.75}, {"key": "mask", "type": str, "default": None}, {"key": "include_init_images", "type": bool, "default": False, "exclude" : True}, {"key": "script_name", "type": str, "default": None}, {"key": "script_args", "type": list, "default": []}]
+    [
+        {"key": "sampler_index", "type": str, "default": "Euler"},
+        {"key": "init_images", "type": list, "default": None},
+        {"key": "denoising_strength", "type": float, "default": 0.75},
+        {"key": "mask", "type": str, "default": None},
+        {"key": "include_init_images", "type": bool, "default": False, "exclude" : True},
+        {"key": "script_name", "type": str, "default": None},
+        {"key": "script_args", "type": list, "default": []},
+        {"key": "send_images", "type": bool, "default": True},
+        {"key": "save_images", "type": bool, "default": False},
+        {"key": "alwayson_scripts", "type": dict, "default": {}},
+    ]
 ).generate_model()
 
 class TextToImageResponse(BaseModel):
@@ -267,3 +285,7 @@ class EmbeddingsResponse(BaseModel):
 class MemoryResponse(BaseModel):
     ram: dict = Field(title="RAM", description="System memory stats")
     cuda: dict = Field(title="CUDA", description="nVidia CUDA memory stats")
+
+class ScriptsList(BaseModel):
+    txt2img: list = Field(default=None,title="Txt2img", description="Titles of scripts (txt2img)")
+    img2img: list = Field(default=None,title="Img2img", description="Titles of scripts (img2img)")
\ No newline at end of file
diff --git a/modules/cmd_args.py b/modules/cmd_args.py
new file mode 100644
index 00000000..81c0b82a
--- /dev/null
+++ b/modules/cmd_args.py
@@ -0,0 +1,103 @@
+import argparse

+import os

+from modules.paths_internal import models_path, script_path, data_path, extensions_dir, extensions_builtin_dir, sd_default_config, sd_model_file

+

+parser = argparse.ArgumentParser()

+

+parser.add_argument("-f", action='store_true', help=argparse.SUPPRESS)  # allows running as root; implemented outside of webui

+parser.add_argument("--update-all-extensions", action='store_true', help="launch.py argument: download updates for all extensions when starting the program")

+parser.add_argument("--skip-python-version-check", action='store_true', help="launch.py argument: do not check python version")

+parser.add_argument("--skip-torch-cuda-test", action='store_true', help="launch.py argument: do not check if CUDA is able to work properly")

+parser.add_argument("--reinstall-xformers", action='store_true', help="launch.py argument: install the appropriate version of xformers even if you have some version already installed")

+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: chck for updates at startup")

+parser.add_argument("--tests", type=str, default=None, help="launch.py argument: run tests in the specified directory")

+parser.add_argument("--no-tests", action='store_true', help="launch.py argument: do not run tests even if --tests option is specified")

+parser.add_argument("--skip-install", action='store_true', help="launch.py argument: skip installation of packages")

+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",)

+parser.add_argument("--ckpt-dir", type=str, default=None, help="Path to directory with stable diffusion checkpoints")

+parser.add_argument("--vae-dir", type=str, default=None, help="Path to directory with VAE files")

+parser.add_argument("--gfpgan-dir", type=str, help="GFPGAN directory", default=('./src/gfpgan' if os.path.exists('./src/gfpgan') else './GFPGAN'))

+parser.add_argument("--gfpgan-model", type=str, help="GFPGAN model file name", default=None)

+parser.add_argument("--no-half", action='store_true', help="do not switch the model to 16-bit floats")

+parser.add_argument("--no-half-vae", action='store_true', help="do not switch the VAE model to 16-bit floats")

+parser.add_argument("--no-progressbar-hiding", action='store_true', help="do not hide progressbar in gradio UI (we hide it because it slows down ML if you have hardware acceleration in browser)")

+parser.add_argument("--max-batch-count", type=int, default=16, help="maximum batch count value for the UI")

+parser.add_argument("--embeddings-dir", type=str, default=os.path.join(data_path, 'embeddings'), help="embeddings directory for textual inversion (default: embeddings)")

+parser.add_argument("--textual-inversion-templates-dir", type=str, default=os.path.join(script_path, 'textual_inversion_templates'), help="directory with textual inversion templates")

+parser.add_argument("--hypernetwork-dir", type=str, default=os.path.join(models_path, 'hypernetworks'), help="hypernetwork directory")

+parser.add_argument("--localizations-dir", type=str, default=os.path.join(script_path, 'localizations'), help="localizations directory")

+parser.add_argument("--allow-code", action='store_true', help="allow custom script execution from webui")

+parser.add_argument("--medvram", action='store_true', help="enable stable diffusion model optimizations for sacrificing a little speed for low VRM usage")

+parser.add_argument("--lowvram", action='store_true', help="enable stable diffusion model optimizations for sacrificing a lot of speed for very low VRM usage")

+parser.add_argument("--lowram", action='store_true', help="load stable diffusion checkpoint weights to VRAM instead of RAM")

+parser.add_argument("--always-batch-cond-uncond", action='store_true', help="disables cond/uncond batching that is enabled to save memory with --medvram or --lowvram")

+parser.add_argument("--unload-gfpgan", action='store_true', help="does not do anything.")

+parser.add_argument("--precision", type=str, help="evaluate at this precision", choices=["full", "autocast"], default="autocast")

+parser.add_argument("--upcast-sampling", action='store_true', help="upcast sampling. No effect with --no-half. Usually produces similar results to --no-half with better performance while using less memory.")

+parser.add_argument("--share", action='store_true', help="use share=True for gradio and make the UI accessible through their site")

+parser.add_argument("--ngrok", type=str, help="ngrok authtoken, alternative to gradio --share", default=None)

+parser.add_argument("--ngrok-region", type=str, help="The region in which ngrok should start.", default="us")

+parser.add_argument("--enable-insecure-extension-access", action='store_true', help="enable extensions tab regardless of other options")

+parser.add_argument("--codeformer-models-path", type=str, help="Path to directory with codeformer model file(s).", default=os.path.join(models_path, 'Codeformer'))

+parser.add_argument("--gfpgan-models-path", type=str, help="Path to directory with GFPGAN model file(s).", default=os.path.join(models_path, 'GFPGAN'))

+parser.add_argument("--esrgan-models-path", type=str, help="Path to directory with ESRGAN model file(s).", default=os.path.join(models_path, 'ESRGAN'))

+parser.add_argument("--bsrgan-models-path", type=str, help="Path to directory with BSRGAN model file(s).", default=os.path.join(models_path, 'BSRGAN'))

+parser.add_argument("--realesrgan-models-path", type=str, help="Path to directory with RealESRGAN model file(s).", default=os.path.join(models_path, 'RealESRGAN'))

+parser.add_argument("--clip-models-path", type=str, help="Path to directory with CLIP model file(s).", default=None)

+parser.add_argument("--xformers", action='store_true', help="enable xformers for cross attention layers")

+parser.add_argument("--force-enable-xformers", action='store_true', help="enable xformers for cross attention layers regardless of whether the checking code thinks you can run it; do not make bug reports if this fails to work")

+parser.add_argument("--xformers-flash-attention", action='store_true', help="enable xformers with Flash Attention to improve reproducibility (supported for SD2.x or variant only)")

+parser.add_argument("--deepdanbooru", action='store_true', help="does not do anything")

+parser.add_argument("--opt-split-attention", action='store_true', help="force-enables Doggettx's cross-attention layer optimization. By default, it's on for torch cuda.")

+parser.add_argument("--opt-sub-quad-attention", action='store_true', help="enable memory efficient sub-quadratic cross-attention layer optimization")

+parser.add_argument("--sub-quad-q-chunk-size", type=int, help="query chunk size for the sub-quadratic cross-attention layer optimization to use", default=1024)

+parser.add_argument("--sub-quad-kv-chunk-size", type=int, help="kv chunk size for the sub-quadratic cross-attention layer optimization to use", default=None)

+parser.add_argument("--sub-quad-chunk-threshold", type=int, help="the percentage of VRAM threshold for the sub-quadratic cross-attention layer optimization to use chunking", default=None)

+parser.add_argument("--opt-split-attention-invokeai", action='store_true', help="force-enables InvokeAI's cross-attention layer optimization. By default, it's on when cuda is unavailable.")

+parser.add_argument("--opt-split-attention-v1", action='store_true', help="enable older version of split attention optimization that does not consume all the VRAM it can find")

+parser.add_argument("--opt-sdp-attention", action='store_true', help="enable scaled dot product cross-attention layer optimization; requires PyTorch 2.*")

+parser.add_argument("--opt-sdp-no-mem-attention", action='store_true', help="enable scaled dot product cross-attention layer optimization without memory efficient attention, makes image generation deterministic; requires PyTorch 2.*")

+parser.add_argument("--disable-opt-split-attention", action='store_true', help="force-disables cross-attention layer 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("--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)

+parser.add_argument("--ui-config-file", type=str, help="filename to use for ui configuration", default=os.path.join(data_path, 'ui-config.json'))

+parser.add_argument("--hide-ui-dir-config", action='store_true', help="hide directory configuration from webui", default=False)

+parser.add_argument("--freeze-settings", action='store_true', help="disable editing settings", default=False)

+parser.add_argument("--ui-settings-file", type=str, help="filename to use for ui settings", default=os.path.join(data_path, 'config.json'))

+parser.add_argument("--gradio-debug",  action='store_true', help="launch gradio with --debug option")

+parser.add_argument("--gradio-auth", type=str, help='set gradio authentication like "username:password"; or comma-delimit multiple like "u1:p1,u2:p2,u3:p3"', default=None)

+parser.add_argument("--gradio-auth-path", type=str, help='set gradio authentication file path ex. "/path/to/auth/file" same auth format as --gradio-auth', default=None)

+parser.add_argument("--gradio-img2img-tool", type=str, help='does not do anything')

+parser.add_argument("--gradio-inpaint-tool", type=str, help="does not do anything")

+parser.add_argument("--opt-channelslast", action='store_true', help="change memory type for stable diffusion to channels last")

+parser.add_argument("--styles-file", type=str, help="filename to use for styles", default=os.path.join(data_path, 'styles.csv'))

+parser.add_argument("--autolaunch", action='store_true', help="open the webui URL in the system's default browser upon launch", default=False)

+parser.add_argument("--theme", type=str, help="launches the UI with light or dark theme", default=None)

+parser.add_argument("--use-textbox-seed", action='store_true', help="use textbox for seeds in UI (no up/down, but possible to input long seeds)", default=False)

+parser.add_argument("--disable-console-progressbars", action='store_true', help="do not output progressbars to console", default=False)

+parser.add_argument("--enable-console-prompts", action='store_true', help="print prompts to console when generating with txt2img and img2img", default=False)

+parser.add_argument('--vae-path', type=str, help='Checkpoint to use as VAE; setting this argument disables all settings related to VAE', default=None)

+parser.add_argument("--disable-safe-unpickle", action='store_true', help="disable checking pytorch models for malicious code", default=False)

+parser.add_argument("--api", action='store_true', help="use api=True to launch the API together with the webui (use --nowebui instead for only the API)")

+parser.add_argument("--api-auth", type=str, help='Set authentication for API like "username:password"; or comma-delimit multiple like "u1:p1,u2:p2,u3:p3"', default=None)

+parser.add_argument("--api-log", action='store_true', help="use api-log=True to enable logging of all API requests")

+parser.add_argument("--nowebui", action='store_true', help="use api=True to launch the API instead of the webui")

+parser.add_argument("--ui-debug-mode", action='store_true', help="Don't load model to quickly launch UI")

+parser.add_argument("--device-id", type=str, help="Select the default CUDA device to use (export CUDA_VISIBLE_DEVICES=0,1,etc might be needed before)", default=None)

+parser.add_argument("--administrator", action='store_true', help="Administrator rights", default=False)

+parser.add_argument("--cors-allow-origins", type=str, help="Allowed CORS origin(s) in the form of a comma-separated list (no spaces)", default=None)

+parser.add_argument("--cors-allow-origins-regex", type=str, help="Allowed CORS origin(s) in the form of a single regular expression", default=None)

+parser.add_argument("--tls-keyfile", type=str, help="Partially enables TLS, requires --tls-certfile to fully function", default=None)

+parser.add_argument("--tls-certfile", type=str, help="Partially enables TLS, requires --tls-keyfile to fully function", default=None)

+parser.add_argument("--server-name", type=str, help="Sets hostname of server", default=None)

+parser.add_argument("--gradio-queue", action='store_true', help="does not do anything", default=True)

+parser.add_argument("--no-gradio-queue", action='store_true', help="Disables gradio queue; causes the webpage to use http requests instead of websockets; was the defaul in earlier versions")

+parser.add_argument("--skip-version-check", action='store_true', help="Do not check versions of torch and xformers")

+parser.add_argument("--no-hashing", action='store_true', help="disable sha256 hashing of checkpoints to help loading performance", default=False)

+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)

diff --git a/modules/codeformer_model.py b/modules/codeformer_model.py
index 01fb7bd8..8d84bbc9 100644
--- a/modules/codeformer_model.py
+++ b/modules/codeformer_model.py
@@ -55,7 +55,7 @@ def setup_model(dirname):
                 if self.net is not None and self.face_helper is not None:

                     self.net.to(devices.device_codeformer)

                     return self.net, self.face_helper

-                model_paths = modelloader.load_models(model_path, model_url, self.cmd_dir, download_name='codeformer-v0.1.0.pth')

+                model_paths = modelloader.load_models(model_path, model_url, self.cmd_dir, download_name='codeformer-v0.1.0.pth', ext_filter=['.pth'])

                 if len(model_paths) != 0:

                     ckpt_path = model_paths[0]

                 else:

diff --git a/modules/config_states.py b/modules/config_states.py
new file mode 100644
index 00000000..2ea00929
--- /dev/null
+++ b/modules/config_states.py
@@ -0,0 +1,200 @@
+"""
+Supports saving and restoring webui and extensions from a known working set of commits
+"""
+
+import os
+import sys
+import traceback
+import json
+import time
+import tqdm
+
+from datetime import datetime
+from collections import OrderedDict
+import git
+
+from modules import shared, extensions
+from modules.paths_internal import extensions_dir, extensions_builtin_dir, script_path, config_states_dir
+
+
+all_config_states = OrderedDict()
+
+
+def list_config_states():
+    global all_config_states
+
+    all_config_states.clear()
+    os.makedirs(config_states_dir, exist_ok=True)
+
+    config_states = []
+    for filename in os.listdir(config_states_dir):
+        if filename.endswith(".json"):
+            path = os.path.join(config_states_dir, filename)
+            with open(path, "r", encoding="utf-8") as f:
+                j = json.load(f)
+                j["filepath"] = path
+                config_states.append(j)
+
+    config_states = list(sorted(config_states, key=lambda cs: cs["created_at"], reverse=True))
+
+    for cs in config_states:
+        timestamp = time.asctime(time.gmtime(cs["created_at"]))
+        name = cs.get("name", "Config")
+        full_name = f"{name}: {timestamp}"
+        all_config_states[full_name] = cs
+
+    return all_config_states
+
+
+def get_webui_config():
+    webui_repo = None
+
+    try:
+        if os.path.exists(os.path.join(script_path, ".git")):
+            webui_repo = git.Repo(script_path)
+    except Exception:
+        print(f"Error reading webui git info from {script_path}:", file=sys.stderr)
+        print(traceback.format_exc(), file=sys.stderr)
+
+    webui_remote = None
+    webui_commit_hash = None
+    webui_commit_date = None
+    webui_branch = None
+    if webui_repo and not webui_repo.bare:
+        try:
+            webui_remote = next(webui_repo.remote().urls, None)
+            head = webui_repo.head.commit
+            webui_commit_date = webui_repo.head.commit.committed_date
+            webui_commit_hash = head.hexsha
+            webui_branch = webui_repo.active_branch.name
+
+        except Exception:
+            webui_remote = None
+
+    return {
+        "remote": webui_remote,
+        "commit_hash": webui_commit_hash,
+        "commit_date": webui_commit_date,
+        "branch": webui_branch,
+    }
+
+
+def get_extension_config():
+    ext_config = {}
+
+    for ext in extensions.extensions:
+        entry = {
+            "name": ext.name,
+            "path": ext.path,
+            "enabled": ext.enabled,
+            "is_builtin": ext.is_builtin,
+            "remote": ext.remote,
+            "commit_hash": ext.commit_hash,
+            "commit_date": ext.commit_date,
+            "branch": ext.branch,
+            "have_info_from_repo": ext.have_info_from_repo
+        }
+
+        ext_config[ext.name] = entry
+
+    return ext_config
+
+
+def get_config():
+    creation_time = datetime.now().timestamp()
+    webui_config = get_webui_config()
+    ext_config = get_extension_config()
+
+    return {
+        "created_at": creation_time,
+        "webui": webui_config,
+        "extensions": ext_config
+    }
+
+
+def restore_webui_config(config):
+    print("* Restoring webui state...")
+
+    if "webui" not in config:
+        print("Error: No webui data saved to config")
+        return
+
+    webui_config = config["webui"]
+
+    if "commit_hash" not in webui_config:
+        print("Error: No commit saved to webui config")
+        return
+
+    webui_commit_hash = webui_config.get("commit_hash", None)
+    webui_repo = None
+
+    try:
+        if os.path.exists(os.path.join(script_path, ".git")):
+            webui_repo = git.Repo(script_path)
+    except Exception:
+        print(f"Error reading webui git info from {script_path}:", file=sys.stderr)
+        print(traceback.format_exc(), file=sys.stderr)
+        return
+
+    try:
+        webui_repo.git.fetch(all=True)
+        webui_repo.git.reset(webui_commit_hash, hard=True)
+        print(f"* Restored webui to commit {webui_commit_hash}.")
+    except Exception:
+        print(f"Error restoring webui to commit {webui_commit_hash}:", file=sys.stderr)
+        print(traceback.format_exc(), file=sys.stderr)
+
+
+def restore_extension_config(config):
+    print("* Restoring extension state...")
+
+    if "extensions" not in config:
+        print("Error: No extension data saved to config")
+        return
+
+    ext_config = config["extensions"]
+
+    results = []
+    disabled = []
+
+    for ext in tqdm.tqdm(extensions.extensions):
+        if ext.is_builtin:
+            continue
+
+        ext.read_info_from_repo()
+        current_commit = ext.commit_hash
+
+        if ext.name not in ext_config:
+            ext.disabled = True
+            disabled.append(ext.name)
+            results.append((ext, current_commit[:8], False, "Saved extension state not found in config, marking as disabled"))
+            continue
+
+        entry = ext_config[ext.name]
+
+        if "commit_hash" in entry and entry["commit_hash"]:
+            try:
+                ext.fetch_and_reset_hard(entry["commit_hash"])
+                ext.read_info_from_repo()
+                if current_commit != entry["commit_hash"]:
+                    results.append((ext, current_commit[:8], True, entry["commit_hash"][:8]))
+            except Exception as ex:
+                results.append((ext, current_commit[:8], False, ex))
+        else:
+            results.append((ext, current_commit[:8], False, "No commit hash found in config"))
+
+        if not entry.get("enabled", False):
+            ext.disabled = True
+            disabled.append(ext.name)
+        else:
+            ext.disabled = False
+
+    shared.opts.disabled_extensions = disabled
+    shared.opts.save(shared.config_filename)
+
+    print("* Finished restoring extensions. Results:")
+    for ext, prev_commit, success, result in results:
+        if success:
+            print(f"  + {ext.name}: {prev_commit} -> {result}")
+        else:
+            print(f"  ! {ext.name}: FAILURE ({result})")
diff --git a/modules/devices.py b/modules/devices.py
index 919048d0..c705a3cb 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -1,22 +1,17 @@
-import sys, os, shlex
+import sys
 import contextlib
 import torch
 from modules import errors
-from modules.sd_hijack_utils import CondFunc
-from packaging import version
+
+if sys.platform == "darwin":
+    from modules import mac_specific
 
 
-# has_mps is only available in nightly pytorch (for now) and macOS 12.3+.
-# check `getattr` and try it for compatibility
 def has_mps() -> bool:
-    if not getattr(torch, 'has_mps', False):
-        return False
-    try:
-        torch.zeros(1).to(torch.device("mps"))
-        return True
-    except Exception:
+    if sys.platform != "darwin":
         return False
-
+    else:
+        return mac_specific.has_mps
 
 def extract_device_id(args, name):
     for x in range(len(args)):
@@ -97,14 +92,18 @@ def cond_cast_float(input):
 
 
 def randn(seed, shape):
+    from modules.shared import opts
+
     torch.manual_seed(seed)
-    if device.type == 'mps':
+    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_without_seed(shape):
-    if device.type == 'mps':
+    from modules.shared import opts
+
+    if opts.randn_source == "CPU" or device.type == 'mps':
         return torch.randn(shape, device=cpu).to(device)
     return torch.randn(shape, device=device)
 
@@ -155,36 +154,3 @@ def test_for_nans(x, where):
     message += " Use --disable-nan-check commandline argument to disable this check."
 
     raise NansException(message)
-
-
-# MPS workaround for https://github.com/pytorch/pytorch/issues/89784
-def cumsum_fix(input, cumsum_func, *args, **kwargs):
-    if input.device.type == 'mps':
-        output_dtype = kwargs.get('dtype', input.dtype)
-        if output_dtype == torch.int64:
-            return cumsum_func(input.cpu(), *args, **kwargs).to(input.device)
-        elif cumsum_needs_bool_fix and output_dtype == torch.bool or cumsum_needs_int_fix and (output_dtype == torch.int8 or output_dtype == torch.int16):
-            return cumsum_func(input.to(torch.int32), *args, **kwargs).to(torch.int64)
-    return cumsum_func(input, *args, **kwargs)
-
-
-if has_mps():
-    if version.parse(torch.__version__) < version.parse("1.13"):
-        # PyTorch 1.13 doesn't need these fixes but unfortunately is slower and has regressions that prevent training from working
-
-        # MPS workaround for https://github.com/pytorch/pytorch/issues/79383
-        CondFunc('torch.Tensor.to', lambda orig_func, self, *args, **kwargs: orig_func(self.contiguous(), *args, **kwargs),
-                                                          lambda _, self, *args, **kwargs: self.device.type != 'mps' and (args and isinstance(args[0], torch.device) and args[0].type == 'mps' or isinstance(kwargs.get('device'), torch.device) and kwargs['device'].type == 'mps'))
-        # MPS workaround for https://github.com/pytorch/pytorch/issues/80800 
-        CondFunc('torch.nn.functional.layer_norm', lambda orig_func, *args, **kwargs: orig_func(*([args[0].contiguous()] + list(args[1:])), **kwargs),
-                                                                                        lambda _, *args, **kwargs: args and isinstance(args[0], torch.Tensor) and args[0].device.type == 'mps')
-        # MPS workaround for https://github.com/pytorch/pytorch/issues/90532
-        CondFunc('torch.Tensor.numpy', lambda orig_func, self, *args, **kwargs: orig_func(self.detach(), *args, **kwargs), lambda _, self, *args, **kwargs: self.requires_grad)
-    elif version.parse(torch.__version__) > version.parse("1.13.1"):
-        cumsum_needs_int_fix = not torch.Tensor([1,2]).to(torch.device("mps")).equal(torch.ShortTensor([1,1]).to(torch.device("mps")).cumsum(0))
-        cumsum_needs_bool_fix = not torch.BoolTensor([True,True]).to(device=torch.device("mps"), dtype=torch.int64).equal(torch.BoolTensor([True,False]).to(torch.device("mps")).cumsum(0))
-        cumsum_fix_func = lambda orig_func, input, *args, **kwargs: cumsum_fix(input, orig_func, *args, **kwargs)
-        CondFunc('torch.cumsum', cumsum_fix_func, None)
-        CondFunc('torch.Tensor.cumsum', cumsum_fix_func, None)
-        CondFunc('torch.narrow', lambda orig_func, *args, **kwargs: orig_func(*args, **kwargs).clone(), None)
-
diff --git a/modules/esrgan_model_arch.py b/modules/esrgan_model_arch.py
index bc9ceb2a..1b52b0f5 100644
--- a/modules/esrgan_model_arch.py
+++ b/modules/esrgan_model_arch.py
@@ -1,5 +1,6 @@
 # this file is adapted from https://github.com/victorca25/iNNfer

 

+from collections import OrderedDict

 import math

 import functools

 import torch

diff --git a/modules/extensions.py b/modules/extensions.py
index 5e12b1aa..34d9d654 100644
--- a/modules/extensions.py
+++ b/modules/extensions.py
@@ -2,19 +2,26 @@ import os
 import sys

 import traceback

 

+import time

+from datetime import datetime

 import git

 

-from modules import paths, shared

+from modules import shared

+from modules.paths_internal import extensions_dir, extensions_builtin_dir, script_path

 

 extensions = []

-extensions_dir = os.path.join(paths.data_path, "extensions")

-extensions_builtin_dir = os.path.join(paths.script_path, "extensions-builtin")

 

 if not os.path.exists(extensions_dir):

     os.makedirs(extensions_dir)

 

+

 def active():

-    return [x for x in extensions if x.enabled]

+    if shared.opts.disable_all_extensions == "all":

+        return []

+    elif 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]

 

 

 class Extension:

@@ -25,22 +32,43 @@ class Extension:
         self.status = ''

         self.can_update = False

         self.is_builtin = is_builtin

+        self.commit_hash = ''

+        self.commit_date = None

+        self.version = ''

+        self.branch = None

+        self.remote = None

+        self.have_info_from_repo = False

+

+    def read_info_from_repo(self):

+        if self.is_builtin or self.have_info_from_repo:

+            return

+

+        self.have_info_from_repo = True

 

         repo = None

         try:

-            if os.path.exists(os.path.join(path, ".git")):

-                repo = git.Repo(path)

+            if os.path.exists(os.path.join(self.path, ".git")):

+                repo = git.Repo(self.path)

         except Exception:

-            print(f"Error reading github repository info from {path}:", file=sys.stderr)

+            print(f"Error reading github repository info from {self.path}:", file=sys.stderr)

             print(traceback.format_exc(), file=sys.stderr)

 

         if repo is None or repo.bare:

             self.remote = None

         else:

             try:

-                self.remote = next(repo.remote().urls, None)

                 self.status = 'unknown'

-            except Exception:

+                self.remote = next(repo.remote().urls, None)

+                head = repo.head.commit

+                self.commit_date = repo.head.commit.committed_date

+                ts = time.asctime(time.gmtime(self.commit_date))

+                if repo.active_branch:

+                    self.branch = repo.active_branch.name

+                self.commit_hash = head.hexsha

+                self.version = f'{self.commit_hash[:8]} ({ts})'

+

+            except Exception as ex:

+                print(f"Failed reading extension data from Git repository ({self.name}): {ex}", file=sys.stderr)

                 self.remote = None

 

     def list_files(self, subdir, extension):

@@ -60,21 +88,33 @@ class Extension:
 

     def check_updates(self):

         repo = git.Repo(self.path)

-        for fetch in repo.remote().fetch("--dry-run"):

+        for fetch in repo.remote().fetch(dry_run=True):

             if fetch.flags != fetch.HEAD_UPTODATE:

                 self.can_update = True

-                self.status = "behind"

+                self.status = "new commits"

                 return

 

+        try:

+            origin = repo.rev_parse('origin')

+            if repo.head.commit != origin:

+                self.can_update = True

+                self.status = "behind HEAD"

+                return

+        except Exception:

+            self.can_update = False

+            self.status = "unknown (remote error)"

+            return

+

         self.can_update = False

         self.status = "latest"

 

-    def fetch_and_reset_hard(self):

+    def fetch_and_reset_hard(self, commit='origin'):

         repo = git.Repo(self.path)

         # Fix: `error: Your local changes to the following files would be overwritten by merge`,

         # because WSL2 Docker set 755 file permissions instead of 644, this results to the error.

-        repo.git.fetch('--all')

-        repo.git.reset('--hard', 'origin')

+        repo.git.fetch(all=True)

+        repo.git.reset(commit, hard=True)

+        self.have_info_from_repo = False

 

 

 def list_extensions():

@@ -83,7 +123,12 @@ def list_extensions():
     if not os.path.isdir(extensions_dir):

         return

 

-    paths = []

+    if shared.opts.disable_all_extensions == "all":

+        print("*** \"Disable all extensions\" option was set, will not load any extensions ***")

+    elif shared.opts.disable_all_extensions == "extra":

+        print("*** \"Disable all extensions\" option was set, will only load built-in extensions ***")

+

+    extension_paths = []

     for dirname in [extensions_dir, extensions_builtin_dir]:

         if not os.path.isdir(dirname):

             return

@@ -93,9 +138,8 @@ def list_extensions():
             if not os.path.isdir(path):

                 continue

 

-            paths.append((extension_dirname, path, dirname == extensions_builtin_dir))

+            extension_paths.append((extension_dirname, path, dirname == extensions_builtin_dir))

 

-    for dirname, path, is_builtin in paths:

+    for dirname, path, is_builtin in extension_paths:

         extension = Extension(name=dirname, path=path, enabled=dirname not in shared.opts.disabled_extensions, is_builtin=is_builtin)

         extensions.append(extension)

-

diff --git a/modules/extra_networks_hypernet.py b/modules/extra_networks_hypernet.py
index d3a4d7ad..33d100dd 100644
--- a/modules/extra_networks_hypernet.py
+++ b/modules/extra_networks_hypernet.py
@@ -9,7 +9,7 @@ class ExtraNetworkHypernet(extra_networks.ExtraNetwork):
     def activate(self, p, params_list):

         additional = shared.opts.sd_hypernetwork

 

-        if additional != "" and additional in shared.hypernetworks and len([x for x in params_list if x.items[0] == additional]) == 0:

+        if additional != "None" and additional in shared.hypernetworks and len([x for x in params_list if x.items[0] == additional]) == 0:

             p.all_prompts = [x + f"<hypernet:{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]))

 

diff --git a/modules/extras.py b/modules/extras.py
index d8ece955..ff4e9c4e 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -1,6 +1,7 @@
 import os

 import re

 import shutil

+import json

 

 

 import torch

@@ -71,7 +72,7 @@ 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):

+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):

     shared.state.begin()

     shared.state.job = 'model-merge'

 

@@ -241,13 +242,54 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
     shared.state.textinfo = "Saving"

     print(f"Saving to {output_modelname}...")

 

+    metadata = {"format": "pt", "sd_merge_models": {}, "sd_merge_recipe": None}

+

+    if save_metadata:

+        merge_recipe = {

+            "type": "webui", # indicate this model was merged with webui's built-in merger

+            "primary_model_hash": primary_model_info.sha256,

+            "secondary_model_hash": secondary_model_info.sha256 if secondary_model_info else None,

+            "tertiary_model_hash": tertiary_model_info.sha256 if tertiary_model_info else None,

+            "interp_method": interp_method,

+            "multiplier": multiplier,

+            "save_as_half": save_as_half,

+            "custom_name": custom_name,

+            "config_source": config_source,

+            "bake_in_vae": bake_in_vae,

+            "discard_weights": discard_weights,

+            "is_inpainting": result_is_inpainting_model,

+            "is_instruct_pix2pix": result_is_instruct_pix2pix_model

+        }

+        metadata["sd_merge_recipe"] = json.dumps(merge_recipe)

+

+        def add_model_metadata(checkpoint_info):

+            checkpoint_info.calculate_shorthash()

+            metadata["sd_merge_models"][checkpoint_info.sha256] = {

+                "name": checkpoint_info.name,

+                "legacy_hash": checkpoint_info.hash,

+                "sd_merge_recipe": checkpoint_info.metadata.get("sd_merge_recipe", None)

+            }

+

+            metadata["sd_merge_models"].update(checkpoint_info.metadata.get("sd_merge_models", {}))

+

+        add_model_metadata(primary_model_info)

+        if secondary_model_info:

+            add_model_metadata(secondary_model_info)

+        if tertiary_model_info:

+            add_model_metadata(tertiary_model_info)

+

+        metadata["sd_merge_models"] = json.dumps(metadata["sd_merge_models"])

+

     _, extension = os.path.splitext(output_modelname)

     if extension.lower() == ".safetensors":

-        safetensors.torch.save_file(theta_0, output_modelname, metadata={"format": "pt"})

+        safetensors.torch.save_file(theta_0, output_modelname, metadata=metadata)

     else:

         torch.save(theta_0, output_modelname)

 

     sd_models.list_models()

+    created_model = next((ckpt for ckpt in sd_models.checkpoints_list.values() if ckpt.name == filename), None)

+    if created_model:

+        created_model.calculate_shorthash()

 

     create_config(output_modelname, config_source, primary_model_info, secondary_model_info, tertiary_model_info)

 

diff --git a/modules/generation_parameters_copypaste.py b/modules/generation_parameters_copypaste.py
index fc9e17aa..99f1a0d3 100644
--- a/modules/generation_parameters_copypaste.py
+++ b/modules/generation_parameters_copypaste.py
@@ -23,13 +23,14 @@ registered_param_bindings = []
 

 

 class ParamBinding:

-    def __init__(self, paste_button, tabname, source_text_component=None, source_image_component=None, source_tabname=None, override_settings_component=None):

+    def __init__(self, paste_button, tabname, source_text_component=None, source_image_component=None, source_tabname=None, override_settings_component=None, paste_field_names=[]):

         self.paste_button = paste_button

         self.tabname = tabname

         self.source_text_component = source_text_component

         self.source_image_component = source_image_component

         self.source_tabname = source_tabname

         self.override_settings_component = override_settings_component

+        self.paste_field_names = paste_field_names

 

 

 def reset():

@@ -74,8 +75,8 @@ def image_from_url_text(filedata):
     return image

 

 

-def add_paste_fields(tabname, init_img, fields):

-    paste_fields[tabname] = {"init_img": init_img, "fields": fields}

+def add_paste_fields(tabname, init_img, fields, override_settings_component=None):

+    paste_fields[tabname] = {"init_img": init_img, "fields": fields, "override_settings_component": override_settings_component}

 

     # backwards compatibility for existing extensions

     import modules.ui

@@ -110,6 +111,7 @@ def connect_paste_params_buttons():
     for binding in registered_param_bindings:

         destination_image_component = paste_fields[binding.tabname]["init_img"]

         fields = paste_fields[binding.tabname]["fields"]

+        override_settings_component = binding.override_settings_component or paste_fields[binding.tabname]["override_settings_component"]

 

         destination_width_component = next(iter([field for field, name in fields if name == "Size-1"] if fields else []), None)

         destination_height_component = next(iter([field for field, name in fields if name == "Size-2"] if fields else []), None)

@@ -130,10 +132,10 @@ def connect_paste_params_buttons():
             )

 

         if binding.source_text_component is not None and fields is not None:

-            connect_paste(binding.paste_button, fields, binding.source_text_component, binding.override_settings_component, binding.tabname)

+            connect_paste(binding.paste_button, fields, binding.source_text_component, override_settings_component, binding.tabname)

 

         if binding.source_tabname is not None and fields is not None:

-            paste_field_names = ['Prompt', 'Negative prompt', 'Steps', 'Face restoration'] + (["Seed"] if shared.opts.send_seed else [])

+            paste_field_names = ['Prompt', 'Negative prompt', 'Steps', 'Face restoration'] + (["Seed"] if shared.opts.send_seed else []) + binding.paste_field_names

             binding.paste_button.click(

                 fn=lambda *x: x,

                 inputs=[field for field, name in paste_fields[binding.source_tabname]["fields"] if name in paste_field_names],

@@ -282,11 +284,17 @@ Steps: 20, Sampler: Euler a, CFG scale: 7, Seed: 965400086, Size: 512x512, Model
 

     restore_old_hires_fix_params(res)

 

+    # Missing RNG means the default was set, which is GPU RNG

+    if "RNG" not in res:

+        res["RNG"] = "GPU"

+

     return res

 

 

 settings_map = {}

 

+

+

 infotext_to_setting_name_mapping = [

     ('Clip skip', 'CLIP_stop_at_last_layers', ),

     ('Conditional mask weight', 'inpainting_mask_weight'),

@@ -295,7 +303,13 @@ infotext_to_setting_name_mapping = [
     ('Noise multiplier', 'initial_noise_multiplier'),

     ('Eta', 'eta_ancestral'),

     ('Eta DDIM', 'eta_ddim'),

-    ('Discard penultimate sigma', 'always_discard_next_to_last_sigma')

+    ('Discard penultimate sigma', 'always_discard_next_to_last_sigma'),

+    ('UniPC variant', 'uni_pc_variant'),

+    ('UniPC skip type', 'uni_pc_skip_type'),

+    ('UniPC order', 'uni_pc_order'),

+    ('UniPC lower order final', 'uni_pc_lower_order_final'),

+    ('RNG', 'randn_source'),

+    ('NGMS', 's_min_uncond'),

 ]

 

 

@@ -393,9 +407,14 @@ def connect_paste(button, paste_fields, input_comp, override_settings_component,
 

     button.click(

         fn=paste_func,

-        _js=f"recalculate_prompts_{tabname}",

         inputs=[input_comp],

         outputs=[x[0] for x in paste_fields],

     )

+    button.click(

+        fn=None,

+        _js=f"recalculate_prompts_{tabname}",

+        inputs=[],

+        outputs=[],

+    )

 

 

diff --git a/modules/hashes.py b/modules/hashes.py
index 819362a3..83272a07 100644
--- a/modules/hashes.py
+++ b/modules/hashes.py
@@ -4,6 +4,7 @@ import os.path
 

 import filelock

 

+from modules import shared

 from modules.paths import data_path

 

 

@@ -68,6 +69,9 @@ def sha256(filename, title):
     if sha256_value is not None:

         return sha256_value

 

+    if shared.cmd_opts.no_hashing:

+        return None

+

     print(f"Calculating sha256 for {filename}: ", end='')

     sha256_value = calculate_sha256(filename)

     print(f"{sha256_value}")

diff --git a/modules/hypernetworks/hypernetwork.py b/modules/hypernetworks/hypernetwork.py
index 503534e2..1fc49537 100644
--- a/modules/hypernetworks/hypernetwork.py
+++ b/modules/hypernetworks/hypernetwork.py
@@ -307,12 +307,12 @@ class Hypernetwork:
     def shorthash(self):

         sha256 = hashes.sha256(self.filename, f'hypernet/{self.name}')

 

-        return sha256[0:10]

+        return sha256[0:10] if sha256 else None

 

 

 def list_hypernetworks(path):

     res = {}

-    for filename in sorted(glob.iglob(os.path.join(path, '**/*.pt'), recursive=True)):

+    for filename in sorted(glob.iglob(os.path.join(path, '**/*.pt'), recursive=True), key=str.lower):

         name = os.path.splitext(os.path.basename(filename))[0]

         # Prevent a hypothetical "None.pt" from being listed.

         if name != "None":

@@ -380,8 +380,8 @@ def apply_single_hypernetwork(hypernetwork, context_k, context_v, layer=None):
         layer.hyper_k = hypernetwork_layers[0]

         layer.hyper_v = hypernetwork_layers[1]

 

-    context_k = hypernetwork_layers[0](context_k)

-    context_v = hypernetwork_layers[1](context_v)

+    context_k = devices.cond_cast_unet(hypernetwork_layers[0](devices.cond_cast_float(context_k)))

+    context_v = devices.cond_cast_unet(hypernetwork_layers[1](devices.cond_cast_float(context_v)))

     return context_k, context_v

 

 

@@ -496,7 +496,7 @@ def create_hypernetwork(name, enable_sizes, overwrite_old, layer_structure=None,
     shared.reload_hypernetworks()

 

 

-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, 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):

+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

 

@@ -554,7 +554,7 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi
 

     pin_memory = shared.opts.pin_memory

 

-    ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=training_width, height=training_height, repeats=shared.opts.training_image_repeats_per_epoch, placeholder_token=hypernetwork_name, model=shared.sd_model, cond_model=shared.sd_model.cond_stage_model, device=devices.device, template_file=template_file, include_cond=True, batch_size=batch_size, gradient_step=gradient_step, shuffle_tags=shuffle_tags, tag_drop_out=tag_drop_out, latent_sampling_method=latent_sampling_method, varsize=varsize)

+    ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=training_width, height=training_height, repeats=shared.opts.training_image_repeats_per_epoch, placeholder_token=hypernetwork_name, model=shared.sd_model, cond_model=shared.sd_model.cond_stage_model, device=devices.device, template_file=template_file, include_cond=True, batch_size=batch_size, gradient_step=gradient_step, shuffle_tags=shuffle_tags, tag_drop_out=tag_drop_out, latent_sampling_method=latent_sampling_method, varsize=varsize, use_weight=use_weight)

 

     if shared.opts.save_training_settings_to_txt:

         saved_params = dict(

@@ -640,13 +640,19 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi
                 

                 with devices.autocast():

                     x = batch.latent_sample.to(devices.device, non_blocking=pin_memory)

+                    if use_weight:

+                        w = batch.weight.to(devices.device, non_blocking=pin_memory)

                     if tag_drop_out != 0 or shuffle_tags:

                         shared.sd_model.cond_stage_model.to(devices.device)

                         c = shared.sd_model.cond_stage_model(batch.cond_text).to(devices.device, non_blocking=pin_memory)

                         shared.sd_model.cond_stage_model.to(devices.cpu)

                     else:

                         c = stack_conds(batch.cond).to(devices.device, non_blocking=pin_memory)

-                    loss = shared.sd_model(x, c)[0] / gradient_step

+                    if use_weight:

+                        loss = shared.sd_model.weighted_forward(x, c, w)[0] / gradient_step

+                        del w

+                    else:

+                        loss = shared.sd_model.forward(x, c)[0] / gradient_step

                     del x

                     del c

 

diff --git a/modules/images.py b/modules/images.py
index ae3cdaf4..fd173829 100644
--- a/modules/images.py
+++ b/modules/images.py
@@ -16,8 +16,9 @@ from PIL import Image, ImageFont, ImageDraw, PngImagePlugin
 from fonts.ttf import Roboto

 import string

 import json

+import hashlib

 

-from modules import sd_samplers, shared, script_callbacks

+from modules import sd_samplers, shared, script_callbacks, errors

 from modules.shared import opts, cmd_opts

 

 LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)

@@ -130,7 +131,7 @@ class GridAnnotation:
         self.size = None

 

 

-def draw_grid_annotations(im, width, height, hor_texts, ver_texts):

+def draw_grid_annotations(im, width, height, hor_texts, ver_texts, margin=0):

     def wrap(drawing, text, font, line_length):

         lines = ['']

         for word in text.split():

@@ -194,32 +195,35 @@ def draw_grid_annotations(im, width, height, hor_texts, ver_texts):
             line.allowed_width = allowed_width

 

     hor_text_heights = [sum([line.size[1] + line_spacing for line in lines]) - line_spacing for lines in hor_texts]

-    ver_text_heights = [sum([line.size[1] + line_spacing for line in lines]) - line_spacing * len(lines) for lines in

-                        ver_texts]

+    ver_text_heights = [sum([line.size[1] + line_spacing for line in lines]) - line_spacing * len(lines) for lines in ver_texts]

 

     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, im.height + pad_top), "white")

-    result.paste(im, (pad_left, pad_top))

+    result = Image.new("RGB", (im.width + pad_left + margin * (cols-1), im.height + pad_top + margin * (rows-1)), "white")

+

+    for row in range(rows):

+        for col in range(cols):

+            cell = im.crop((width * col, height * row, width * (col+1), height * (row+1)))

+            result.paste(cell, (pad_left + (width + margin) * col, pad_top + (height + margin) * row))

 

     d = ImageDraw.Draw(result)

 

     for col in range(cols):

-        x = pad_left + width * col + width / 2

+        x = pad_left + (width + margin) * col + width / 2

         y = pad_top / 2 - hor_text_heights[col] / 2

 

         draw_texts(d, x, y, hor_texts[col], fnt, fontsize)

 

     for row in range(rows):

         x = pad_left / 2

-        y = pad_top + height * row + height / 2 - ver_text_heights[row] / 2

+        y = pad_top + (height + margin) * row + height / 2 - ver_text_heights[row] / 2

 

         draw_texts(d, x, y, ver_texts[row], fnt, fontsize)

 

     return result

 

 

-def draw_prompt_matrix(im, width, height, all_prompts):

+def draw_prompt_matrix(im, width, height, all_prompts, margin=0):

     prompts = all_prompts[1:]

     boundary = math.ceil(len(prompts) / 2)

 

@@ -229,7 +233,7 @@ def draw_prompt_matrix(im, width, height, all_prompts):
     hor_texts = [[GridAnnotation(x, is_active=pos & (1 << i) != 0) for i, x in enumerate(prompts_horiz)] for pos in range(1 << len(prompts_horiz))]

     ver_texts = [[GridAnnotation(x, is_active=pos & (1 << i) != 0) for i, x in enumerate(prompts_vert)] for pos in range(1 << len(prompts_vert))]

 

-    return draw_grid_annotations(im, width, height, hor_texts, ver_texts)

+    return draw_grid_annotations(im, width, height, hor_texts, ver_texts, margin)

 

 

 def resize_image(resize_mode, im, width, height, upscaler_name=None):

@@ -257,9 +261,12 @@ def resize_image(resize_mode, im, width, height, upscaler_name=None):
 

         if scale > 1.0:

             upscalers = [x for x in shared.sd_upscalers if x.name == upscaler_name]

-            assert len(upscalers) > 0, f"could not find upscaler named {upscaler_name}"

+            if len(upscalers) == 0:

+                upscaler = shared.sd_upscalers[0]

+                print(f"could not find upscaler named {upscaler_name or '<empty string>'}, using {upscaler.name} as a fallback")

+            else:

+                upscaler = upscalers[0]

 

-            upscaler = upscalers[0]

             im = upscaler.scaler.upscale(im, scale, upscaler.data_path)

 

         if im.width != w or im.height != h:

@@ -311,6 +318,7 @@ re_nonletters = re.compile(r'[\s' + string.punctuation + ']+')
 re_pattern = re.compile(r"(.*?)(?:\[([^\[\]]+)\]|$)")

 re_pattern_arg = re.compile(r"(.*)<([^>]*)>$")

 max_filename_part_length = 128

+NOTHING_AND_SKIP_PREVIOUS_TEXT = object()

 

 

 def sanitize_filename_part(text, replace_spaces=True):

@@ -340,10 +348,15 @@ class FilenameGenerator:
         '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),

+        'prompt_hash': lambda self: hashlib.sha256(self.prompt.encode()).hexdigest()[0:8],

         'prompt': lambda self: sanitize_filename_part(self.prompt),

         'prompt_no_styles': lambda self: self.prompt_no_style(),

         'prompt_spaces': lambda self: sanitize_filename_part(self.prompt, replace_spaces=False),

         'prompt_words': lambda self: self.prompt_words(),

+        'batch_number': lambda self: NOTHING_AND_SKIP_PREVIOUS_TEXT if self.p.batch_size == 1 else self.p.batch_index + 1,

+        'generation_number': lambda self: NOTHING_AND_SKIP_PREVIOUS_TEXT if self.p.n_iter == 1 and self.p.batch_size == 1 else self.p.iteration * self.p.batch_size + self.p.batch_index + 1,

+        'hasprompt': lambda self, *args: self.hasprompt(*args),  # accepts formats:[hasprompt<prompt1|default><prompt2>..]

+        'clip_skip': lambda self: opts.data["CLIP_stop_at_last_layers"],

     }

     default_time_format = '%Y%m%d%H%M%S'

 

@@ -352,6 +365,22 @@ class FilenameGenerator:
         self.seed = seed

         self.prompt = prompt

         self.image = image

+        

+    def hasprompt(self, *args):

+        lower = self.prompt.lower()

+        if self.p is None or self.prompt is None:

+            return None

+        outres = ""

+        for arg in args:

+            if arg != "":

+                division = arg.split("|")

+                expected = division[0].lower()

+                default = division[1] if len(division) > 1 else ""

+                if lower.find(expected) >= 0:

+                    outres = f'{outres}{expected}'

+                else:

+                    outres = outres if default == "" else f'{outres}{default}'

+        return sanitize_filename_part(outres)

 

     def prompt_no_style(self):

         if self.p is None or self.prompt is None:

@@ -395,9 +424,9 @@ class FilenameGenerator:
 

         for m in re_pattern.finditer(x):

             text, pattern = m.groups()

-            res += text

 

             if pattern is None:

+                res += text

                 continue

 

             pattern_args = []

@@ -418,11 +447,13 @@ class FilenameGenerator:
                     print(f"Error adding [{pattern}] to filename", file=sys.stderr)

                     print(traceback.format_exc(), file=sys.stderr)

 

-                if replacement is not None:

-                    res += str(replacement)

+                if replacement == NOTHING_AND_SKIP_PREVIOUS_TEXT:

+                    continue

+                elif replacement is not None:

+                    res += text + str(replacement)

                     continue

 

-            res += f'[{pattern}]'

+            res += f'{text}[{pattern}]'

 

         return res

 

@@ -548,8 +579,10 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
         elif extension.lower() in (".jpg", ".jpeg", ".webp"):

             if image_to_save.mode == 'RGBA':

                 image_to_save = image_to_save.convert("RGB")

+            elif image_to_save.mode == 'I;16':

+                image_to_save = image_to_save.point(lambda p: p * 0.0038910505836576).convert("RGB" if extension.lower() == ".webp" else "L")

 

-            image_to_save.save(temp_file_path, format=image_format, quality=opts.jpeg_quality)

+            image_to_save.save(temp_file_path, format=image_format, quality=opts.jpeg_quality, lossless=opts.webp_lossless)

 

             if opts.enable_pnginfo and info is not None:

                 exif_bytes = piexif.dump({

@@ -566,21 +599,28 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
         os.replace(temp_file_path, filename_without_extension + extension)

 

     fullfn_without_extension, extension = os.path.splitext(params.filename)

+    if hasattr(os, 'statvfs'):

+        max_name_len = os.statvfs(path).f_namemax

+        fullfn_without_extension = fullfn_without_extension[:max_name_len - max(4, len(extension))]

+        params.filename = fullfn_without_extension + extension

+        fullfn = params.filename

     _atomically_save_image(image, fullfn_without_extension, extension)

 

     image.already_saved_as = fullfn

 

-    target_side_length = 4000

-    oversize = image.width > target_side_length or image.height > target_side_length

-    if opts.export_for_4chan and (oversize or os.stat(fullfn).st_size > 4 * 1024 * 1024):

+    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

 

         if oversize and ratio > 1:

-            image = image.resize((target_side_length, image.height * target_side_length // image.width), LANCZOS)

+            image = image.resize((round(opts.target_side_length), round(image.height * opts.target_side_length / image.width)), LANCZOS)

         elif oversize:

-            image = image.resize((image.width * target_side_length // image.height, target_side_length), LANCZOS)

+            image = image.resize((round(image.width * opts.target_side_length / image.height), round(opts.target_side_length)), LANCZOS)

 

-        _atomically_save_image(image, fullfn_without_extension, ".jpg")

+        try:

+            _atomically_save_image(image, fullfn_without_extension, ".jpg")

+        except Exception as e:

+            errors.display(e, "saving image as downscaled JPG")

 

     if opts.save_txt and info is not None:

         txt_fullfn = f"{fullfn_without_extension}.txt"

@@ -631,6 +671,8 @@ Steps: {json_info["steps"]}, Sampler: {sampler}, CFG scale: {json_info["scale"]}
 

 

 def image_data(data):

+    import gradio as gr

+

     try:

         image = Image.open(io.BytesIO(data))

         textinfo, _ = read_info_from_image(image)

@@ -646,7 +688,7 @@ def image_data(data):
     except Exception:

         pass

 

-    return '', None

+    return gr.update(), None

 

 

 def flatten(img, bgcolor):

diff --git a/modules/img2img.py b/modules/img2img.py
index f813299c..56c846d6 100644
--- a/modules/img2img.py
+++ b/modules/img2img.py
@@ -4,7 +4,7 @@ import sys
 import traceback

 

 import numpy as np

-from PIL import Image, ImageOps, ImageFilter, ImageEnhance, ImageChops

+from PIL import Image, ImageOps, ImageFilter, ImageEnhance, ImageChops, UnidentifiedImageError

 

 from modules import devices, sd_samplers

 from modules.generation_parameters_copypaste import create_override_settings_dict

@@ -46,7 +46,10 @@ def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args):
         if state.interrupted:

             break

 

-        img = Image.open(image)

+        try:

+            img = Image.open(image)

+        except UnidentifiedImageError:

+            continue

         # Use the EXIF orientation of photos taken by smartphones.

         img = ImageOps.exif_transpose(img)

         p.init_images = [img] * p.batch_size

@@ -73,10 +76,12 @@ def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args):
 

             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, denoising_strength: 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, 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, *args):

     override_settings = create_override_settings_dict(override_settings_texts)

 

     is_batch = mode == 5

@@ -112,6 +117,12 @@ def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_s
     if image is not None:

         image = ImageOps.exif_transpose(image)

 

+    if selected_scale_tab == 1:

+        assert image, "Can't scale by because no image is selected"

+

+        width = int(image.width * scale_by)

+        height = int(image.height * scale_by)

+

     assert 0. <= denoising_strength <= 1., 'can only work with strength in [0.0, 1.0]'

 

     p = StableDiffusionProcessingImg2Img(

@@ -142,19 +153,21 @@ def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_s
         inpainting_fill=inpainting_fill,

         resize_mode=resize_mode,

         denoising_strength=denoising_strength,

+        image_cfg_scale=image_cfg_scale,

         inpaint_full_res=inpaint_full_res,

         inpaint_full_res_padding=inpaint_full_res_padding,

         inpainting_mask_invert=inpainting_mask_invert,

         override_settings=override_settings,

     )

 

-    p.scripts = modules.scripts.scripts_txt2img

+    p.scripts = modules.scripts.scripts_img2img

     p.script_args = args

 

     if shared.cmd_opts.enable_console_prompts:

         print(f"\nimg2img: {prompt}", file=shared.progress_print_out)

 

-    p.extra_generation_params["Mask blur"] = mask_blur

+    if mask:

+        p.extra_generation_params["Mask blur"] = mask_blur

 

     if is_batch:

         assert not shared.cmd_opts.hide_ui_dir_config, "Launched with --hide-ui-dir-config, batch img2img disabled"

diff --git a/modules/interrogate.py b/modules/interrogate.py
index cbb80683..e1665708 100644
--- a/modules/interrogate.py
+++ b/modules/interrogate.py
@@ -32,7 +32,7 @@ def download_default_clip_interrogate_categories(content_dir):
     category_types = ["artists", "flavors", "mediums", "movements"]

 

     try:

-        os.makedirs(tmpdir)

+        os.makedirs(tmpdir, exist_ok=True)

         for category_type in category_types:

             torch.hub.download_url_to_file(f"https://raw.githubusercontent.com/pharmapsychotic/clip-interrogator/main/clip_interrogator/data/{category_type}.txt", os.path.join(tmpdir, f"{category_type}.txt"))

         os.rename(tmpdir, content_dir)

@@ -41,7 +41,7 @@ def download_default_clip_interrogate_categories(content_dir):
         errors.display(e, "downloading default CLIP interrogate categories")

     finally:

         if os.path.exists(tmpdir):

-            os.remove(tmpdir)

+            os.removedirs(tmpdir)

 

 

 class InterrogateModels:

diff --git a/modules/lowvram.py b/modules/lowvram.py
index 042a0254..e254cc13 100644
--- a/modules/lowvram.py
+++ b/modules/lowvram.py
@@ -55,12 +55,12 @@ def setup_for_low_vram(sd_model, use_medvram):
     if hasattr(sd_model.cond_stage_model, 'model'):

         sd_model.cond_stage_model.transformer = sd_model.cond_stage_model.model

 

-    # remove four big modules, cond, first_stage, depth (if applicable), and unet from the model and then

+    # remove several big modules: cond, first_stage, depth/embedder (if applicable), and unet from the model and then

     # send the model to GPU. Then put modules back. the modules will be in CPU.

-    stored = sd_model.cond_stage_model.transformer, sd_model.first_stage_model, getattr(sd_model, 'depth_model', None), sd_model.model

-    sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.depth_model, sd_model.model = None, None, None, None

+    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

     sd_model.to(devices.device)

-    sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.depth_model, sd_model.model = stored

+    sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.depth_model, sd_model.embedder, sd_model.model = stored

 

     # register hooks for those the first three models

     sd_model.cond_stage_model.transformer.register_forward_pre_hook(send_me_to_gpu)

@@ -69,6 +69,8 @@ def setup_for_low_vram(sd_model, use_medvram):
     sd_model.first_stage_model.decode = first_stage_model_decode_wrap

     if sd_model.depth_model:

         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'):

diff --git a/modules/mac_specific.py b/modules/mac_specific.py
new file mode 100644
index 00000000..6fe8dea0
--- /dev/null
+++ b/modules/mac_specific.py
@@ -0,0 +1,59 @@
+import torch
+import platform
+from modules import paths
+from modules.sd_hijack_utils import CondFunc
+from packaging import version
+
+
+# has_mps is only available in nightly pytorch (for now) and macOS 12.3+.
+# check `getattr` and try it for compatibility
+def check_for_mps() -> bool:
+    if not getattr(torch, 'has_mps', False):
+        return False
+    try:
+        torch.zeros(1).to(torch.device("mps"))
+        return True
+    except Exception:
+        return False
+has_mps = check_for_mps()
+
+
+# MPS workaround for https://github.com/pytorch/pytorch/issues/89784
+def cumsum_fix(input, cumsum_func, *args, **kwargs):
+    if input.device.type == 'mps':
+        output_dtype = kwargs.get('dtype', input.dtype)
+        if output_dtype == torch.int64:
+            return cumsum_func(input.cpu(), *args, **kwargs).to(input.device)
+        elif output_dtype == torch.bool or cumsum_needs_int_fix and (output_dtype == torch.int8 or output_dtype == torch.int16):
+            return cumsum_func(input.to(torch.int32), *args, **kwargs).to(torch.int64)
+    return cumsum_func(input, *args, **kwargs)
+
+
+if has_mps:
+    # MPS fix for randn in torchsde
+    CondFunc('torchsde._brownian.brownian_interval._randn', lambda _, size, dtype, device, seed: torch.randn(size, dtype=dtype, device=torch.device("cpu"), generator=torch.Generator(torch.device("cpu")).manual_seed(int(seed))).to(device), lambda _, size, dtype, device, seed: device.type == 'mps')
+
+    if platform.mac_ver()[0].startswith("13.2."):
+        # MPS workaround for https://github.com/pytorch/pytorch/issues/95188, thanks to danieldk (https://github.com/explosion/curated-transformers/pull/124)
+        CondFunc('torch.nn.functional.linear', lambda _, input, weight, bias: (torch.matmul(input, weight.t()) + bias) if bias is not None else torch.matmul(input, weight.t()), lambda _, input, weight, bias: input.numel() > 10485760)
+
+    if version.parse(torch.__version__) < version.parse("1.13"):
+        # PyTorch 1.13 doesn't need these fixes but unfortunately is slower and has regressions that prevent training from working
+
+        # MPS workaround for https://github.com/pytorch/pytorch/issues/79383
+        CondFunc('torch.Tensor.to', lambda orig_func, self, *args, **kwargs: orig_func(self.contiguous(), *args, **kwargs),
+                                                          lambda _, self, *args, **kwargs: self.device.type != 'mps' and (args and isinstance(args[0], torch.device) and args[0].type == 'mps' or isinstance(kwargs.get('device'), torch.device) and kwargs['device'].type == 'mps'))
+        # MPS workaround for https://github.com/pytorch/pytorch/issues/80800 
+        CondFunc('torch.nn.functional.layer_norm', lambda orig_func, *args, **kwargs: orig_func(*([args[0].contiguous()] + list(args[1:])), **kwargs),
+                                                                                        lambda _, *args, **kwargs: args and isinstance(args[0], torch.Tensor) and args[0].device.type == 'mps')
+        # MPS workaround for https://github.com/pytorch/pytorch/issues/90532
+        CondFunc('torch.Tensor.numpy', lambda orig_func, self, *args, **kwargs: orig_func(self.detach(), *args, **kwargs), lambda _, self, *args, **kwargs: self.requires_grad)
+    elif version.parse(torch.__version__) > version.parse("1.13.1"):
+        cumsum_needs_int_fix = not torch.Tensor([1,2]).to(torch.device("mps")).equal(torch.ShortTensor([1,1]).to(torch.device("mps")).cumsum(0))
+        cumsum_fix_func = lambda orig_func, input, *args, **kwargs: cumsum_fix(input, orig_func, *args, **kwargs)
+        CondFunc('torch.cumsum', cumsum_fix_func, None)
+        CondFunc('torch.Tensor.cumsum', cumsum_fix_func, None)
+        CondFunc('torch.narrow', lambda orig_func, *args, **kwargs: orig_func(*args, **kwargs).clone(), None)
+    if version.parse(torch.__version__) == version.parse("2.0"):
+        # MPS workaround for https://github.com/pytorch/pytorch/issues/96113
+        CondFunc('torch.nn.functional.layer_norm', lambda orig_func, x, normalized_shape, weight, bias, eps, **kwargs: orig_func(x.float(), normalized_shape, weight.float() if weight is not None else None, bias.float() if bias is not None else bias, eps).to(x.dtype), lambda *args, **kwargs: len(args) == 6)
diff --git a/modules/memmon.py b/modules/memmon.py
index a7060f58..4018edcc 100644
--- a/modules/memmon.py
+++ b/modules/memmon.py
@@ -23,12 +23,16 @@ class MemUsageMonitor(threading.Thread):
         self.data = defaultdict(int)
 
         try:
-            torch.cuda.mem_get_info()
+            self.cuda_mem_get_info()
             torch.cuda.memory_stats(self.device)
         except Exception as e:  # AMD or whatever
             print(f"Warning: caught exception '{e}', memory monitor disabled")
             self.disabled = True
 
+    def cuda_mem_get_info(self):
+        index = self.device.index if self.device.index is not None else torch.cuda.current_device()
+        return torch.cuda.mem_get_info(index)
+
     def run(self):
         if self.disabled:
             return
@@ -43,10 +47,10 @@ class MemUsageMonitor(threading.Thread):
                 self.run_flag.clear()
                 continue
 
-            self.data["min_free"] = torch.cuda.mem_get_info()[0]
+            self.data["min_free"] = self.cuda_mem_get_info()[0]
 
             while self.run_flag.is_set():
-                free, total = torch.cuda.mem_get_info()  # calling with self.device errors, torch bug?
+                free, total = self.cuda_mem_get_info()
                 self.data["min_free"] = min(self.data["min_free"], free)
 
                 time.sleep(1 / self.opts.memmon_poll_rate)
@@ -70,7 +74,7 @@ class MemUsageMonitor(threading.Thread):
 
     def read(self):
         if not self.disabled:
-            free, total = torch.cuda.mem_get_info()
+            free, total = self.cuda_mem_get_info()
             self.data["free"] = free
             self.data["total"] = total
 
diff --git a/modules/modelloader.py b/modules/modelloader.py
index e9aa514e..522affc6 100644
--- a/modules/modelloader.py
+++ b/modules/modelloader.py
@@ -4,9 +4,8 @@ import shutil
 import importlib
 from urllib.parse import urlparse
 
-from basicsr.utils.download_util import load_file_from_url
 from modules import shared
-from modules.upscaler import Upscaler
+from modules.upscaler import Upscaler, UpscalerLanczos, UpscalerNearest, UpscalerNone
 from modules.paths import script_path, models_path
 
 
@@ -45,6 +44,9 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None
                     full_path = file
                     if os.path.isdir(full_path):
                         continue
+                    if os.path.islink(full_path) and not os.path.exists(full_path):
+                        print(f"Skipping broken symlink: {full_path}")
+                        continue
                     if ext_blacklist is not None and any([full_path.endswith(x) for x in ext_blacklist]):
                         continue
                     if len(ext_filter) != 0:
@@ -56,6 +58,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
                 dl = load_file_from_url(model_url, model_path, True, download_name)
                 output.append(dl)
             else:
@@ -166,4 +169,8 @@ def load_upscalers():
         scaler = cls(commandline_options.get(cmd_name, None))
         datas += scaler.scalers
 
-    shared.sd_upscalers = datas
+    shared.sd_upscalers = sorted(
+        datas,
+        # Special case for UpscalerNone keeps it at the beginning of the list.
+        key=lambda x: x.name.lower() if not isinstance(x.scaler, (UpscalerNone, UpscalerLanczos, UpscalerNearest)) else ""
+    )
diff --git a/modules/models/diffusion/uni_pc/__init__.py b/modules/models/diffusion/uni_pc/__init__.py
new file mode 100644
index 00000000..e1265e3f
--- /dev/null
+++ b/modules/models/diffusion/uni_pc/__init__.py
@@ -0,0 +1 @@
+from .sampler import UniPCSampler
diff --git a/modules/models/diffusion/uni_pc/sampler.py b/modules/models/diffusion/uni_pc/sampler.py
new file mode 100644
index 00000000..a241c8a7
--- /dev/null
+++ b/modules/models/diffusion/uni_pc/sampler.py
@@ -0,0 +1,100 @@
+"""SAMPLING ONLY."""
+
+import torch
+
+from .uni_pc import NoiseScheduleVP, model_wrapper, UniPC
+from modules import shared, devices
+
+
+class UniPCSampler(object):
+    def __init__(self, model, **kwargs):
+        super().__init__()
+        self.model = model
+        to_torch = lambda x: x.clone().detach().to(torch.float32).to(model.device)
+        self.before_sample = None
+        self.after_sample = None
+        self.register_buffer('alphas_cumprod', to_torch(model.alphas_cumprod))
+
+    def register_buffer(self, name, attr):
+        if type(attr) == torch.Tensor:
+            if attr.device != devices.device:
+                attr = attr.to(devices.device)
+        setattr(self, name, attr)
+
+    def set_hooks(self, before_sample, after_sample, after_update):
+        self.before_sample = before_sample
+        self.after_sample = after_sample
+        self.after_update = after_update
+
+    @torch.no_grad()
+    def sample(self,
+               S,
+               batch_size,
+               shape,
+               conditioning=None,
+               callback=None,
+               normals_sequence=None,
+               img_callback=None,
+               quantize_x0=False,
+               eta=0.,
+               mask=None,
+               x0=None,
+               temperature=1.,
+               noise_dropout=0.,
+               score_corrector=None,
+               corrector_kwargs=None,
+               verbose=True,
+               x_T=None,
+               log_every_t=100,
+               unconditional_guidance_scale=1.,
+               unconditional_conditioning=None,
+               # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
+               **kwargs
+               ):
+        if conditioning is not None:
+            if isinstance(conditioning, dict):
+                ctmp = conditioning[list(conditioning.keys())[0]]
+                while isinstance(ctmp, list): ctmp = ctmp[0]
+                cbs = ctmp.shape[0]
+                if cbs != batch_size:
+                    print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
+
+            elif isinstance(conditioning, list):
+                for ctmp in conditioning:
+                    if ctmp.shape[0] != batch_size:
+                        print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
+
+            else:
+                if conditioning.shape[0] != batch_size:
+                    print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
+
+        # sampling
+        C, H, W = shape
+        size = (batch_size, C, H, W)
+        # print(f'Data shape for UniPC sampling is {size}')
+
+        device = self.model.betas.device
+        if x_T is None:
+            img = torch.randn(size, device=device)
+        else:
+            img = x_T
+
+        ns = NoiseScheduleVP('discrete', alphas_cumprod=self.alphas_cumprod)
+
+        # SD 1.X is "noise", SD 2.X is "v"
+        model_type = "v" if self.model.parameterization == "v" else "noise"
+
+        model_fn = model_wrapper(
+            lambda x, t, c: self.model.apply_model(x, t, c),
+            ns,
+            model_type=model_type,
+            guidance_type="classifier-free",
+            #condition=conditioning,
+            #unconditional_condition=unconditional_conditioning,
+            guidance_scale=unconditional_guidance_scale,
+        )
+
+        uni_pc = UniPC(model_fn, ns, predict_x0=True, thresholding=False, variant=shared.opts.uni_pc_variant, condition=conditioning, unconditional_condition=unconditional_conditioning, before_sample=self.before_sample, after_sample=self.after_sample, after_update=self.after_update)
+        x = uni_pc.sample(img, steps=S, skip_type=shared.opts.uni_pc_skip_type, method="multistep", order=shared.opts.uni_pc_order, lower_order_final=shared.opts.uni_pc_lower_order_final)
+
+        return x.to(device), None
diff --git a/modules/models/diffusion/uni_pc/uni_pc.py b/modules/models/diffusion/uni_pc/uni_pc.py
new file mode 100644
index 00000000..eb5f4e76
--- /dev/null
+++ b/modules/models/diffusion/uni_pc/uni_pc.py
@@ -0,0 +1,857 @@
+import torch
+import torch.nn.functional as F
+import math
+from tqdm.auto import trange
+
+
+class NoiseScheduleVP:
+    def __init__(
+            self,
+            schedule='discrete',
+            betas=None,
+            alphas_cumprod=None,
+            continuous_beta_0=0.1,
+            continuous_beta_1=20.,
+        ):
+        """Create a wrapper class for the forward SDE (VP type).
+
+        ***
+        Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t.
+                We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for high-resolution images.
+        ***
+
+        The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ).
+        We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper).
+        Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have:
+
+            log_alpha_t = self.marginal_log_mean_coeff(t)
+            sigma_t = self.marginal_std(t)
+            lambda_t = self.marginal_lambda(t)
+
+        Moreover, as lambda(t) is an invertible function, we also support its inverse function:
+
+            t = self.inverse_lambda(lambda_t)
+
+        ===============================================================
+
+        We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and continuous-time DPMs (trained on t in [t_0, T]).
+
+        1. For discrete-time DPMs:
+
+            For discrete-time DPMs trained on n = 0, 1, ..., N-1, we convert the discrete steps to continuous time steps by:
+                t_i = (i + 1) / N
+            e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1.
+            We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3.
+
+            Args:
+                betas: A `torch.Tensor`. The beta array for the discrete-time DPM. (See the original DDPM paper for details)
+                alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. (See the original DDPM paper for details)
+
+            Note that we always have alphas_cumprod = cumprod(betas). Therefore, we only need to set one of `betas` and `alphas_cumprod`.
+
+            **Important**:  Please pay special attention for the args for `alphas_cumprod`:
+                The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. Specifically, DDPMs assume that
+                    q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ).
+                Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. In fact, we have
+                    alpha_{t_n} = \sqrt{\hat{alpha_n}},
+                and
+                    log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}).
+
+
+        2. For continuous-time DPMs:
+
+            We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). The hyperparameters for the noise
+            schedule are the default settings in DDPM and improved-DDPM:
+
+            Args:
+                beta_min: A `float` number. The smallest beta for the linear schedule.
+                beta_max: A `float` number. The largest beta for the linear schedule.
+                cosine_s: A `float` number. The hyperparameter in the cosine schedule.
+                cosine_beta_max: A `float` number. The hyperparameter in the cosine schedule.
+                T: A `float` number. The ending time of the forward process.
+
+        ===============================================================
+
+        Args:
+            schedule: A `str`. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs,
+                    'linear' or 'cosine' for continuous-time DPMs.
+        Returns:
+            A wrapper object of the forward SDE (VP type).
+
+        ===============================================================
+
+        Example:
+
+        # For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1):
+        >>> ns = NoiseScheduleVP('discrete', betas=betas)
+
+        # For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array for n = 0, 1, ..., N - 1):
+        >>> ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod)
+
+        # For continuous-time DPMs (VPSDE), linear schedule:
+        >>> ns = NoiseScheduleVP('linear', continuous_beta_0=0.1, continuous_beta_1=20.)
+
+        """
+
+        if schedule not in ['discrete', 'linear', 'cosine']:
+            raise ValueError("Unsupported noise schedule {}. The schedule needs to be 'discrete' or 'linear' or 'cosine'".format(schedule))
+
+        self.schedule = schedule
+        if schedule == 'discrete':
+            if betas is not None:
+                log_alphas = 0.5 * torch.log(1 - betas).cumsum(dim=0)
+            else:
+                assert alphas_cumprod is not None
+                log_alphas = 0.5 * torch.log(alphas_cumprod)
+            self.total_N = len(log_alphas)
+            self.T = 1.
+            self.t_array = torch.linspace(0., 1., self.total_N + 1)[1:].reshape((1, -1))
+            self.log_alpha_array = log_alphas.reshape((1, -1,))
+        else:
+            self.total_N = 1000
+            self.beta_0 = continuous_beta_0
+            self.beta_1 = continuous_beta_1
+            self.cosine_s = 0.008
+            self.cosine_beta_max = 999.
+            self.cosine_t_max = math.atan(self.cosine_beta_max * (1. + self.cosine_s) / math.pi) * 2. * (1. + self.cosine_s) / math.pi - self.cosine_s
+            self.cosine_log_alpha_0 = math.log(math.cos(self.cosine_s / (1. + self.cosine_s) * math.pi / 2.))
+            self.schedule = schedule
+            if schedule == 'cosine':
+                # For the cosine schedule, T = 1 will have numerical issues. So we manually set the ending time T.
+                # Note that T = 0.9946 may be not the optimal setting. However, we find it works well.
+                self.T = 0.9946
+            else:
+                self.T = 1.
+
+    def marginal_log_mean_coeff(self, t):
+        """
+        Compute log(alpha_t) of a given continuous-time label t in [0, T].
+        """
+        if self.schedule == 'discrete':
+            return interpolate_fn(t.reshape((-1, 1)), self.t_array.to(t.device), self.log_alpha_array.to(t.device)).reshape((-1))
+        elif self.schedule == 'linear':
+            return -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
+        elif self.schedule == 'cosine':
+            log_alpha_fn = lambda s: torch.log(torch.cos((s + self.cosine_s) / (1. + self.cosine_s) * math.pi / 2.))
+            log_alpha_t =  log_alpha_fn(t) - self.cosine_log_alpha_0
+            return log_alpha_t
+
+    def marginal_alpha(self, t):
+        """
+        Compute alpha_t of a given continuous-time label t in [0, T].
+        """
+        return torch.exp(self.marginal_log_mean_coeff(t))
+
+    def marginal_std(self, t):
+        """
+        Compute sigma_t of a given continuous-time label t in [0, T].
+        """
+        return torch.sqrt(1. - torch.exp(2. * self.marginal_log_mean_coeff(t)))
+
+    def marginal_lambda(self, t):
+        """
+        Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T].
+        """
+        log_mean_coeff = self.marginal_log_mean_coeff(t)
+        log_std = 0.5 * torch.log(1. - torch.exp(2. * log_mean_coeff))
+        return log_mean_coeff - log_std
+
+    def inverse_lambda(self, lamb):
+        """
+        Compute the continuous-time label t in [0, T] of a given half-logSNR lambda_t.
+        """
+        if self.schedule == 'linear':
+            tmp = 2. * (self.beta_1 - self.beta_0) * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
+            Delta = self.beta_0**2 + tmp
+            return tmp / (torch.sqrt(Delta) + self.beta_0) / (self.beta_1 - self.beta_0)
+        elif self.schedule == 'discrete':
+            log_alpha = -0.5 * torch.logaddexp(torch.zeros((1,)).to(lamb.device), -2. * lamb)
+            t = interpolate_fn(log_alpha.reshape((-1, 1)), torch.flip(self.log_alpha_array.to(lamb.device), [1]), torch.flip(self.t_array.to(lamb.device), [1]))
+            return t.reshape((-1,))
+        else:
+            log_alpha = -0.5 * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
+            t_fn = lambda log_alpha_t: torch.arccos(torch.exp(log_alpha_t + self.cosine_log_alpha_0)) * 2. * (1. + self.cosine_s) / math.pi - self.cosine_s
+            t = t_fn(log_alpha)
+            return t
+
+
+def model_wrapper(
+    model,
+    noise_schedule,
+    model_type="noise",
+    model_kwargs={},
+    guidance_type="uncond",
+    #condition=None,
+    #unconditional_condition=None,
+    guidance_scale=1.,
+    classifier_fn=None,
+    classifier_kwargs={},
+):
+    """Create a wrapper function for the noise prediction model.
+
+    DPM-Solver needs to solve the continuous-time diffusion ODEs. For DPMs trained on discrete-time labels, we need to
+    firstly wrap the model function to a noise prediction model that accepts the continuous time as the input.
+
+    We support four types of the diffusion model by setting `model_type`:
+
+        1. "noise": noise prediction model. (Trained by predicting noise).
+
+        2. "x_start": data prediction model. (Trained by predicting the data x_0 at time 0).
+
+        3. "v": velocity prediction model. (Trained by predicting the velocity).
+            The "v" prediction is derivation detailed in Appendix D of [1], and is used in Imagen-Video [2].
+
+            [1] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models."
+                arXiv preprint arXiv:2202.00512 (2022).
+            [2] Ho, Jonathan, et al. "Imagen Video: High Definition Video Generation with Diffusion Models."
+                arXiv preprint arXiv:2210.02303 (2022).
+
+        4. "score": marginal score function. (Trained by denoising score matching).
+            Note that the score function and the noise prediction model follows a simple relationship:
+            ```
+                noise(x_t, t) = -sigma_t * score(x_t, t)
+            ```
+
+    We support three types of guided sampling by DPMs by setting `guidance_type`:
+        1. "uncond": unconditional sampling by DPMs.
+            The input `model` has the following format:
+            ``
+                model(x, t_input, **model_kwargs) -> noise | x_start | v | score
+            ``
+
+        2. "classifier": classifier guidance sampling [3] by DPMs and another classifier.
+            The input `model` has the following format:
+            ``
+                model(x, t_input, **model_kwargs) -> noise | x_start | v | score
+            ``
+
+            The input `classifier_fn` has the following format:
+            ``
+                classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond)
+            ``
+
+            [3] P. Dhariwal and A. Q. Nichol, "Diffusion models beat GANs on image synthesis,"
+                in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 8780-8794.
+
+        3. "classifier-free": classifier-free guidance sampling by conditional DPMs.
+            The input `model` has the following format:
+            ``
+                model(x, t_input, cond, **model_kwargs) -> noise | x_start | v | score
+            ``
+            And if cond == `unconditional_condition`, the model output is the unconditional DPM output.
+
+            [4] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance."
+                arXiv preprint arXiv:2207.12598 (2022).
+
+
+    The `t_input` is the time label of the model, which may be discrete-time labels (i.e. 0 to 999)
+    or continuous-time labels (i.e. epsilon to T).
+
+    We wrap the model function to accept only `x` and `t_continuous` as inputs, and outputs the predicted noise:
+    ``
+        def model_fn(x, t_continuous) -> noise:
+            t_input = get_model_input_time(t_continuous)
+            return noise_pred(model, x, t_input, **model_kwargs)
+    ``
+    where `t_continuous` is the continuous time labels (i.e. epsilon to T). And we use `model_fn` for DPM-Solver.
+
+    ===============================================================
+
+    Args:
+        model: A diffusion model with the corresponding format described above.
+        noise_schedule: A noise schedule object, such as NoiseScheduleVP.
+        model_type: A `str`. The parameterization type of the diffusion model.
+                    "noise" or "x_start" or "v" or "score".
+        model_kwargs: A `dict`. A dict for the other inputs of the model function.
+        guidance_type: A `str`. The type of the guidance for sampling.
+                    "uncond" or "classifier" or "classifier-free".
+        condition: A pytorch tensor. The condition for the guided sampling.
+                    Only used for "classifier" or "classifier-free" guidance type.
+        unconditional_condition: A pytorch tensor. The condition for the unconditional sampling.
+                    Only used for "classifier-free" guidance type.
+        guidance_scale: A `float`. The scale for the guided sampling.
+        classifier_fn: A classifier function. Only used for the classifier guidance.
+        classifier_kwargs: A `dict`. A dict for the other inputs of the classifier function.
+    Returns:
+        A noise prediction model that accepts the noised data and the continuous time as the inputs.
+    """
+
+    def get_model_input_time(t_continuous):
+        """
+        Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time.
+        For discrete-time DPMs, we convert `t_continuous` in [1 / N, 1] to `t_input` in [0, 1000 * (N - 1) / N].
+        For continuous-time DPMs, we just use `t_continuous`.
+        """
+        if noise_schedule.schedule == 'discrete':
+            return (t_continuous - 1. / noise_schedule.total_N) * 1000.
+        else:
+            return t_continuous
+
+    def noise_pred_fn(x, t_continuous, cond=None):
+        if t_continuous.reshape((-1,)).shape[0] == 1:
+            t_continuous = t_continuous.expand((x.shape[0]))
+        t_input = get_model_input_time(t_continuous)
+        if cond is None:
+            output = model(x, t_input, None, **model_kwargs)
+        else:
+            output = model(x, t_input, cond, **model_kwargs)
+        if model_type == "noise":
+            return output
+        elif model_type == "x_start":
+            alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
+            dims = x.dim()
+            return (x - expand_dims(alpha_t, dims) * output) / expand_dims(sigma_t, dims)
+        elif model_type == "v":
+            alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
+            dims = x.dim()
+            return expand_dims(alpha_t, dims) * output + expand_dims(sigma_t, dims) * x
+        elif model_type == "score":
+            sigma_t = noise_schedule.marginal_std(t_continuous)
+            dims = x.dim()
+            return -expand_dims(sigma_t, dims) * output
+
+    def cond_grad_fn(x, t_input, condition):
+        """
+        Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t).
+        """
+        with torch.enable_grad():
+            x_in = x.detach().requires_grad_(True)
+            log_prob = classifier_fn(x_in, t_input, condition, **classifier_kwargs)
+            return torch.autograd.grad(log_prob.sum(), x_in)[0]
+
+    def model_fn(x, t_continuous, condition, unconditional_condition):
+        """
+        The noise predicition model function that is used for DPM-Solver.
+        """
+        if t_continuous.reshape((-1,)).shape[0] == 1:
+            t_continuous = t_continuous.expand((x.shape[0]))
+        if guidance_type == "uncond":
+            return noise_pred_fn(x, t_continuous)
+        elif guidance_type == "classifier":
+            assert classifier_fn is not None
+            t_input = get_model_input_time(t_continuous)
+            cond_grad = cond_grad_fn(x, t_input, condition)
+            sigma_t = noise_schedule.marginal_std(t_continuous)
+            noise = noise_pred_fn(x, t_continuous)
+            return noise - guidance_scale * expand_dims(sigma_t, dims=cond_grad.dim()) * cond_grad
+        elif guidance_type == "classifier-free":
+            if guidance_scale == 1. or unconditional_condition is None:
+                return noise_pred_fn(x, t_continuous, cond=condition)
+            else:
+                x_in = torch.cat([x] * 2)
+                t_in = torch.cat([t_continuous] * 2)
+                if isinstance(condition, dict):
+                    assert isinstance(unconditional_condition, dict)
+                    c_in = dict()
+                    for k in condition:
+                        if isinstance(condition[k], list):
+                            c_in[k] = [torch.cat([
+                                unconditional_condition[k][i],
+                                condition[k][i]]) for i in range(len(condition[k]))]
+                        else:
+                            c_in[k] = torch.cat([
+                                unconditional_condition[k],
+                                condition[k]])
+                elif isinstance(condition, list):
+                    c_in = list()
+                    assert isinstance(unconditional_condition, list)
+                    for i in range(len(condition)):
+                        c_in.append(torch.cat([unconditional_condition[i], condition[i]]))
+                else:
+                    c_in = torch.cat([unconditional_condition, condition])
+                noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
+                return noise_uncond + guidance_scale * (noise - noise_uncond)
+
+    assert model_type in ["noise", "x_start", "v"]
+    assert guidance_type in ["uncond", "classifier", "classifier-free"]
+    return model_fn
+
+
+class UniPC:
+    def __init__(
+        self,
+        model_fn,
+        noise_schedule,
+        predict_x0=True,
+        thresholding=False,
+        max_val=1.,
+        variant='bh1',
+        condition=None,
+        unconditional_condition=None,
+        before_sample=None,
+        after_sample=None,
+        after_update=None
+    ):
+        """Construct a UniPC.
+
+        We support both data_prediction and noise_prediction.
+        """
+        self.model_fn_ = model_fn
+        self.noise_schedule = noise_schedule
+        self.variant = variant
+        self.predict_x0 = predict_x0
+        self.thresholding = thresholding
+        self.max_val = max_val
+        self.condition = condition
+        self.unconditional_condition = unconditional_condition
+        self.before_sample = before_sample
+        self.after_sample = after_sample
+        self.after_update = after_update
+
+    def dynamic_thresholding_fn(self, x0, t=None):
+        """
+        The dynamic thresholding method.
+        """
+        dims = x0.dim()
+        p = self.dynamic_thresholding_ratio
+        s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
+        s = expand_dims(torch.maximum(s, self.thresholding_max_val * torch.ones_like(s).to(s.device)), dims)
+        x0 = torch.clamp(x0, -s, s) / s
+        return x0
+
+    def model(self, x, t):
+        cond = self.condition
+        uncond = self.unconditional_condition
+        if self.before_sample is not None:
+            x, t, cond, uncond = self.before_sample(x, t, cond, uncond)
+        res = self.model_fn_(x, t, cond, uncond)
+        if self.after_sample is not None:
+            x, t, cond, uncond, res = self.after_sample(x, t, cond, uncond, res)
+
+        if isinstance(res, tuple):
+            # (None, pred_x0)
+            res = res[1]
+
+        return res
+
+    def noise_prediction_fn(self, x, t):
+        """
+        Return the noise prediction model.
+        """
+        return self.model(x, t)
+
+    def data_prediction_fn(self, x, t):
+        """
+        Return the data prediction model (with thresholding).
+        """
+        noise = self.noise_prediction_fn(x, t)
+        dims = x.dim()
+        alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t)
+        x0 = (x - expand_dims(sigma_t, dims) * noise) / expand_dims(alpha_t, dims)
+        if self.thresholding:
+            p = 0.995   # A hyperparameter in the paper of "Imagen" [1].
+            s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
+            s = expand_dims(torch.maximum(s, self.max_val * torch.ones_like(s).to(s.device)), dims)
+            x0 = torch.clamp(x0, -s, s) / s
+        return x0
+
+    def model_fn(self, x, t):
+        """
+        Convert the model to the noise prediction model or the data prediction model.
+        """
+        if self.predict_x0:
+            return self.data_prediction_fn(x, t)
+        else:
+            return self.noise_prediction_fn(x, t)
+
+    def get_time_steps(self, skip_type, t_T, t_0, N, device):
+        """Compute the intermediate time steps for sampling.
+        """
+        if skip_type == 'logSNR':
+            lambda_T = self.noise_schedule.marginal_lambda(torch.tensor(t_T).to(device))
+            lambda_0 = self.noise_schedule.marginal_lambda(torch.tensor(t_0).to(device))
+            logSNR_steps = torch.linspace(lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1).to(device)
+            return self.noise_schedule.inverse_lambda(logSNR_steps)
+        elif skip_type == 'time_uniform':
+            return torch.linspace(t_T, t_0, N + 1).to(device)
+        elif skip_type == 'time_quadratic':
+            t_order = 2
+            t = torch.linspace(t_T**(1. / t_order), t_0**(1. / t_order), N + 1).pow(t_order).to(device)
+            return t
+        else:
+            raise ValueError("Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(skip_type))
+
+    def get_orders_and_timesteps_for_singlestep_solver(self, steps, order, skip_type, t_T, t_0, device):
+        """
+        Get the order of each step for sampling by the singlestep DPM-Solver.
+        """
+        if order == 3:
+            K = steps // 3 + 1
+            if steps % 3 == 0:
+                orders = [3,] * (K - 2) + [2, 1]
+            elif steps % 3 == 1:
+                orders = [3,] * (K - 1) + [1]
+            else:
+                orders = [3,] * (K - 1) + [2]
+        elif order == 2:
+            if steps % 2 == 0:
+                K = steps // 2
+                orders = [2,] * K
+            else:
+                K = steps // 2 + 1
+                orders = [2,] * (K - 1) + [1]
+        elif order == 1:
+            K = steps
+            orders = [1,] * steps
+        else:
+            raise ValueError("'order' must be '1' or '2' or '3'.")
+        if skip_type == 'logSNR':
+            # To reproduce the results in DPM-Solver paper
+            timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, K, device)
+        else:
+            timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, steps, device)[torch.cumsum(torch.tensor([0,] + orders), 0).to(device)]
+        return timesteps_outer, orders
+
+    def denoise_to_zero_fn(self, x, s):
+        """
+        Denoise at the final step, which is equivalent to solve the ODE from lambda_s to infty by first-order discretization.
+        """
+        return self.data_prediction_fn(x, s)
+
+    def multistep_uni_pc_update(self, x, model_prev_list, t_prev_list, t, order, **kwargs):
+        if len(t.shape) == 0:
+            t = t.view(-1)
+        if 'bh' in self.variant:
+            return self.multistep_uni_pc_bh_update(x, model_prev_list, t_prev_list, t, order, **kwargs)
+        else:
+            assert self.variant == 'vary_coeff'
+            return self.multistep_uni_pc_vary_update(x, model_prev_list, t_prev_list, t, order, **kwargs)
+
+    def multistep_uni_pc_vary_update(self, x, model_prev_list, t_prev_list, t, order, use_corrector=True):
+        #print(f'using unified predictor-corrector with order {order} (solver type: vary coeff)')
+        ns = self.noise_schedule
+        assert order <= len(model_prev_list)
+
+        # first compute rks
+        t_prev_0 = t_prev_list[-1]
+        lambda_prev_0 = ns.marginal_lambda(t_prev_0)
+        lambda_t = ns.marginal_lambda(t)
+        model_prev_0 = model_prev_list[-1]
+        sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
+        log_alpha_t = ns.marginal_log_mean_coeff(t)
+        alpha_t = torch.exp(log_alpha_t)
+
+        h = lambda_t - lambda_prev_0
+
+        rks = []
+        D1s = []
+        for i in range(1, order):
+            t_prev_i = t_prev_list[-(i + 1)]
+            model_prev_i = model_prev_list[-(i + 1)]
+            lambda_prev_i = ns.marginal_lambda(t_prev_i)
+            rk = (lambda_prev_i - lambda_prev_0) / h
+            rks.append(rk)
+            D1s.append((model_prev_i - model_prev_0) / rk)
+
+        rks.append(1.)
+        rks = torch.tensor(rks, device=x.device)
+
+        K = len(rks)
+        # build C matrix
+        C = []
+
+        col = torch.ones_like(rks)
+        for k in range(1, K + 1):
+            C.append(col)
+            col = col * rks / (k + 1)
+        C = torch.stack(C, dim=1)
+
+        if len(D1s) > 0:
+            D1s = torch.stack(D1s, dim=1) # (B, K)
+            C_inv_p = torch.linalg.inv(C[:-1, :-1])
+            A_p = C_inv_p
+
+        if use_corrector:
+            #print('using corrector')
+            C_inv = torch.linalg.inv(C)
+            A_c = C_inv
+
+        hh = -h if self.predict_x0 else h
+        h_phi_1 = torch.expm1(hh)
+        h_phi_ks = []
+        factorial_k = 1
+        h_phi_k = h_phi_1
+        for k in range(1, K + 2):
+            h_phi_ks.append(h_phi_k)
+            h_phi_k = h_phi_k / hh - 1 / factorial_k
+            factorial_k *= (k + 1)
+
+        model_t = None
+        if self.predict_x0:
+            x_t_ = (
+                sigma_t / sigma_prev_0 * x
+                - alpha_t * h_phi_1 * model_prev_0
+            )
+            # now predictor
+            x_t = x_t_
+            if len(D1s) > 0:
+                # compute the residuals for predictor
+                for k in range(K - 1):
+                    x_t = x_t - alpha_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_p[k])
+            # now corrector
+            if use_corrector:
+                model_t = self.model_fn(x_t, t)
+                D1_t = (model_t - model_prev_0)
+                x_t = x_t_
+                k = 0
+                for k in range(K - 1):
+                    x_t = x_t - alpha_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_c[k][:-1])
+                x_t = x_t - alpha_t * h_phi_ks[K] * (D1_t * A_c[k][-1])
+        else:
+            log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
+            x_t_ = (
+                (torch.exp(log_alpha_t - log_alpha_prev_0)) * x
+                - (sigma_t * h_phi_1) * model_prev_0
+            )
+            # now predictor
+            x_t = x_t_
+            if len(D1s) > 0:
+                # compute the residuals for predictor
+                for k in range(K - 1):
+                    x_t = x_t - sigma_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_p[k])
+            # now corrector
+            if use_corrector:
+                model_t = self.model_fn(x_t, t)
+                D1_t = (model_t - model_prev_0)
+                x_t = x_t_
+                k = 0
+                for k in range(K - 1):
+                    x_t = x_t - sigma_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_c[k][:-1])
+                x_t = x_t - sigma_t * h_phi_ks[K] * (D1_t * A_c[k][-1])
+        return x_t, model_t
+
+    def multistep_uni_pc_bh_update(self, x, model_prev_list, t_prev_list, t, order, x_t=None, use_corrector=True):
+        #print(f'using unified predictor-corrector with order {order} (solver type: B(h))')
+        ns = self.noise_schedule
+        assert order <= len(model_prev_list)
+        dims = x.dim()
+
+        # first compute rks
+        t_prev_0 = t_prev_list[-1]
+        lambda_prev_0 = ns.marginal_lambda(t_prev_0)
+        lambda_t = ns.marginal_lambda(t)
+        model_prev_0 = model_prev_list[-1]
+        sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
+        log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
+        alpha_t = torch.exp(log_alpha_t)
+
+        h = lambda_t - lambda_prev_0
+
+        rks = []
+        D1s = []
+        for i in range(1, order):
+            t_prev_i = t_prev_list[-(i + 1)]
+            model_prev_i = model_prev_list[-(i + 1)]
+            lambda_prev_i = ns.marginal_lambda(t_prev_i)
+            rk = ((lambda_prev_i - lambda_prev_0) / h)[0]
+            rks.append(rk)
+            D1s.append((model_prev_i - model_prev_0) / rk)
+
+        rks.append(1.)
+        rks = torch.tensor(rks, device=x.device)
+
+        R = []
+        b = []
+
+        hh = -h[0] if self.predict_x0 else h[0]
+        h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1
+        h_phi_k = h_phi_1 / hh - 1
+
+        factorial_i = 1
+
+        if self.variant == 'bh1':
+            B_h = hh
+        elif self.variant == 'bh2':
+            B_h = torch.expm1(hh)
+        else:
+            raise NotImplementedError()
+
+        for i in range(1, order + 1):
+            R.append(torch.pow(rks, i - 1))
+            b.append(h_phi_k * factorial_i / B_h)
+            factorial_i *= (i + 1)
+            h_phi_k = h_phi_k / hh - 1 / factorial_i
+
+        R = torch.stack(R)
+        b = torch.tensor(b, device=x.device)
+
+        # now predictor
+        use_predictor = len(D1s) > 0 and x_t is None
+        if len(D1s) > 0:
+            D1s = torch.stack(D1s, dim=1) # (B, K)
+            if x_t is None:
+                # for order 2, we use a simplified version
+                if order == 2:
+                    rhos_p = torch.tensor([0.5], device=b.device)
+                else:
+                    rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1])
+        else:
+            D1s = None
+
+        if use_corrector:
+            #print('using corrector')
+            # for order 1, we use a simplified version
+            if order == 1:
+                rhos_c = torch.tensor([0.5], device=b.device)
+            else:
+                rhos_c = torch.linalg.solve(R, b)
+
+        model_t = None
+        if self.predict_x0:
+            x_t_ = (
+                expand_dims(sigma_t / sigma_prev_0, dims) * x
+                - expand_dims(alpha_t * h_phi_1, dims)* model_prev_0
+            )
+
+            if x_t is None:
+                if use_predictor:
+                    pred_res = torch.einsum('k,bkchw->bchw', rhos_p, D1s)
+                else:
+                    pred_res = 0
+                x_t = x_t_ - expand_dims(alpha_t * B_h, dims) * pred_res
+
+            if use_corrector:
+                model_t = self.model_fn(x_t, t)
+                if D1s is not None:
+                    corr_res = torch.einsum('k,bkchw->bchw', rhos_c[:-1], D1s)
+                else:
+                    corr_res = 0
+                D1_t = (model_t - model_prev_0)
+                x_t = x_t_ - expand_dims(alpha_t * B_h, dims) * (corr_res + rhos_c[-1] * D1_t)
+        else:
+            x_t_ = (
+                expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
+                - expand_dims(sigma_t * h_phi_1, dims) * model_prev_0
+            )
+            if x_t is None:
+                if use_predictor:
+                    pred_res = torch.einsum('k,bkchw->bchw', rhos_p, D1s)
+                else:
+                    pred_res = 0
+                x_t = x_t_ - expand_dims(sigma_t * B_h, dims) * pred_res
+
+            if use_corrector:
+                model_t = self.model_fn(x_t, t)
+                if D1s is not None:
+                    corr_res = torch.einsum('k,bkchw->bchw', rhos_c[:-1], D1s)
+                else:
+                    corr_res = 0
+                D1_t = (model_t - model_prev_0)
+                x_t = x_t_ - expand_dims(sigma_t * B_h, dims) * (corr_res + rhos_c[-1] * D1_t)
+        return x_t, model_t
+
+
+    def sample(self, x, steps=20, t_start=None, t_end=None, order=3, skip_type='time_uniform',
+        method='singlestep', lower_order_final=True, denoise_to_zero=False, solver_type='dpm_solver',
+        atol=0.0078, rtol=0.05, corrector=False,
+    ):
+        t_0 = 1. / self.noise_schedule.total_N if t_end is None else t_end
+        t_T = self.noise_schedule.T if t_start is None else t_start
+        device = x.device
+        if method == 'multistep':
+            assert steps >= order, "UniPC order must be < sampling steps"
+            timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device)
+            #print(f"Running UniPC Sampling with {timesteps.shape[0]} timesteps, order {order}")
+            assert timesteps.shape[0] - 1 == steps
+            with torch.no_grad():
+                vec_t = timesteps[0].expand((x.shape[0]))
+                model_prev_list = [self.model_fn(x, vec_t)]
+                t_prev_list = [vec_t]
+                # Init the first `order` values by lower order multistep DPM-Solver.
+                for init_order in range(1, order):
+                    vec_t = timesteps[init_order].expand(x.shape[0])
+                    x, model_x = self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, vec_t, init_order, use_corrector=True)
+                    if model_x is None:
+                        model_x = self.model_fn(x, vec_t)
+                    if self.after_update is not None:
+                        self.after_update(x, model_x)
+                    model_prev_list.append(model_x)
+                    t_prev_list.append(vec_t)
+                for step in trange(order, steps + 1):
+                    vec_t = timesteps[step].expand(x.shape[0])
+                    if lower_order_final:
+                        step_order = min(order, steps + 1 - step)
+                    else:
+                        step_order = order
+                    #print('this step order:', step_order)
+                    if step == steps:
+                        #print('do not run corrector at the last step')
+                        use_corrector = False
+                    else:
+                        use_corrector = True
+                    x, model_x =  self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, vec_t, step_order, use_corrector=use_corrector)
+                    if self.after_update is not None:
+                        self.after_update(x, model_x)
+                    for i in range(order - 1):
+                        t_prev_list[i] = t_prev_list[i + 1]
+                        model_prev_list[i] = model_prev_list[i + 1]
+                    t_prev_list[-1] = vec_t
+                    # We do not need to evaluate the final model value.
+                    if step < steps:
+                        if model_x is None:
+                            model_x = self.model_fn(x, vec_t)
+                        model_prev_list[-1] = model_x
+        else:
+            raise NotImplementedError()
+        if denoise_to_zero:
+            x = self.denoise_to_zero_fn(x, torch.ones((x.shape[0],)).to(device) * t_0)
+        return x
+
+
+#############################################################
+# other utility functions
+#############################################################
+
+def interpolate_fn(x, xp, yp):
+    """
+    A piecewise linear function y = f(x), using xp and yp as keypoints.
+    We implement f(x) in a differentiable way (i.e. applicable for autograd).
+    The function f(x) is well-defined for all x-axis. (For x beyond the bounds of xp, we use the outmost points of xp to define the linear function.)
+
+    Args:
+        x: PyTorch tensor with shape [N, C], where N is the batch size, C is the number of channels (we use C = 1 for DPM-Solver).
+        xp: PyTorch tensor with shape [C, K], where K is the number of keypoints.
+        yp: PyTorch tensor with shape [C, K].
+    Returns:
+        The function values f(x), with shape [N, C].
+    """
+    N, K = x.shape[0], xp.shape[1]
+    all_x = torch.cat([x.unsqueeze(2), xp.unsqueeze(0).repeat((N, 1, 1))], dim=2)
+    sorted_all_x, x_indices = torch.sort(all_x, dim=2)
+    x_idx = torch.argmin(x_indices, dim=2)
+    cand_start_idx = x_idx - 1
+    start_idx = torch.where(
+        torch.eq(x_idx, 0),
+        torch.tensor(1, device=x.device),
+        torch.where(
+            torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
+        ),
+    )
+    end_idx = torch.where(torch.eq(start_idx, cand_start_idx), start_idx + 2, start_idx + 1)
+    start_x = torch.gather(sorted_all_x, dim=2, index=start_idx.unsqueeze(2)).squeeze(2)
+    end_x = torch.gather(sorted_all_x, dim=2, index=end_idx.unsqueeze(2)).squeeze(2)
+    start_idx2 = torch.where(
+        torch.eq(x_idx, 0),
+        torch.tensor(0, device=x.device),
+        torch.where(
+            torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
+        ),
+    )
+    y_positions_expanded = yp.unsqueeze(0).expand(N, -1, -1)
+    start_y = torch.gather(y_positions_expanded, dim=2, index=start_idx2.unsqueeze(2)).squeeze(2)
+    end_y = torch.gather(y_positions_expanded, dim=2, index=(start_idx2 + 1).unsqueeze(2)).squeeze(2)
+    cand = start_y + (x - start_x) * (end_y - start_y) / (end_x - start_x)
+    return cand
+
+
+def expand_dims(v, dims):
+    """
+    Expand the tensor `v` to the dim `dims`.
+
+    Args:
+        `v`: a PyTorch tensor with shape [N].
+        `dim`: a `int`.
+    Returns:
+        a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`.
+    """
+    return v[(...,) + (None,)*(dims - 1)]
diff --git a/modules/ngrok.py b/modules/ngrok.py
index 3df2c06b..1ad7989b 100644
--- a/modules/ngrok.py
+++ b/modules/ngrok.py
@@ -13,6 +13,18 @@ def connect(token, port, region):
     config = conf.PyngrokConfig(
         auth_token=token, region=region
     )
+    
+    # Guard for existing tunnels
+    existing = ngrok.get_tunnels(pyngrok_config=config)
+    if existing:
+        for established in existing:
+            # Extra configuration in the case that the user is also using ngrok for other tunnels
+            if established.config['addr'][-4:] == str(port):
+                public_url = existing[0].public_url
+                print(f'ngrok has already been connected to localhost:{port}! URL: {public_url}\n'
+                    'You can use this link after the launch is complete.')
+                return
+    
     try:
         if account is None:
             public_url = ngrok.connect(port, pyngrok_config=config, bind_tls=True).public_url
diff --git a/modules/paths.py b/modules/paths.py
index d991cc71..0e1e00e7 100644
--- a/modules/paths.py
+++ b/modules/paths.py
@@ -1,16 +1,9 @@
-import argparse

 import os

 import sys

-import modules.safe

+from modules.paths_internal import models_path, script_path, data_path, extensions_dir, extensions_builtin_dir

 

-script_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))

+import modules.safe

 

-# Parse the --data-dir flag first so we can use it as a base for our other argument default values

-parser = argparse.ArgumentParser(add_help=False)

-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",)

-cmd_opts_pre = parser.parse_known_args()[0]

-data_path = cmd_opts_pre.data_dir

-models_path = os.path.join(data_path, "models")

 

 # data_path = cmd_opts_pre.data

 sys.path.insert(0, script_path)

diff --git a/modules/paths_internal.py b/modules/paths_internal.py
new file mode 100644
index 00000000..6765bafe
--- /dev/null
+++ b/modules/paths_internal.py
@@ -0,0 +1,23 @@
+"""this module defines internal paths used by program and is safe to import before dependencies are installed in launch.py"""

+

+import argparse

+import os

+

+script_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))

+

+sd_configs_path = os.path.join(script_path, "configs")

+sd_default_config = os.path.join(sd_configs_path, "v1-inference.yaml")

+sd_model_file = os.path.join(script_path, 'model.ckpt')

+default_sd_model_file = sd_model_file

+

+# Parse the --data-dir flag first so we can use it as a base for our other argument default values

+parser_pre = argparse.ArgumentParser(add_help=False)

+parser_pre.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",)

+cmd_opts_pre = parser_pre.parse_known_args()[0]

+

+data_path = cmd_opts_pre.data_dir

+

+models_path = os.path.join(data_path, "models")

+extensions_dir = os.path.join(data_path, "extensions")

+extensions_builtin_dir = os.path.join(script_path, "extensions-builtin")

+config_states_dir = os.path.join(script_path, "config_states")

diff --git a/modules/postprocessing.py b/modules/postprocessing.py
index 09d8e605..736315e2 100644
--- a/modules/postprocessing.py
+++ b/modules/postprocessing.py
@@ -18,9 +18,14 @@ def run_postprocessing(extras_mode, image, image_folder, input_dir, output_dir,
 

     if extras_mode == 1:

         for img in image_folder:

-            image = Image.open(img)

+            if isinstance(img, Image.Image):

+                image = img

+                fn = ''

+            else:

+                image = Image.open(os.path.abspath(img.name))

+                fn = os.path.splitext(img.orig_name)[0]

             image_data.append(image)

-            image_names.append(os.path.splitext(img.orig_name)[0])

+            image_names.append(fn)

     elif extras_mode == 2:

         assert not shared.cmd_opts.hide_ui_dir_config, '--hide-ui-dir-config option must be disabled'

         assert input_dir, 'input directory not selected'

diff --git a/modules/processing.py b/modules/processing.py
index e544c2e1..a48fff99 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -3,6 +3,7 @@ import math
 import os

 import sys

 import warnings

+import hashlib

 

 import torch

 import numpy as np

@@ -78,28 +79,34 @@ def apply_overlay(image, paste_loc, index, overlays):
 

 

 def txt2img_image_conditioning(sd_model, x, width, height):

-    if sd_model.model.conditioning_key not in {'hybrid', 'concat'}:

-        # Dummy zero conditioning if we're not using inpainting model.

-        # Still takes up a bit of memory, but no encoder call.

-        # Pretty sure we can just make this a 1x1 image since its not going to be used besides its batch size.

-        return x.new_zeros(x.shape[0], 5, 1, 1, dtype=x.dtype, device=x.device)

+    if sd_model.model.conditioning_key in {'hybrid', 'concat'}: # Inpainting models

+

+        # 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))

 

-    # 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))

+        # 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)

+        image_conditioning = image_conditioning.to(x.dtype)

 

-    # 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)

-    image_conditioning = image_conditioning.to(x.dtype)

+        return image_conditioning

+

+    elif sd_model.model.conditioning_key == "crossattn-adm": # UnCLIP models

 

-    return image_conditioning

+        return x.new_zeros(x.shape[0], 2*sd_model.noise_augmentor.time_embed.dim, dtype=x.dtype, device=x.device)

+

+    else:

+        # Dummy zero conditioning if we're not using inpainting or unclip models.

+        # Still takes up a bit of memory, but no encoder call.

+        # Pretty sure we can just make this a 1x1 image since its not going to be used besides its batch size.

+        return x.new_zeros(x.shape[0], 5, 1, 1, dtype=x.dtype, device=x.device)

 

 

 class StableDiffusionProcessing:

     """

     The first set of paramaters: sd_models -> do_not_reload_embeddings represent the minimum required to create a StableDiffusionProcessing

     """

-    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_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 = 1.0, 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)

 

@@ -134,6 +141,7 @@ class StableDiffusionProcessing:
         self.denoising_strength: float = denoising_strength

         self.sampler_noise_scheduler_override = None

         self.ddim_discretize = ddim_discretize or opts.ddim_discretize

+        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

@@ -156,6 +164,8 @@ class StableDiffusionProcessing:
         self.all_seeds = None

         self.all_subseeds = None

         self.iteration = 0

+        self.is_hr_pass = False

+        

 

     @property

     def sd_model(self):

@@ -186,10 +196,18 @@ class StableDiffusionProcessing:
         return conditioning

 

     def edit_image_conditioning(self, source_image):

-        conditioning_image = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(source_image))

+        conditioning_image = self.sd_model.encode_first_stage(source_image).mode()

 

         return conditioning_image

 

+    def unclip_image_conditioning(self, source_image):

+        c_adm = self.sd_model.embedder(source_image)

+        if self.sd_model.noise_augmentor is not None:

+            noise_level = 0 # TODO: Allow other noise levels?

+            c_adm, noise_level_emb = self.sd_model.noise_augmentor(c_adm, noise_level=repeat(torch.tensor([noise_level]).to(c_adm.device), '1 -> b', b=c_adm.shape[0]))

+            c_adm = torch.cat((c_adm, noise_level_emb), 1)

+        return c_adm

+

     def inpainting_image_conditioning(self, source_image, latent_image, image_mask=None):

         self.is_using_inpainting_conditioning = True

 

@@ -241,6 +259,9 @@ class StableDiffusionProcessing:
         if self.sampler.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":

+            return self.unclip_image_conditioning(source_image)

+

         # Dummy zero conditioning if we're not using inpainting or depth model.

         return latent_image.new_zeros(latent_image.shape[0], 5, 1, 1)

 

@@ -268,6 +289,7 @@ class Processed:
         self.height = p.height

         self.sampler_name = p.sampler_name

         self.cfg_scale = p.cfg_scale

+        self.image_cfg_scale = getattr(p, 'image_cfg_scale', None)

         self.steps = p.steps

         self.batch_size = p.batch_size

         self.restore_faces = p.restore_faces

@@ -445,6 +467,7 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iter
         "Steps": p.steps,

         "Sampler": p.sampler_name,

         "CFG scale": p.cfg_scale,

+        "Image CFG scale": getattr(p, 'image_cfg_scale', None),

         "Seed": all_seeds[index],

         "Face restoration": (opts.face_restoration_model if p.restore_faces else None),

         "Size": f"{p.width}x{p.height}",

@@ -457,6 +480,9 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iter
         "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,

         "ENSD": None if opts.eta_noise_seed_delta == 0 else opts.eta_noise_seed_delta,

+        "Init image hash": getattr(p, 'init_img_hash', None),

+        "RNG": opts.randn_source if opts.randn_source != "GPU" else None,

+        "NGMS": None if p.s_min_uncond == 0 else p.s_min_uncond,

     }

 

     generation_params.update(p.extra_generation_params)

@@ -541,8 +567,6 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
     if os.path.exists(cmd_opts.embeddings_dir) and not p.do_not_reload_embeddings:

         model_hijack.embedding_db.load_textual_inversion_embeddings()

 

-    _, extra_network_data = extra_networks.parse_prompts(p.all_prompts[0:1])

-

     if p.scripts is not None:

         p.scripts.process(p)

 

@@ -580,16 +604,10 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             if shared.opts.live_previews_enable and opts.show_progress_type == "Approx NN":

                 sd_vae_approx.model()

 

-            if not p.disable_extra_networks:

-                extra_networks.activate(p, extra_network_data)

-

-        with open(os.path.join(paths.data_path, "params.txt"), "w", encoding="utf8") as file:

-            processed = Processed(p, [], p.seed, "")

-            file.write(processed.infotext(p, 0))

-

         if state.job_count == -1:

             state.job_count = p.n_iter

 

+        extra_network_data = None

         for n in range(p.n_iter):

             p.iteration = n

 

@@ -604,16 +622,38 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             seeds = p.all_seeds[n * p.batch_size:(n + 1) * p.batch_size]

             subseeds = p.all_subseeds[n * p.batch_size:(n + 1) * p.batch_size]

 

+            if p.scripts is not None:

+                p.scripts.before_process_batch(p, batch_number=n, prompts=prompts, seeds=seeds, subseeds=subseeds)

+

             if len(prompts) == 0:

                 break

 

-            prompts, _ = extra_networks.parse_prompts(prompts)

+            prompts, extra_network_data = extra_networks.parse_prompts(prompts)

+

+            if not p.disable_extra_networks:

+                with devices.autocast():

+                    extra_networks.activate(p, extra_network_data)

 

             if p.scripts is not None:

                 p.scripts.process_batch(p, batch_number=n, prompts=prompts, seeds=seeds, subseeds=subseeds)

 

-            uc = get_conds_with_caching(prompt_parser.get_learned_conditioning, negative_prompts, p.steps, cached_uc)

-            c = get_conds_with_caching(prompt_parser.get_multicond_learned_conditioning, prompts, p.steps, cached_c)

+            # params.txt should be saved after scripts.process_batch, since the

+            # infotext could be modified by that callback

+            # Example: a wildcard processed by process_batch sets an extra model

+            # strength, which is saved as "Model Strength: 1.0" in the infotext

+            if n == 0:

+                with open(os.path.join(paths.data_path, "params.txt"), "w", encoding="utf8") as file:

+                    processed = Processed(p, [], p.seed, "")

+                    file.write(processed.infotext(p, 0))

+

+            step_multiplier = 1

+            if not shared.opts.dont_fix_second_order_samplers_schedule:

+                try:

+                    step_multiplier = 2 if sd_samplers.all_samplers_map.get(p.sampler_name).aliases[0] in ['k_dpmpp_2s_a', 'k_dpmpp_2s_a_ka', 'k_dpmpp_sde', 'k_dpmpp_sde_ka', 'k_dpm_2', 'k_dpm_2_a', 'k_heun'] else 1

+                except:

+                    pass

+            uc = get_conds_with_caching(prompt_parser.get_learned_conditioning, negative_prompts, p.steps * step_multiplier, cached_uc)

+            c = get_conds_with_caching(prompt_parser.get_multicond_learned_conditioning, prompts, p.steps * step_multiplier, cached_c)

 

             if len(model_hijack.comments) > 0:

                 for comment in model_hijack.comments:

@@ -643,6 +683,8 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                 p.scripts.postprocess_batch(p, x_samples_ddim, batch_number=n)

 

             for i, x_sample in enumerate(x_samples_ddim):

+                p.batch_index = i

+

                 x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)

                 x_sample = x_sample.astype(np.uint8)

 

@@ -679,6 +721,22 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                     image.info["parameters"] = text

                 output_images.append(image)

 

+                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, "", seeds[i], prompts[i], opts.samples_format, info=infotext(n, i), p=p, suffix="-mask")

+

+                    if opts.save_mask_composite:

+                        images.save_image(image_mask_composite, p.outpath_samples, "", seeds[i], prompts[i], opts.samples_format, info=infotext(n, i), p=p, suffix="-mask-composite")

+

+                    if opts.return_mask:

+                        output_images.append(image_mask)

+

+                    if opts.return_mask_composite:

+                        output_images.append(image_mask_composite)

+

             del x_samples_ddim

 

             devices.torch_gc()

@@ -703,7 +761,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             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)

 

-    if not p.disable_extra_networks:

+    if not p.disable_extra_networks and extra_network_data:

         extra_networks.deactivate(p, extra_network_data)

 

     devices.torch_gc()

@@ -828,6 +886,8 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
         if not self.enable_hr:

             return samples

 

+        self.is_hr_pass = True

+

         target_width = self.hr_upscale_to_x

         target_height = self.hr_upscale_to_y

 

@@ -882,7 +942,9 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 

         shared.state.nextjob()

 

-        img2img_sampler_name = self.sampler_name if self.sampler_name != 'PLMS' else 'DDIM'  # PLMS does not support img2img so we just silently switch ot DDIM

+        img2img_sampler_name = 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]

@@ -895,18 +957,21 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 

         samples = self.sampler.sample_img2img(self, samples, noise, conditioning, unconditional_conditioning, steps=self.hr_second_pass_steps or self.steps, image_conditioning=image_conditioning)

 

+        self.is_hr_pass = False

+

         return samples

 

 

 class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):

     sampler = None

 

-    def __init__(self, init_images: list = None, resize_mode: int = 0, denoising_strength: float = 0.75, mask: Any = None, mask_blur: int = 4, inpainting_fill: int = 0, inpaint_full_res: bool = True, inpaint_full_res_padding: int = 0, inpainting_mask_invert: int = 0, initial_noise_multiplier: float = None, **kwargs):

+    def __init__(self, init_images: list = None, resize_mode: int = 0, denoising_strength: float = 0.75, image_cfg_scale: float = None, mask: Any = None, mask_blur: int = 4, inpainting_fill: int = 0, inpaint_full_res: bool = True, inpaint_full_res_padding: int = 0, inpainting_mask_invert: int = 0, initial_noise_multiplier: float = None, **kwargs):

         super().__init__(**kwargs)

 

         self.init_images = init_images

         self.resize_mode: int = resize_mode

         self.denoising_strength: float = denoising_strength

+        self.image_cfg_scale: float = image_cfg_scale if shared.sd_model.cond_stage_key == "edit" else None

         self.init_latent = None

         self.image_mask = mask

         self.latent_mask = None

@@ -961,6 +1026,12 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
             self.color_corrections = []

         imgs = []

         for img in self.init_images:

+

+            # Save init image

+            if opts.save_init_img:

+                self.init_img_hash = hashlib.md5(img.tobytes()).hexdigest()

+                images.save_image(img, path=opts.outdir_init_images, basename=None, forced_filename=self.init_img_hash, save_to_dirs=False)

+

             image = images.flatten(img, opts.img2img_background_color)

 

             if crop_region is None and self.resize_mode != 3:

diff --git a/modules/realesrgan_model.py b/modules/realesrgan_model.py
index aad4a629..d6079433 100644
--- a/modules/realesrgan_model.py
+++ b/modules/realesrgan_model.py
@@ -9,7 +9,7 @@ from realesrgan import RealESRGANer
 

 from modules.upscaler import Upscaler, UpscalerData

 from modules.shared import cmd_opts, opts

-

+from modules import modelloader

 

 class UpscalerRealESRGAN(Upscaler):

     def __init__(self, path):

@@ -23,7 +23,15 @@ class UpscalerRealESRGAN(Upscaler):
             self.enable = True

             self.scalers = []

             scalers = self.load_models(path)

+

+            local_model_paths = self.find_models(ext_filter=[".pth"])

             for scaler in scalers:

+                if scaler.local_data_path.startswith("http"):

+                    filename = modelloader.friendly_name(scaler.local_data_path)

+                    local = next(iter([local_model for local_model in local_model_paths if local_model.endswith(filename + '.pth')]), None)

+                    if local:

+                        scaler.local_data_path = local

+

                 if scaler.name in opts.realesrgan_enabled_models:

                     self.scalers.append(scaler)

 

@@ -64,7 +72,9 @@ class UpscalerRealESRGAN(Upscaler):
                 print(f"Unable to find model info: {path}")

                 return None

 

-            info.local_data_path = load_file_from_url(url=info.data_path, model_dir=self.model_path, progress=True)

+            if info.local_data_path.startswith("http"):

+                info.local_data_path = load_file_from_url(url=info.data_path, model_dir=self.model_path, progress=True)

+

             return info

         except Exception as e:

             print(f"Error making Real-ESRGAN models list: {e}", file=sys.stderr)

diff --git a/modules/safe.py b/modules/safe.py
index 82d44be3..dadf319c 100644
--- a/modules/safe.py
+++ b/modules/safe.py
@@ -1,6 +1,5 @@
 # this code is adapted from the script contributed by anon from /h/

 

-import io

 import pickle

 import collections

 import sys

@@ -12,11 +11,9 @@ import _codecs
 import zipfile

 import re

 

-

 # PyTorch 1.13 and later have _TypedStorage renamed to TypedStorage

 TypedStorage = torch.storage.TypedStorage if hasattr(torch.storage, 'TypedStorage') else torch.storage._TypedStorage

 

-

 def encode(*args):

     out = _codecs.encode(*args)

     return out

@@ -27,7 +24,7 @@ class RestrictedUnpickler(pickle.Unpickler):
 

     def persistent_load(self, saved_id):

         assert saved_id[0] == 'storage'

-        return TypedStorage()

+        return TypedStorage(_internal=True)

 

     def find_class(self, module, name):

         if self.extra_handler is not None:

diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py
index 8e80dd8d..cede1cb0 100644
--- a/modules/script_callbacks.py
+++ b/modules/script_callbacks.py
@@ -29,7 +29,7 @@ class ImageSaveParams:
 

 

 class CFGDenoiserParams:

-    def __init__(self, x, image_cond, sigma, sampling_step, total_sampling_steps):

+    def __init__(self, x, image_cond, sigma, sampling_step, total_sampling_steps, text_cond, text_uncond):

         self.x = x

         """Latent image representation in the process of being denoised"""

         

@@ -44,6 +44,24 @@ class CFGDenoiserParams:
         

         self.total_sampling_steps = total_sampling_steps

         """Total number of sampling steps planned"""

+        

+        self.text_cond = text_cond

+        """ Encoder hidden states of text conditioning from prompt"""

+        

+        self.text_uncond = text_uncond

+        """ Encoder hidden states of text conditioning from negative prompt"""

+

+

+class CFGDenoisedParams:

+    def __init__(self, x, sampling_step, total_sampling_steps):

+        self.x = x

+        """Latent image representation in the process of being denoised"""

+

+        self.sampling_step = sampling_step

+        """Current Sampling step number"""

+

+        self.total_sampling_steps = total_sampling_steps

+        """Total number of sampling steps planned"""

 

 

 class UiTrainTabParams:

@@ -68,6 +86,7 @@ callback_map = dict(
     callbacks_before_image_saved=[],

     callbacks_image_saved=[],

     callbacks_cfg_denoiser=[],

+    callbacks_cfg_denoised=[],

     callbacks_before_component=[],

     callbacks_after_component=[],

     callbacks_image_grid=[],

@@ -165,6 +184,14 @@ def cfg_denoiser_callback(params: CFGDenoiserParams):
             report_exception(c, 'cfg_denoiser_callback')

 

 

+def cfg_denoised_callback(params: CFGDenoisedParams):

+    for c in callback_map['callbacks_cfg_denoised']:

+        try:

+            c.callback(params)

+        except Exception:

+            report_exception(c, 'cfg_denoised_callback')

+

+

 def before_component_callback(component, **kwargs):

     for c in callback_map['callbacks_before_component']:

         try:

@@ -306,6 +333,14 @@ def on_cfg_denoiser(callback):
     add_callback(callback_map['callbacks_cfg_denoiser'], callback)

 

 

+def on_cfg_denoised(callback):

+    """register a function to be called in the kdiffussion cfg_denoiser method after building the inner model inputs.

+    The callback is called with one argument:

+        - params: CFGDenoisedParams - parameters to be passed to the inner model and sampling state details.

+    """

+    add_callback(callback_map['callbacks_cfg_denoised'], callback)

+

+

 def on_before_component(callback):

     """register a function to be called before a component is created.

     The callback is called with arguments:

diff --git a/modules/scripts.py b/modules/scripts.py
index 24056a12..4d0bbd66 100644
--- a/modules/scripts.py
+++ b/modules/scripts.py
@@ -33,6 +33,11 @@ class Script:
     parsing infotext to set the value for the component; see ui.py's txt2img_paste_fields for an example

     """

 

+    paste_field_names = None

+    """if set in ui(), this is a list of names of infotext fields; the fields will be sent through the

+    various "Send to <X>" buttons when clicked

+    """

+

     def title(self):

         """this function should return the title of the script. This is what will be displayed in the dropdown menu."""

 

@@ -80,6 +85,20 @@ class Script:
 

         pass

 

+    def before_process_batch(self, p, *args, **kwargs):

+        """

+        Called before extra networks are parsed from the prompt, so you can add

+        new extra network keywords to the prompt with this callback.

+

+        **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

+        """

+

+        pass

+

     def process_batch(self, p, *args, **kwargs):

         """

         Same as process(), but called for every batch.

@@ -220,7 +239,15 @@ def load_scripts():
             elif issubclass(script_class, scripts_postprocessing.ScriptPostprocessing):

                 postprocessing_scripts_data.append(ScriptClassData(script_class, scriptfile.path, scriptfile.basedir, module))

 

-    for scriptfile in sorted(scripts_list):

+    def orderby(basedir):

+        # 1st webui, 2nd extensions-builtin, 3rd extensions

+        priority = {os.path.join(paths.script_path, "extensions-builtin"):1, paths.script_path:0}

+        for key in priority:

+            if basedir.startswith(key):

+                return priority[key]

+        return 9999

+

+    for scriptfile in sorted(scripts_list, key=lambda x: [orderby(x.basedir), x]):

         try:

             if scriptfile.basedir != paths.script_path:

                 sys.path = [scriptfile.basedir] + sys.path

@@ -256,6 +283,7 @@ class ScriptRunner:
         self.alwayson_scripts = []

         self.titles = []

         self.infotext_fields = []

+        self.paste_field_names = []

 

     def initialize_scripts(self, is_img2img):

         from modules import scripts_auto_postprocessing

@@ -304,6 +332,9 @@ class ScriptRunner:
             if script.infotext_fields is not None:

                 self.infotext_fields += script.infotext_fields

 

+            if script.paste_field_names is not None:

+                self.paste_field_names += script.paste_field_names

+

             inputs += controls

             inputs_alwayson += [script.alwayson for _ in controls]

             script.args_to = len(inputs)

@@ -388,6 +419,15 @@ class ScriptRunner:
                 print(f"Error running process: {script.filename}", file=sys.stderr)

                 print(traceback.format_exc(), file=sys.stderr)

 

+    def before_process_batch(self, p, **kwargs):

+        for script in self.alwayson_scripts:

+            try:

+                script_args = p.script_args[script.args_from:script.args_to]

+                script.before_process_batch(p, *script_args, **kwargs)

+            except Exception:

+                print(f"Error running before_process_batch: {script.filename}", file=sys.stderr)

+                print(traceback.format_exc(), file=sys.stderr)

+

     def process_batch(self, p, **kwargs):

         for script in self.alwayson_scripts:

             try:

@@ -481,6 +521,18 @@ def reload_scripts():
     scripts_postproc = scripts_postprocessing.ScriptPostprocessingRunner()

 

 

+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 = ["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)

@@ -489,6 +541,8 @@ def IOComponent_init(self, *args, **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:

@@ -499,3 +553,15 @@ def IOComponent_init(self, *args, **kwargs):
 

 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

diff --git a/modules/scripts_postprocessing.py b/modules/scripts_postprocessing.py
index ce0ebb61..b11568c0 100644
--- a/modules/scripts_postprocessing.py
+++ b/modules/scripts_postprocessing.py
@@ -109,7 +109,7 @@ class ScriptPostprocessingRunner:
         inputs = []

 

         for script in self.scripts_in_preferred_order():

-            with gr.Box() as group:

+            with gr.Row() as group:

                 self.create_script_ui(script, inputs)

 

             script.group = group

diff --git a/modules/sd_disable_initialization.py b/modules/sd_disable_initialization.py
index e90aa9fe..c4a09d15 100644
--- a/modules/sd_disable_initialization.py
+++ b/modules/sd_disable_initialization.py
@@ -20,8 +20,9 @@ class DisableInitialization:
     ```

     """

 

-    def __init__(self):

+    def __init__(self, disable_clip=True):

         self.replaced = []

+        self.disable_clip = disable_clip

 

     def replace(self, obj, field, func):

         original = getattr(obj, field, None)

@@ -75,12 +76,14 @@ class DisableInitialization:
         self.replace(torch.nn.init, 'kaiming_uniform_', do_nothing)

         self.replace(torch.nn.init, '_no_grad_normal_', do_nothing)

         self.replace(torch.nn.init, '_no_grad_uniform_', do_nothing)

-        self.create_model_and_transforms = self.replace(open_clip, 'create_model_and_transforms', create_model_and_transforms_without_pretrained)

-        self.CLIPTextModel_from_pretrained = self.replace(ldm.modules.encoders.modules.CLIPTextModel, 'from_pretrained', CLIPTextModel_from_pretrained)

-        self.transformers_modeling_utils_load_pretrained_model = self.replace(transformers.modeling_utils.PreTrainedModel, '_load_pretrained_model', transformers_modeling_utils_load_pretrained_model)

-        self.transformers_tokenization_utils_base_cached_file = self.replace(transformers.tokenization_utils_base, 'cached_file', transformers_tokenization_utils_base_cached_file)

-        self.transformers_configuration_utils_cached_file = self.replace(transformers.configuration_utils, 'cached_file', transformers_configuration_utils_cached_file)

-        self.transformers_utils_hub_get_from_cache = self.replace(transformers.utils.hub, 'get_from_cache', transformers_utils_hub_get_from_cache)

+

+        if self.disable_clip:

+            self.create_model_and_transforms = self.replace(open_clip, 'create_model_and_transforms', create_model_and_transforms_without_pretrained)

+            self.CLIPTextModel_from_pretrained = self.replace(ldm.modules.encoders.modules.CLIPTextModel, 'from_pretrained', CLIPTextModel_from_pretrained)

+            self.transformers_modeling_utils_load_pretrained_model = self.replace(transformers.modeling_utils.PreTrainedModel, '_load_pretrained_model', transformers_modeling_utils_load_pretrained_model)

+            self.transformers_tokenization_utils_base_cached_file = self.replace(transformers.tokenization_utils_base, 'cached_file', transformers_tokenization_utils_base_cached_file)

+            self.transformers_configuration_utils_cached_file = self.replace(transformers.configuration_utils, 'cached_file', transformers_configuration_utils_cached_file)

+            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:

diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index 8fdc5990..f4bb0266 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -1,5 +1,6 @@
 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, sd_hijack_checkpoint

@@ -36,11 +37,23 @@ def apply_optimizations():
     

     optimization_method = None

 

+    can_use_sdp = hasattr(torch.nn.functional, "scaled_dot_product_attention") and callable(getattr(torch.nn.functional, "scaled_dot_product_attention")) # not everyone has torch 2.x to use sdp

+

     if cmd_opts.force_enable_xformers or (cmd_opts.xformers and shared.xformers_available and torch.version.cuda and (6, 0) <= torch.cuda.get_device_capability(shared.device) <= (9, 0)):

         print("Applying xformers cross attention optimization.")

         ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.xformers_attention_forward

         ldm.modules.diffusionmodules.model.AttnBlock.forward = sd_hijack_optimizations.xformers_attnblock_forward

         optimization_method = 'xformers'

+    elif cmd_opts.opt_sdp_no_mem_attention and can_use_sdp:

+        print("Applying scaled dot product cross attention optimization (without memory efficient attention).")

+        ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.scaled_dot_product_no_mem_attention_forward

+        ldm.modules.diffusionmodules.model.AttnBlock.forward = sd_hijack_optimizations.sdp_no_mem_attnblock_forward

+        optimization_method = 'sdp-no-mem'

+    elif cmd_opts.opt_sdp_attention and can_use_sdp:

+        print("Applying scaled dot product cross attention optimization.")

+        ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.scaled_dot_product_attention_forward

+        ldm.modules.diffusionmodules.model.AttnBlock.forward = sd_hijack_optimizations.sdp_attnblock_forward

+        optimization_method = 'sdp'

     elif cmd_opts.opt_sub_quad_attention:

         print("Applying sub-quadratic cross attention optimization.")

         ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.sub_quad_attention_forward

@@ -76,6 +89,54 @@ def fix_checkpoint():
     pass

 

 

+def weighted_loss(sd_model, pred, target, mean=True):

+    #Calculate the weight normally, but ignore the mean

+    loss = sd_model._old_get_loss(pred, target, mean=False)

+    

+    #Check if we have weights available

+    weight = getattr(sd_model, '_custom_loss_weight', None)

+    if weight is not None:

+        loss *= weight

+    

+    #Return the loss, as mean if specified

+    return loss.mean() if mean else loss

+

+def weighted_forward(sd_model, x, c, w, *args, **kwargs):

+    try:

+        #Temporarily append weights to a place accessible during loss calc

+        sd_model._custom_loss_weight = w

+        

+        #Replace 'get_loss' with a weight-aware one. Otherwise we need to reimplement 'forward' completely

+        #Keep 'get_loss', but don't overwrite the previous old_get_loss if it's already set

+        if not hasattr(sd_model, '_old_get_loss'):

+            sd_model._old_get_loss = sd_model.get_loss

+        sd_model.get_loss = MethodType(weighted_loss, sd_model)

+

+        #Run the standard forward function, but with the patched 'get_loss'

+        return sd_model.forward(x, c, *args, **kwargs)

+    finally:

+        try:

+            #Delete temporary weights if appended

+            del sd_model._custom_loss_weight

+        except AttributeError as e:

+            pass

+            

+        #If we have an old loss function, reset the loss function to the original one

+        if hasattr(sd_model, '_old_get_loss'):

+            sd_model.get_loss = sd_model._old_get_loss

+            del sd_model._old_get_loss

+

+def apply_weighted_forward(sd_model):

+    #Add new function 'weighted_forward' that can be called to calc weighted loss

+    sd_model.weighted_forward = MethodType(weighted_forward, sd_model)

+

+def undo_weighted_forward(sd_model):

+    try:

+        del sd_model.weighted_forward

+    except AttributeError as e:

+        pass

+

+

 class StableDiffusionModelHijack:

     fixes = None

     comments = []

@@ -104,6 +165,10 @@ class StableDiffusionModelHijack:
             m.cond_stage_model.model.token_embedding = EmbeddingsWithFixes(m.cond_stage_model.model.token_embedding, self)

             m.cond_stage_model = sd_hijack_open_clip.FrozenOpenCLIPEmbedderWithCustomWords(m.cond_stage_model, self)

 

+        apply_weighted_forward(m)

+        if m.cond_stage_key == "edit":

+            sd_hijack_unet.hijack_ddpm_edit()

+

         self.optimization_method = apply_optimizations()

 

         self.clip = m.cond_stage_model

@@ -132,6 +197,7 @@ class StableDiffusionModelHijack:
             m.cond_stage_model = m.cond_stage_model.wrapped

 

         undo_optimizations()

+        undo_weighted_forward(m)

 

         self.apply_circular(False)

         self.layers = None

diff --git a/modules/sd_hijack_inpainting.py b/modules/sd_hijack_inpainting.py
index 478cd499..55a2ce4d 100644
--- a/modules/sd_hijack_inpainting.py
+++ b/modules/sd_hijack_inpainting.py
@@ -11,6 +11,7 @@ import ldm.models.diffusion.plms
 from ldm.models.diffusion.ddpm import LatentDiffusion
 from ldm.models.diffusion.plms import PLMSSampler
 from ldm.models.diffusion.ddim import DDIMSampler, noise_like
+from ldm.models.diffusion.sampling_util import norm_thresholding
 
 
 @torch.no_grad()
diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index c02d954c..372555ff 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -337,7 +337,7 @@ def xformers_attention_forward(self, x, context=None, mask=None):
 

     dtype = q.dtype

     if shared.opts.upcast_attn:

-        q, k = q.float(), k.float()

+        q, k, v = q.float(), k.float(), v.float()

 

     out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=get_xformers_flash_attention_op(q, k, v))

 

@@ -346,6 +346,52 @@ 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):

+    batch_size, sequence_length, inner_dim = x.shape

+

+    if mask is not None:

+        mask = self.prepare_attention_mask(mask, sequence_length, batch_size)

+        mask = mask.view(batch_size, self.heads, -1, mask.shape[-1])

+

+    h = self.heads

+    q_in = self.to_q(x)

+    context = default(context, x)

+

+    context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context)

+    k_in = self.to_k(context_k)

+    v_in = self.to_v(context_v)

+

+    head_dim = inner_dim // h

+    q = q_in.view(batch_size, -1, h, head_dim).transpose(1, 2)

+    k = k_in.view(batch_size, -1, h, head_dim).transpose(1, 2)

+    v = v_in.view(batch_size, -1, h, head_dim).transpose(1, 2)

+    

+    del q_in, k_in, v_in

+

+    dtype = q.dtype

+    if shared.opts.upcast_attn:

+        q, k, v = q.float(), k.float(), v.float()

+

+    # the output of sdp = (batch, num_heads, seq_len, head_dim)

+    hidden_states = torch.nn.functional.scaled_dot_product_attention(

+        q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False

+    )

+

+    hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, h * head_dim)

+    hidden_states = hidden_states.to(dtype)

+

+    # linear proj

+    hidden_states = self.to_out[0](hidden_states)

+    # dropout

+    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):

+    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_)

@@ -427,6 +473,30 @@ 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_)

+    q = self.q(h_)

+    k = self.k(h_)

+    v = self.v(h_)

+    b, c, h, w = q.shape

+    q, k, v = map(lambda t: rearrange(t, 'b c h w -> b (h w) c'), (q, k, v))

+    dtype = q.dtype

+    if shared.opts.upcast_attn:

+        q, k = q.float(), k.float()

+    q = q.contiguous()

+    k = k.contiguous()

+    v = v.contiguous()

+    out = torch.nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False)

+    out = out.to(dtype)

+    out = rearrange(out, 'b (h w) c -> b c h w', h=h)

+    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_)

diff --git a/modules/sd_hijack_unet.py b/modules/sd_hijack_unet.py
index 45cf2b18..15858263 100644
--- a/modules/sd_hijack_unet.py
+++ b/modules/sd_hijack_unet.py
@@ -44,6 +44,7 @@ def apply_model(orig_func, self, x_noisy, t, cond, **kwargs):
     with devices.autocast():

         return orig_func(self, x_noisy.to(devices.dtype_unet), t.to(devices.dtype_unet), cond, **kwargs).float()

 

+

 class GELUHijack(torch.nn.GELU, torch.nn.Module):

     def __init__(self, *args, **kwargs):

         torch.nn.GELU.__init__(self, *args, **kwargs)

@@ -53,10 +54,20 @@ class GELUHijack(torch.nn.GELU, torch.nn.Module):
         else:

             return torch.nn.GELU.forward(self, x)

 

+

+ddpm_edit_hijack = None

+def hijack_ddpm_edit():

+    global ddpm_edit_hijack

+    if not ddpm_edit_hijack:

+        CondFunc('modules.models.diffusion.ddpm_edit.LatentDiffusion.decode_first_stage', first_stage_sub, first_stage_cond)

+        CondFunc('modules.models.diffusion.ddpm_edit.LatentDiffusion.encode_first_stage', first_stage_sub, first_stage_cond)

+        ddpm_edit_hijack = CondFunc('modules.models.diffusion.ddpm_edit.LatentDiffusion.apply_model', apply_model, unet_needs_upcast)

+

+

 unet_needs_upcast = lambda *args, **kwargs: devices.unet_needs_upcast

 CondFunc('ldm.models.diffusion.ddpm.LatentDiffusion.apply_model', apply_model, unet_needs_upcast)

 CondFunc('ldm.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)

-if version.parse(torch.__version__) <= version.parse("1.13.1"):

+if version.parse(torch.__version__) <= version.parse("1.13.2") or torch.cuda.is_available():

     CondFunc('ldm.modules.diffusionmodules.util.GroupNorm32.forward', lambda orig_func, self, *args, **kwargs: orig_func(self.float(), *args, **kwargs), unet_needs_upcast)

     CondFunc('ldm.modules.attention.GEGLU.forward', lambda orig_func, self, x: orig_func(self.float(), x.float()).to(devices.dtype_unet), unet_needs_upcast)

     CondFunc('open_clip.transformer.ResidualAttentionBlock.__init__', lambda orig_func, *args, **kwargs: kwargs.update({'act_layer': GELUHijack}) and False or orig_func(*args, **kwargs), lambda _, *args, **kwargs: kwargs.get('act_layer') is None or kwargs['act_layer'] == torch.nn.GELU)

diff --git a/modules/sd_models.py b/modules/sd_models.py
index 300387a9..4f7613a1 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -52,6 +52,15 @@ class CheckpointInfo:
 

         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.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:

@@ -59,13 +68,17 @@ class CheckpointInfo:
 

     def calculate_shorthash(self):

         self.sha256 = hashes.sha256(self.filename, "checkpoint/" + self.name)

+        if self.sha256 is None:

+            return

+

         self.shorthash = self.sha256[0:10]

 

         if self.shorthash not in self.ids:

-            self.ids += [self.shorthash, self.sha256]

-            self.register()

+            self.ids += [self.shorthash, self.sha256, f'{self.name} [{self.shorthash}]']

 

+        checkpoints_list.pop(self.title)

         self.title = f'{self.name} [{self.shorthash}]'

+        self.register()

 

         return self.shorthash

 

@@ -101,9 +114,15 @@ def checkpoint_tiles():
 def list_models():

     checkpoints_list.clear()

     checkpoint_alisases.clear()

-    model_list = modelloader.load_models(model_path=model_path, command_path=shared.cmd_opts.ckpt_dir, ext_filter=[".ckpt", ".safetensors"], ext_blacklist=[".vae.safetensors"])

 

     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):

+        model_url = None

+    else:

+        model_url = "https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors"

+

+    model_list = modelloader.load_models(model_path=model_path, model_url=model_url, command_path=shared.cmd_opts.ckpt_dir, ext_filter=[".ckpt", ".safetensors"], download_name="v1-5-pruned-emaonly.safetensors", ext_blacklist=[".vae.ckpt", ".vae.safetensors"])

+

     if os.path.exists(cmd_ckpt):

         checkpoint_info = CheckpointInfo(cmd_ckpt)

         checkpoint_info.register()

@@ -112,7 +131,7 @@ 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 model_list:

+    for filename in sorted(model_list, key=str.lower):

         checkpoint_info = CheckpointInfo(filename)

         checkpoint_info.register()

 

@@ -158,7 +177,7 @@ def select_checkpoint():
         print(f" - directory {model_path}", file=sys.stderr)

         if shared.cmd_opts.ckpt_dir is not None:

             print(f" - directory {os.path.abspath(shared.cmd_opts.ckpt_dir)}", file=sys.stderr)

-        print("Can't run without a checkpoint. Find and place a .ckpt file into any of those locations. The program will exit.", file=sys.stderr)

+        print("Can't run without a checkpoint. Find and place a .ckpt or .safetensors file into any of those locations. The program will exit.", file=sys.stderr)

         exit(1)

 

     checkpoint_info = next(iter(checkpoints_list.values()))

@@ -168,7 +187,7 @@ def select_checkpoint():
     return checkpoint_info

 

 

-chckpoint_dict_replacements = {

+checkpoint_dict_replacements = {

     'cond_stage_model.transformer.embeddings.': 'cond_stage_model.transformer.text_model.embeddings.',

     'cond_stage_model.transformer.encoder.': 'cond_stage_model.transformer.text_model.encoder.',

     'cond_stage_model.transformer.final_layer_norm.': 'cond_stage_model.transformer.text_model.final_layer_norm.',

@@ -176,7 +195,7 @@ chckpoint_dict_replacements = {
 

 

 def transform_checkpoint_dict_key(k):

-    for text, replacement in chckpoint_dict_replacements.items():

+    for text, replacement in checkpoint_dict_replacements.items():

         if k.startswith(text):

             k = replacement + k[len(text):]

 

@@ -200,6 +219,30 @@ def get_state_dict_from_checkpoint(pl_sd):
     return pl_sd

 

 

+def read_metadata_from_safetensors(filename):

+    import json

+

+    with open(filename, mode="rb") as file:

+        metadata_len = file.read(8)

+        metadata_len = int.from_bytes(metadata_len, "little")

+        json_start = file.read(2)

+

+        assert metadata_len > 2 and json_start in (b'{"', b"{'"), f"{filename} is not a safetensors file"

+        json_data = json_start + file.read(metadata_len-2)

+        json_obj = json.loads(json_data)

+

+        res = {}

+        for k, v in json_obj.get("__metadata__", {}).items():

+            res[k] = v

+            if isinstance(v, str) and v[0:1] == '{':

+                try:

+                    res[k] = json.loads(v)

+                except Exception as e:

+                    pass

+

+        return res

+

+

 def read_state_dict(checkpoint_file, print_global_state=False, map_location=None):

     _, extension = os.path.splitext(checkpoint_file)

     if extension.lower() == ".safetensors":

@@ -349,6 +392,17 @@ def repair_config(sd_config):
     elif shared.cmd_opts.upcast_sampling:

         sd_config.model.params.unet_config.params.use_fp16 = True

 

+    if getattr(sd_config.model.params.first_stage_config.params.ddconfig, "attn_type", None) == "vanilla-xformers" and not shared.xformers_available:

+        sd_config.model.params.first_stage_config.params.ddconfig.attn_type = "vanilla"

+

+    # For UnCLIP-L, override the hardcoded karlo directory

+    if hasattr(sd_config.model.params, "noise_aug_config") and hasattr(sd_config.model.params.noise_aug_config.params, "clip_stats_path"):

+        karlo_path = os.path.join(paths.models_path, 'karlo')

+        sd_config.model.params.noise_aug_config.params.clip_stats_path = sd_config.model.params.noise_aug_config.params.clip_stats_path.replace("checkpoints/karlo_models", karlo_path)

+

+

+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'

 

 def load_model(checkpoint_info=None, already_loaded_state_dict=None, time_taken_to_load_state_dict=None):

     from modules import lowvram, sd_hijack

@@ -370,6 +424,7 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, time_taken_
         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

 

     timer.record("find config")

 

@@ -382,7 +437,7 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, time_taken_
 

     sd_model = None

     try:

-        with sd_disable_initialization.DisableInitialization():

+        with sd_disable_initialization.DisableInitialization(disable_clip=clip_is_included_into_sd):

             sd_model = instantiate_from_config(sd_config.model)

     except Exception as e:

         pass

@@ -456,7 +511,7 @@ def reload_model_weights(sd_model=None, info=None):
     if sd_model is None or checkpoint_config != sd_model.used_config:

         del sd_model

         checkpoints_loaded.clear()

-        load_model(checkpoint_info, already_loaded_state_dict=state_dict, time_taken_to_load_state_dict=timer.records["load weights from disk"])

+        load_model(checkpoint_info, already_loaded_state_dict=state_dict)

         return shared.sd_model

 

     try:

@@ -479,3 +534,23 @@ def reload_model_weights(sd_model=None, info=None):
     print(f"Weights loaded in {timer.summary()}.")

 

     return sd_model

+

+def unload_model_weights(sd_model=None, info=None):

+    from modules import lowvram, devices, sd_hijack

+    timer = Timer()

+

+    if shared.sd_model:

+

+        # shared.sd_model.cond_stage_model.to(devices.cpu)

+        # shared.sd_model.first_stage_model.to(devices.cpu)

+        shared.sd_model.to(devices.cpu)

+        sd_hijack.model_hijack.undo_hijack(shared.sd_model)

+        shared.sd_model = None

+        sd_model = None

+        gc.collect()

+        devices.torch_gc()

+        torch.cuda.empty_cache()

+

+    print(f"Unloaded weights {timer.summary()}.")

+

+    return sd_model

diff --git a/modules/sd_models_config.py b/modules/sd_models_config.py
index 91c21700..9398f528 100644
--- a/modules/sd_models_config.py
+++ b/modules/sd_models_config.py
@@ -14,6 +14,8 @@ 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_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")

 config_inpainting = os.path.join(sd_configs_path, "v1-inpainting-inference.yaml")

 config_instruct_pix2pix = os.path.join(sd_configs_path, "instruct-pix2pix.yaml")

 config_alt_diffusion = os.path.join(sd_configs_path, "alt-diffusion-inference.yaml")

@@ -65,9 +67,14 @@ def is_using_v_parameterization_for_sd2(state_dict):
 def guess_model_config_from_state_dict(sd, filename):

     sd2_cond_proj_weight = sd.get('cond_stage_model.model.transformer.resblocks.0.attn.in_proj_weight', None)

     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:

         return config_depth_model

+    elif sd2_variations_weight is not None and sd2_variations_weight.shape[0] == 768:

+        return config_unclip

+    elif sd2_variations_weight is not None and sd2_variations_weight.shape[0] == 1024:

+        return config_unopenclip

 

     if sd2_cond_proj_weight is not None and sd2_cond_proj_weight.shape[1] == 1024:

         if diffusion_model_input.shape[1] == 9:

diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index 28c2136f..ff361f22 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -32,7 +32,7 @@ def set_samplers():
     global samplers, samplers_for_img2img

 

     hidden = set(shared.opts.hide_samplers)

-    hidden_img2img = set(shared.opts.hide_samplers + ['PLMS'])

+    hidden_img2img = set(shared.opts.hide_samplers + ['PLMS', 'UniPC'])

 

     samplers = [x for x in all_samplers if x.name not in hidden]

     samplers_for_img2img = [x for x in all_samplers if x.name not in hidden_img2img]

diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py
index 3c03d442..bc074238 100644
--- a/modules/sd_samplers_common.py
+++ b/modules/sd_samplers_common.py
@@ -2,7 +2,6 @@ from collections import namedtuple
 import numpy as np

 import torch

 from PIL import Image

-import torchsde._brownian.brownian_interval

 from modules import devices, processing, images, sd_vae_approx

 

 from modules.shared import opts, state

@@ -63,16 +62,11 @@ class InterruptedException(BaseException):
     pass

 

 

-# MPS fix for randn in torchsde

-# XXX move this to separate file for MPS

-def torchsde_randn(size, dtype, device, seed):

-    if device.type == 'mps':

+if opts.randn_source == "CPU":

+    import torchsde._brownian.brownian_interval

+

+    def torchsde_randn(size, dtype, device, seed):

         generator = torch.Generator(devices.cpu).manual_seed(int(seed))

         return torch.randn(size, dtype=dtype, device=devices.cpu, generator=generator).to(device)

-    else:

-        generator = torch.Generator(device).manual_seed(int(seed))

-        return torch.randn(size, dtype=dtype, device=device, generator=generator)

-

-

-torchsde._brownian.brownian_interval._randn = torchsde_randn

 

+    torchsde._brownian.brownian_interval._randn = torchsde_randn

diff --git a/modules/sd_samplers_compvis.py b/modules/sd_samplers_compvis.py
index d03131cd..bfcc5574 100644
--- a/modules/sd_samplers_compvis.py
+++ b/modules/sd_samplers_compvis.py
@@ -7,19 +7,27 @@ import torch
 

 from modules.shared import state

 from modules import sd_samplers_common, prompt_parser, shared

+import modules.models.diffusion.uni_pc

 

 

 samplers_data_compvis = [

     sd_samplers_common.SamplerData('DDIM', lambda model: VanillaStableDiffusionSampler(ldm.models.diffusion.ddim.DDIMSampler, model), [], {}),

     sd_samplers_common.SamplerData('PLMS', lambda model: VanillaStableDiffusionSampler(ldm.models.diffusion.plms.PLMSSampler, model), [], {}),

+    sd_samplers_common.SamplerData('UniPC', lambda model: VanillaStableDiffusionSampler(modules.models.diffusion.uni_pc.UniPCSampler, model), [], {}),

 ]

 

 

 class VanillaStableDiffusionSampler:

     def __init__(self, constructor, sd_model):

         self.sampler = constructor(sd_model)

+        self.is_ddim = hasattr(self.sampler, 'p_sample_ddim')

         self.is_plms = hasattr(self.sampler, 'p_sample_plms')

-        self.orig_p_sample_ddim = self.sampler.p_sample_plms if self.is_plms else self.sampler.p_sample_ddim

+        self.is_unipc = isinstance(self.sampler, modules.models.diffusion.uni_pc.UniPCSampler)

+        self.orig_p_sample_ddim = None

+        if self.is_plms:

+            self.orig_p_sample_ddim = self.sampler.p_sample_plms

+        elif self.is_ddim:

+            self.orig_p_sample_ddim = self.sampler.p_sample_ddim

         self.mask = None

         self.nmask = None

         self.init_latent = None

@@ -45,6 +53,15 @@ class VanillaStableDiffusionSampler:
             return self.last_latent

 

     def p_sample_ddim_hook(self, x_dec, cond, ts, unconditional_conditioning, *args, **kwargs):

+        x_dec, ts, cond, unconditional_conditioning = self.before_sample(x_dec, ts, cond, unconditional_conditioning)

+

+        res = self.orig_p_sample_ddim(x_dec, cond, ts, unconditional_conditioning=unconditional_conditioning, *args, **kwargs)

+

+        x_dec, ts, cond, unconditional_conditioning, res = self.after_sample(x_dec, ts, cond, unconditional_conditioning, res)

+

+        return res

+

+    def before_sample(self, x, ts, cond, unconditional_conditioning):

         if state.interrupted or state.skipped:

             raise sd_samplers_common.InterruptedException

 

@@ -53,8 +70,13 @@ class VanillaStableDiffusionSampler:
 

         # Have to unwrap the inpainting conditioning here to perform pre-processing

         image_conditioning = None

+        uc_image_conditioning = None

         if isinstance(cond, dict):

-            image_conditioning = cond["c_concat"][0]

+            if self.conditioning_key == "crossattn-adm":

+                image_conditioning = cond["c_adm"]

+                uc_image_conditioning = unconditional_conditioning["c_adm"]

+            else:

+                image_conditioning = cond["c_concat"][0]

             cond = cond["c_crossattn"][0]

             unconditional_conditioning = unconditional_conditioning["c_crossattn"][0]

 

@@ -76,20 +98,25 @@ class VanillaStableDiffusionSampler:
 

         if self.mask is not None:

             img_orig = self.sampler.model.q_sample(self.init_latent, ts)

-            x_dec = img_orig * self.mask + self.nmask * x_dec

+            x = img_orig * self.mask + self.nmask * x

 

         # Wrap the image conditioning back up since the DDIM code can accept the dict directly.

         # Note that they need to be lists because it just concatenates them later.

         if image_conditioning is not None:

-            cond = {"c_concat": [image_conditioning], "c_crossattn": [cond]}

-            unconditional_conditioning = {"c_concat": [image_conditioning], "c_crossattn": [unconditional_conditioning]}

+            if self.conditioning_key == "crossattn-adm":

+                cond = {"c_adm": image_conditioning, "c_crossattn": [cond]}

+                unconditional_conditioning = {"c_adm": uc_image_conditioning, "c_crossattn": [unconditional_conditioning]}

+            else:

+                cond = {"c_concat": [image_conditioning], "c_crossattn": [cond]}

+                unconditional_conditioning = {"c_concat": [image_conditioning], "c_crossattn": [unconditional_conditioning]}

 

-        res = self.orig_p_sample_ddim(x_dec, cond, ts, unconditional_conditioning=unconditional_conditioning, *args, **kwargs)

+        return x, ts, cond, unconditional_conditioning

 

+    def update_step(self, last_latent):

         if self.mask is not None:

-            self.last_latent = self.init_latent * self.mask + self.nmask * res[1]

+            self.last_latent = self.init_latent * self.mask + self.nmask * last_latent

         else:

-            self.last_latent = res[1]

+            self.last_latent = last_latent

 

         sd_samplers_common.store_latent(self.last_latent)

 

@@ -97,26 +124,51 @@ class VanillaStableDiffusionSampler:
         state.sampling_step = self.step

         shared.total_tqdm.update()

 

-        return res

+    def after_sample(self, x, ts, cond, uncond, res):

+        if not self.is_unipc:

+            self.update_step(res[1])

+

+        return x, ts, cond, uncond, res

+

+    def unipc_after_update(self, x, model_x):

+        self.update_step(x)

 

     def initialize(self, p):

         self.eta = p.eta if p.eta is not None else shared.opts.eta_ddim

         if self.eta != 0.0:

             p.extra_generation_params["Eta DDIM"] = self.eta

 

+        if self.is_unipc:

+            keys = [

+                ('UniPC variant', 'uni_pc_variant'),

+                ('UniPC skip type', 'uni_pc_skip_type'),

+                ('UniPC order', 'uni_pc_order'),

+                ('UniPC lower order final', 'uni_pc_lower_order_final'),

+            ]

+

+            for name, key in keys:

+                v = getattr(shared.opts, key)

+                if v != shared.opts.get_default(key):

+                    p.extra_generation_params[name] = v

+

         for fieldname in ['p_sample_ddim', 'p_sample_plms']:

             if hasattr(self.sampler, fieldname):

                 setattr(self.sampler, fieldname, self.p_sample_ddim_hook)

+        if self.is_unipc:

+            self.sampler.set_hooks(lambda x, t, c, u: self.before_sample(x, t, c, u), lambda x, t, c, u, r: self.after_sample(x, t, c, u, r), lambda x, mx: self.unipc_after_update(x, mx))

 

         self.mask = p.mask if hasattr(p, 'mask') else None

         self.nmask = p.nmask if hasattr(p, 'nmask') else None

 

+

     def adjust_steps_if_invalid(self, p, num_steps):

-        if (self.config.name == 'DDIM' and p.ddim_discretize == 'uniform') or (self.config.name == 'PLMS'):

+        if ((self.config.name == 'DDIM') and p.ddim_discretize == 'uniform') or (self.config.name == 'PLMS') or (self.config.name == 'UniPC'):

+            if self.config.name == 'UniPC' and num_steps < shared.opts.uni_pc_order:

+                num_steps = shared.opts.uni_pc_order

             valid_step = 999 / (1000 // num_steps)

             if valid_step == math.floor(valid_step):

                 return int(valid_step) + 1

-        

+

         return num_steps

 

     def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning, steps=None, image_conditioning=None):

@@ -133,8 +185,12 @@ class VanillaStableDiffusionSampler:
 

         # Wrap the conditioning models with additional image conditioning for inpainting model

         if image_conditioning is not None:

-            conditioning = {"c_concat": [image_conditioning], "c_crossattn": [conditioning]}

-            unconditional_conditioning = {"c_concat": [image_conditioning], "c_crossattn": [unconditional_conditioning]}

+            if self.conditioning_key == "crossattn-adm":

+                conditioning = {"c_adm": image_conditioning, "c_crossattn": [conditioning]}

+                unconditional_conditioning = {"c_adm": torch.zeros_like(image_conditioning), "c_crossattn": [unconditional_conditioning]}

+            else:

+                conditioning = {"c_concat": [image_conditioning], "c_crossattn": [conditioning]}

+                unconditional_conditioning = {"c_concat": [image_conditioning], "c_crossattn": [unconditional_conditioning]}

 

         samples = self.launch_sampling(t_enc + 1, lambda: self.sampler.decode(x1, conditioning, t_enc, unconditional_guidance_scale=p.cfg_scale, unconditional_conditioning=unconditional_conditioning))

 

@@ -152,8 +208,12 @@ class VanillaStableDiffusionSampler:
         # Wrap the conditioning models with additional image conditioning for inpainting model

         # dummy_for_plms is needed because PLMS code checks the first item in the dict to have the right shape

         if image_conditioning is not None:

-            conditioning = {"dummy_for_plms": np.zeros((conditioning.shape[0],)), "c_crossattn": [conditioning], "c_concat": [image_conditioning]}

-            unconditional_conditioning = {"c_crossattn": [unconditional_conditioning], "c_concat": [image_conditioning]}

+            if self.conditioning_key == "crossattn-adm":

+                conditioning = {"dummy_for_plms": np.zeros((conditioning.shape[0],)), "c_crossattn": [conditioning], "c_adm": image_conditioning}

+                unconditional_conditioning = {"c_crossattn": [unconditional_conditioning], "c_adm": torch.zeros_like(image_conditioning)}

+            else:

+                conditioning = {"dummy_for_plms": np.zeros((conditioning.shape[0],)), "c_crossattn": [conditioning], "c_concat": [image_conditioning]}

+                unconditional_conditioning = {"c_crossattn": [unconditional_conditioning], "c_concat": [image_conditioning]}

 

         samples_ddim = self.launch_sampling(steps, lambda: self.sampler.sample(S=steps, conditioning=conditioning, batch_size=int(x.shape[0]), shape=x[0].shape, verbose=False, unconditional_guidance_scale=p.cfg_scale, unconditional_conditioning=unconditional_conditioning, x_T=x, eta=self.eta)[0])

 

diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py
index aa7f106b..eb98e599 100644
--- a/modules/sd_samplers_kdiffusion.py
+++ b/modules/sd_samplers_kdiffusion.py
@@ -1,12 +1,14 @@
 from collections import deque

 import torch

 import inspect

+import einops

 import k_diffusion.sampling

 from modules import prompt_parser, devices, sd_samplers_common

 

 from modules.shared import opts, state

 import modules.shared as shared

 from modules.script_callbacks import CFGDenoiserParams, cfg_denoiser_callback

+from modules.script_callbacks import CFGDenoisedParams, cfg_denoised_callback

 

 samplers_k_diffusion = [

     ('Euler a', 'sample_euler_ancestral', ['k_euler_a', 'k_euler_ancestral'], {}),

@@ -56,6 +58,7 @@ class CFGDenoiser(torch.nn.Module):
         self.nmask = None

         self.init_latent = None

         self.step = 0

+        self.image_cfg_scale = None

 

     def combine_denoised(self, x_out, conds_list, uncond, cond_scale):

         denoised_uncond = x_out[-uncond.shape[0]:]

@@ -67,61 +70,118 @@ class CFGDenoiser(torch.nn.Module):
 

         return denoised

 

-    def forward(self, x, sigma, uncond, cond, cond_scale, image_cond):

+    def combine_denoised_for_edit_model(self, x_out, cond_scale):

+        out_cond, out_img_cond, out_uncond = x_out.chunk(3)

+        denoised = out_uncond + cond_scale * (out_cond - out_img_cond) + self.image_cfg_scale * (out_img_cond - out_uncond)

+

+        return denoised

+

+    def forward(self, x, sigma, uncond, cond, cond_scale, s_min_uncond, image_cond):

         if state.interrupted or state.skipped:

             raise sd_samplers_common.InterruptedException

 

+        # at self.image_cfg_scale == 1.0 produced results for edit model are the same as with normal sampling,

+        # so is_edit_model is set to False to support AND composition.

+        is_edit_model = shared.sd_model.cond_stage_key == "edit" and self.image_cfg_scale is not None and self.image_cfg_scale != 1.0

+

         conds_list, tensor = prompt_parser.reconstruct_multicond_batch(cond, self.step)

         uncond = prompt_parser.reconstruct_cond_batch(uncond, self.step)

 

+        assert not is_edit_model or all([len(conds) == 1 for conds in conds_list]), "AND is not supported for InstructPix2Pix checkpoint (unless using Image CFG scale = 1.0)"

+

         batch_size = len(conds_list)

         repeats = [len(conds_list[i]) for i in range(batch_size)]

 

-        x_in = torch.cat([torch.stack([x[i] for _ in range(n)]) for i, n in enumerate(repeats)] + [x])

-        image_cond_in = torch.cat([torch.stack([image_cond[i] for _ in range(n)]) for i, n in enumerate(repeats)] + [image_cond])

-        sigma_in = torch.cat([torch.stack([sigma[i] for _ in range(n)]) for i, n in enumerate(repeats)] + [sigma])

+        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} 

+        else:

+            image_uncond = image_cond

+            make_condition_dict = lambda c_crossattn, c_concat: {"c_crossattn": c_crossattn, "c_concat": [c_concat]} 

 

-        denoiser_params = CFGDenoiserParams(x_in, image_cond_in, sigma_in, state.sampling_step, state.sampling_steps)

+        if not is_edit_model:

+            x_in = torch.cat([torch.stack([x[i] for _ in range(n)]) for i, n in enumerate(repeats)] + [x])

+            sigma_in = torch.cat([torch.stack([sigma[i] for _ in range(n)]) for i, n in enumerate(repeats)] + [sigma])

+            image_cond_in = torch.cat([torch.stack([image_cond[i] for _ in range(n)]) for i, n in enumerate(repeats)] + [image_uncond])

+        else:

+            x_in = torch.cat([torch.stack([x[i] for _ in range(n)]) for i, n in enumerate(repeats)] + [x] + [x])

+            sigma_in = torch.cat([torch.stack([sigma[i] for _ in range(n)]) for i, n in enumerate(repeats)] + [sigma] + [sigma])

+            image_cond_in = torch.cat([torch.stack([image_cond[i] for _ in range(n)]) for i, n in enumerate(repeats)] + [image_uncond] + [torch.zeros_like(self.init_latent)])

+

+        denoiser_params = CFGDenoiserParams(x_in, image_cond_in, sigma_in, state.sampling_step, state.sampling_steps, tensor, uncond)

         cfg_denoiser_callback(denoiser_params)

         x_in = denoiser_params.x

         image_cond_in = denoiser_params.image_cond

         sigma_in = denoiser_params.sigma

-

-        if tensor.shape[1] == uncond.shape[1]:

-            cond_in = torch.cat([tensor, uncond])

+        tensor = denoiser_params.text_cond

+        uncond = denoiser_params.text_uncond

+        skip_uncond = False

+

+        # alternating uncond allows for higher thresholds without the quality loss normally expected from raising it

+        if self.step % 2 and s_min_uncond > 0 and sigma[0] < s_min_uncond and not is_edit_model:

+            skip_uncond = True

+            x_in = x_in[:-batch_size]

+            sigma_in = sigma_in[:-batch_size]

+

+        if tensor.shape[1] == uncond.shape[1] or skip_uncond:

+            if is_edit_model:

+                cond_in = torch.cat([tensor, uncond, uncond])

+            elif skip_uncond:

+                cond_in = tensor

+            else:

+                cond_in = torch.cat([tensor, uncond])

 

             if shared.batch_cond_uncond:

-                x_out = self.inner_model(x_in, sigma_in, cond={"c_crossattn": [cond_in], "c_concat": [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={"c_crossattn": [cond_in[a:b]], "c_concat": [image_cond_in[a:b]]})

+                    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]))

         else:

             x_out = torch.zeros_like(x_in)

             batch_size = batch_size*2 if shared.batch_cond_uncond else batch_size

             for batch_offset in range(0, tensor.shape[0], batch_size):

                 a = batch_offset

                 b = min(a + batch_size, tensor.shape[0])

-                x_out[a:b] = self.inner_model(x_in[a:b], sigma_in[a:b], cond={"c_crossattn": [tensor[a:b]], "c_concat": [image_cond_in[a:b]]})

 

-            x_out[-uncond.shape[0]:] = self.inner_model(x_in[-uncond.shape[0]:], sigma_in[-uncond.shape[0]:], cond={"c_crossattn": [uncond], "c_concat": [image_cond_in[-uncond.shape[0]:]]})

+                if not is_edit_model:

+                    c_crossattn = [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]:]))

+

+        denoised_image_indexes = [x[0][0] for x in conds_list]

+        if skip_uncond:

+            fake_uncond = torch.cat([x_out[i:i+1] for i in denoised_image_indexes])

+            x_out = torch.cat([x_out, fake_uncond])  # we skipped uncond denoising, so we put cond-denoised image to where the uncond-denoised image should be

+

+        denoised_params = CFGDenoisedParams(x_out, state.sampling_step, state.sampling_steps)

+        cfg_denoised_callback(denoised_params)

 

         devices.test_for_nans(x_out, "unet")

 

         if opts.live_preview_content == "Prompt":

-            sd_samplers_common.store_latent(x_out[0:uncond.shape[0]])

+            sd_samplers_common.store_latent(torch.cat([x_out[i:i+1] for i in denoised_image_indexes]))

         elif opts.live_preview_content == "Negative prompt":

             sd_samplers_common.store_latent(x_out[-uncond.shape[0]:])

 

-        denoised = self.combine_denoised(x_out, conds_list, uncond, cond_scale)

+        if is_edit_model:

+            denoised = self.combine_denoised_for_edit_model(x_out, cond_scale)

+        elif skip_uncond:

+            denoised = self.combine_denoised(x_out, conds_list, uncond, 1.0)

+        else:

+            denoised = self.combine_denoised(x_out, conds_list, uncond, cond_scale)

 

         if self.mask is not None:

             denoised = self.init_latent * self.mask + self.nmask * denoised

 

         self.step += 1

-

         return denoised

 

 

@@ -146,7 +206,7 @@ class TorchHijack:
             if noise.shape == x.shape:

                 return noise

 

-        if x.device.type == 'mps':

+        if opts.randn_source == "CPU" or x.device.type == 'mps':

             return torch.randn_like(x, device=devices.cpu).to(x.device)

         else:

             return torch.randn_like(x)

@@ -166,6 +226,7 @@ class KDiffusionSampler:
         self.eta = None

         self.config = None

         self.last_latent = None

+        self.s_min_uncond = None

 

         self.conditioning_key = sd_model.model.conditioning_key

 

@@ -198,7 +259,9 @@ class KDiffusionSampler:
         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

+        self.model_wrap_cfg.image_cfg_scale = getattr(p, 'image_cfg_scale', None)

         self.eta = p.eta if p.eta is not None else opts.eta_ancestral

+        self.s_min_uncond = getattr(p, 's_min_uncond', 0.0)

 

         k_diffusion.sampling.torch = TorchHijack(self.sampler_noises if self.sampler_noises is not None else [])

 

@@ -237,6 +300,16 @@ class KDiffusionSampler:
 

         return sigmas

 

+    def create_noise_sampler(self, x, sigmas, p):

+        """For DPM++ SDE: manually create noise sampler to enable deterministic results across different batch sizes"""

+        if shared.opts.no_dpmpp_sde_batch_determinism:

+            return None

+

+        from k_diffusion.sampling import BrownianTreeNoiseSampler

+        sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()

+        current_iter_seeds = p.all_seeds[p.iteration * p.batch_size:(p.iteration + 1) * p.batch_size]

+        return BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=current_iter_seeds)

+

     def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning, steps=None, image_conditioning=None):

         steps, t_enc = sd_samplers_common.setup_img2img_steps(p, steps)

 

@@ -246,31 +319,39 @@ class KDiffusionSampler:
         xi = x + noise * sigma_sched[0]

         

         extra_params_kwargs = self.initialize(p)

-        if 'sigma_min' in inspect.signature(self.func).parameters:

+        parameters = inspect.signature(self.func).parameters

+

+        if 'sigma_min' in parameters:

             ## last sigma is zero which isn't allowed by DPM Fast & Adaptive so taking value before last

             extra_params_kwargs['sigma_min'] = sigma_sched[-2]

-        if 'sigma_max' in inspect.signature(self.func).parameters:

+        if 'sigma_max' in parameters:

             extra_params_kwargs['sigma_max'] = sigma_sched[0]

-        if 'n' in inspect.signature(self.func).parameters:

+        if 'n' in parameters:

             extra_params_kwargs['n'] = len(sigma_sched) - 1

-        if 'sigma_sched' in inspect.signature(self.func).parameters:

+        if 'sigma_sched' in parameters:

             extra_params_kwargs['sigma_sched'] = sigma_sched

-        if 'sigmas' in inspect.signature(self.func).parameters:

+        if 'sigmas' in parameters:

             extra_params_kwargs['sigmas'] = sigma_sched

 

+        if self.funcname == 'sample_dpmpp_sde':

+            noise_sampler = self.create_noise_sampler(x, sigmas, p)

+            extra_params_kwargs['noise_sampler'] = noise_sampler

+

         self.model_wrap_cfg.init_latent = x

         self.last_latent = x

-

-        samples = self.launch_sampling(t_enc + 1, lambda: self.func(self.model_wrap_cfg, xi, extra_args={

+        extra_args={

             'cond': conditioning, 

             'image_cond': image_conditioning, 

             'uncond': unconditional_conditioning, 

-            'cond_scale': p.cfg_scale

-        }, disable=False, callback=self.callback_state, **extra_params_kwargs))

+            'cond_scale': p.cfg_scale,

+            's_min_uncond': self.s_min_uncond

+        }

+

+        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))

 

         return samples

 

-    def sample(self, p, x, conditioning, unconditional_conditioning, steps=None, image_conditioning = None):

+    def sample(self, p, x, conditioning, unconditional_conditioning, steps=None, image_conditioning=None):

         steps = steps or p.steps

 

         sigmas = self.get_sigmas(p, steps)

@@ -278,20 +359,27 @@ class KDiffusionSampler:
         x = x * sigmas[0]

 

         extra_params_kwargs = self.initialize(p)

-        if 'sigma_min' in inspect.signature(self.func).parameters:

+        parameters = inspect.signature(self.func).parameters

+

+        if 'sigma_min' in parameters:

             extra_params_kwargs['sigma_min'] = self.model_wrap.sigmas[0].item()

             extra_params_kwargs['sigma_max'] = self.model_wrap.sigmas[-1].item()

-            if 'n' in inspect.signature(self.func).parameters:

+            if 'n' in parameters:

                 extra_params_kwargs['n'] = steps

         else:

             extra_params_kwargs['sigmas'] = sigmas

 

+        if self.funcname == 'sample_dpmpp_sde':

+            noise_sampler = self.create_noise_sampler(x, sigmas, p)

+            extra_params_kwargs['noise_sampler'] = noise_sampler

+

         self.last_latent = x

         samples = self.launch_sampling(steps, lambda: self.func(self.model_wrap_cfg, x, extra_args={

             'cond': conditioning, 

             'image_cond': image_conditioning, 

             'uncond': unconditional_conditioning, 

-            'cond_scale': p.cfg_scale

+            'cond_scale': p.cfg_scale,

+            's_min_uncond': self.s_min_uncond

         }, disable=False, callback=self.callback_state, **extra_params_kwargs))

 

         return samples

diff --git a/modules/sd_vae_approx.py b/modules/sd_vae_approx.py
index 0027343a..e2f00468 100644
--- a/modules/sd_vae_approx.py
+++ b/modules/sd_vae_approx.py
@@ -35,8 +35,11 @@ def model():
     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()

-        sd_vae_approx_model.load_state_dict(torch.load(os.path.join(paths.models_path, "VAE-approx", "model.pt"), map_location='cpu' if devices.device.type != 'cuda' else None))

+        if not os.path.exists(model_path):

+            model_path = os.path.join(paths.script_path, "models", "VAE-approx", "model.pt")

+        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)

 

diff --git a/modules/shared.py b/modules/shared.py
index 5600d480..6a2b3c2b 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -4,6 +4,7 @@ import json
 import os

 import sys

 import time

+import requests

 

 from PIL import Image

 import gradio as gr

@@ -13,106 +14,21 @@ import modules.interrogate
 import modules.memmon

 import modules.styles

 import modules.devices as devices

-from modules import localization, extensions, script_loading, errors, ui_components, shared_items

-from modules.paths import models_path, script_path, data_path

-

+from modules import localization, script_loading, errors, ui_components, shared_items, cmd_args

+from modules.paths_internal import models_path, script_path, data_path, sd_configs_path, sd_default_config, sd_model_file, default_sd_model_file, extensions_dir, extensions_builtin_dir

 

 demo = None

 

-sd_configs_path = os.path.join(script_path, "configs")

-sd_default_config = os.path.join(sd_configs_path, "v1-inference.yaml")

-sd_model_file = os.path.join(script_path, 'model.ckpt')

-default_sd_model_file = sd_model_file

-

-parser = argparse.ArgumentParser()

-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",)

-parser.add_argument("--ckpt-dir", type=str, default=None, help="Path to directory with stable diffusion checkpoints")

-parser.add_argument("--vae-dir", type=str, default=None, help="Path to directory with VAE files")

-parser.add_argument("--gfpgan-dir", type=str, help="GFPGAN directory", default=('./src/gfpgan' if os.path.exists('./src/gfpgan') else './GFPGAN'))

-parser.add_argument("--gfpgan-model", type=str, help="GFPGAN model file name", default=None)

-parser.add_argument("--no-half", action='store_true', help="do not switch the model to 16-bit floats")

-parser.add_argument("--no-half-vae", action='store_true', help="do not switch the VAE model to 16-bit floats")

-parser.add_argument("--no-progressbar-hiding", action='store_true', help="do not hide progressbar in gradio UI (we hide it because it slows down ML if you have hardware acceleration in browser)")

-parser.add_argument("--max-batch-count", type=int, default=16, help="maximum batch count value for the UI")

-parser.add_argument("--embeddings-dir", type=str, default=os.path.join(data_path, 'embeddings'), help="embeddings directory for textual inversion (default: embeddings)")

-parser.add_argument("--textual-inversion-templates-dir", type=str, default=os.path.join(script_path, 'textual_inversion_templates'), help="directory with textual inversion templates")

-parser.add_argument("--hypernetwork-dir", type=str, default=os.path.join(models_path, 'hypernetworks'), help="hypernetwork directory")

-parser.add_argument("--localizations-dir", type=str, default=os.path.join(script_path, 'localizations'), help="localizations directory")

-parser.add_argument("--allow-code", action='store_true', help="allow custom script execution from webui")

-parser.add_argument("--medvram", action='store_true', help="enable stable diffusion model optimizations for sacrificing a little speed for low VRM usage")

-parser.add_argument("--lowvram", action='store_true', help="enable stable diffusion model optimizations for sacrificing a lot of speed for very low VRM usage")

-parser.add_argument("--lowram", action='store_true', help="load stable diffusion checkpoint weights to VRAM instead of RAM")

-parser.add_argument("--always-batch-cond-uncond", action='store_true', help="disables cond/uncond batching that is enabled to save memory with --medvram or --lowvram")

-parser.add_argument("--unload-gfpgan", action='store_true', help="does not do anything.")

-parser.add_argument("--precision", type=str, help="evaluate at this precision", choices=["full", "autocast"], default="autocast")

-parser.add_argument("--upcast-sampling", action='store_true', help="upcast sampling. No effect with --no-half. Usually produces similar results to --no-half with better performance while using less memory.")

-parser.add_argument("--share", action='store_true', help="use share=True for gradio and make the UI accessible through their site")

-parser.add_argument("--ngrok", type=str, help="ngrok authtoken, alternative to gradio --share", default=None)

-parser.add_argument("--ngrok-region", type=str, help="The region in which ngrok should start.", default="us")

-parser.add_argument("--enable-insecure-extension-access", action='store_true', help="enable extensions tab regardless of other options")

-parser.add_argument("--codeformer-models-path", type=str, help="Path to directory with codeformer model file(s).", default=os.path.join(models_path, 'Codeformer'))

-parser.add_argument("--gfpgan-models-path", type=str, help="Path to directory with GFPGAN model file(s).", default=os.path.join(models_path, 'GFPGAN'))

-parser.add_argument("--esrgan-models-path", type=str, help="Path to directory with ESRGAN model file(s).", default=os.path.join(models_path, 'ESRGAN'))

-parser.add_argument("--bsrgan-models-path", type=str, help="Path to directory with BSRGAN model file(s).", default=os.path.join(models_path, 'BSRGAN'))

-parser.add_argument("--realesrgan-models-path", type=str, help="Path to directory with RealESRGAN model file(s).", default=os.path.join(models_path, 'RealESRGAN'))

-parser.add_argument("--clip-models-path", type=str, help="Path to directory with CLIP model file(s).", default=None)

-parser.add_argument("--xformers", action='store_true', help="enable xformers for cross attention layers")

-parser.add_argument("--force-enable-xformers", action='store_true', help="enable xformers for cross attention layers regardless of whether the checking code thinks you can run it; do not make bug reports if this fails to work")

-parser.add_argument("--xformers-flash-attention", action='store_true', help="enable xformers with Flash Attention to improve reproducibility (supported for SD2.x or variant only)")

-parser.add_argument("--deepdanbooru", action='store_true', help="does not do anything")

-parser.add_argument("--opt-split-attention", action='store_true', help="force-enables Doggettx's cross-attention layer optimization. By default, it's on for torch cuda.")

-parser.add_argument("--opt-sub-quad-attention", action='store_true', help="enable memory efficient sub-quadratic cross-attention layer optimization")

-parser.add_argument("--sub-quad-q-chunk-size", type=int, help="query chunk size for the sub-quadratic cross-attention layer optimization to use", default=1024)

-parser.add_argument("--sub-quad-kv-chunk-size", type=int, help="kv chunk size for the sub-quadratic cross-attention layer optimization to use", default=None)

-parser.add_argument("--sub-quad-chunk-threshold", type=int, help="the percentage of VRAM threshold for the sub-quadratic cross-attention layer optimization to use chunking", default=None)

-parser.add_argument("--opt-split-attention-invokeai", action='store_true', help="force-enables InvokeAI's cross-attention layer optimization. By default, it's on when cuda is unavailable.")

-parser.add_argument("--opt-split-attention-v1", action='store_true', help="enable older version of split attention optimization that does not consume all the VRAM it can find")

-parser.add_argument("--disable-opt-split-attention", action='store_true', help="force-disables cross-attention layer 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("--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)

-parser.add_argument("--ui-config-file", type=str, help="filename to use for ui configuration", default=os.path.join(data_path, 'ui-config.json'))

-parser.add_argument("--hide-ui-dir-config", action='store_true', help="hide directory configuration from webui", default=False)

-parser.add_argument("--freeze-settings", action='store_true', help="disable editing settings", default=False)

-parser.add_argument("--ui-settings-file", type=str, help="filename to use for ui settings", default=os.path.join(data_path, 'config.json'))

-parser.add_argument("--gradio-debug",  action='store_true', help="launch gradio with --debug option")

-parser.add_argument("--gradio-auth", type=str, help='set gradio authentication like "username:password"; or comma-delimit multiple like "u1:p1,u2:p2,u3:p3"', default=None)

-parser.add_argument("--gradio-img2img-tool", type=str, help='does not do anything')

-parser.add_argument("--gradio-inpaint-tool", type=str, help="does not do anything")

-parser.add_argument("--opt-channelslast", action='store_true', help="change memory type for stable diffusion to channels last")

-parser.add_argument("--styles-file", type=str, help="filename to use for styles", default=os.path.join(data_path, 'styles.csv'))

-parser.add_argument("--autolaunch", action='store_true', help="open the webui URL in the system's default browser upon launch", default=False)

-parser.add_argument("--theme", type=str, help="launches the UI with light or dark theme", default=None)

-parser.add_argument("--use-textbox-seed", action='store_true', help="use textbox for seeds in UI (no up/down, but possible to input long seeds)", default=False)

-parser.add_argument("--disable-console-progressbars", action='store_true', help="do not output progressbars to console", default=False)

-parser.add_argument("--enable-console-prompts", action='store_true', help="print prompts to console when generating with txt2img and img2img", default=False)

-parser.add_argument('--vae-path', type=str, help='Checkpoint to use as VAE; setting this argument disables all settings related to VAE', default=None)

-parser.add_argument("--disable-safe-unpickle", action='store_true', help="disable checking pytorch models for malicious code", default=False)

-parser.add_argument("--api", action='store_true', help="use api=True to launch the API together with the webui (use --nowebui instead for only the API)")

-parser.add_argument("--api-auth", type=str, help='Set authentication for API like "username:password"; or comma-delimit multiple like "u1:p1,u2:p2,u3:p3"', default=None)

-parser.add_argument("--api-log", action='store_true', help="use api-log=True to enable logging of all API requests")

-parser.add_argument("--nowebui", action='store_true', help="use api=True to launch the API instead of the webui")

-parser.add_argument("--ui-debug-mode", action='store_true', help="Don't load model to quickly launch UI")

-parser.add_argument("--device-id", type=str, help="Select the default CUDA device to use (export CUDA_VISIBLE_DEVICES=0,1,etc might be needed before)", default=None)

-parser.add_argument("--administrator", action='store_true', help="Administrator rights", default=False)

-parser.add_argument("--cors-allow-origins", type=str, help="Allowed CORS origin(s) in the form of a comma-separated list (no spaces)", default=None)

-parser.add_argument("--cors-allow-origins-regex", type=str, help="Allowed CORS origin(s) in the form of a single regular expression", default=None)

-parser.add_argument("--tls-keyfile", type=str, help="Partially enables TLS, requires --tls-certfile to fully function", default=None)

-parser.add_argument("--tls-certfile", type=str, help="Partially enables TLS, requires --tls-keyfile to fully function", default=None)

-parser.add_argument("--server-name", type=str, help="Sets hostname of server", default=None)

-parser.add_argument("--gradio-queue", action='store_true', help="Uses gradio queue; experimental option; breaks restart UI button")

-parser.add_argument("--skip-version-check", action='store_true', help="Do not check versions of torch and xformers")

-

-

-

-script_loading.preload_extensions(extensions.extensions_dir, parser)

-script_loading.preload_extensions(extensions.extensions_builtin_dir, parser)

-

-cmd_opts = parser.parse_args()

+parser = cmd_args.parser

+

+script_loading.preload_extensions(extensions_dir, parser)

+script_loading.preload_extensions(extensions_builtin_dir, parser)

+

+if os.environ.get('IGNORE_CMD_ARGS_ERRORS', None) is None:

+    cmd_opts = parser.parse_args()

+else:

+    cmd_opts, _ = parser.parse_known_args()

+

 

 restricted_opts = {

     "samples_filename_pattern",

@@ -124,6 +40,7 @@ restricted_opts = {
     "outdir_grids",

     "outdir_txt2img_grids",

     "outdir_save",

+    "outdir_init_images"

 }

 

 ui_reorder_categories = [

@@ -139,6 +56,21 @@ ui_reorder_categories = [
     "scripts",

 ]

 

+# https://huggingface.co/datasets/freddyaboulton/gradio-theme-subdomains/resolve/main/subdomains.json

+gradio_hf_hub_themes = [

+    "gradio/glass",

+    "gradio/monochrome",

+    "gradio/seafoam",

+    "gradio/soft",

+    "freddyaboulton/dracula_revamped",

+    "gradio/dracula_test",

+    "abidlabs/dracula_test",

+    "abidlabs/pakistan",

+    "dawood/microsoft_windows",

+    "ysharma/steampunk"

+]

+

+

 cmd_opts.disable_extension_access = (cmd_opts.share or cmd_opts.listen or cmd_opts.server_name) and not cmd_opts.enable_insecure_extension_access

 

 devices.device, devices.device_interrogate, devices.device_gfpgan, devices.device_esrgan, devices.device_codeformer = \

@@ -303,6 +235,7 @@ def list_samplers():
 

 

 hide_dirs = {"visible": not cmd_opts.hide_ui_dir_config}

+tab_names = []

 

 options_templates = {}

 

@@ -324,13 +257,19 @@ options_templates.update(options_section(('saving-images', "Saving images/grids"
     "save_images_before_face_restoration": OptionInfo(False, "Save a copy of image before doing face restoration."),

     "save_images_before_highres_fix": OptionInfo(False, "Save a copy of image before applying highres fix."),

     "save_images_before_color_correction": OptionInfo(False, "Save a copy of image before applying color correction to img2img results"),

+    "save_mask": OptionInfo(False, "For inpainting, save a copy of the greyscale mask"),

+    "save_mask_composite": OptionInfo(False, "For inpainting, save a masked composite"),

     "jpeg_quality": OptionInfo(80, "Quality for saved jpeg images", gr.Slider, {"minimum": 1, "maximum": 100, "step": 1}),

-    "export_for_4chan": OptionInfo(True, "If PNG image is larger than 4MB or any dimension is larger than 4000, downscale and save copy as JPG"),

+    "webp_lossless": OptionInfo(False, "Use lossless compression for webp images"),

+    "export_for_4chan": OptionInfo(True, "If the saved image file size is above the limit, or its either width or height are above the limit, save a downscaled copy as JPG"),

+    "img_downscale_threshold": OptionInfo(4.0, "File size limit for the above option, MB", gr.Number),

+    "target_side_length": OptionInfo(4000, "Width/height limit for the above option, in pixels", gr.Number),

+    "img_max_size_mp": OptionInfo(200, "Maximum image size, in megapixels", gr.Number),

 

     "use_original_name_batch": OptionInfo(True, "Use original name for output filename during batch process in extras tab"),

     "use_upscaler_name_as_suffix": OptionInfo(False, "Use upscaler name as filename suffix in the extras tab"),

     "save_selected_only": OptionInfo(True, "When using 'Save' button, only save a single selected image"),

-    "do_not_add_watermark": OptionInfo(False, "Do not add watermark to images"),

+    "save_init_img": OptionInfo(False, "Save init images when using img2img"),

 

     "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"),

@@ -346,6 +285,7 @@ options_templates.update(options_section(('saving-paths', "Paths for saving"), {
     "outdir_txt2img_grids": OptionInfo("outputs/txt2img-grids", 'Output directory for txt2img grids', component_args=hide_dirs),

     "outdir_img2img_grids": OptionInfo("outputs/img2img-grids", 'Output directory for img2img grids', component_args=hide_dirs),

     "outdir_save": OptionInfo("log/images", "Directory for saving images using the Save button", component_args=hide_dirs),

+    "outdir_init_images": OptionInfo("outputs/init-images", "Directory for saving init images when using img2img", component_args=hide_dirs),

 }))

 

 options_templates.update(options_section(('saving-to-dirs', "Saving to a directory"), {

@@ -361,10 +301,12 @@ options_templates.update(options_section(('upscaling', "Upscaling"), {
     "ESRGAN_tile_overlap": OptionInfo(8, "Tile overlap, in pixels for ESRGAN upscalers. Low values = visible seam.", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}),

     "realesrgan_enabled_models": OptionInfo(["R-ESRGAN 4x+", "R-ESRGAN 4x+ Anime6B"], "Select which Real-ESRGAN models to show in the web UI. (Requires restart)", gr.CheckboxGroup, lambda: {"choices": shared_items.realesrgan_models_names()}),

     "upscaler_for_img2img": OptionInfo(None, "Upscaler for img2img", gr.Dropdown, lambda: {"choices": [x.name for x in sd_upscalers]}),

+    "SCUNET_tile": OptionInfo(256, "Tile size for SCUNET upscalers. 0 = no tiling.", gr.Slider, {"minimum": 0, "maximum": 512, "step": 16}),

+    "SCUNET_tile_overlap": OptionInfo(8, "Tile overlap, in pixels for SCUNET upscalers. Low values = visible seam.", gr.Slider, {"minimum": 0, "maximum": 64, "step": 1}),

 }))

 

 options_templates.update(options_section(('face-restoration', "Face restoration"), {

-    "face_restoration_model": OptionInfo(None, "Face restoration model", gr.Radio, lambda: {"choices": [x.name() for x in face_restorers]}),

+    "face_restoration_model": OptionInfo("CodeFormer", "Face restoration model", gr.Radio, lambda: {"choices": [x.name() for x in face_restorers]}),

     "code_former_weight": OptionInfo(0.5, "CodeFormer weight parameter; 0 = maximum effect; 1 = minimum effect", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01}),

     "face_restoration_unload": OptionInfo(False, "Move face restoration model from VRAM into RAM after processing"),

 }))

@@ -409,12 +351,15 @@ options_templates.update(options_section(('sd', "Stable Diffusion"), {
     "comma_padding_backtrack": OptionInfo(20, "Increase coherency by padding from the last comma within n tokens when using more than 75 tokens", gr.Slider, {"minimum": 0, "maximum": 74, "step": 1 }),

     "CLIP_stop_at_last_layers": OptionInfo(1, "Clip skip", gr.Slider, {"minimum": 1, "maximum": 12, "step": 1}),

     "upcast_attn": OptionInfo(False, "Upcast cross attention layer to float32"),

+    "randn_source": OptionInfo("GPU", "Random number generator source. Changes seeds drastically. Use CPU to produce the same picture across different vidocard vendors.", gr.Radio, {"choices": ["GPU", "CPU"]}),

 }))

 

 options_templates.update(options_section(('compatibility', "Compatibility"), {

     "use_old_emphasis_implementation": OptionInfo(False, "Use old emphasis implementation. Can be useful to reproduce old seeds."),

     "use_old_karras_scheduler_sigmas": OptionInfo(False, "Use old karras scheduler sigmas (0.1 to 10)."),

+    "no_dpmpp_sde_batch_determinism": OptionInfo(False, "Do not make DPM++ SDE deterministic across different batch sizes."),

     "use_old_hires_fix_width_height": OptionInfo(False, "For hires fix, use width/height sliders to set final resolution rather than first pass (disables Upscale by, Resize width/height to)."),

+    "dont_fix_second_order_samplers_schedule": OptionInfo(False, "Do not fix prompt schedule for second order samplers."),

 }))

 

 options_templates.update(options_section(('interrogate', "Interrogate Options"), {

@@ -435,11 +380,16 @@ options_templates.update(options_section(('interrogate', "Interrogate Options"),
 options_templates.update(options_section(('extra_networks', "Extra Networks"), {

     "extra_networks_default_view": OptionInfo("cards", "Default view for Extra Networks", gr.Dropdown, {"choices": ["cards", "thumbs"]}),

     "extra_networks_default_multiplier": OptionInfo(1.0, "Multiplier for extra networks", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),

-    "sd_hypernetwork": OptionInfo("None", "Add hypernetwork to prompt", gr.Dropdown, lambda: {"choices": [""] + [x for x in hypernetworks.keys()]}, refresh=reload_hypernetworks),

+    "extra_networks_card_width": OptionInfo(0, "Card width for Extra Networks (px)"),

+    "extra_networks_card_height": OptionInfo(0, "Card height for Extra Networks (px)"),

+    "extra_networks_add_text_separator": OptionInfo(" ", "Extra text to add before <...> when adding extra network to prompt"),

+    "sd_hypernetwork": OptionInfo("None", "Add hypernetwork to prompt", gr.Dropdown, lambda: {"choices": ["None"] + [x for x in hypernetworks.keys()]}, refresh=reload_hypernetworks),

 }))

 

 options_templates.update(options_section(('ui', "User interface"), {

     "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"),

     "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"),

@@ -454,10 +404,13 @@ options_templates.update(options_section(('ui', "User interface"), {
     "dimensions_and_batch_together": OptionInfo(True, "Show Width/Height and Batch sliders in same row"),

     "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": OptionInfo("sd_model_checkpoint", "Quicksettings list"),

+    "hidden_tabs": OptionInfo([], "Hidden UI tabs (requires restart)", ui_components.DropdownMulti, lambda: {"choices": [x for x in tab_names]}),

     "ui_reorder": OptionInfo(", ".join(ui_reorder_categories), "txt2img/img2img UI item order"),

     "ui_extra_networks_tab_reorder": OptionInfo("", "Extra networks tab order"),

     "localization": OptionInfo("None", "Localization (requires restart)", gr.Dropdown, lambda: {"choices": ["None"] + list(localization.localizations.keys())}, refresh=lambda: localization.list_localizations(cmd_opts.localizations_dir)),

+    "gradio_theme": OptionInfo("Default", "Gradio theme (requires restart)", ui_components.DropdownEditable, lambda: {"choices": ["Default"] + gradio_hf_hub_themes})

 }))

 

 options_templates.update(options_section(('ui', "Live previews"), {

@@ -476,10 +429,15 @@ options_templates.update(options_section(('sampler-params', "Sampler parameters"
     "eta_ancestral": OptionInfo(1.0, "eta (noise multiplier) for ancestral samplers", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),

     "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_min_uncond': OptionInfo(0, "Negative Guidance minimum sigma", gr.Slider, {"minimum": 0.0, "maximum": 4.0, "step": 0.01}),

     's_tmin':  OptionInfo(0.0, "sigma tmin",  gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),

     's_noise': OptionInfo(1.0, "sigma noise", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),

     'eta_noise_seed_delta': OptionInfo(0, "Eta noise seed delta", gr.Number, {"precision": 0}),

     'always_discard_next_to_last_sigma': OptionInfo(False, "Always discard next-to-last sigma"),

+    'uni_pc_variant': OptionInfo("bh1", "UniPC variant", gr.Radio, {"choices": ["bh1", "bh2", "vary_coeff"]}),

+    'uni_pc_skip_type': OptionInfo("time_uniform", "UniPC skip type", gr.Radio, {"choices": ["time_uniform", "time_quadratic", "logSNR"]}),

+    'uni_pc_order': OptionInfo(3, "UniPC order (must be < sampling steps)", gr.Slider, {"minimum": 1, "maximum": 50, "step": 1}),

+    'uni_pc_lower_order_final': OptionInfo(True, "UniPC lower order final"),

 }))

 

 options_templates.update(options_section(('postprocessing', "Postprocessing"), {

@@ -489,7 +447,9 @@ options_templates.update(options_section(('postprocessing', "Postprocessing"), {
 }))

 

 options_templates.update(options_section((None, "Hidden options"), {

-    "disabled_extensions": OptionInfo([], "Disable those extensions"),

+    "disabled_extensions": OptionInfo([], "Disable these extensions"),

+    "disable_all_extensions": OptionInfo("none", "Disable all extensions (preserves the list of disabled extensions)", gr.Radio, {"choices": ["none", "extra", "all"]}),

+    "restore_config_state_file": OptionInfo("", "Config state file to restore from, under 'config-states/' folder"),

     "sd_checkpoint_hash": OptionInfo("", "SHA256 hash of the current checkpoint"),

 }))

 

@@ -554,6 +514,15 @@ class Options:
 

         return True

 

+    def get_default(self, key):

+        """returns the default value for the key"""

+

+        data_label = self.data_labels.get(key)

+        if data_label is None:

+            return None

+

+        return data_label.default

+

     def save(self, filename):

         assert not cmd_opts.freeze_settings, "saving settings is disabled"

 

@@ -657,6 +626,24 @@ clip_model = None
 

 progress_print_out = sys.stdout

 

+gradio_theme = gr.themes.Base()

+

+

+def reload_gradio_theme(theme_name=None):

+    global gradio_theme

+    if not theme_name:

+        theme_name = opts.gradio_theme

+

+    if theme_name == "Default":

+        gradio_theme = gr.themes.Default()

+    else:

+        try:

+            gradio_theme = gr.themes.ThemeClass.from_hub(theme_name)

+        except requests.exceptions.ConnectionError:

+            print("Can't access HuggingFace Hub, falling back to default Gradio theme")

+            gradio_theme = gr.themes.Default()

+

+

 

 class TotalTQDM:

     def __init__(self):

@@ -686,6 +673,7 @@ class TotalTQDM:
 

     def clear(self):

         if self._tqdm is not None:

+            self._tqdm.refresh()

             self._tqdm.close()

             self._tqdm = None

 

@@ -697,7 +685,7 @@ mem_mon.start()
 

 

 def listfiles(dirname):

-    filenames = [os.path.join(dirname, x) for x in sorted(os.listdir(dirname)) if not x.startswith(".")]

+    filenames = [os.path.join(dirname, x) for x in sorted(os.listdir(dirname), key=str.lower) if not x.startswith(".")]

     return [file for file in filenames if os.path.isfile(file)]

 

 

diff --git a/modules/shared_items.py b/modules/shared_items.py
index 8b5ec96d..e792a134 100644
--- a/modules/shared_items.py
+++ b/modules/shared_items.py
@@ -20,4 +20,4 @@ def sd_vae_items():
 def refresh_vae_list():

     import modules.sd_vae

 

-    return modules.sd_vae.refresh_vae_list

+    modules.sd_vae.refresh_vae_list()

diff --git a/modules/styles.py b/modules/styles.py
index 990d5623..9ed85991 100644
--- a/modules/styles.py
+++ b/modules/styles.py
@@ -72,16 +72,14 @@ class StyleDatabase:
         return apply_styles_to_prompt(prompt, [self.styles.get(x, self.no_style).negative_prompt for x in styles])

 

     def save_styles(self, path: str) -> None:

-        # Write to temporary file first, so we don't nuke the file if something goes wrong

-        fd, temp_path = tempfile.mkstemp(".csv")

+        # Always keep a backup file around

+        if os.path.exists(path):

+            shutil.copy(path, path + ".bak")

+

+        fd = os.open(path, os.O_RDWR|os.O_CREAT)

         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())

-

-        # Always keep a backup file around

-        if os.path.exists(path):

-            shutil.move(path, path + ".bak")

-        shutil.move(temp_path, path)

diff --git a/modules/textual_inversion/dataset.py b/modules/textual_inversion/dataset.py
index d31963d4..af9fbcf2 100644
--- a/modules/textual_inversion/dataset.py
+++ b/modules/textual_inversion/dataset.py
@@ -19,9 +19,10 @@ re_numbers_at_start = re.compile(r"^[-\d]+\s*")
 

 

 class DatasetEntry:

-    def __init__(self, filename=None, filename_text=None, latent_dist=None, latent_sample=None, cond=None, cond_text=None, pixel_values=None):

+    def __init__(self, filename=None, filename_text=None, latent_dist=None, latent_sample=None, cond=None, cond_text=None, pixel_values=None, weight=None):

         self.filename = filename

         self.filename_text = filename_text

+        self.weight = weight

         self.latent_dist = latent_dist

         self.latent_sample = latent_sample

         self.cond = cond

@@ -30,7 +31,7 @@ class DatasetEntry:
 

 

 class PersonalizedBase(Dataset):

-    def __init__(self, data_root, width, height, repeats, flip_p=0.5, placeholder_token="*", model=None, cond_model=None, device=None, template_file=None, include_cond=False, batch_size=1, gradient_step=1, shuffle_tags=False, tag_drop_out=0, latent_sampling_method='once', varsize=False):

+    def __init__(self, data_root, width, height, repeats, flip_p=0.5, placeholder_token="*", model=None, cond_model=None, device=None, template_file=None, include_cond=False, batch_size=1, gradient_step=1, shuffle_tags=False, tag_drop_out=0, latent_sampling_method='once', varsize=False, use_weight=False):

         re_word = re.compile(shared.opts.dataset_filename_word_regex) if len(shared.opts.dataset_filename_word_regex) > 0 else None

 

         self.placeholder_token = placeholder_token

@@ -56,10 +57,16 @@ class PersonalizedBase(Dataset):
 

         print("Preparing dataset...")

         for path in tqdm.tqdm(self.image_paths):

+            alpha_channel = None

             if shared.state.interrupted:

                 raise Exception("interrupted")

             try:

-                image = Image.open(path).convert('RGB')

+                image = Image.open(path)

+                #Currently does not work for single color transparency

+                #We would need to read image.info['transparency'] for that

+                if use_weight and 'A' in image.getbands():

+                    alpha_channel = image.getchannel('A')

+                image = image.convert('RGB')

                 if not varsize:

                     image = image.resize((width, height), PIL.Image.BICUBIC)

             except Exception:

@@ -87,17 +94,35 @@ class PersonalizedBase(Dataset):
             with devices.autocast():

                 latent_dist = model.encode_first_stage(torchdata.unsqueeze(dim=0))

 

-            if latent_sampling_method == "once" or (latent_sampling_method == "deterministic" and not isinstance(latent_dist, DiagonalGaussianDistribution)):

-                latent_sample = model.get_first_stage_encoding(latent_dist).squeeze().to(devices.cpu)

-                latent_sampling_method = "once"

-                entry = DatasetEntry(filename=path, filename_text=filename_text, latent_sample=latent_sample)

-            elif latent_sampling_method == "deterministic":

-                # Works only for DiagonalGaussianDistribution

-                latent_dist.std = 0

-                latent_sample = model.get_first_stage_encoding(latent_dist).squeeze().to(devices.cpu)

-                entry = DatasetEntry(filename=path, filename_text=filename_text, latent_sample=latent_sample)

-            elif latent_sampling_method == "random":

-                entry = DatasetEntry(filename=path, filename_text=filename_text, latent_dist=latent_dist)

+            #Perform latent sampling, even for random sampling.

+            #We need the sample dimensions for the weights

+            if latent_sampling_method == "deterministic":

+                if isinstance(latent_dist, DiagonalGaussianDistribution):

+                    # Works only for DiagonalGaussianDistribution

+                    latent_dist.std = 0

+                else:

+                    latent_sampling_method = "once"

+            latent_sample = model.get_first_stage_encoding(latent_dist).squeeze().to(devices.cpu)

+

+            if use_weight and alpha_channel is not None:

+                channels, *latent_size = latent_sample.shape

+                weight_img = alpha_channel.resize(latent_size)

+                npweight = np.array(weight_img).astype(np.float32)

+                #Repeat for every channel in the latent sample

+                weight = torch.tensor([npweight] * channels).reshape([channels] + latent_size)

+                #Normalize the weight to a minimum of 0 and a mean of 1, that way the loss will be comparable to default.

+                weight -= weight.min()

+                weight /= weight.mean()

+            elif use_weight:

+                #If an image does not have a alpha channel, add a ones weight map anyway so we can stack it later

+                weight = torch.ones(latent_sample.shape)

+            else:

+                weight = None

+            

+            if latent_sampling_method == "random":

+                entry = DatasetEntry(filename=path, filename_text=filename_text, latent_dist=latent_dist, weight=weight)

+            else:

+                entry = DatasetEntry(filename=path, filename_text=filename_text, latent_sample=latent_sample, weight=weight)

 

             if not (self.tag_drop_out != 0 or self.shuffle_tags):

                 entry.cond_text = self.create_text(filename_text)

@@ -110,6 +135,7 @@ class PersonalizedBase(Dataset):
             del torchdata

             del latent_dist

             del latent_sample

+            del weight

 

         self.length = len(self.dataset)

         self.groups = list(groups.values())

@@ -195,6 +221,10 @@ class BatchLoader:
         self.cond_text = [entry.cond_text for entry in data]

         self.cond = [entry.cond for entry in data]

         self.latent_sample = torch.stack([entry.latent_sample for entry in data]).squeeze(1)

+        if all(entry.weight is not None for entry in data):

+            self.weight = torch.stack([entry.weight for entry in data]).squeeze(1)

+        else:

+            self.weight = None

         #self.emb_index = [entry.emb_index for entry in data]

         #print(self.latent_sample.device)

 

diff --git a/modules/textual_inversion/preprocess.py b/modules/textual_inversion/preprocess.py
index 2239cb84..4a29151d 100644
--- a/modules/textual_inversion/preprocess.py
+++ b/modules/textual_inversion/preprocess.py
@@ -11,7 +11,7 @@ from modules.shared import opts, cmd_opts
 from modules.textual_inversion import autocrop

 

 

-def preprocess(id_task, process_src, process_dst, process_width, process_height, preprocess_txt_action, 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.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):

     try:

         if process_caption:

             shared.interrogator.load()

@@ -19,7 +19,7 @@ def preprocess(id_task, process_src, process_dst, process_width, process_height,
         if process_caption_deepbooru:

             deepbooru.model.start()

 

-        preprocess_work(process_src, process_dst, process_width, process_height, preprocess_txt_action, process_flip, process_split, process_caption, process_caption_deepbooru, split_threshold, overlap_ratio, process_focal_crop, process_focal_crop_face_weight, process_focal_crop_entropy_weight, process_focal_crop_edges_weight, process_focal_crop_debug, process_multicrop, process_multicrop_mindim, process_multicrop_maxdim, process_multicrop_minarea, process_multicrop_maxarea, process_multicrop_objective, process_multicrop_threshold)

+        preprocess_work(process_src, process_dst, process_width, process_height, preprocess_txt_action, process_keep_original_size, process_flip, process_split, process_caption, process_caption_deepbooru, split_threshold, overlap_ratio, process_focal_crop, process_focal_crop_face_weight, process_focal_crop_entropy_weight, process_focal_crop_edges_weight, process_focal_crop_debug, process_multicrop, process_multicrop_mindim, process_multicrop_maxdim, process_multicrop_minarea, process_multicrop_maxarea, process_multicrop_objective, process_multicrop_threshold)

 

     finally:

 

@@ -131,7 +131,7 @@ def multicrop_pic(image: Image, mindim, maxdim, minarea, maxarea, objective, thr
     return wh and center_crop(image, *wh)

     

 

-def preprocess_work(process_src, process_dst, process_width, process_height, preprocess_txt_action, 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_work(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):

     width = process_width

     height = process_height

     src = os.path.abspath(process_src)

@@ -161,7 +161,9 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
         params.subindex = 0

         filename = os.path.join(src, imagefile)

         try:

-            img = Image.open(filename).convert("RGB")

+            img = Image.open(filename)

+            img = ImageOps.exif_transpose(img)

+            img = img.convert("RGB")

         except Exception:

             continue

 

@@ -223,6 +225,10 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
                 print(f"skipped {img.width}x{img.height} image {filename} (can't find suitable size within error threshold)")

             process_default_resize = False

 

+        if process_keep_original_size:

+            save_pic(img, index, params, existing_caption=existing_caption)

+            process_default_resize = False

+

         if process_default_resize:

             img = images.resize_image(1, img, width, height)

             save_pic(img, index, params, existing_caption=existing_caption)

diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py
index a1a406c2..379df243 100644
--- a/modules/textual_inversion/textual_inversion.py
+++ b/modules/textual_inversion/textual_inversion.py
@@ -152,7 +152,11 @@ class EmbeddingDatabase:
                 name = data.get('name', name)

             else:

                 data = extract_image_data_embed(embed_image)

-                name = data.get('name', name)

+                if data:

+                    name = data.get('name', name)

+                else:

+                    # if data is None, means this is not an embeding, just a preview image

+                    return

         elif ext in ['.BIN', '.PT']:

             data = torch.load(path, map_location="cpu")

         elif ext in ['.SAFETENSORS']:

@@ -229,6 +233,12 @@ class EmbeddingDatabase:
             self.load_from_dir(embdir)

             embdir.update()

 

+        # re-sort word_embeddings because load_from_dir may not load in alphabetic order.

+        # using a temporary copy so we don't reinitialize self.word_embeddings in case other objects have a reference to it.

+        sorted_word_embeddings = {e.name: e for e in sorted(self.word_embeddings.values(), key=lambda e: e.name.lower())}

+        self.word_embeddings.clear()

+        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:

             self.previously_displayed_embeddings = displayed_embeddings

@@ -351,7 +361,7 @@ def validate_train_inputs(model_name, learn_rate, batch_size, gradient_step, dat
         assert log_directory, "Log directory is empty"

 

 

-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, 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):

+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):

     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)

@@ -410,7 +420,7 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st
 

     pin_memory = shared.opts.pin_memory

 

-    ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=training_width, height=training_height, repeats=shared.opts.training_image_repeats_per_epoch, placeholder_token=embedding_name, model=shared.sd_model, cond_model=shared.sd_model.cond_stage_model, device=devices.device, template_file=template_file, batch_size=batch_size, gradient_step=gradient_step, shuffle_tags=shuffle_tags, tag_drop_out=tag_drop_out, latent_sampling_method=latent_sampling_method, varsize=varsize)

+    ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=training_width, height=training_height, repeats=shared.opts.training_image_repeats_per_epoch, placeholder_token=embedding_name, model=shared.sd_model, cond_model=shared.sd_model.cond_stage_model, device=devices.device, template_file=template_file, batch_size=batch_size, gradient_step=gradient_step, shuffle_tags=shuffle_tags, tag_drop_out=tag_drop_out, latent_sampling_method=latent_sampling_method, varsize=varsize, use_weight=use_weight)

 

     if shared.opts.save_training_settings_to_txt:

         save_settings_to_file(log_directory, {**dict(model_name=checkpoint.model_name, model_hash=checkpoint.shorthash, num_of_dataset_images=len(ds), num_vectors_per_token=len(embedding.vec)), **locals()})

@@ -480,6 +490,8 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st
             

                 with devices.autocast():

                     x = batch.latent_sample.to(devices.device, non_blocking=pin_memory)

+                    if use_weight:

+                        w = batch.weight.to(devices.device, non_blocking=pin_memory)

                     c = shared.sd_model.cond_stage_model(batch.cond_text)

 

                     if is_training_inpainting_model:

@@ -490,7 +502,11 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st
                     else:

                         cond = c

 

-                    loss = shared.sd_model(x, cond)[0] / gradient_step

+                    if use_weight:

+                        loss = shared.sd_model.weighted_forward(x, cond, w)[0] / gradient_step

+                        del w

+                    else:

+                        loss = shared.sd_model.forward(x, cond)[0] / gradient_step

                     del x

 

                     _loss_step += loss.item()

diff --git a/modules/timer.py b/modules/timer.py
index 57a4f17a..ba92be33 100644
--- a/modules/timer.py
+++ b/modules/timer.py
@@ -33,3 +33,6 @@ class Timer:
         res += ")"

 

         return res

+

+    def reset(self):

+        self.__init__()

diff --git a/modules/ui.py b/modules/ui.py
index 5e34fb07..1130345c 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -20,7 +20,7 @@ from PIL import Image, PngImagePlugin
 from modules.call_queue import wrap_gradio_gpu_call, wrap_queued_call, wrap_gradio_call

 

 from modules import sd_hijack, sd_models, localization, script_callbacks, ui_extensions, deepbooru, sd_vae, extra_networks, postprocessing, ui_components, ui_common, ui_postprocessing

-from modules.ui_components import FormRow, FormGroup, ToolButton, FormHTML

+from modules.ui_components import FormRow, FormColumn, FormGroup, ToolButton, FormHTML

 from modules.paths import script_path, data_path

 

 from modules.shared import opts, cmd_opts, restricted_opts

@@ -70,17 +70,6 @@ def gr_show(visible=True):
 sample_img2img = "assets/stable-samples/img2img/sketch-mountains-input.jpg"

 sample_img2img = sample_img2img if os.path.exists(sample_img2img) else None

 

-css_hide_progressbar = """

-.wrap .m-12 svg { display:none!important; }

-.wrap .m-12::before { content:"Loading..." }

-.wrap .z-20 svg { display:none!important; }

-.wrap .z-20::before { content:"Loading..." }

-.wrap.cover-bg .z-20::before { content:"" }

-.progress-bar { display:none!important; }

-.meta-text { display:none!important; }

-.meta-text-center { display:none!important; }

-"""

-

 # Using constants for these since the variation selector isn't visible.

 # Important that they exactly match script.js for tooltip to work.

 random_symbol = '\U0001f3b2\ufe0f'  # 🎲️

@@ -89,7 +78,7 @@ paste_symbol = '\u2199\ufe0f'  # ↙
 refresh_symbol = '\U0001f504'  # 🔄

 save_style_symbol = '\U0001f4be'  # 💾

 apply_style_symbol = '\U0001f4cb'  # 📋

-clear_prompt_symbol = '\U0001F5D1'  # 🗑️

+clear_prompt_symbol = '\U0001f5d1\ufe0f'  # 🗑️

 extra_networks_symbol = '\U0001F3B4'  # 🎴

 switch_values_symbol = '\U000021C5' # ⇅

 

@@ -105,6 +94,9 @@ def send_gradio_gallery_to_image(x):
 

 def visit(x, func, path=""):

     if hasattr(x, 'children'):

+        if isinstance(x, gr.Tabs) and x.elem_id is not None:

+            # Tabs element can't have a label, have to use elem_id instead

+            func(f"{path}/Tabs@{x.elem_id}", x)

         for c in x.children:

             visit(c, func, path)

     elif x.label is not None:

@@ -138,6 +130,16 @@ def calc_resolution_hires(enable, width, height, hr_scale, hr_resize_x, hr_resiz
     return f"resize: from <span class='resolution'>{p.width}x{p.height}</span> to <span class='resolution'>{p.hr_resize_x or p.hr_upscale_to_x}x{p.hr_resize_y or p.hr_upscale_to_y}</span>"

 

 

+def resize_from_to_html(width, height, scale_by):

+    target_width = int(width * scale_by)

+    target_height = int(height * scale_by)

+

+    if not target_width or not target_height:

+        return "no image selected"

+

+    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)

@@ -179,14 +181,13 @@ def interrogate_deepbooru(image):
 

 

 def create_seed_inputs(target_interface):

-    with FormRow(elem_id=target_interface + '_seed_row'):

+    with FormRow(elem_id=target_interface + '_seed_row', variant="compact"):

         seed = (gr.Textbox if cmd_opts.use_textbox_seed else gr.Number)(label='Seed', value=-1, elem_id=target_interface + '_seed')

         seed.style(container=False)

-        random_seed = gr.Button(random_symbol, elem_id=target_interface + '_random_seed')

-        reuse_seed = gr.Button(reuse_symbol, elem_id=target_interface + '_reuse_seed')

+        random_seed = ToolButton(random_symbol, elem_id=target_interface + '_random_seed', label='Random seed')

+        reuse_seed = ToolButton(reuse_symbol, elem_id=target_interface + '_reuse_seed', label='Reuse seed')

 

-        with gr.Group(elem_id=target_interface + '_subseed_show_box'):

-            seed_checkbox = gr.Checkbox(label='Extra', elem_id=target_interface + '_subseed_show', value=False)

+        seed_checkbox = gr.Checkbox(label='Extra', elem_id=target_interface + '_subseed_show', value=False)

 

     # Components to show/hide based on the 'Extra' checkbox

     seed_extras = []

@@ -195,8 +196,8 @@ def create_seed_inputs(target_interface):
         seed_extras.append(seed_extra_row_1)

         subseed = gr.Number(label='Variation seed', value=-1, elem_id=target_interface + '_subseed')

         subseed.style(container=False)

-        random_subseed = gr.Button(random_symbol, elem_id=target_interface + '_random_subseed')

-        reuse_subseed = gr.Button(reuse_symbol, elem_id=target_interface + '_reuse_subseed')

+        random_subseed = ToolButton(random_symbol, elem_id=target_interface + '_random_subseed')

+        reuse_subseed = ToolButton(reuse_symbol, elem_id=target_interface + '_reuse_subseed')

         subseed_strength = gr.Slider(label='Variation strength', value=0.0, minimum=0, maximum=1, step=0.01, elem_id=target_interface + '_subseed_strength')

 

     with FormRow(visible=False) as seed_extra_row_2:

@@ -291,19 +292,19 @@ def create_toprow(is_img2img):
             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=2, placeholder="Negative prompt (press Ctrl+Enter or Alt+Enter to generate)")

+                        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)")

 

         button_interrogate = None

         button_deepbooru = None

         if is_img2img:

-            with gr.Column(scale=1, elem_id="interrogate_col"):

+            with gr.Column(scale=1, elem_classes="interrogate-col"):

                 button_interrogate = gr.Button('Interrogate\nCLIP', elem_id="interrogate")

                 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"):

-                interrupt = gr.Button('Interrupt', elem_id=f"{id_part}_interrupt")

-                skip = gr.Button('Skip', elem_id=f"{id_part}_skip")

+            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")

+                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')

 

                 skip.click(

@@ -325,9 +326,9 @@ def create_toprow(is_img2img):
                 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")

 

-                token_counter = gr.HTML(value="<span></span>", elem_id=f"{id_part}_token_counter")

+                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")

-                negative_token_counter = gr.HTML(value="<span></span>", elem_id=f"{id_part}_negative_token_counter")

+                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")

 

                 clear_prompt_button.click(

@@ -479,7 +480,9 @@ def create_ui():
                                 width = gr.Slider(minimum=64, maximum=2048, step=8, label="Width", value=512, elem_id="txt2img_width")

                                 height = gr.Slider(minimum=64, maximum=2048, step=8, label="Height", value=512, elem_id="txt2img_height")

 

-                            res_switch_btn = ToolButton(value=switch_values_symbol, elem_id="txt2img_res_switch_btn")

+                            with gr.Column(elem_id="txt2img_dimensions_row", scale=1, elem_classes="dimensions-tools"):

+                                res_switch_btn = ToolButton(value=switch_values_symbol, elem_id="txt2img_res_switch_btn", label="Switch dims")

+

                             if opts.dimensions_and_batch_together:

                                 with gr.Column(elem_id="txt2img_column_batch"):

                                     batch_count = gr.Slider(minimum=1, step=1, label='Batch count', value=1, elem_id="txt2img_batch_count")

@@ -492,7 +495,7 @@ def create_ui():
                         seed, reuse_seed, subseed, reuse_subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w, seed_checkbox = create_seed_inputs('txt2img')

 

                     elif category == "checkboxes":

-                        with FormRow(elem_id="txt2img_checkboxes", variant="compact"):

+                        with FormRow(elem_classes="checkboxes-row", variant="compact"):

                             restore_faces = gr.Checkbox(label='Restore faces', value=False, visible=len(shared.face_restorers) > 1, elem_id="txt2img_restore_faces")

                             tiling = gr.Checkbox(label='Tiling', value=False, elem_id="txt2img_tiling")

                             enable_hr = gr.Checkbox(label='Hires. fix', value=False, elem_id="txt2img_enable_hr")

@@ -586,7 +589,7 @@ def create_ui():
             txt2img_prompt.submit(**txt2img_args)

             submit.click(**txt2img_args)

 

-            res_switch_btn.click(lambda w, h: (h, w), inputs=[width, height], outputs=[width, height])

+            res_switch_btn.click(lambda w, h: (h, w), inputs=[width, height], outputs=[width, height], show_progress=False)

 

             txt_prompt_img.change(

                 fn=modules.images.image_data,

@@ -631,9 +634,9 @@ def create_ui():
                 (hr_resize_y, "Hires resize-2"),

                 *modules.scripts.scripts_txt2img.infotext_fields

             ]

-            parameters_copypaste.add_paste_fields("txt2img", None, txt2img_paste_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, override_settings_component=override_settings,

+                paste_button=txt2img_paste, tabname="txt2img", source_text_component=txt2img_prompt, source_image_component=None,

             ))

 

             txt2img_preview_params = [

@@ -683,6 +686,8 @@ def create_ui():
                             copy_image_buttons.append((button, name, elem))

 

                 with gr.Tabs(elem_id="mode_img2img"):

+                    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=480)

                         add_copy_image_controls('img2img', init_img)

@@ -725,6 +730,12 @@ def create_ui():
                         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")

 

+                    img2img_tabs = [tab_img2img, tab_sketch, tab_inpaint, tab_inpaint_color, tab_inpaint_upload, tab_batch]

+                    img2img_image_inputs = [init_img, sketch, init_img_with_mask, inpaint_color_sketch]

+

+                    for i, tab in enumerate(img2img_tabs):

+                        tab.select(fn=lambda tabnum=i: tabnum, inputs=[], outputs=[img2img_selected_tab])

+

                 def copy_image(img):

                     if isinstance(img, dict) and 'image' in img:

                         return img['image']

@@ -754,10 +765,35 @@ def create_ui():
                     elif category == "dimensions":

                         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")

-                                height = gr.Slider(minimum=64, maximum=2048, step=8, label="Height", value=512, elem_id="img2img_height")

+                                selected_scale_tab = gr.State(value=0)

+

+                                with gr.Tabs():

+                                    with gr.Tab(label="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")

+                                                height = gr.Slider(minimum=64, maximum=2048, step=8, label="Height", value=512, elem_id="img2img_height")

+                                            with gr.Column(elem_id="img2img_dimensions_row", scale=1, elem_classes="dimensions-tools"):

+                                                res_switch_btn = ToolButton(value=switch_values_symbol, elem_id="img2img_res_switch_btn")

+

+                                    with gr.Tab(label="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():

+                                            scale_by_html = FormHTML(resize_from_to_html(0, 0, 0.0), elem_id="img2img_scale_resolution_preview")

+                                            gr.Slider(label="Unused", elem_id="img2img_unused_scale_by_slider")

+

+                                    scale_by.change(

+                                        fn=resize_from_to_html,

+                                        _js="currentImg2imgSourceResolution",

+                                        inputs=[dummy_component, dummy_component, scale_by],

+                                        outputs=scale_by_html,

+                                        show_progress=False,

+                                    )

+

+                            tab_scale_to.select(fn=lambda: 0, inputs=[], outputs=[selected_scale_tab])

+                            tab_scale_by.select(fn=lambda: 1, inputs=[], outputs=[selected_scale_tab])

 

-                            res_switch_btn = ToolButton(value=switch_values_symbol, elem_id="img2img_res_switch_btn")

                             if opts.dimensions_and_batch_together:

                                 with gr.Column(elem_id="img2img_column_batch"):

                                     batch_count = gr.Slider(minimum=1, step=1, label='Batch count', value=1, elem_id="img2img_batch_count")

@@ -765,14 +801,16 @@ def create_ui():
 

                     elif category == "cfg":

                         with FormGroup():

-                            cfg_scale = gr.Slider(minimum=1.0, maximum=30.0, step=0.5, label='CFG Scale', value=7.0, elem_id="img2img_cfg_scale")

+                            with FormRow():

+                                cfg_scale = gr.Slider(minimum=1.0, maximum=30.0, step=0.5, label='CFG Scale', value=7.0, elem_id="img2img_cfg_scale")

+                                image_cfg_scale = gr.Slider(minimum=0, maximum=3.0, step=0.05, label='Image CFG Scale', value=1.5, elem_id="img2img_image_cfg_scale", visible=shared.sd_model and shared.sd_model.cond_stage_key == "edit")

                             denoising_strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Denoising strength', value=0.75, elem_id="img2img_denoising_strength")

 

                     elif category == "seed":

                         seed, reuse_seed, subseed, reuse_subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w, seed_checkbox = create_seed_inputs('img2img')

 

                     elif category == "checkboxes":

-                        with FormRow(elem_id="img2img_checkboxes", variant="compact"):

+                        with FormRow(elem_classes="checkboxes-row", variant="compact"):

                             restore_faces = gr.Checkbox(label='Restore faces', value=False, visible=len(shared.face_restorers) > 1, elem_id="img2img_restore_faces")

                             tiling = gr.Checkbox(label='Tiling', value=False, elem_id="img2img_tiling")

 

@@ -812,7 +850,7 @@ def create_ui():
                             def select_img2img_tab(tab):

                                 return gr.update(visible=tab in [2, 3, 4]), gr.update(visible=tab == 3),

 

-                            for i, elem in enumerate([tab_img2img, tab_sketch, tab_inpaint, tab_inpaint_color, tab_inpaint_upload, tab_batch]):

+                            for i, elem in enumerate(img2img_tabs):

                                 elem.select(

                                     fn=lambda tab=i: select_img2img_tab(tab),

                                     inputs=[],

@@ -861,11 +899,14 @@ def create_ui():
                     batch_count,

                     batch_size,

                     cfg_scale,

+                    image_cfg_scale,

                     denoising_strength,

                     seed,

                     subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w, seed_checkbox,

+                    selected_scale_tab,

                     height,

                     width,

+                    scale_by,

                     resize_mode,

                     inpaint_full_res,

                     inpaint_full_res_padding,

@@ -901,7 +942,7 @@ def create_ui():
 

             img2img_prompt.submit(**img2img_args)

             submit.click(**img2img_args)

-            res_switch_btn.click(lambda w, h: (h, w), inputs=[width, height], outputs=[width, height])

+            res_switch_btn.click(lambda w, h: (h, w), inputs=[width, height], outputs=[width, height], show_progress=False)

 

             img2img_interrogate.click(

                 fn=lambda *args: process_interrogate(interrogate, *args),

@@ -936,7 +977,7 @@ def create_ui():
                 )

 

             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=[txt2img_negative_prompt, steps], outputs=[negative_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)

 

@@ -947,6 +988,7 @@ def create_ui():
                 (sampler_index, "Sampler"),

                 (restore_faces, "Face restoration"),

                 (cfg_scale, "CFG scale"),

+                (image_cfg_scale, "Image CFG scale"),

                 (seed, "Seed"),

                 (width, "Size-1"),

                 (height, "Size-2"),

@@ -959,10 +1001,10 @@ def create_ui():
                 (mask_blur, "Mask blur"),

                 *modules.scripts.scripts_img2img.infotext_fields

             ]

-            parameters_copypaste.add_paste_fields("img2img", init_img, img2img_paste_fields)

-            parameters_copypaste.add_paste_fields("inpaint", init_img_with_mask, img2img_paste_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, override_settings_component=override_settings,

+                paste_button=img2img_paste, tabname="img2img", source_text_component=img2img_prompt, source_image_component=None,

             ))

 

     modules.scripts.scripts_current = None

@@ -1023,8 +1065,9 @@ def create_ui():
                 interp_method.change(fn=update_interp_description, inputs=[interp_method], outputs=[interp_description])

 

                 with FormRow():

-                    checkpoint_format = gr.Radio(choices=["ckpt", "safetensors"], value="ckpt", label="Checkpoint format", elem_id="modelmerger_checkpoint_format")

+                    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():

@@ -1052,7 +1095,7 @@ def create_ui():
         with gr.Row(variant="compact").style(equal_height=False):

             with gr.Tabs(elem_id="train_tabs"):

 

-                with gr.Tab(label="Create embedding"):

+                with gr.Tab(label="Create embedding", id="create_embedding"):

                     new_embedding_name = gr.Textbox(label="Name", elem_id="train_new_embedding_name")

                     initialization_text = gr.Textbox(label="Initialization text", value="*", elem_id="train_initialization_text")

                     nvpt = gr.Slider(label="Number of vectors per token", minimum=1, maximum=75, step=1, value=1, elem_id="train_nvpt")

@@ -1065,7 +1108,7 @@ def create_ui():
                         with gr.Column():

                             create_embedding = gr.Button(value="Create embedding", variant='primary', elem_id="train_create_embedding")

 

-                with gr.Tab(label="Create hypernetwork"):

+                with gr.Tab(label="Create hypernetwork", id="create_hypernetwork"):

                     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")

@@ -1083,7 +1126,7 @@ def create_ui():
                         with gr.Column():

                             create_hypernetwork = gr.Button(value="Create hypernetwork", variant='primary', elem_id="train_create_hypernetwork")

 

-                with gr.Tab(label="Preprocess images"):

+                with gr.Tab(label="Preprocess images", id="preprocess_images"):

                     process_src = gr.Textbox(label='Source directory', elem_id="train_process_src")

                     process_dst = gr.Textbox(label='Destination directory', elem_id="train_process_dst")

                     process_width = gr.Slider(minimum=64, maximum=2048, step=8, label="Width", value=512, elem_id="train_process_width")

@@ -1091,6 +1134,7 @@ def create_ui():
                     preprocess_txt_action = gr.Dropdown(label='Existing Caption txt Action', value="ignore", choices=["ignore", "copy", "prepend", "append"], elem_id="train_preprocess_txt_action")

 

                     with gr.Row():

+                        process_keep_original_size = gr.Checkbox(label='Keep original size', elem_id="train_process_keep_original_size")

                         process_flip = gr.Checkbox(label='Create flipped copies', elem_id="train_process_flip")

                         process_split = gr.Checkbox(label='Split oversized images', elem_id="train_process_split")

                         process_focal_crop = gr.Checkbox(label='Auto focal point crop', elem_id="train_process_focal_crop")

@@ -1150,7 +1194,7 @@ def create_ui():
                 def get_textual_inversion_template_names():

                     return sorted([x for x in textual_inversion.textual_inversion_templates])

 

-                with gr.Tab(label="Train"):

+                with gr.Tab(label="Train", id="train"):

                     gr.HTML(value="<p style='margin-bottom: 0.7em'>Train an embedding or Hypernetwork; you must specify a directory with a set of 1:1 ratio images <a href=\"https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Textual-Inversion\" style=\"font-weight:bold;\">[wiki]</a></p>")

                     with FormRow():

                         train_embedding_name = gr.Dropdown(label='Embedding', elem_id="train_embedding", choices=sorted(sd_hijack.model_hijack.embedding_db.word_embeddings.keys()))

@@ -1187,6 +1231,8 @@ def create_ui():
                         create_image_every = gr.Number(label='Save an image to log directory every N steps, 0 to disable', value=500, precision=0, elem_id="train_create_image_every")

                         save_embedding_every = gr.Number(label='Save a copy of embedding to log directory every N steps, 0 to disable', value=500, precision=0, elem_id="train_save_embedding_every")

 

+                    use_weight = gr.Checkbox(label="Use PNG alpha channel as loss weight", value=False, elem_id="use_weight")

+

                     save_image_with_stored_embedding = gr.Checkbox(label='Save images with embedding in PNG chunks', value=True, elem_id="train_save_image_with_stored_embedding")

                     preview_from_txt2img = gr.Checkbox(label='Read parameters (prompt, etc...) from txt2img tab when making previews', value=False, elem_id="train_preview_from_txt2img")

 

@@ -1206,7 +1252,7 @@ def create_ui():
 

             with gr.Column(elem_id='ti_gallery_container'):

                 ti_output = gr.Text(elem_id="ti_output", value="", show_label=False)

-                ti_gallery = gr.Gallery(label='Output', show_label=False, elem_id='ti_gallery').style(grid=4)

+                ti_gallery = gr.Gallery(label='Output', show_label=False, elem_id='ti_gallery').style(columns=4)

                 ti_progress = gr.HTML(elem_id="ti_progress", value="")

                 ti_outcome = gr.HTML(elem_id="ti_error", value="")

 

@@ -1255,6 +1301,7 @@ def create_ui():
                 process_width,

                 process_height,

                 preprocess_txt_action,

+                process_keep_original_size,

                 process_flip,

                 process_split,

                 process_caption,

@@ -1300,6 +1347,7 @@ def create_ui():
                 shuffle_tags,

                 tag_drop_out,

                 latent_sampling_method,

+                use_weight,

                 create_image_every,

                 save_embedding_every,

                 template_file,

@@ -1333,6 +1381,7 @@ def create_ui():
                 shuffle_tags,

                 tag_drop_out,

                 latent_sampling_method,

+                use_weight,

                 create_image_every,

                 save_embedding_every,

                 template_file,

@@ -1479,15 +1528,37 @@ def create_ui():
                 current_row.__exit__()

                 current_tab.__exit__()

 

-            with gr.TabItem("Actions"):

+            with gr.TabItem("Actions", id="actions"):

                 request_notifications = gr.Button(value='Request browser notifications', elem_id="request_notifications")

                 download_localization = gr.Button(value='Download localization template', elem_id="download_localization")

                 reload_script_bodies = gr.Button(value='Reload custom script bodies (No ui updates, No restart)', variant='secondary', elem_id="settings_reload_script_bodies")

+                with gr.Row():

+                    unload_sd_model = gr.Button(value='Unload SD checkpoint to free VRAM', elem_id="sett_unload_sd_model")

+                    reload_sd_model = gr.Button(value='Reload the last SD checkpoint back into VRAM', elem_id="sett_reload_sd_model")

 

-            with gr.TabItem("Licenses"):

+            with gr.TabItem("Licenses", id="licenses"):

                 gr.HTML(shared.html("licenses.html"), elem_id="licenses")

 

             gr.Button(value="Show all pages", elem_id="settings_show_all_pages")

+            

+

+        def unload_sd_weights():

+            modules.sd_models.unload_model_weights()

+

+        def reload_sd_weights():

+            modules.sd_models.reload_model_weights()

+

+        unload_sd_model.click(

+            fn=unload_sd_weights,

+            inputs=[],

+            outputs=[]

+        )

+

+        reload_sd_model.click(

+            fn=reload_sd_weights,

+            inputs=[],

+            outputs=[]

+        )

 

         request_notifications.click(

             fn=lambda: None,

@@ -1533,29 +1604,17 @@ def create_ui():
         (train_interface, "Train", "ti"),

     ]

 

-    css = ""

-

-    for cssfile in modules.scripts.list_files_with_name("style.css"):

-        if not os.path.isfile(cssfile):

-            continue

-

-        with open(cssfile, "r", encoding="utf8") as file:

-            css += file.read() + "\n"

-

-    if os.path.exists(os.path.join(data_path, "user.css")):

-        with open(os.path.join(data_path, "user.css"), "r", encoding="utf8") as file:

-            css += file.read() + "\n"

-

-    if not cmd_opts.no_progressbar_hiding:

-        css += css_hide_progressbar

-

     interfaces += script_callbacks.ui_tabs_callback()

     interfaces += [(settings_interface, "Settings", "settings")]

 

     extensions_interface = ui_extensions.create_ui()

     interfaces += [(extensions_interface, "Extensions", "extensions")]

 

-    with gr.Blocks(css=css, analytics_enabled=False, title="Stable Diffusion") as demo:

+    shared.tab_names = []

+    for _interface, label, _ifid in interfaces:

+        shared.tab_names.append(label)

+

+    with gr.Blocks(theme=shared.gradio_theme, analytics_enabled=False, title="Stable Diffusion") as demo:

         with gr.Row(elem_id="quicksettings", variant="compact"):

             for i, k, item in sorted(quicksettings_list, key=lambda x: quicksettings_names.get(x[1], x[0])):

                 component = create_setting_component(k, is_quicksettings=True)

@@ -1565,6 +1624,8 @@ def create_ui():
 

         with gr.Tabs(elem_id="tabs") as tabs:

             for interface, label, ifid in interfaces:

+                if label in shared.opts.hidden_tabs:

+                    continue

                 with gr.TabItem(label, id=ifid, elem_id='tab_' + ifid):

                     interface.render()

 

@@ -1584,13 +1645,21 @@ def create_ui():
 

         for i, k, item in quicksettings_list:

             component = component_dict[k]

+            info = opts.data_labels[k]

 

             component.change(

                 fn=lambda value, k=k: run_settings_single(value, key=k),

                 inputs=[component],

                 outputs=[component, text_settings],

+                show_progress=info.refresh is not None,

             )

 

+        text_settings.change(

+            fn=lambda: gr.update(visible=shared.sd_model and shared.sd_model.cond_stage_key == "edit"),

+            inputs=[],

+            outputs=[image_cfg_scale],

+        )

+

         button_set_checkpoint = gr.Button('Change checkpoint', elem_id='change_checkpoint', visible=False)

         button_set_checkpoint.click(

             fn=lambda value, _: run_settings_single(value, key='sd_model_checkpoint'),

@@ -1608,6 +1677,7 @@ def create_ui():
             fn=get_settings_values,

             inputs=[],

             outputs=[component_dict[k] for k in component_keys],

+            queue=False,

         )

 

         def modelmerger(*args):

@@ -1637,6 +1707,7 @@ def create_ui():
                 config_source,

                 bake_in_vae,

                 discard_weights,

+                save_metadata,

             ],

             outputs=[

                 primary_model_name,

@@ -1684,7 +1755,7 @@ def create_ui():
                 if init_field is not None:

                     init_field(saved_value)

 

-        if type(x) in [gr.Slider, gr.Radio, gr.Checkbox, gr.Textbox, gr.Number, gr.Dropdown] and x.visible:

+        if type(x) in [gr.Slider, gr.Radio, gr.Checkbox, gr.Textbox, gr.Number, gr.Dropdown, ToolButton] and x.visible:

             apply_field(x, 'visible')

 

         if type(x) == gr.Slider:

@@ -1714,12 +1785,27 @@ def create_ui():
 

             apply_field(x, 'value', check_dropdown, getattr(x, 'init_field', None))

 

+        def check_tab_id(tab_id):

+            tab_items = list(filter(lambda e: isinstance(e, gr.TabItem), x.children))

+            if type(tab_id) == str:

+                tab_ids = [t.id for t in tab_items]

+                return tab_id in tab_ids

+            elif type(tab_id) == int:

+                return tab_id >= 0 and tab_id < len(tab_items)

+            else:

+                return False

+

+        if type(x) == gr.Tabs:

+            apply_field(x, 'selected', check_tab_id)

+

     visit(txt2img_interface, loadsave, "txt2img")

     visit(img2img_interface, loadsave, "img2img")

     visit(extras_interface, loadsave, "extras")

     visit(modelmerger_interface, loadsave, "modelmerger")

     visit(train_interface, loadsave, "train")

 

+    loadsave(f"webui/Tabs@{tabs.elem_id}", tabs)

+

     if not error_loading and (not os.path.exists(ui_config_file) or settings_count != len(ui_settings)):

         with open(ui_config_file, "w", encoding="utf8") as file:

             json.dump(ui_settings, file, indent=4)

@@ -1730,21 +1816,60 @@ def create_ui():
     return demo

 

 

-def reload_javascript():

-    head = f'<script type="text/javascript" src="file={os.path.abspath("script.js")}?{os.path.getmtime("script.js")}"></script>\n'

+def webpath(fn):

+    if fn.startswith(script_path):

+        web_path = os.path.relpath(fn, script_path).replace('\\', '/')

+    else:

+        web_path = os.path.abspath(fn)

+

+    return f'file={web_path}?{os.path.getmtime(fn)}'

+

+

+def javascript_html():

+    script_js = os.path.join(script_path, "script.js")

+    head = f'<script type="text/javascript" src="{webpath(script_js)}"></script>\n'

 

     inline = f"{localization.localization_js(shared.opts.localization)};"

     if cmd_opts.theme is not None:

         inline += f"set_theme('{cmd_opts.theme}');"

 

     for script in modules.scripts.list_scripts("javascript", ".js"):

-        head += f'<script type="text/javascript" src="file={script.path}?{os.path.getmtime(script.path)}"></script>\n'

+        head += f'<script type="text/javascript" src="{webpath(script.path)}"></script>\n'

+

+    for script in modules.scripts.list_scripts("javascript", ".mjs"):

+        head += f'<script type="module" src="{webpath(script.path)}"></script>\n'

 

     head += f'<script type="text/javascript">{inline}</script>\n'

 

+    return head

+

+

+def css_html():

+    head = ""

+

+    def stylesheet(fn):

+        return f'<link rel="stylesheet" property="stylesheet" href="{webpath(fn)}">'

+

+    for cssfile in modules.scripts.list_files_with_name("style.css"):

+        if not os.path.isfile(cssfile):

+            continue

+

+        head += stylesheet(cssfile)

+

+    if os.path.exists(os.path.join(data_path, "user.css")):

+        head += stylesheet(os.path.join(data_path, "user.css"))

+

+    return head

+

+

+def reload_javascript():

+    js = javascript_html()

+    css = css_html()

+

     def template_response(*args, **kwargs):

         res = shared.GradioTemplateResponseOriginal(*args, **kwargs)

-        res.body = res.body.replace(b'</head>', f'{head}</head>'.encode("utf8"))

+        res.body = res.body.replace(b'</head>', f'{js}</head>'.encode("utf8"))

+        res.body = res.body.replace(b'</body>', f'{css}</body>'.encode("utf8"))

         res.init_headers()

         return res

 

@@ -1772,7 +1897,7 @@ def versions_html():
     return f"""

 python: <span title="{sys.version}">{python_version}</span>

  • 

-torch: {torch.__version__}

+torch: {getattr(torch, '__long_version__',torch.__version__)}

  • 

 xformers: {xformers_version}

  • 

diff --git a/modules/ui_common.py b/modules/ui_common.py
index fd047f31..27ab3ebb 100644
--- a/modules/ui_common.py
+++ b/modules/ui_common.py
@@ -125,12 +125,12 @@ Requested path was: {f}
 

     with gr.Column(variant='panel', elem_id=f"{tabname}_results"):

         with gr.Group(elem_id=f"{tabname}_gallery_container"):

-            result_gallery = gr.Gallery(label='Output', show_label=False, elem_id=f"{tabname}_gallery").style(grid=4)

+            result_gallery = gr.Gallery(label='Output', show_label=False, elem_id=f"{tabname}_gallery").style(columns=4)

 

         generation_info = None

         with gr.Column():

-            with gr.Row(elem_id=f"image_buttons_{tabname}"):

-                open_folder_button = gr.Button(folder_symbol, elem_id="hidden_element" if shared.cmd_opts.hide_ui_dir_config else f'open_folder_{tabname}')

+            with gr.Row(elem_id=f"image_buttons_{tabname}", elem_classes="image-buttons"):

+                open_folder_button = gr.Button(folder_symbol, visible=not shared.cmd_opts.hide_ui_dir_config)

 

                 if tabname != "extras":

                     save = gr.Button('Save', elem_id=f'save_{tabname}')

@@ -145,11 +145,10 @@ Requested path was: {f}
             )

 

             if tabname != "extras":

-                with gr.Row():

-                    download_files = gr.File(None, file_count="multiple", interactive=False, show_label=False, visible=False, elem_id=f'download_files_{tabname}')

+                download_files = gr.File(None, file_count="multiple", interactive=False, show_label=False, visible=False, elem_id=f'download_files_{tabname}')

 

                 with gr.Group():

-                    html_info = gr.HTML(elem_id=f'html_info_{tabname}')

+                    html_info = gr.HTML(elem_id=f'html_info_{tabname}', elem_classes="infotext")

                     html_log = gr.HTML(elem_id=f'html_log_{tabname}')

 

                     generation_info = gr.Textbox(visible=False, elem_id=f'generation_info_{tabname}')

@@ -160,6 +159,7 @@ Requested path was: {f}
                             _js="function(x, y, z){ return [x, y, selected_gallery_index()] }",

                             inputs=[generation_info, html_info, html_info],

                             outputs=[html_info, html_info],

+                            show_progress=False,

                         )

 

                     save.click(

@@ -195,12 +195,19 @@ Requested path was: {f}
 

             else:

                 html_info_x = gr.HTML(elem_id=f'html_info_x_{tabname}')

-                html_info = gr.HTML(elem_id=f'html_info_{tabname}')

+                html_info = gr.HTML(elem_id=f'html_info_{tabname}', elem_classes="infotext")

                 html_log = gr.HTML(elem_id=f'html_log_{tabname}')

 

+            paste_field_names = []

+            if tabname == "txt2img":

+                paste_field_names = modules.scripts.scripts_txt2img.paste_field_names

+            elif tabname == "img2img":

+                paste_field_names = modules.scripts.scripts_img2img.paste_field_names

+

             for paste_tabname, paste_button in buttons.items():

                 parameters_copypaste.register_paste_params_button(parameters_copypaste.ParamBinding(

-                    paste_button=paste_button, tabname=paste_tabname, source_tabname="txt2img" if tabname == "txt2img" else None, source_image_component=result_gallery

+                    paste_button=paste_button, tabname=paste_tabname, source_tabname="txt2img" if tabname == "txt2img" else None, source_image_component=result_gallery,

+                    paste_field_names=paste_field_names

                 ))

 

             return result_gallery, generation_info if tabname != "extras" else html_info_x, html_info, html_log

diff --git a/modules/ui_components.py b/modules/ui_components.py
index 284ca0cf..64451df7 100644
--- a/modules/ui_components.py
+++ b/modules/ui_components.py
@@ -1,58 +1,74 @@
 import gradio as gr

 

 

-class ToolButton(gr.Button, gr.components.FormComponent):

-    """Small button with single emoji as text, fits inside gradio forms"""

+class FormComponent:

+    def get_expected_parent(self):

+        return gr.components.Form

 

-    def __init__(self, **kwargs):

-        super().__init__(variant="tool", **kwargs)

 

-    def get_block_name(self):

-        return "button"

+gr.Dropdown.get_expected_parent = FormComponent.get_expected_parent

 

 

-class ToolButtonTop(gr.Button, gr.components.FormComponent):

-    """Small button with single emoji as text, with extra margin at top, fits inside gradio forms"""

+class ToolButton(FormComponent, gr.Button):

+    """Small button with single emoji as text, fits inside gradio forms"""

 

-    def __init__(self, **kwargs):

-        super().__init__(variant="tool-top", **kwargs)

+    def __init__(self, *args, **kwargs):

+        classes = kwargs.pop("elem_classes", [])

+        super().__init__(*args, elem_classes=["tool", *classes], **kwargs)

 

     def get_block_name(self):

         return "button"

 

 

-class FormRow(gr.Row, gr.components.FormComponent):

+class FormRow(FormComponent, gr.Row):

     """Same as gr.Row but fits inside gradio forms"""

 

     def get_block_name(self):

         return "row"

 

 

-class FormGroup(gr.Group, gr.components.FormComponent):

+class FormColumn(FormComponent, gr.Column):

+    """Same as gr.Column but fits inside gradio forms"""

+

+    def get_block_name(self):

+        return "column"

+

+

+class FormGroup(FormComponent, gr.Group):

     """Same as gr.Row but fits inside gradio forms"""

 

     def get_block_name(self):

         return "group"

 

 

-class FormHTML(gr.HTML, gr.components.FormComponent):

+class FormHTML(FormComponent, gr.HTML):

     """Same as gr.HTML but fits inside gradio forms"""

 

     def get_block_name(self):

         return "html"

 

 

-class FormColorPicker(gr.ColorPicker, gr.components.FormComponent):

+class FormColorPicker(FormComponent, gr.ColorPicker):

     """Same as gr.ColorPicker but fits inside gradio forms"""

 

     def get_block_name(self):

         return "colorpicker"

 

 

-class DropdownMulti(gr.Dropdown):

+class DropdownMulti(FormComponent, gr.Dropdown):

     """Same as gr.Dropdown but always multiselect"""

     def __init__(self, **kwargs):

         super().__init__(multiselect=True, **kwargs)

 

     def get_block_name(self):

         return "dropdown"

+

+

+class DropdownEditable(FormComponent, gr.Dropdown):

+    """Same as gr.Dropdown but allows editing value"""

+    def __init__(self, **kwargs):

+        super().__init__(allow_custom_value=True, **kwargs)

+

+    def get_block_name(self):

+        return "dropdown"

+

diff --git a/modules/ui_extensions.py b/modules/ui_extensions.py
index 37d30e1f..99ac8756 100644
--- a/modules/ui_extensions.py
+++ b/modules/ui_extensions.py
@@ -1,8 +1,8 @@
 import json

 import os.path

-import shutil

 import sys

 import time

+from datetime import datetime

 import traceback

 

 import git

@@ -12,17 +12,19 @@ import html
 import shutil

 import errno

 

-from modules import extensions, shared, paths

+from modules import extensions, shared, paths, config_states

+from modules.paths_internal import config_states_dir

 from modules.call_queue import wrap_gradio_gpu_call

 

 available_extensions = {"extensions": []}

+STYLE_PRIMARY = ' style="color: var(--primary-400)"'

 

 

 def check_access():

     assert not shared.cmd_opts.disable_extension_access, "extension access disabled because of command line flags"

 

 

-def apply_and_restart(disable_list, update_list):

+def apply_and_restart(disable_list, update_list, disable_all):

     check_access()

 

     disabled = json.loads(disable_list)

@@ -31,6 +33,9 @@ def apply_and_restart(disable_list, update_list):
     update = json.loads(update_list)

     assert type(update) == list, f"wrong update_list data for apply_and_restart: {update_list}"

 

+    if update:

+        save_config_state("Backup (pre-update)")

+

     update = set(update)

 

     for ext in extensions.extensions:

@@ -44,12 +49,53 @@ def apply_and_restart(disable_list, update_list):
             print(traceback.format_exc(), file=sys.stderr)

 

     shared.opts.disabled_extensions = disabled

+    shared.opts.disable_all_extensions = disable_all

     shared.opts.save(shared.config_filename)

 

     shared.state.interrupt()

     shared.state.need_restart = True

 

 

+def save_config_state(name):

+    current_config_state = config_states.get_config()

+    if not name:

+        name = "Config"

+    current_config_state["name"] = name

+    filename = os.path.join(config_states_dir, datetime.now().strftime("%Y_%m_%d-%H_%M_%S") + "_" + name + ".json")

+    print(f"Saving backup of webui/extension state to {filename}.")

+    with open(filename, "w", encoding="utf-8") as f:

+        json.dump(current_config_state, f)

+    config_states.list_config_states()

+    new_value = next(iter(config_states.all_config_states.keys()), "Current")

+    new_choices = ["Current"] + list(config_states.all_config_states.keys())

+    return gr.Dropdown.update(value=new_value, choices=new_choices), f"<span>Saved current webui/extension state to \"{filename}\"</span>"

+

+

+def restore_config_state(confirmed, config_state_name, restore_type):

+    if config_state_name == "Current":

+        return "<span>Select a config to restore from.</span>"

+    if not confirmed:

+        return "<span>Cancelled.</span>"

+

+    check_access()

+

+    config_state = config_states.all_config_states[config_state_name]

+

+    print(f"*** Restoring webui state from backup: {restore_type} ***")

+

+    if restore_type == "extensions" or restore_type == "both":

+        shared.opts.restore_config_state_file = config_state["filepath"]

+        shared.opts.save(shared.config_filename)

+

+    if restore_type == "webui" or restore_type == "both":

+        config_states.restore_webui_config(config_state)

+

+    shared.state.interrupt()

+    shared.state.need_restart = True

+

+    return ""

+

+

 def check_updates(id_task, disable_list):

     check_access()

 

@@ -64,6 +110,9 @@ def check_updates(id_task, disable_list):
 

         try:

             ext.check_updates()

+        except FileNotFoundError as e:

+            if 'FETCH_HEAD' not in str(e):

+                raise

         except Exception:

             print(f"Error checking updates for {ext.name}:", file=sys.stderr)

             print(traceback.format_exc(), file=sys.stderr)

@@ -73,6 +122,16 @@ def check_updates(id_task, disable_list):
     return extension_table(), ""

 

 

+def make_commit_link(commit_hash, remote, text=None):

+    if text is None:

+        text = commit_hash[:8]

+    if remote.startswith("https://github.com/"):

+        href = os.path.join(remote, "commit", commit_hash)

+        return f'<a href="{href}" target="_blank">{text}</a>'

+    else:

+        return text

+

+

 def extension_table():

     code = f"""<!-- {time.time()} -->

     <table id="extensions">

@@ -80,6 +139,7 @@ def extension_table():
             <tr>

                 <th><abbr title="Use checkbox to enable the extension; it will be enabled or disabled when you click apply button">Extension</abbr></th>

                 <th>URL</th>

+                <th><abbr title="Extension version">Version</abbr></th>

                 <th><abbr title="Use checkbox to mark the extension for update; it will be updated when you click apply button">Update</abbr></th>

             </tr>

         </thead>

@@ -87,21 +147,28 @@ def extension_table():
     """

 

     for ext in extensions.extensions:

-        remote = ""

-        if ext.is_builtin:

-            remote = "built-in"

-        elif ext.remote:

-            remote = f"""<a href="{html.escape(ext.remote or '')}" target="_blank">{html.escape("built-in" if ext.is_builtin else ext.remote or '')}</a>"""

+        ext.read_info_from_repo()

+

+        remote = f"""<a href="{html.escape(ext.remote or '')}" target="_blank">{html.escape("built-in" if ext.is_builtin else ext.remote or '')}</a>"""

 

         if ext.can_update:

             ext_status = f"""<label><input class="gr-check-radio gr-checkbox" name="update_{html.escape(ext.name)}" checked="checked" type="checkbox">{html.escape(ext.status)}</label>"""

         else:

             ext_status = ext.status

 

+        style = ""

+        if shared.opts.disable_all_extensions == "extra" and not ext.is_builtin or shared.opts.disable_all_extensions == "all":

+            style = STYLE_PRIMARY

+

+        version_link = ext.version

+        if ext.commit_hash and ext.remote:

+            version_link = make_commit_link(ext.commit_hash, ext.remote, ext.version)

+

         code += f"""

             <tr>

-                <td><label><input class="gr-check-radio gr-checkbox" name="enable_{html.escape(ext.name)}" type="checkbox" {'checked="checked"' if ext.enabled else ''}>{html.escape(ext.name)}</label></td>

+                <td><label{style}><input class="gr-check-radio gr-checkbox" name="enable_{html.escape(ext.name)}" type="checkbox" {'checked="checked"' if ext.enabled else ''}>{html.escape(ext.name)}</label></td>

                 <td>{remote}</td>

+                <td>{version_link}</td>

                 <td{' class="extension_status"' if ext.remote is not None else ''}>{ext_status}</td>

             </tr>

     """

@@ -114,6 +181,133 @@ def extension_table():
     return code

 

 

+def update_config_states_table(state_name):

+    if state_name == "Current":

+        config_state = config_states.get_config()

+    else:

+        config_state = config_states.all_config_states[state_name]

+

+    config_name = config_state.get("name", "Config")

+    created_date = time.asctime(time.gmtime(config_state["created_at"]))

+    filepath = config_state.get("filepath", "<unknown>")

+

+    code = f"""<!-- {time.time()} -->"""

+

+    webui_remote = config_state["webui"]["remote"] or ""

+    webui_branch = config_state["webui"]["branch"]

+    webui_commit_hash = config_state["webui"]["commit_hash"] or "<unknown>"

+    webui_commit_date = config_state["webui"]["commit_date"]

+    if webui_commit_date:

+        webui_commit_date = time.asctime(time.gmtime(webui_commit_date))

+    else:

+        webui_commit_date = "<unknown>"

+

+    remote = f"""<a href="{html.escape(webui_remote)}" target="_blank">{html.escape(webui_remote or '')}</a>"""

+    commit_link = make_commit_link(webui_commit_hash, webui_remote)

+    date_link = make_commit_link(webui_commit_hash, webui_remote, webui_commit_date)

+

+    current_webui = config_states.get_webui_config()

+

+    style_remote = ""

+    style_branch = ""

+    style_commit = ""

+    if current_webui["remote"] != webui_remote:

+        style_remote = STYLE_PRIMARY

+    if current_webui["branch"] != webui_branch:

+        style_branch = STYLE_PRIMARY

+    if current_webui["commit_hash"] != webui_commit_hash:

+        style_commit = STYLE_PRIMARY

+

+    code += f"""<h2>Config Backup: {config_name}</h2>

+      <div><b>Filepath:</b> {filepath}</div>

+      <div><b>Created at:</b> {created_date}</div>"""

+

+    code += f"""<h2>WebUI State</h2>

+      <table id="config_state_webui">

+        <thead>

+            <tr>

+                <th>URL</th>

+                <th>Branch</th>

+                <th>Commit</th>

+                <th>Date</th>

+            </tr>

+        </thead>

+        <tbody>

+            <tr>

+                <td><label{style_remote}>{remote}</label></td>

+                <td><label{style_branch}>{webui_branch}</label></td>

+                <td><label{style_commit}>{commit_link}</label></td>

+                <td><label{style_commit}>{date_link}</label></td>

+            </tr>

+        </tbody>

+      </table>

+    """

+

+    code += """<h2>Extension State</h2>

+      <table id="config_state_extensions">

+        <thead>

+            <tr>

+                <th>Extension</th>

+                <th>URL</th>

+                <th>Branch</th>

+                <th>Commit</th>

+                <th>Date</th>

+            </tr>

+        </thead>

+        <tbody>

+    """

+

+    ext_map = {ext.name: ext for ext in extensions.extensions}

+

+    for ext_name, ext_conf in config_state["extensions"].items():

+        ext_remote = ext_conf["remote"] or ""

+        ext_branch = ext_conf["branch"] or "<unknown>"

+        ext_enabled = ext_conf["enabled"]

+        ext_commit_hash = ext_conf["commit_hash"] or "<unknown>"

+        ext_commit_date = ext_conf["commit_date"]

+        if ext_commit_date:

+            ext_commit_date = time.asctime(time.gmtime(ext_commit_date))

+        else:

+            ext_commit_date = "<unknown>"

+

+        remote = f"""<a href="{html.escape(ext_remote)}" target="_blank">{html.escape(ext_remote or '')}</a>"""

+        commit_link = make_commit_link(ext_commit_hash, ext_remote)

+        date_link = make_commit_link(ext_commit_hash, ext_remote, ext_commit_date)

+

+        style_enabled = ""

+        style_remote = ""

+        style_branch = ""

+        style_commit = ""

+        if ext_name in ext_map:

+            current_ext = ext_map[ext_name]

+            current_ext.read_info_from_repo()

+            if current_ext.enabled != ext_enabled:

+                style_enabled = STYLE_PRIMARY

+            if current_ext.remote != ext_remote:

+                style_remote = STYLE_PRIMARY

+            if current_ext.branch != ext_branch:

+                style_branch = STYLE_PRIMARY

+            if current_ext.commit_hash != ext_commit_hash:

+                style_commit = STYLE_PRIMARY

+

+        code += f"""

+            <tr>

+                <td><label{style_enabled}><input class="gr-check-radio gr-checkbox" type="checkbox" disabled="true" {'checked="checked"' if ext_enabled else ''}>{html.escape(ext_name)}</label></td>

+                <td><label{style_remote}>{remote}</label></td>

+                <td><label{style_branch}>{ext_branch}</label></td>

+                <td><label{style_commit}>{commit_link}</label></td>

+                <td><label{style_commit}>{date_link}</label></td>

+            </tr>

+    """

+

+    code += """

+        </tbody>

+    </table>

+    """

+

+    return code

+

+

 def normalize_git_url(url):

     if url is None:

         return ""

@@ -122,7 +316,7 @@ def normalize_git_url(url):
     return url

 

 

-def install_extension_from_url(dirname, url):

+def install_extension_from_url(dirname, url, branch_name=None):

     check_access()

 

     assert url, 'No URL specified'

@@ -143,22 +337,27 @@ def install_extension_from_url(dirname, url):
 

     try:

         shutil.rmtree(tmpdir, True)

-

-        repo = git.Repo.clone_from(url, tmpdir)

-        repo.remote().fetch()

-

+        if not branch_name:

+            # if no branch is specified, use the default branch

+            with git.Repo.clone_from(url, tmpdir) as repo:

+                repo.remote().fetch()

+                for submodule in repo.submodules:

+                    submodule.update()

+        else:

+            with git.Repo.clone_from(url, tmpdir, branch=branch_name) as repo:

+                repo.remote().fetch()

+                for submodule in repo.submodules:

+                    submodule.update()

         try:

             os.rename(tmpdir, target_dir)

         except OSError as err:

-            # TODO what does this do on windows? I think it'll be a different error code but I don't have a system to check it

-            # Shouldn't cause any new issues at least but we probably want to handle it there too.

             if err.errno == errno.EXDEV:

                 # Cross device link, typical in docker or when tmp/ and extensions/ are on different file systems

                 # Since we can't use a rename, do the slower but more versitile shutil.move()

                 shutil.move(tmpdir, target_dir)

             else:

                 # Something else, not enough free space, permissions, etc.  rethrow it so that it gets handled.

-                raise(err)

+                raise err

 

         import launch

         launch.run_extension_installer(target_dir)

@@ -169,12 +368,12 @@ def install_extension_from_url(dirname, url):
         shutil.rmtree(tmpdir, True)

 

 

-def install_extension_from_index(url, hide_tags, sort_column):

+def install_extension_from_index(url, hide_tags, sort_column, filter_text):

     ext_table, message = install_extension_from_url(None, url)

 

-    code, _ = refresh_available_extensions_from_data(hide_tags, sort_column)

+    code, _ = refresh_available_extensions_from_data(hide_tags, sort_column, filter_text)

 

-    return code, ext_table, message

+    return code, ext_table, message, ''

 

 

 def refresh_available_extensions(url, hide_tags, sort_column):

@@ -188,11 +387,17 @@ def refresh_available_extensions(url, hide_tags, sort_column):
 

     code, tags = refresh_available_extensions_from_data(hide_tags, sort_column)

 

-    return url, code, gr.CheckboxGroup.update(choices=tags), ''

+    return url, code, gr.CheckboxGroup.update(choices=tags), '', ''

 

 

-def refresh_available_extensions_for_tags(hide_tags, sort_column):

-    code, _ = refresh_available_extensions_from_data(hide_tags, sort_column)

+def refresh_available_extensions_for_tags(hide_tags, sort_column, filter_text):

+    code, _ = refresh_available_extensions_from_data(hide_tags, sort_column, filter_text)

+

+    return code, ''

+

+

+def search_extensions(filter_text, hide_tags, sort_column):

+    code, _ = refresh_available_extensions_from_data(hide_tags, sort_column, filter_text)

 

     return code, ''

 

@@ -207,7 +412,7 @@ sort_ordering = [
 ]

 

 

-def refresh_available_extensions_from_data(hide_tags, sort_column):

+def refresh_available_extensions_from_data(hide_tags, sort_column, filter_text=""):

     extlist = available_extensions["extensions"]

     installed_extension_urls = {normalize_git_url(extension.remote): extension.name for extension in extensions.extensions}

 

@@ -246,7 +451,12 @@ def refresh_available_extensions_from_data(hide_tags, sort_column):
             hidden += 1

             continue

 

-        install_code = f"""<input onclick="install_extension_from_index(this, '{html.escape(url)}')" type="button" value="{"Install" if not existing else "Installed"}" {"disabled=disabled" if existing else ""} class="gr-button gr-button-lg gr-button-secondary">"""

+        if filter_text and filter_text.strip():

+            if filter_text.lower() not in html.escape(name).lower() and filter_text.lower() not in html.escape(description).lower():

+                hidden += 1

+                continue

+

+        install_code = f"""<button onclick="install_extension_from_index(this, '{html.escape(url)}')" {"disabled=disabled" if existing else ""} class="lg secondary gradio-button custom-button">{"Install" if not existing else "Installed"}</button>"""

 

         tags_text = ", ".join([f"<span class='extension-tag' title='{tags.get(x, '')}'>{x}</span>" for x in extension_tags])

 

@@ -276,23 +486,33 @@ def refresh_available_extensions_from_data(hide_tags, sort_column):
 def create_ui():

     import modules.ui

 

+    config_states.list_config_states()

+

     with gr.Blocks(analytics_enabled=False) as ui:

         with gr.Tabs(elem_id="tabs_extensions") as tabs:

-            with gr.TabItem("Installed"):

+            with gr.TabItem("Installed", id="installed"):

 

                 with gr.Row(elem_id="extensions_installed_top"):

                     apply = gr.Button(value="Apply and restart UI", variant="primary")

                     check = gr.Button(value="Check for updates")

+                    extensions_disable_all = gr.Radio(label="Disable all extensions", choices=["none", "extra", "all"], value=shared.opts.disable_all_extensions, elem_id="extensions_disable_all")

                     extensions_disabled_list = gr.Text(elem_id="extensions_disabled_list", visible=False).style(container=False)

                     extensions_update_list = gr.Text(elem_id="extensions_update_list", visible=False).style(container=False)

 

-                info = gr.HTML()

+                html = ""

+                if shared.opts.disable_all_extensions != "none":

+                    html = """

+<span style="color: var(--primary-400);">

+    "Disable all extensions" was set, change it to "none" to load all extensions again

+</span>

+                    """

+                info = gr.HTML(html)

                 extensions_table = gr.HTML(lambda: extension_table())

 

                 apply.click(

                     fn=apply_and_restart,

                     _js="extensions_apply",

-                    inputs=[extensions_disabled_list, extensions_update_list],

+                    inputs=[extensions_disabled_list, extensions_update_list, extensions_disable_all],

                     outputs=[],

                 )

 

@@ -303,10 +523,10 @@ def create_ui():
                     outputs=[extensions_table, info],

                 )

 

-            with gr.TabItem("Available"):

+            with gr.TabItem("Available", id="available"):

                 with gr.Row():

                     refresh_available_extensions_button = gr.Button(value="Load from:", variant="primary")

-                    available_extensions_index = gr.Text(value="https://raw.githubusercontent.com/wiki/AUTOMATIC1111/stable-diffusion-webui/Extensions-index.md", label="Extension index URL").style(container=False)

+                    available_extensions_index = gr.Text(value="https://raw.githubusercontent.com/AUTOMATIC1111/stable-diffusion-webui-extensions/master/index.json", label="Extension index URL").style(container=False)

                     extension_to_install = gr.Text(elem_id="extension_to_install", visible=False)

                     install_extension_button = gr.Button(elem_id="install_extension_button", visible=False)

 

@@ -314,43 +534,77 @@ def create_ui():
                     hide_tags = gr.CheckboxGroup(value=["ads", "localization", "installed"], label="Hide extensions with tags", choices=["script", "ads", "localization", "installed"])

                     sort_column = gr.Radio(value="newest first", label="Order", choices=["newest first", "oldest first", "a-z", "z-a", "internal order", ], type="index")

 

+                with gr.Row(): 

+                    search_extensions_text = gr.Text(label="Search").style(container=False)

+                   

                 install_result = gr.HTML()

                 available_extensions_table = gr.HTML()

 

                 refresh_available_extensions_button.click(

                     fn=modules.ui.wrap_gradio_call(refresh_available_extensions, extra_outputs=[gr.update(), gr.update(), gr.update()]),

                     inputs=[available_extensions_index, hide_tags, sort_column],

-                    outputs=[available_extensions_index, available_extensions_table, hide_tags, install_result],

+                    outputs=[available_extensions_index, available_extensions_table, hide_tags, install_result, search_extensions_text],

                 )

 

                 install_extension_button.click(

                     fn=modules.ui.wrap_gradio_call(install_extension_from_index, extra_outputs=[gr.update(), gr.update()]),

-                    inputs=[extension_to_install, hide_tags, sort_column],

+                    inputs=[extension_to_install, hide_tags, sort_column, search_extensions_text],

                     outputs=[available_extensions_table, extensions_table, install_result],

                 )

 

+                search_extensions_text.change(

+                    fn=modules.ui.wrap_gradio_call(search_extensions, extra_outputs=[gr.update()]),

+                    inputs=[search_extensions_text, hide_tags, sort_column],

+                    outputs=[available_extensions_table, install_result],

+                )

+

                 hide_tags.change(

                     fn=modules.ui.wrap_gradio_call(refresh_available_extensions_for_tags, extra_outputs=[gr.update()]),

-                    inputs=[hide_tags, sort_column],

+                    inputs=[hide_tags, sort_column, search_extensions_text],

                     outputs=[available_extensions_table, install_result]

                 )

 

                 sort_column.change(

                     fn=modules.ui.wrap_gradio_call(refresh_available_extensions_for_tags, extra_outputs=[gr.update()]),

-                    inputs=[hide_tags, sort_column],

+                    inputs=[hide_tags, sort_column, search_extensions_text],

                     outputs=[available_extensions_table, install_result]

                 )

 

-            with gr.TabItem("Install from URL"):

+            with gr.TabItem("Install from URL", id="install_from_url"):

                 install_url = gr.Text(label="URL for extension's git repository")

+                install_branch = gr.Text(label="Specific branch name", placeholder="Leave empty for default main branch")

                 install_dirname = gr.Text(label="Local directory name", placeholder="Leave empty for auto")

                 install_button = gr.Button(value="Install", variant="primary")

                 install_result = gr.HTML(elem_id="extension_install_result")

 

                 install_button.click(

                     fn=modules.ui.wrap_gradio_call(install_extension_from_url, extra_outputs=[gr.update()]),

-                    inputs=[install_dirname, install_url],

+                    inputs=[install_dirname, install_url, install_branch],

                     outputs=[extensions_table, install_result],

                 )

 

+            with gr.TabItem("Backup/Restore"):

+                with gr.Row(elem_id="extensions_backup_top_row"):

+                    config_states_list = gr.Dropdown(label="Saved Configs", elem_id="extension_backup_saved_configs", value="Current", choices=["Current"] + list(config_states.all_config_states.keys()))

+                    modules.ui.create_refresh_button(config_states_list, config_states.list_config_states, lambda: {"choices": ["Current"] + list(config_states.all_config_states.keys())}, "refresh_config_states")

+                    config_restore_type = gr.Radio(label="State to restore", choices=["extensions", "webui", "both"], value="extensions", elem_id="extension_backup_restore_type")

+                    config_restore_button = gr.Button(value="Restore Selected Config", variant="primary", elem_id="extension_backup_restore")

+                with gr.Row(elem_id="extensions_backup_top_row2"):

+                    config_save_name = gr.Textbox("", placeholder="Config Name", show_label=False)

+                    config_save_button = gr.Button(value="Save Current Config")

+

+                config_states_info = gr.HTML("")

+                config_states_table = gr.HTML(lambda: update_config_states_table("Current"))

+

+                config_save_button.click(fn=save_config_state, inputs=[config_save_name], outputs=[config_states_list, config_states_info])

+

+                dummy_component = gr.Label(visible=False)

+                config_restore_button.click(fn=restore_config_state, _js="config_state_confirm_restore", inputs=[dummy_component, config_states_list, config_restore_type], outputs=[config_states_info])

+

+                config_states_list.change(

+                    fn=update_config_states_table,

+                    inputs=[config_states_list],

+                    outputs=[config_states_table],

+                )

+

     return ui

diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py
index 83367968..aa2f5d1b 100644
--- a/modules/ui_extra_networks.py
+++ b/modules/ui_extra_networks.py
@@ -2,8 +2,10 @@ import glob
 import os.path

 import urllib.parse

 from pathlib import Path

+from PIL import PngImagePlugin

 

 from modules import shared

+from modules.images import read_info_from_image

 import gradio as gr

 import json

 import html

@@ -22,20 +24,37 @@ def register_page(page):
     allowed_dirs.update(set(sum([x.allowed_directories_for_previews() for x in extra_pages], [])))

 

 

-def add_pages_to_demo(app):

-    def fetch_file(filename: str = ""):

-        from starlette.responses import FileResponse

+def fetch_file(filename: str = ""):

+    from starlette.responses import FileResponse

+

+    if not any([Path(x).absolute() in Path(filename).absolute().parents for x in allowed_dirs]):

+        raise ValueError(f"File cannot be fetched: {filename}. Must be in one of directories registered by extra pages.")

+

+    ext = os.path.splitext(filename)[1].lower()

+    if ext not in (".png", ".jpg", ".webp"):

+        raise ValueError(f"File cannot be fetched: {filename}. Only png and jpg and webp.")

+

+    # would profit from returning 304

+    return FileResponse(filename, headers={"Accept-Ranges": "bytes"})

+

 

-        if not any([Path(x).resolve() in Path(filename).resolve().parents for x in allowed_dirs]):

-            raise ValueError(f"File cannot be fetched: {filename}. Must be in one of directories registered by extra pages.")

+def get_metadata(page: str = "", item: str = ""):

+    from starlette.responses import JSONResponse

 

-        if os.path.splitext(filename)[1].lower() != ".png":

-            raise ValueError(f"File cannot be fetched: {filename}. Only png.")

+    page = next(iter([x for x in extra_pages if x.name == page]), None)

+    if page is None:

+        return JSONResponse({})

 

-        # would profit from returning 304

-        return FileResponse(filename, headers={"Accept-Ranges": "bytes"})

+    metadata = page.metadata.get(item)

+    if metadata is None:

+        return JSONResponse({})

 

+    return JSONResponse({"metadata": metadata})

+

+

+def add_pages_to_demo(app):

     app.add_api_route("/sd_extra_networks/thumb", fetch_file, methods=["GET"])

+    app.add_api_route("/sd_extra_networks/metadata", get_metadata, methods=["GET"])

 

 

 class ExtraNetworksPage:

@@ -44,6 +63,7 @@ class ExtraNetworksPage:
         self.name = title.lower()

         self.card_page = shared.html("extra-networks-card.html")

         self.allow_negative_prompt = False

+        self.metadata = {}

 

     def refresh(self):

         pass

@@ -65,6 +85,8 @@ class ExtraNetworksPage:
         view = shared.opts.extra_networks_default_view

         items_html = ''

 

+        self.metadata = {}

+

         subdirs = {}

         for parentdir in [os.path.abspath(x) for x in self.allowed_directories_for_previews()]:

             for x in glob.glob(os.path.join(parentdir, '**/*'), recursive=True):

@@ -75,29 +97,39 @@ class ExtraNetworksPage:
                 while subdir.startswith("/"):

                     subdir = subdir[1:]

 

+                is_empty = len(os.listdir(x)) == 0

+                if not is_empty and not subdir.endswith("/"):

+                    subdir = subdir + "/"

+

                 subdirs[subdir] = 1

 

         if subdirs:

             subdirs = {"": 1, **subdirs}

 

         subdirs_html = "".join([f"""

-<button class='gr-button gr-button-lg gr-button-secondary{" search-all" if subdir=="" else ""}' onclick='extraNetworksSearchButton("{tabname}_extra_tabs", event)'>

+<button class='lg secondary gradio-button custom-button{" search-all" if subdir=="" else ""}' onclick='extraNetworksSearchButton("{tabname}_extra_tabs", event)'>

 {html.escape(subdir if subdir!="" else "all")}

 </button>

 """ for subdir in subdirs])

 

         for item in self.list_items():

+            metadata = item.get("metadata")

+            if metadata:

+                self.metadata[item["name"]] = metadata

+

             items_html += self.create_html_for_item(item, tabname)

 

         if items_html == '':

             dirs = "".join([f"<li>{x}</li>" for x in self.allowed_directories_for_previews()])

             items_html = shared.html("extra-networks-no-cards.html").format(dirs=dirs)

 

+        self_name_id = self.name.replace(" ", "_")

+

         res = f"""

-<div id='{tabname}_{self.name}_subdirs' class='extra-network-subdirs extra-network-subdirs-{view}'>

+<div id='{tabname}_{self_name_id}_subdirs' class='extra-network-subdirs extra-network-subdirs-{view}'>

 {subdirs_html}

 </div>

-<div id='{tabname}_{self.name}_cards' class='extra-network-{view}'>

+<div id='{tabname}_{self_name_id}_cards' class='extra-network-{view}'>

 {items_html}

 </div>

 """

@@ -117,19 +149,58 @@ class ExtraNetworksPage:
         if onclick is None:

             onclick = '"' + html.escape(f"""return cardClicked({json.dumps(tabname)}, {item["prompt"]}, {"true" if self.allow_negative_prompt else "false"})""") + '"'

 

+        height = f"height: {shared.opts.extra_networks_card_height}px;" if shared.opts.extra_networks_card_height else ''

+        width = f"width: {shared.opts.extra_networks_card_width}px;" if shared.opts.extra_networks_card_width else ''

+        background_image = f"background-image: url(\"{html.escape(preview)}\");" if preview else ''

+        metadata_button = ""

+        metadata = item.get("metadata")

+        if metadata:

+            metadata_button = f"<div class='metadata-button' title='Show metadata' onclick='extraNetworksRequestMetadata(event, {json.dumps(self.name)}, {json.dumps(item['name'])})'></div>"

+

         args = {

-            "preview_html": "style='background-image: url(\"" + html.escape(preview) + "\")'" if preview else '',

+            "style": f"'{height}{width}{background_image}'",

             "prompt": item.get("prompt", None),

             "tabname": json.dumps(tabname),

             "local_preview": json.dumps(item["local_preview"]),

             "name": item["name"],

+            "description": (item.get("description") or ""),

             "card_clicked": onclick,

             "save_card_preview": '"' + html.escape(f"""return saveCardPreview(event, {json.dumps(tabname)}, {json.dumps(item["local_preview"])})""") + '"',

             "search_term": item.get("search_term", ""),

+            "metadata_button": metadata_button,

         }

 

         return self.card_page.format(**args)

 

+    def find_preview(self, path):

+        """

+        Find a preview PNG for a given path (without extension) and call link_preview on it.

+        """

+

+        preview_extensions = ["png", "jpg", "webp"]

+        if shared.opts.samples_format not in preview_extensions:

+            preview_extensions.append(shared.opts.samples_format)

+

+        potential_files = sum([[path + "." + ext, path + ".preview." + ext] for ext in preview_extensions], [])

+

+        for file in potential_files:

+            if os.path.isfile(file):

+                return self.link_preview(file)

+

+        return None

+

+    def find_description(self, path):

+        """

+        Find and read a description file for a given path (without extension).

+        """

+        for file in [f"{path}.txt", f"{path}.description.txt"]:

+            try:

+                with open(file, "r", encoding="utf-8", errors="replace") as f:

+                    return f.read()

+            except OSError:

+                pass

+        return None

+

 

 def intialize():

     extra_pages.clear()

@@ -170,24 +241,23 @@ def create_ui(container, button, tabname):
 

     with gr.Tabs(elem_id=tabname+"_extra_tabs") as tabs:

         for page in ui.stored_extra_pages:

-            with gr.Tab(page.title):

+            with gr.Tab(page.title, id=page.title.lower().replace(" ", "_")):

+

                 page_elem = gr.HTML(page.create_html(ui.tabname))

                 ui.pages.append(page_elem)

 

     filter = gr.Textbox('', show_label=False, elem_id=tabname+"_extra_search", placeholder="Search...", visible=False)

     button_refresh = gr.Button('Refresh', elem_id=tabname+"_extra_refresh")

-    button_close = gr.Button('Close', elem_id=tabname+"_extra_close")

 

     ui.button_save_preview = gr.Button('Save preview', elem_id=tabname+"_save_preview", visible=False)

     ui.preview_target_filename = gr.Textbox('Preview save filename', elem_id=tabname+"_preview_filename", visible=False)

 

     def toggle_visibility(is_visible):

         is_visible = not is_visible

-        return is_visible, gr.update(visible=is_visible)

+        return is_visible, gr.update(visible=is_visible), gr.update(variant=("secondary-down" if is_visible else "secondary"))

 

     state_visible = gr.State(value=False)

-    button.click(fn=toggle_visibility, inputs=[state_visible], outputs=[state_visible, container])

-    button_close.click(fn=toggle_visibility, inputs=[state_visible], outputs=[state_visible, container])

+    button.click(fn=toggle_visibility, inputs=[state_visible], outputs=[state_visible, container, button])

 

     def refresh():

         res = []

@@ -222,6 +292,7 @@ def setup_ui(ui, gallery):
 

         img_info = images[index if index >= 0 else 0]

         image = image_from_url_text(img_info)

+        geninfo, items = read_info_from_image(image)

 

         is_allowed = False

         for extra_page in ui.stored_extra_pages:

@@ -231,7 +302,12 @@ def setup_ui(ui, gallery):
 

         assert is_allowed, f'writing to {filename} is not allowed'

 

-        image.save(filename)

+        if geninfo:

+            pnginfo_data = PngImagePlugin.PngInfo()

+            pnginfo_data.add_text('parameters', geninfo)

+            image.save(filename, pnginfo=pnginfo_data)

+        else:

+            image.save(filename)

 

         return [page.create_html(ui.tabname) for page in ui.stored_extra_pages]

 

diff --git a/modules/ui_extra_networks_checkpoints.py b/modules/ui_extra_networks_checkpoints.py
index 04097a79..a17aa9c9 100644
--- a/modules/ui_extra_networks_checkpoints.py
+++ b/modules/ui_extra_networks_checkpoints.py
@@ -1,7 +1,6 @@
 import html

 import json

 import os

-import urllib.parse

 

 from modules import shared, ui_extra_networks, sd_models

 

@@ -17,21 +16,14 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
         checkpoint: sd_models.CheckpointInfo

         for name, checkpoint in sd_models.checkpoints_list.items():

             path, ext = os.path.splitext(checkpoint.filename)

-            previews = [path + ".png", path + ".preview.png"]

-

-            preview = None

-            for file in previews:

-                if os.path.isfile(file):

-                    preview = self.link_preview(file)

-                    break

-

             yield {

                 "name": checkpoint.name_for_extra,

                 "filename": path,

-                "preview": preview,

+                "preview": self.find_preview(path),

+                "description": self.find_description(path),

                 "search_term": self.search_terms_from_path(checkpoint.filename) + " " + (checkpoint.sha256 or ""),

                 "onclick": '"' + html.escape(f"""return selectCheckpoint({json.dumps(name)})""") + '"',

-                "local_preview": path + ".png",

+                "local_preview": f"{path}.{shared.opts.samples_format}",

             }

 

     def allowed_directories_for_previews(self):

diff --git a/modules/ui_extra_networks_hypernets.py b/modules/ui_extra_networks_hypernets.py
index 57851088..6187e000 100644
--- a/modules/ui_extra_networks_hypernets.py
+++ b/modules/ui_extra_networks_hypernets.py
@@ -14,21 +14,15 @@ class ExtraNetworksPageHypernetworks(ui_extra_networks.ExtraNetworksPage):
     def list_items(self):

         for name, path in shared.hypernetworks.items():

             path, ext = os.path.splitext(path)

-            previews = [path + ".png", path + ".preview.png"]

-

-            preview = None

-            for file in previews:

-                if os.path.isfile(file):

-                    preview = self.link_preview(file)

-                    break

 

             yield {

                 "name": name,

                 "filename": path,

-                "preview": preview,

+                "preview": self.find_preview(path),

+                "description": self.find_description(path),

                 "search_term": self.search_terms_from_path(path),

                 "prompt": json.dumps(f"<hypernet:{name}:") + " + opts.extra_networks_default_multiplier + " + json.dumps(">"),

-                "local_preview": path + ".png",

+                "local_preview": f"{path}.preview.{shared.opts.samples_format}",

             }

 

     def allowed_directories_for_previews(self):

diff --git a/modules/ui_extra_networks_textual_inversion.py b/modules/ui_extra_networks_textual_inversion.py
index bb64eb81..6944d559 100644
--- a/modules/ui_extra_networks_textual_inversion.py
+++ b/modules/ui_extra_networks_textual_inversion.py
@@ -1,7 +1,7 @@
 import json

 import os

 

-from modules import ui_extra_networks, sd_hijack

+from modules import ui_extra_networks, sd_hijack, shared

 

 

 class ExtraNetworksPageTextualInversion(ui_extra_networks.ExtraNetworksPage):

@@ -15,19 +15,14 @@ class ExtraNetworksPageTextualInversion(ui_extra_networks.ExtraNetworksPage):
     def list_items(self):

         for embedding in sd_hijack.model_hijack.embedding_db.word_embeddings.values():

             path, ext = os.path.splitext(embedding.filename)

-            preview_file = path + ".preview.png"

-

-            preview = None

-            if os.path.isfile(preview_file):

-                preview = self.link_preview(preview_file)

-

             yield {

                 "name": embedding.name,

                 "filename": embedding.filename,

-                "preview": preview,

+                "preview": self.find_preview(path),

+                "description": self.find_description(path),

                 "search_term": self.search_terms_from_path(embedding.filename),

                 "prompt": json.dumps(embedding.name),

-                "local_preview": path + ".preview.png",

+                "local_preview": f"{path}.preview.{shared.opts.samples_format}",

             }

 

     def allowed_directories_for_previews(self):

diff --git a/modules/ui_postprocessing.py b/modules/ui_postprocessing.py
index b418d955..f25639e5 100644
--- a/modules/ui_postprocessing.py
+++ b/modules/ui_postprocessing.py
@@ -9,13 +9,13 @@ def create_ui():
     with gr.Row().style(equal_height=False, variant='compact'):

         with gr.Column(variant='compact'):

             with gr.Tabs(elem_id="mode_extras"):

-                with gr.TabItem('Single Image', elem_id="extras_single_tab") as tab_single:

+                with gr.TabItem('Single Image', id="single_image", elem_id="extras_single_tab") as tab_single:

                     extras_image = gr.Image(label="Source", source="upload", interactive=True, type="pil", elem_id="extras_image")

 

-                with gr.TabItem('Batch Process', elem_id="extras_batch_process_tab") as tab_batch:

-                    image_batch = gr.File(label="Batch Process", file_count="multiple", interactive=True, type="file", elem_id="extras_image_batch")

+                with gr.TabItem('Batch Process', id="batch_process", elem_id="extras_batch_process_tab") as tab_batch:

+                    image_batch = gr.Files(label="Batch Process", interactive=True, elem_id="extras_image_batch")

 

-                with gr.TabItem('Batch from Directory', elem_id="extras_batch_directory_tab") as tab_batch_dir:

+                with gr.TabItem('Batch from Directory', id="batch_from_directory", elem_id="extras_batch_directory_tab") as tab_batch_dir:

                     extras_batch_input_dir = gr.Textbox(label="Input directory", **shared.hide_dirs, placeholder="A directory on the same machine where the server is running.", elem_id="extras_batch_input_dir")

                     extras_batch_output_dir = gr.Textbox(label="Output directory", **shared.hide_dirs, placeholder="Leave blank to save images to the default path.", elem_id="extras_batch_output_dir")

                     show_extras_results = gr.Checkbox(label='Show result images', value=True, elem_id="extras_show_extras_results")