73 files changed, 6985 insertions, 2355 deletions

diff --git a/modules/api/api.py b/modules/api/api.py
index 1ceba75d..eb7b1da5 100644
--- a/modules/api/api.py
+++ b/modules/api/api.py
@@ -1,34 +1,42 @@
 import base64
 import io
 import time
+import datetime
 import uvicorn
 from threading import Lock
 from io import BytesIO
 from gradio.processing_utils import decode_base64_to_file
-from fastapi import APIRouter, Depends, FastAPI, HTTPException
+from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request, Response
 from fastapi.security import HTTPBasic, HTTPBasicCredentials
 from secrets import compare_digest
 
 import modules.shared as shared
-from modules import sd_samplers, deepbooru, sd_hijack
+from modules import sd_samplers, deepbooru, sd_hijack, images, scripts, ui, postprocessing
 from modules.api.models import *
 from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, process_images
-from modules.extras import run_extras, run_pnginfo
 from modules.textual_inversion.textual_inversion import create_embedding, train_embedding
 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_config import find_checkpoint_config_near_filename
 from modules.realesrgan_model import get_realesrgan_models
 from modules import devices
 from typing import List
+import piexif
+import piexif.helper
 
 def upscaler_to_index(name: str):
     try:
         return [x.name.lower() for x in shared.sd_upscalers].index(name.lower())
     except:
-        raise HTTPException(status_code=400, detail=f"Invalid upscaler, needs to be on of these: {' , '.join([x.name for x in sd_upscalers])}")
+        raise HTTPException(status_code=400, detail=f"Invalid upscaler, needs to be one of these: {' , '.join([x.name for x in sd_upscalers])}")
 
+def script_name_to_index(name, scripts):
+    try:
+        return [script.title().lower() for script in scripts].index(name.lower())
+    except:
+        raise HTTPException(status_code=422, detail=f"Script '{name}' not found")
 
 def validate_sampler_name(name):
     config = sd_samplers.all_samplers_map.get(name, None)
@@ -39,34 +47,69 @@ def validate_sampler_name(name):
 
 def setUpscalers(req: dict):
     reqDict = vars(req)
-    reqDict['extras_upscaler_1'] = upscaler_to_index(req.upscaler_1)
-    reqDict['extras_upscaler_2'] = upscaler_to_index(req.upscaler_2)
-    reqDict.pop('upscaler_1')
-    reqDict.pop('upscaler_2')
+    reqDict['extras_upscaler_1'] = reqDict.pop('upscaler_1', None)
+    reqDict['extras_upscaler_2'] = reqDict.pop('upscaler_2', None)
     return reqDict
 
 def decode_base64_to_image(encoding):
     if encoding.startswith("data:image/"):
         encoding = encoding.split(";")[1].split(",")[1]
-    return Image.open(BytesIO(base64.b64decode(encoding)))
+    try:
+        image = Image.open(BytesIO(base64.b64decode(encoding)))
+        return image
+    except Exception as err:
+        raise HTTPException(status_code=500, detail="Invalid encoded image")
 
 def encode_pil_to_base64(image):
     with io.BytesIO() as output_bytes:
 
-        # Copy any text-only metadata
-        use_metadata = False
-        metadata = PngImagePlugin.PngInfo()
-        for key, value in image.info.items():
-            if isinstance(key, str) and isinstance(value, str):
-                metadata.add_text(key, value)
-                use_metadata = True
+        if opts.samples_format.lower() == 'png':
+            use_metadata = False
+            metadata = PngImagePlugin.PngInfo()
+            for key, value in image.info.items():
+                if isinstance(key, str) and isinstance(value, str):
+                    metadata.add_text(key, value)
+                    use_metadata = True
+            image.save(output_bytes, format="PNG", pnginfo=(metadata if use_metadata else None), quality=opts.jpeg_quality)
+
+        elif opts.samples_format.lower() in ("jpg", "jpeg", "webp"):
+            parameters = image.info.get('parameters', None)
+            exif_bytes = piexif.dump({
+                "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(parameters or "", encoding="unicode") }
+            })
+            if opts.samples_format.lower() in ("jpg", "jpeg"):
+                image.save(output_bytes, format="JPEG", exif = exif_bytes, quality=opts.jpeg_quality)
+            else:
+                image.save(output_bytes, format="WEBP", exif = exif_bytes, quality=opts.jpeg_quality)
+
+        else:
+            raise HTTPException(status_code=500, detail="Invalid image format")
 
-        image.save(
-            output_bytes, "PNG", pnginfo=(metadata if use_metadata else None)
-        )
         bytes_data = output_bytes.getvalue()
+
     return base64.b64encode(bytes_data)
 
+def api_middleware(app: FastAPI):
+    @app.middleware("http")
+    async def log_and_time(req: Request, call_next):
+        ts = time.time()
+        res: Response = await call_next(req)
+        duration = str(round(time.time() - ts, 4))
+        res.headers["X-Process-Time"] = duration
+        endpoint = req.scope.get('path', 'err')
+        if shared.cmd_opts.api_log and endpoint.startswith('/sdapi'):
+            print('API {t} {code} {prot}/{ver} {method} {endpoint} {cli} {duration}'.format(
+                t = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"),
+                code = res.status_code,
+                ver = req.scope.get('http_version', '0.0'),
+                cli = req.scope.get('client', ('0:0.0.0', 0))[0],
+                prot = req.scope.get('scheme', 'err'),
+                method = req.scope.get('method', 'err'),
+                endpoint = endpoint,
+                duration = duration,
+            ))
+        return res
+
 
 class Api:
     def __init__(self, app: FastAPI, queue_lock: Lock):
@@ -79,6 +122,7 @@ class Api:
         self.router = APIRouter()
         self.app = app
         self.queue_lock = queue_lock
+        api_middleware(self.app)
         self.add_api_route("/sdapi/v1/txt2img", self.text2imgapi, methods=["POST"], response_model=TextToImageResponse)
         self.add_api_route("/sdapi/v1/img2img", self.img2imgapi, methods=["POST"], response_model=ImageToImageResponse)
         self.add_api_route("/sdapi/v1/extra-single-image", self.extras_single_image_api, methods=["POST"], response_model=ExtrasSingleImageResponse)
@@ -98,14 +142,14 @@ class Api:
         self.add_api_route("/sdapi/v1/face-restorers", self.get_face_restorers, methods=["GET"], response_model=List[FaceRestorerItem])
         self.add_api_route("/sdapi/v1/realesrgan-models", self.get_realesrgan_models, methods=["GET"], response_model=List[RealesrganItem])
         self.add_api_route("/sdapi/v1/prompt-styles", self.get_prompt_styles, methods=["GET"], response_model=List[PromptStyleItem])
-        self.add_api_route("/sdapi/v1/artist-categories", self.get_artists_categories, methods=["GET"], response_model=List[str])
-        self.add_api_route("/sdapi/v1/artists", self.get_artists, methods=["GET"], response_model=List[ArtistItem])
+        self.add_api_route("/sdapi/v1/embeddings", self.get_embeddings, methods=["GET"], response_model=EmbeddingsResponse)
         self.add_api_route("/sdapi/v1/refresh-checkpoints", self.refresh_checkpoints, methods=["POST"])
         self.add_api_route("/sdapi/v1/create/embedding", self.create_embedding, methods=["POST"], response_model=CreateResponse)
         self.add_api_route("/sdapi/v1/create/hypernetwork", self.create_hypernetwork, methods=["POST"], response_model=CreateResponse)
         self.add_api_route("/sdapi/v1/preprocess", self.preprocess, methods=["POST"], response_model=PreprocessResponse)
         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)
 
     def add_api_route(self, path: str, endpoint, **kwargs):
         if shared.cmd_opts.api_auth:
@@ -119,9 +163,22 @@ 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:
+            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 text2imgapi(self, txt2imgreq: StableDiffusionTxt2ImgProcessingAPI):
+        script, script_idx = self.get_script(txt2imgreq.script_name, scripts.scripts_txt2img)
+
         populate = txt2imgreq.copy(update={ # Override __init__ params
-            "sd_model": shared.sd_model,
             "sampler_name": validate_sampler_name(txt2imgreq.sampler_name or txt2imgreq.sampler_index),
             "do_not_save_samples": True,
             "do_not_save_grid": True
@@ -129,15 +186,22 @@ class Api:
         )
         if populate.sampler_name:
             populate.sampler_index = None  # prevent a warning later on
-        p = StableDiffusionProcessingTxt2Img(**vars(populate))
-        # Override object param
 
-        shared.state.begin()
+        args = vars(populate)
+        args.pop('script_name', None)
 
         with self.queue_lock:
-            processed = process_images(p)
+            p = StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **args)
 
-        shared.state.end()
+            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)
+            else:
+                processed = process_images(p)
+            shared.state.end()
 
         b64images = list(map(encode_pil_to_base64, processed.images))
 
@@ -148,12 +212,13 @@ 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
-            "sd_model": shared.sd_model,
             "sampler_name": validate_sampler_name(img2imgreq.sampler_name or img2imgreq.sampler_index),
             "do_not_save_samples": True,
             "do_not_save_grid": True,
@@ -165,16 +230,21 @@ class Api:
 
         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.
-        p = StableDiffusionProcessingImg2Img(**args)
-
-        p.init_images = [decode_base64_to_image(x) for x in init_images]
-
-        shared.state.begin()
+        args.pop('script_name', None)
 
         with self.queue_lock:
-            processed = process_images(p)
+            p = StableDiffusionProcessingImg2Img(sd_model=shared.sd_model, **args)
+            p.init_images = [decode_base64_to_image(x) for x in init_images]
 
-        shared.state.end()
+            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)
+            else:
+                processed = process_images(p)
+            shared.state.end()
 
         b64images = list(map(encode_pil_to_base64, processed.images))
 
@@ -190,7 +260,7 @@ class Api:
         reqDict['image'] = decode_base64_to_image(reqDict['image'])
 
         with self.queue_lock:
-            result = run_extras(extras_mode=0, image_folder="", input_dir="", output_dir="", save_output=False, **reqDict)
+            result = postprocessing.run_extras(extras_mode=0, image_folder="", input_dir="", output_dir="", save_output=False, **reqDict)
 
         return ExtrasSingleImageResponse(image=encode_pil_to_base64(result[0][0]), html_info=result[1])
 
@@ -206,7 +276,7 @@ class Api:
         reqDict.pop('imageList')
 
         with self.queue_lock:
-            result = run_extras(extras_mode=1, image="", input_dir="", output_dir="", save_output=False, **reqDict)
+            result = postprocessing.run_extras(extras_mode=1, image="", input_dir="", output_dir="", save_output=False, **reqDict)
 
         return ExtrasBatchImagesResponse(images=list(map(encode_pil_to_base64, result[0])), html_info=result[1])
 
@@ -214,15 +284,23 @@ class Api:
         if(not req.image.strip()):
             return PNGInfoResponse(info="")
 
-        result = run_pnginfo(decode_base64_to_image(req.image.strip()))
+        image = decode_base64_to_image(req.image.strip())
+        if image is None:
+            return PNGInfoResponse(info="")
 
-        return PNGInfoResponse(info=result[1])
+        geninfo, items = images.read_info_from_image(image)
+        if geninfo is None:
+            geninfo = ""
+
+        items = {**{'parameters': geninfo}, **items}
+
+        return PNGInfoResponse(info=geninfo, items=items)
 
     def progressapi(self, req: ProgressRequest = Depends()):
         # copy from check_progress_call of ui.py
 
         if shared.state.job_count == 0:
-            return ProgressResponse(progress=0, eta_relative=0, state=shared.state.dict())
+            return ProgressResponse(progress=0, eta_relative=0, state=shared.state.dict(), textinfo=shared.state.textinfo)
 
         # avoid dividing zero
         progress = 0.01
@@ -244,7 +322,7 @@ class Api:
         if shared.state.current_image and not req.skip_current_image:
             current_image = encode_pil_to_base64(shared.state.current_image)
 
-        return ProgressResponse(progress=progress, eta_relative=eta_relative, state=shared.state.dict(), current_image=current_image)
+        return ProgressResponse(progress=progress, eta_relative=eta_relative, state=shared.state.dict(), current_image=current_image, textinfo=shared.state.textinfo)
 
     def interrogateapi(self, interrogatereq: InterrogateRequest):
         image_b64 = interrogatereq.image
@@ -298,16 +376,19 @@ class Api:
         return [{"name": sampler[0], "aliases":sampler[2], "options":sampler[3]} for sampler in sd_samplers.all_samplers]
 
     def get_upscalers(self):
-        upscalers = []
-
-        for upscaler in shared.sd_upscalers:
-            u = upscaler.scaler
-            upscalers.append({"name":u.name, "model_name":u.model_name, "model_path":u.model_path, "model_url":u.model_url})
-
-        return upscalers
+        return [
+            {
+                "name": upscaler.name,
+                "model_name": upscaler.scaler.model_name,
+                "model_path": upscaler.data_path,
+                "model_url": None,
+                "scale": upscaler.scale,
+            }
+            for upscaler in shared.sd_upscalers
+        ]
 
     def get_sd_models(self):
-        return [{"title":x.title, "model_name":x.model_name, "hash":x.hash, "filename": x.filename, "config": x.config} for x in checkpoints_list.values()]
+        return [{"title": x.title, "model_name": x.model_name, "hash": x.shorthash, "sha256": x.sha256, "filename": x.filename, "config": find_checkpoint_config_near_filename(x)} for x in checkpoints_list.values()]
 
     def get_hypernetworks(self):
         return [{"name": name, "path": shared.hypernetworks[name]} for name in shared.hypernetworks]
@@ -326,11 +407,25 @@ class Api:
 
         return styleList
 
-    def get_artists_categories(self):
-        return shared.artist_db.cats
+    def get_embeddings(self):
+        db = sd_hijack.model_hijack.embedding_db
+
+        def convert_embedding(embedding):
+            return {
+                "step": embedding.step,
+                "sd_checkpoint": embedding.sd_checkpoint,
+                "sd_checkpoint_name": embedding.sd_checkpoint_name,
+                "shape": embedding.shape,
+                "vectors": embedding.vectors,
+            }
 
-    def get_artists(self):
-        return [{"name":x[0], "score":x[1], "category":x[2]} for x in shared.artist_db.artists]
+        def convert_embeddings(embeddings):
+            return {embedding.name: convert_embedding(embedding) for embedding in embeddings.values()}
+
+        return {
+            "loaded": convert_embeddings(db.word_embeddings),
+            "skipped": convert_embeddings(db.skipped_embeddings),
+        }
 
     def refresh_checkpoints(self):
         shared.refresh_checkpoints()
@@ -396,7 +491,7 @@ class Api:
     def train_hypernetwork(self, args: dict):
         try:
             shared.state.begin()
-            initial_hypernetwork = shared.loaded_hypernetwork
+            shared.loaded_hypernetworks = []
             apply_optimizations = shared.opts.training_xattention_optimizations
             error = None
             filename = ''
@@ -407,16 +502,49 @@ class Api:
             except Exception as e:
                 error = e
             finally:
-                shared.loaded_hypernetwork = initial_hypernetwork
                 shared.sd_model.cond_stage_model.to(devices.device)
                 shared.sd_model.first_stage_model.to(devices.device)
                 if not apply_optimizations:
                     sd_hijack.apply_optimizations()
                 shared.state.end()
-            return TrainResponse(info = "train embedding complete: filename: {filename} error: {error}".format(filename = filename, error = error))
+            return TrainResponse(info="train embedding complete: filename: {filename} error: {error}".format(filename=filename, error=error))
         except AssertionError as msg:
             shared.state.end()
-            return TrainResponse(info = "train embedding error: {error}".format(error = error))
+            return TrainResponse(info="train embedding error: {error}".format(error=error))
+
+    def get_memory(self):
+        try:
+            import os, psutil
+            process = psutil.Process(os.getpid())
+            res = process.memory_info() # only rss is cross-platform guaranteed so we dont rely on other values
+            ram_total = 100 * res.rss / process.memory_percent() # and total memory is calculated as actual value is not cross-platform safe
+            ram = { 'free': ram_total - res.rss, 'used': res.rss, 'total': ram_total }
+        except Exception as err:
+            ram = { 'error': f'{err}' }
+        try:
+            import torch
+            if torch.cuda.is_available():
+                s = torch.cuda.mem_get_info()
+                system = { 'free': s[0], 'used': s[1] - s[0], 'total': s[1] }
+                s = dict(torch.cuda.memory_stats(shared.device))
+                allocated = { 'current': s['allocated_bytes.all.current'], 'peak': s['allocated_bytes.all.peak'] }
+                reserved = { 'current': s['reserved_bytes.all.current'], 'peak': s['reserved_bytes.all.peak'] }
+                active = { 'current': s['active_bytes.all.current'], 'peak': s['active_bytes.all.peak'] }
+                inactive = { 'current': s['inactive_split_bytes.all.current'], 'peak': s['inactive_split_bytes.all.peak'] }
+                warnings = { 'retries': s['num_alloc_retries'], 'oom': s['num_ooms'] }
+                cuda = {
+                    'system': system,
+                    'active': active,
+                    'allocated': allocated,
+                    'reserved': reserved,
+                    'inactive': inactive,
+                    'events': warnings,
+                }
+            else:
+                cuda = { 'error': 'unavailable' }
+        except Exception as err:
+            cuda = { 'error': f'{err}' }
+        return MemoryResponse(ram = ram, cuda = cuda)
 
     def launch(self, server_name, port):
         self.app.include_router(self.router)
diff --git a/modules/api/models.py b/modules/api/models.py
index c446ce7a..cba43d3b 100644
--- a/modules/api/models.py
+++ b/modules/api/models.py
@@ -100,13 +100,13 @@ class PydanticModelGenerator:
 StableDiffusionTxt2ImgProcessingAPI = PydanticModelGenerator(
     "StableDiffusionProcessingTxt2Img",
     StableDiffusionProcessingTxt2Img,
-    [{"key": "sampler_index", "type": str, "default": "Euler"}]
+    [{"key": "sampler_index", "type": str, "default": "Euler"}, {"key": "script_name", "type": str, "default": None}, {"key": "script_args", "type": list, "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": "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": []}]
 ).generate_model()
 
 class TextToImageResponse(BaseModel):
@@ -125,7 +125,7 @@ class ExtrasBaseRequest(BaseModel):
     gfpgan_visibility: float = Field(default=0, title="GFPGAN Visibility", ge=0, le=1, allow_inf_nan=False, description="Sets the visibility of GFPGAN, values should be between 0 and 1.")
     codeformer_visibility: float = Field(default=0, title="CodeFormer Visibility", ge=0, le=1, allow_inf_nan=False, description="Sets the visibility of CodeFormer, values should be between 0 and 1.")
     codeformer_weight: float = Field(default=0, title="CodeFormer Weight", ge=0, le=1, allow_inf_nan=False, description="Sets the weight of CodeFormer, values should be between 0 and 1.")
-    upscaling_resize: float = Field(default=2, title="Upscaling Factor", ge=1, le=4, description="By how much to upscale the image, only used when resize_mode=0.")
+    upscaling_resize: float = Field(default=2, title="Upscaling Factor", ge=1, le=8, description="By how much to upscale the image, only used when resize_mode=0.")
     upscaling_resize_w: int = Field(default=512, title="Target Width", ge=1, description="Target width for the upscaler to hit. Only used when resize_mode=1.")
     upscaling_resize_h: int = Field(default=512, title="Target Height", ge=1, description="Target height for the upscaler to hit. Only used when resize_mode=1.")
     upscaling_crop: bool = Field(default=True, title="Crop to fit", description="Should the upscaler crop the image to fit in the chosen size?")
@@ -157,7 +157,8 @@ class PNGInfoRequest(BaseModel):
     image: str = Field(title="Image", description="The base64 encoded PNG image")
 
 class PNGInfoResponse(BaseModel):
-    info: str = Field(title="Image info", description="A string with all the info the image had")
+    info: str = Field(title="Image info", description="A string with the parameters used to generate the image")
+    items: dict = Field(title="Items", description="An object containing all the info the image had")
 
 class ProgressRequest(BaseModel):
     skip_current_image: bool = Field(default=False, title="Skip current image", description="Skip current image serialization")
@@ -167,6 +168,7 @@ class ProgressResponse(BaseModel):
     eta_relative: float = Field(title="ETA in secs")
     state: dict = Field(title="State", description="The current state snapshot")
     current_image: str = Field(default=None, title="Current image", description="The current image in base64 format. opts.show_progress_every_n_steps is required for this to work.")
+    textinfo: str = Field(default=None, title="Info text", description="Info text used by WebUI.")
 
 class InterrogateRequest(BaseModel):
     image: str = Field(default="", title="Image", description="Image to work on, must be a Base64 string containing the image's data.")
@@ -218,13 +220,15 @@ class UpscalerItem(BaseModel):
     model_name: Optional[str] = Field(title="Model Name")
     model_path: Optional[str] = Field(title="Path")
     model_url: Optional[str] = Field(title="URL")
+    scale: Optional[float] = Field(title="Scale")
 
 class SDModelItem(BaseModel):
     title: str = Field(title="Title")
     model_name: str = Field(title="Model Name")
-    hash: str = Field(title="Hash")
+    hash: Optional[str] = Field(title="Short hash")
+    sha256: Optional[str] = Field(title="sha256 hash")
     filename: str = Field(title="Filename")
-    config: str = Field(title="Config file")
+    config: Optional[str] = Field(title="Config file")
 
 class HypernetworkItem(BaseModel):
     name: str = Field(title="Name")
@@ -249,3 +253,17 @@ class ArtistItem(BaseModel):
     score: float = Field(title="Score")
     category: str = Field(title="Category")
 
+class EmbeddingItem(BaseModel):
+    step: Optional[int] = Field(title="Step", description="The number of steps that were used to train this embedding, if available")
+    sd_checkpoint: Optional[str] = Field(title="SD Checkpoint", description="The hash of the checkpoint this embedding was trained on, if available")
+    sd_checkpoint_name: Optional[str] = Field(title="SD Checkpoint Name", description="The name of the checkpoint this embedding was trained on, if available. Note that this is the name that was used by the trainer; for a stable identifier, use `sd_checkpoint` instead")
+    shape: int = Field(title="Shape", description="The length of each individual vector in the embedding")
+    vectors: int = Field(title="Vectors", description="The number of vectors in the embedding")
+
+class EmbeddingsResponse(BaseModel):
+    loaded: Dict[str, EmbeddingItem] = Field(title="Loaded", description="Embeddings loaded for the current model")
+    skipped: Dict[str, EmbeddingItem] = Field(title="Skipped", description="Embeddings skipped for the current model (likely due to architecture incompatibility)")
+
+class MemoryResponse(BaseModel):
+    ram: dict = Field(title="RAM", description="System memory stats")
+    cuda: dict = Field(title="CUDA", description="nVidia CUDA memory stats")
diff --git a/modules/artists.py b/modules/artists.py
deleted file mode 100644
index 3612758b..00000000
--- a/modules/artists.py
+++ /dev/null
@@ -1,25 +0,0 @@
-import os.path

-import csv

-from collections import namedtuple

-

-Artist = namedtuple("Artist", ['name', 'weight', 'category'])

-

-

-class ArtistsDatabase:

-    def __init__(self, filename):

-        self.cats = set()

-        self.artists = []

-

-        if not os.path.exists(filename):

-            return

-

-        with open(filename, "r", newline='', encoding="utf8") as file:

-            reader = csv.DictReader(file)

-

-            for row in reader:

-                artist = Artist(row["artist"], float(row["score"]), row["category"])

-                self.artists.append(artist)

-                self.cats.add(artist.category)

-

-    def categories(self):

-        return sorted(self.cats)

diff --git a/modules/call_queue.py b/modules/call_queue.py
index 4cd49533..92097c15 100644
--- a/modules/call_queue.py
+++ b/modules/call_queue.py
@@ -4,7 +4,7 @@ import threading
 import traceback

 import time

 

-from modules import shared

+from modules import shared, progress

 

 queue_lock = threading.Lock()

 

@@ -22,12 +22,23 @@ def wrap_queued_call(func):
 def wrap_gradio_gpu_call(func, extra_outputs=None):

     def f(*args, **kwargs):

 

-        shared.state.begin()

+        # if the first argument is a string that says "task(...)", it is treated as a job id

+        if len(args) > 0 and type(args[0]) == str and args[0][0:5] == "task(" and args[0][-1] == ")":

+            id_task = args[0]

+            progress.add_task_to_queue(id_task)

+        else:

+            id_task = None

 

         with queue_lock:

-            res = func(*args, **kwargs)

+            shared.state.begin()

+            progress.start_task(id_task)

+

+            try:

+                res = func(*args, **kwargs)

+            finally:

+                progress.finish_task(id_task)

 

-        shared.state.end()

+            shared.state.end()

 

         return res

 

diff --git a/modules/codeformer_model.py b/modules/codeformer_model.py
index ab40d842..01fb7bd8 100644
--- a/modules/codeformer_model.py
+++ b/modules/codeformer_model.py
@@ -8,7 +8,7 @@ import torch
 import modules.face_restoration

 import modules.shared

 from modules import shared, devices, modelloader

-from modules.paths import script_path, models_path

+from modules.paths import models_path

 

 # codeformer people made a choice to include modified basicsr library to their project which makes

 # it utterly impossible to use it alongside with other libraries that also use basicsr, like GFPGAN.

diff --git a/modules/deepbooru_model.py b/modules/deepbooru_model.py
index edd40c81..83d2ff09 100644
--- a/modules/deepbooru_model.py
+++ b/modules/deepbooru_model.py
@@ -2,6 +2,8 @@ import torch
 import torch.nn as nn

 import torch.nn.functional as F

 

+from modules import devices

+

 # see https://github.com/AUTOMATIC1111/TorchDeepDanbooru for more

 

 

@@ -196,7 +198,7 @@ class DeepDanbooruModel(nn.Module):
         t_358, = inputs

         t_359 = t_358.permute(*[0, 3, 1, 2])

         t_359_padded = F.pad(t_359, [2, 3, 2, 3], value=0)

-        t_360 = self.n_Conv_0(t_359_padded)

+        t_360 = self.n_Conv_0(t_359_padded.to(self.n_Conv_0.bias.dtype) if devices.unet_needs_upcast else t_359_padded)

         t_361 = F.relu(t_360)

         t_361 = F.pad(t_361, [0, 1, 0, 1], value=float('-inf'))

         t_362 = self.n_MaxPool_0(t_361)

diff --git a/modules/devices.py b/modules/devices.py
index 800510b7..4687944e 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -34,14 +34,18 @@ def get_cuda_device_string():
     return "cuda"
 
 
-def get_optimal_device():
+def get_optimal_device_name():
     if torch.cuda.is_available():
-        return torch.device(get_cuda_device_string())
+        return get_cuda_device_string()
 
     if has_mps():
-        return torch.device("mps")
+        return "mps"
+
+    return "cpu"
 
-    return cpu
+
+def get_optimal_device():
+    return torch.device(get_optimal_device_name())
 
 
 def get_device_for(task):
@@ -79,6 +83,8 @@ cpu = torch.device("cpu")
 device = device_interrogate = device_gfpgan = device_esrgan = device_codeformer = None
 dtype = torch.float16
 dtype_vae = torch.float16
+dtype_unet = torch.float16
+unet_needs_upcast = False
 
 
 def randn(seed, shape):
@@ -106,6 +112,42 @@ def autocast(disable=False):
     return torch.autocast("cuda")
 
 
+def without_autocast(disable=False):
+    return torch.autocast("cuda", enabled=False) if torch.is_autocast_enabled() and not disable else contextlib.nullcontext()
+
+
+class NansException(Exception):
+    pass
+
+
+def test_for_nans(x, where):
+    from modules import shared
+
+    if shared.cmd_opts.disable_nan_check:
+        return
+
+    if not torch.all(torch.isnan(x)).item():
+        return
+
+    if where == "unet":
+        message = "A tensor with all NaNs was produced in Unet."
+
+        if not shared.cmd_opts.no_half:
+            message += " This could be either because there's not enough precision to represent the picture, or because your video card does not support half type. Try setting the \"Upcast cross attention layer to float32\" option in Settings > Stable Diffusion or using the --no-half commandline argument to fix this."
+
+    elif where == "vae":
+        message = "A tensor with all NaNs was produced in VAE."
+
+        if not shared.cmd_opts.no_half and not shared.cmd_opts.no_half_vae:
+            message += " This could be because there's not enough precision to represent the picture. Try adding --no-half-vae commandline argument to fix this."
+    else:
+        message = "A tensor with all NaNs was produced."
+
+    message += " Use --disable-nan-check commandline argument to disable this check."
+
+    raise NansException(message)
+
+
 # MPS workaround for https://github.com/pytorch/pytorch/issues/79383
 orig_tensor_to = torch.Tensor.to
 def tensor_to_fix(self, *args, **kwargs):
@@ -133,8 +175,30 @@ def numpy_fix(self, *args, **kwargs):
     return orig_tensor_numpy(self, *args, **kwargs)
 
 
-# PyTorch 1.13 doesn't need these fixes but unfortunately is slower and has regressions that prevent training from working
-if has_mps() and version.parse(torch.__version__) < version.parse("1.13"):
-    torch.Tensor.to = tensor_to_fix
-    torch.nn.functional.layer_norm = layer_norm_fix
-    torch.Tensor.numpy = numpy_fix
+# MPS workaround for https://github.com/pytorch/pytorch/issues/89784
+orig_cumsum = torch.cumsum
+orig_Tensor_cumsum = torch.Tensor.cumsum
+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
+        torch.Tensor.to = tensor_to_fix
+        torch.nn.functional.layer_norm = layer_norm_fix
+        torch.Tensor.numpy = numpy_fix
+    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))
+        torch.cumsum = lambda input, *args, **kwargs: ( cumsum_fix(input, orig_cumsum, *args, **kwargs) )
+        torch.Tensor.cumsum = lambda self, *args, **kwargs: ( cumsum_fix(self, orig_Tensor_cumsum, *args, **kwargs) )
+        orig_narrow = torch.narrow
+        torch.narrow = lambda *args, **kwargs: ( orig_narrow(*args, **kwargs).clone() )
+
diff --git a/modules/errors.py b/modules/errors.py
index 372dc51a..f6b80dbb 100644
--- a/modules/errors.py
+++ b/modules/errors.py
@@ -2,9 +2,42 @@ import sys
 import traceback

 

 

+def print_error_explanation(message):

+    lines = message.strip().split("\n")

+    max_len = max([len(x) for x in lines])

+

+    print('=' * max_len, file=sys.stderr)

+    for line in lines:

+        print(line, file=sys.stderr)

+    print('=' * max_len, file=sys.stderr)

+

+

+def display(e: Exception, task):

+    print(f"{task or 'error'}: {type(e).__name__}", file=sys.stderr)

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

+

+    message = str(e)

+    if "copying a param with shape torch.Size([640, 1024]) from checkpoint, the shape in current model is torch.Size([640, 768])" in message:

+        print_error_explanation("""

+The most likely cause of this is you are trying to load Stable Diffusion 2.0 model without specifying its config file.

+See https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#stable-diffusion-20 for how to solve this.

+        """)

+

+

+already_displayed = {}

+

+

+def display_once(e: Exception, task):

+    if task in already_displayed:

+        return

+

+    display(e, task)

+

+    already_displayed[task] = 1

+

+

 def run(code, task):

     try:

         code()

     except Exception as e:

-        print(f"{task}: {type(e).__name__}", file=sys.stderr)

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

+        display(task, e)

diff --git a/modules/extensions.py b/modules/extensions.py
index b522125c..5e12b1aa 100644
--- a/modules/extensions.py
+++ b/modules/extensions.py
@@ -7,9 +7,11 @@ import git
 from modules import paths, shared

 

 extensions = []

-extensions_dir = os.path.join(paths.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]

diff --git a/modules/extra_networks.py b/modules/extra_networks.py
new file mode 100644
index 00000000..1978673d
--- /dev/null
+++ b/modules/extra_networks.py
@@ -0,0 +1,147 @@
+import re

+from collections import defaultdict

+

+from modules import errors

+

+extra_network_registry = {}

+

+

+def initialize():

+    extra_network_registry.clear()

+

+

+def register_extra_network(extra_network):

+    extra_network_registry[extra_network.name] = extra_network

+

+

+class ExtraNetworkParams:

+    def __init__(self, items=None):

+        self.items = items or []

+

+

+class ExtraNetwork:

+    def __init__(self, name):

+        self.name = name

+

+    def activate(self, p, params_list):

+        """

+        Called by processing on every run. Whatever the extra network is meant to do should be activated here.

+        Passes arguments related to this extra network in params_list.

+        User passes arguments by specifying this in his prompt:

+

+        <name:arg1:arg2:arg3>

+

+        Where name matches the name of this ExtraNetwork object, and arg1:arg2:arg3 are any natural number of text arguments

+        separated by colon.

+

+        Even if the user does not mention this ExtraNetwork in his prompt, the call will stil be made, with empty params_list -

+        in this case, all effects of this extra networks should be disabled.

+

+        Can be called multiple times before deactivate() - each new call should override the previous call completely.

+

+        For example, if this ExtraNetwork's name is 'hypernet' and user's prompt is:

+

+        > "1girl, <hypernet:agm:1.1> <extrasupernet:master:12:13:14> <hypernet:ray>"

+

+        params_list will be:

+

+        [

+            ExtraNetworkParams(items=["agm", "1.1"]),

+            ExtraNetworkParams(items=["ray"])

+        ]

+

+        """

+        raise NotImplementedError

+

+    def deactivate(self, p):

+        """

+        Called at the end of processing for housekeeping. No need to do anything here.

+        """

+

+        raise NotImplementedError

+

+

+def activate(p, extra_network_data):

+    """call activate for extra networks in extra_network_data in specified order, then call

+    activate for all remaining registered networks with an empty argument list"""

+

+    for extra_network_name, extra_network_args in extra_network_data.items():

+        extra_network = extra_network_registry.get(extra_network_name, None)

+        if extra_network is None:

+            print(f"Skipping unknown extra network: {extra_network_name}")

+            continue

+

+        try:

+            extra_network.activate(p, extra_network_args)

+        except Exception as e:

+            errors.display(e, f"activating extra network {extra_network_name} with arguments {extra_network_args}")

+

+    for extra_network_name, extra_network in extra_network_registry.items():

+        args = extra_network_data.get(extra_network_name, None)

+        if args is not None:

+            continue

+

+        try:

+            extra_network.activate(p, [])

+        except Exception as e:

+            errors.display(e, f"activating extra network {extra_network_name}")

+

+

+def deactivate(p, extra_network_data):

+    """call deactivate for extra networks in extra_network_data in specified order, then call

+    deactivate for all remaining registered networks"""

+

+    for extra_network_name, extra_network_args in extra_network_data.items():

+        extra_network = extra_network_registry.get(extra_network_name, None)

+        if extra_network is None:

+            continue

+

+        try:

+            extra_network.deactivate(p)

+        except Exception as e:

+            errors.display(e, f"deactivating extra network {extra_network_name}")

+

+    for extra_network_name, extra_network in extra_network_registry.items():

+        args = extra_network_data.get(extra_network_name, None)

+        if args is not None:

+            continue

+

+        try:

+            extra_network.deactivate(p)

+        except Exception as e:

+            errors.display(e, f"deactivating unmentioned extra network {extra_network_name}")

+

+

+re_extra_net = re.compile(r"<(\w+):([^>]+)>")

+

+

+def parse_prompt(prompt):

+    res = defaultdict(list)

+

+    def found(m):

+        name = m.group(1)

+        args = m.group(2)

+

+        res[name].append(ExtraNetworkParams(items=args.split(":")))

+

+        return ""

+

+    prompt = re.sub(re_extra_net, found, prompt)

+

+    return prompt, res

+

+

+def parse_prompts(prompts):

+    res = []

+    extra_data = None

+

+    for prompt in prompts:

+        updated_prompt, parsed_extra_data = parse_prompt(prompt)

+

+        if extra_data is None:

+            extra_data = parsed_extra_data

+

+        res.append(updated_prompt)

+

+    return res, extra_data

+

diff --git a/modules/extra_networks_hypernet.py b/modules/extra_networks_hypernet.py
new file mode 100644
index 00000000..ff279a1f
--- /dev/null
+++ b/modules/extra_networks_hypernet.py
@@ -0,0 +1,21 @@
+from modules import extra_networks

+from modules.hypernetworks import hypernetwork

+

+

+class ExtraNetworkHypernet(extra_networks.ExtraNetwork):

+    def __init__(self):

+        super().__init__('hypernet')

+

+    def activate(self, p, params_list):

+        names = []

+        multipliers = []

+        for params in params_list:

+            assert len(params.items) > 0

+

+            names.append(params.items[0])

+            multipliers.append(float(params.items[1]) if len(params.items) > 1 else 1.0)

+

+        hypernetwork.load_hypernetworks(names, multipliers)

+

+    def deactivate(self, p):

+        pass

diff --git a/modules/extras.py b/modules/extras.py
index 68939dea..d8ece955 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -1,222 +1,16 @@
-from __future__ import annotations

-import math

 import os

-import sys

-import traceback

+import re

+import shutil

 

-import numpy as np

-from PIL import Image

 

 import torch

 import tqdm

 

-from typing import Callable, List, OrderedDict, Tuple

-from functools import partial

-from dataclasses import dataclass

-

-from modules import processing, shared, images, devices, sd_models, sd_samplers

-from modules.shared import opts

-import modules.gfpgan_model

-from modules.ui import plaintext_to_html

-import modules.codeformer_model

-import piexif

-import piexif.helper

+from modules import shared, images, sd_models, sd_vae, sd_models_config

+from modules.ui_common import plaintext_to_html

 import gradio as gr

 import safetensors.torch

 

-class LruCache(OrderedDict):

-    @dataclass(frozen=True)

-    class Key:

-        image_hash: int

-        info_hash: int

-        args_hash: int

-

-    @dataclass

-    class Value:

-        image: Image.Image

-        info: str

-

-    def __init__(self, max_size: int = 5, *args, **kwargs):

-        super().__init__(*args, **kwargs)

-        self._max_size = max_size

-

-    def get(self, key: LruCache.Key) -> LruCache.Value:

-        ret = super().get(key)

-        if ret is not None:

-            self.move_to_end(key)  # Move to end of eviction list

-        return ret

-

-    def put(self, key: LruCache.Key, value: LruCache.Value) -> None:

-        self[key] = value

-        while len(self) > self._max_size:

-            self.popitem(last=False)

-

-

-cached_images: LruCache = LruCache(max_size=5)

-

-

-def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_dir, show_extras_results, gfpgan_visibility, codeformer_visibility, codeformer_weight, upscaling_resize, upscaling_resize_w, upscaling_resize_h, upscaling_crop, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility, upscale_first: bool, save_output: bool = True):

-    devices.torch_gc()

-

-    imageArr = []

-    # Also keep track of original file names

-    imageNameArr = []

-    outputs = []

-

-    if extras_mode == 1:

-        #convert file to pillow image

-        for img in image_folder:

-            image = Image.open(img)

-            imageArr.append(image)

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

-    elif extras_mode == 2:

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

-

-        if input_dir == '':

-            return outputs, "Please select an input directory.", ''

-        image_list = shared.listfiles(input_dir)

-        for img in image_list:

-            try:

-                image = Image.open(img)

-            except Exception:

-                continue

-            imageArr.append(image)

-            imageNameArr.append(img)

-    else:

-        imageArr.append(image)

-        imageNameArr.append(None)

-

-    if extras_mode == 2 and output_dir != '':

-        outpath = output_dir

-    else:

-        outpath = opts.outdir_samples or opts.outdir_extras_samples

-

-    # Extra operation definitions

-

-    def run_gfpgan(image: Image.Image, info: str) -> Tuple[Image.Image, str]:

-        restored_img = modules.gfpgan_model.gfpgan_fix_faces(np.array(image, dtype=np.uint8))

-        res = Image.fromarray(restored_img)

-

-        if gfpgan_visibility < 1.0:

-            res = Image.blend(image, res, gfpgan_visibility)

-

-        info += f"GFPGAN visibility:{round(gfpgan_visibility, 2)}\n"

-        return (res, info)

-

-    def run_codeformer(image: Image.Image, info: str) -> Tuple[Image.Image, str]:

-        restored_img = modules.codeformer_model.codeformer.restore(np.array(image, dtype=np.uint8), w=codeformer_weight)

-        res = Image.fromarray(restored_img)

-

-        if codeformer_visibility < 1.0:

-            res = Image.blend(image, res, codeformer_visibility)

-

-        info += f"CodeFormer w: {round(codeformer_weight, 2)}, CodeFormer visibility:{round(codeformer_visibility, 2)}\n"

-        return (res, info)

-

-    def upscale(image, scaler_index, resize, mode, resize_w, resize_h, crop):

-        upscaler = shared.sd_upscalers[scaler_index]

-        res = upscaler.scaler.upscale(image, resize, upscaler.data_path)

-        if mode == 1 and crop:

-            cropped = Image.new("RGB", (resize_w, resize_h))

-            cropped.paste(res, box=(resize_w // 2 - res.width // 2, resize_h // 2 - res.height // 2))

-            res = cropped

-        return res

-

-    def run_prepare_crop(image: Image.Image, info: str) -> Tuple[Image.Image, str]:

-        # Actual crop happens in run_upscalers_blend, this just sets upscaling_resize and adds info text

-        nonlocal upscaling_resize

-        if resize_mode == 1:

-            upscaling_resize = max(upscaling_resize_w/image.width, upscaling_resize_h/image.height)

-            crop_info = " (crop)" if upscaling_crop else ""

-            info += f"Resize to: {upscaling_resize_w:g}x{upscaling_resize_h:g}{crop_info}\n"

-        return (image, info)

-

-    @dataclass

-    class UpscaleParams:

-        upscaler_idx: int

-        blend_alpha: float

-

-    def run_upscalers_blend(params: List[UpscaleParams], image: Image.Image, info: str) -> Tuple[Image.Image, str]:

-        blended_result: Image.Image = None

-        image_hash: str = hash(np.array(image.getdata()).tobytes())

-        for upscaler in params:

-            upscale_args = (upscaler.upscaler_idx, upscaling_resize, resize_mode,

-                            upscaling_resize_w, upscaling_resize_h, upscaling_crop)

-            cache_key = LruCache.Key(image_hash=image_hash,

-                                     info_hash=hash(info),

-                                     args_hash=hash(upscale_args))

-            cached_entry = cached_images.get(cache_key)

-            if cached_entry is None:

-                res = upscale(image, *upscale_args)

-                info += f"Upscale: {round(upscaling_resize, 3)}, visibility: {upscaler.blend_alpha}, model:{shared.sd_upscalers[upscaler.upscaler_idx].name}\n"

-                cached_images.put(cache_key, LruCache.Value(image=res, info=info))

-            else:

-                res, info = cached_entry.image, cached_entry.info

-

-            if blended_result is None:

-                blended_result = res

-            else:

-                blended_result = Image.blend(blended_result, res, upscaler.blend_alpha)

-        return (blended_result, info)

-

-    # Build a list of operations to run

-    facefix_ops: List[Callable] = []

-    facefix_ops += [run_gfpgan] if gfpgan_visibility > 0 else []

-    facefix_ops += [run_codeformer] if codeformer_visibility > 0 else []

-

-    upscale_ops: List[Callable] = []

-    upscale_ops += [run_prepare_crop] if resize_mode == 1 else []

-

-    if upscaling_resize != 0:

-        step_params: List[UpscaleParams] = []

-        step_params.append(UpscaleParams(upscaler_idx=extras_upscaler_1, blend_alpha=1.0))

-        if extras_upscaler_2 != 0 and extras_upscaler_2_visibility > 0:

-            step_params.append(UpscaleParams(upscaler_idx=extras_upscaler_2, blend_alpha=extras_upscaler_2_visibility))

-

-        upscale_ops.append(partial(run_upscalers_blend, step_params))

-

-    extras_ops: List[Callable] = (upscale_ops + facefix_ops) if upscale_first else (facefix_ops + upscale_ops)

-

-    for image, image_name in zip(imageArr, imageNameArr):

-        if image is None:

-            return outputs, "Please select an input image.", ''

-        existing_pnginfo = image.info or {}

-

-        image = image.convert("RGB")

-        info = ""

-        # Run each operation on each image

-        for op in extras_ops:

-            image, info = op(image, info)

-

-        if opts.use_original_name_batch and image_name is not None:

-            basename = os.path.splitext(os.path.basename(image_name))[0]

-        else:

-            basename = ''

-

-        if save_output:

-            # Add upscaler name as a suffix.

-            suffix = f"-{shared.sd_upscalers[extras_upscaler_1].name}" if shared.opts.use_upscaler_name_as_suffix else ""

-            # Add second upscaler if applicable.

-            if suffix and extras_upscaler_2 and extras_upscaler_2_visibility:

-                suffix += f"-{shared.sd_upscalers[extras_upscaler_2].name}"

-

-            images.save_image(image, path=outpath, basename=basename, seed=None, prompt=None, extension=opts.samples_format, info=info, short_filename=True,

-                            no_prompt=True, grid=False, pnginfo_section_name="extras", existing_info=existing_pnginfo, forced_filename=None, suffix=suffix)

-

-        if opts.enable_pnginfo:

-            image.info = existing_pnginfo

-            image.info["extras"] = info

-

-        if extras_mode != 2 or show_extras_results :

-            outputs.append(image)

-

-    devices.torch_gc()

-

-    return outputs, plaintext_to_html(info), ''

-

-def clear_cache():

-    cached_images.clear()

-

 

 def run_pnginfo(image):

     if image is None:

@@ -241,7 +35,51 @@ def run_pnginfo(image):
     return '', geninfo, info

 

 

-def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format):

+def create_config(ckpt_result, config_source, a, b, c):

+    def config(x):

+        res = sd_models_config.find_checkpoint_config_near_filename(x) if x else None

+        return res if res != shared.sd_default_config else None

+

+    if config_source == 0:

+        cfg = config(a) or config(b) or config(c)

+    elif config_source == 1:

+        cfg = config(b)

+    elif config_source == 2:

+        cfg = config(c)

+    else:

+        cfg = None

+

+    if cfg is None:

+        return

+

+    filename, _ = os.path.splitext(ckpt_result)

+    checkpoint_filename = filename + ".yaml"

+

+    print("Copying config:")

+    print("   from:", cfg)

+    print("     to:", checkpoint_filename)

+    shutil.copyfile(cfg, checkpoint_filename)

+

+

+checkpoint_dict_skip_on_merge = ["cond_stage_model.transformer.text_model.embeddings.position_ids"]

+

+

+def to_half(tensor, enable):

+    if enable and tensor.dtype == torch.float:

+        return tensor.half()

+

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

+    shared.state.begin()

+    shared.state.job = 'model-merge'

+

+    def fail(message):

+        shared.state.textinfo = message

+        shared.state.end()

+        return [*[gr.update() for _ in range(4)], message]

+

     def weighted_sum(theta0, theta1, alpha):

         return ((1 - alpha) * theta0) + (alpha * theta1)

 

@@ -251,43 +89,94 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
     def add_difference(theta0, theta1_2_diff, alpha):

         return theta0 + (alpha * theta1_2_diff)

 

-    primary_model_info = sd_models.checkpoints_list[primary_model_name]

-    secondary_model_info = sd_models.checkpoints_list[secondary_model_name]

-    tertiary_model_info = sd_models.checkpoints_list.get(tertiary_model_name, None)

-    result_is_inpainting_model = False

+    def filename_weighted_sum():

+        a = primary_model_info.model_name

+        b = secondary_model_info.model_name

+        Ma = round(1 - multiplier, 2)

+        Mb = round(multiplier, 2)

+

+        return f"{Ma}({a}) + {Mb}({b})"

+

+    def filename_add_difference():

+        a = primary_model_info.model_name

+        b = secondary_model_info.model_name

+        c = tertiary_model_info.model_name

+        M = round(multiplier, 2)

+

+        return f"{a} + {M}({b} - {c})"

+

+    def filename_nothing():

+        return primary_model_info.model_name

 

     theta_funcs = {

-        "Weighted sum": (None, weighted_sum),

-        "Add difference": (get_difference, add_difference),

+        "Weighted sum": (filename_weighted_sum, None, weighted_sum),

+        "Add difference": (filename_add_difference, get_difference, add_difference),

+        "No interpolation": (filename_nothing, None, None),

     }

-    theta_func1, theta_func2 = theta_funcs[interp_method]

+    filename_generator, theta_func1, theta_func2 = theta_funcs[interp_method]

+    shared.state.job_count = (1 if theta_func1 else 0) + (1 if theta_func2 else 0)

+

+    if not primary_model_name:

+        return fail("Failed: Merging requires a primary model.")

 

-    if theta_func1 and not tertiary_model_info:

-        return ["Failed: Interpolation method requires a tertiary model."] + [gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)]

+    primary_model_info = sd_models.checkpoints_list[primary_model_name]

+

+    if theta_func2 and not secondary_model_name:

+        return fail("Failed: Merging requires a secondary model.")

+

+    secondary_model_info = sd_models.checkpoints_list[secondary_model_name] if theta_func2 else None

 

-    print(f"Loading {secondary_model_info.filename}...")

-    theta_1 = sd_models.read_state_dict(secondary_model_info.filename, map_location='cpu')

+    if theta_func1 and not tertiary_model_name:

+        return fail(f"Failed: Interpolation method ({interp_method}) requires a tertiary model.")

+

+    tertiary_model_info = sd_models.checkpoints_list[tertiary_model_name] if theta_func1 else None

+

+    result_is_inpainting_model = False

+    result_is_instruct_pix2pix_model = False

+

+    if theta_func2:

+        shared.state.textinfo = f"Loading B"

+        print(f"Loading {secondary_model_info.filename}...")

+        theta_1 = sd_models.read_state_dict(secondary_model_info.filename, map_location='cpu')

+    else:

+        theta_1 = None

 

     if theta_func1:

+        shared.state.textinfo = f"Loading C"

         print(f"Loading {tertiary_model_info.filename}...")

         theta_2 = sd_models.read_state_dict(tertiary_model_info.filename, map_location='cpu')

 

+        shared.state.textinfo = 'Merging B and C'

+        shared.state.sampling_steps = len(theta_1.keys())

         for key in tqdm.tqdm(theta_1.keys()):

+            if key in checkpoint_dict_skip_on_merge:

+                continue

+

             if 'model' in key:

                 if key in theta_2:

                     t2 = theta_2.get(key, torch.zeros_like(theta_1[key]))

                     theta_1[key] = theta_func1(theta_1[key], t2)

                 else:

                     theta_1[key] = torch.zeros_like(theta_1[key])

+

+            shared.state.sampling_step += 1

         del theta_2

 

+        shared.state.nextjob()

+

+    shared.state.textinfo = f"Loading {primary_model_info.filename}..."

     print(f"Loading {primary_model_info.filename}...")

     theta_0 = sd_models.read_state_dict(primary_model_info.filename, map_location='cpu')

 

     print("Merging...")

-

+    shared.state.textinfo = 'Merging A and B'

+    shared.state.sampling_steps = len(theta_0.keys())

     for key in tqdm.tqdm(theta_0.keys()):

-        if 'model' in key and key in theta_1:

+        if theta_1 and 'model' in key and key in theta_1:

+

+            if key in checkpoint_dict_skip_on_merge:

+                continue

+

             a = theta_0[key]

             b = theta_1[key]

 

@@ -297,39 +186,59 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
             if a.shape != b.shape and a.shape[0:1] + a.shape[2:] == b.shape[0:1] + b.shape[2:]:

                 if a.shape[1] == 4 and b.shape[1] == 9:

                     raise RuntimeError("When merging inpainting model with a normal one, A must be the inpainting model.")

+                if a.shape[1] == 4 and b.shape[1] == 8:

+                    raise RuntimeError("When merging instruct-pix2pix model with a normal one, A must be the instruct-pix2pix model.")

 

-                assert a.shape[1] == 9 and b.shape[1] == 4, f"Bad dimensions for merged layer {key}: A={a.shape}, B={b.shape}"

-

-                theta_0[key][:, 0:4, :, :] = theta_func2(a[:, 0:4, :, :], b, multiplier)

-                result_is_inpainting_model = True

+                if a.shape[1] == 8 and b.shape[1] == 4:#If we have an Instruct-Pix2Pix model...

+                    theta_0[key][:, 0:4, :, :] = theta_func2(a[:, 0:4, :, :], b, multiplier)#Merge only the vectors the models have in common.  Otherwise we get an error due to dimension mismatch.

+                    result_is_instruct_pix2pix_model = True

+                else:

+                    assert a.shape[1] == 9 and b.shape[1] == 4, f"Bad dimensions for merged layer {key}: A={a.shape}, B={b.shape}"

+                    theta_0[key][:, 0:4, :, :] = theta_func2(a[:, 0:4, :, :], b, multiplier)

+                    result_is_inpainting_model = True

             else:

                 theta_0[key] = theta_func2(a, b, multiplier)

+            

+            theta_0[key] = to_half(theta_0[key], save_as_half)

 

-            if save_as_half:

-                theta_0[key] = theta_0[key].half()

+        shared.state.sampling_step += 1

 

-    # I believe this part should be discarded, but I'll leave it for now until I am sure

-    for key in theta_1.keys():

-        if 'model' in key and key not in theta_0:

-            theta_0[key] = theta_1[key]

-            if save_as_half:

-                theta_0[key] = theta_0[key].half()

     del theta_1

 

-    ckpt_dir = shared.cmd_opts.ckpt_dir or sd_models.model_path

+    bake_in_vae_filename = sd_vae.vae_dict.get(bake_in_vae, None)

+    if bake_in_vae_filename is not None:

+        print(f"Baking in VAE from {bake_in_vae_filename}")

+        shared.state.textinfo = 'Baking in VAE'

+        vae_dict = sd_vae.load_vae_dict(bake_in_vae_filename, map_location='cpu')

+

+        for key in vae_dict.keys():

+            theta_0_key = 'first_stage_model.' + key

+            if theta_0_key in theta_0:

+                theta_0[theta_0_key] = to_half(vae_dict[key], save_as_half)

+

+        del vae_dict

+

+    if save_as_half and not theta_func2:

+        for key in theta_0.keys():

+            theta_0[key] = to_half(theta_0[key], save_as_half)

 

-    filename = \

-        primary_model_info.model_name + '_' + str(round(1-multiplier, 2)) + '-' + \

-        secondary_model_info.model_name + '_' + str(round(multiplier, 2)) + '-' + \

-        interp_method.replace(" ", "_") + \

-        '-merged.' +  \

-        ("inpainting." if result_is_inpainting_model else "") + \

-        checkpoint_format

+    if discard_weights:

+        regex = re.compile(discard_weights)

+        for key in list(theta_0):

+            if re.search(regex, key):

+                theta_0.pop(key, None)

 

-    filename = filename if custom_name == '' else (custom_name + '.' + checkpoint_format)

+    ckpt_dir = shared.cmd_opts.ckpt_dir or sd_models.model_path

+

+    filename = filename_generator() if custom_name == '' else custom_name

+    filename += ".inpainting" if result_is_inpainting_model else ""

+    filename += ".instruct-pix2pix" if result_is_instruct_pix2pix_model else ""

+    filename += "." + checkpoint_format

 

     output_modelname = os.path.join(ckpt_dir, filename)

 

+    shared.state.nextjob()

+    shared.state.textinfo = "Saving"

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

 

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

@@ -340,5 +249,10 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
 

     sd_models.list_models()

 

-    print("Checkpoint saved.")

-    return ["Checkpoint saved to " + output_modelname] + [gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)]

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

+

+    print(f"Checkpoint saved to {output_modelname}.")

+    shared.state.textinfo = "Checkpoint saved"

+    shared.state.end()

+

+    return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)], "Checkpoint saved to " + output_modelname]

diff --git a/modules/generation_parameters_copypaste.py b/modules/generation_parameters_copypaste.py
index fbd91300..773c5c0e 100644
--- a/modules/generation_parameters_copypaste.py
+++ b/modules/generation_parameters_copypaste.py
@@ -1,12 +1,13 @@
 import base64

 import io

+import math

 import os

 import re

 from pathlib import Path

 

 import gradio as gr

-from modules.shared import script_path

-from modules import shared

+from modules.paths import data_path

+from modules import shared, ui_tempdir, script_callbacks

 import tempfile

 from PIL import Image

 

@@ -36,9 +37,15 @@ def quote(text):
 

 

 def image_from_url_text(filedata):

-    if type(filedata) == dict and filedata["is_file"]:

+    if filedata is None:

+        return None

+

+    if type(filedata) == list and len(filedata) > 0 and type(filedata[0]) == dict and filedata[0].get("is_file", False):

+        filedata = filedata[0]

+

+    if type(filedata) == dict and filedata.get("is_file", False):

         filename = filedata["name"]

-        is_in_right_dir = any(Path(temp_dir).resolve() in Path(filename).resolve().parents for temp_dir in shared.demo.temp_dirs)

+        is_in_right_dir = ui_tempdir.check_tmp_file(shared.demo, filename)

         assert is_in_right_dir, 'trying to open image file outside of allowed directories'

 

         return Image.open(filename)

@@ -72,8 +79,6 @@ def integrate_settings_paste_fields(component_dict):
     from modules import ui

 

     settings_map = {

-        'sd_hypernetwork': 'Hypernet',

-        'sd_hypernetwork_strength': 'Hypernet strength',

         'CLIP_stop_at_last_layers': 'Clip skip',

         'inpainting_mask_weight': 'Conditional mask weight',

         'sd_model_checkpoint': 'Model hash',

@@ -93,7 +98,7 @@ def integrate_settings_paste_fields(component_dict):
 def create_buttons(tabs_list):

     buttons = {}

     for tab in tabs_list:

-        buttons[tab] = gr.Button(f"Send to {tab}")

+        buttons[tab] = gr.Button(f"Send to {tab}", elem_id=f"{tab}_tab")

     return buttons

 

 

@@ -102,35 +107,57 @@ def bind_buttons(buttons, send_image, send_generate_info):
     bind_list.append([buttons, send_image, send_generate_info])

 

 

+def send_image_and_dimensions(x):

+    if isinstance(x, Image.Image):

+        img = x

+    else:

+        img = image_from_url_text(x)

+

+    if shared.opts.send_size and isinstance(img, Image.Image):

+        w = img.width

+        h = img.height

+    else:

+        w = gr.update()

+        h = gr.update()

+

+    return img, w, h

+

+

 def run_bind():

-    for buttons, send_image, send_generate_info in bind_list:

+    for buttons, source_image_component, send_generate_info in bind_list:

         for tab in buttons:

             button = buttons[tab]

-            if send_image and paste_fields[tab]["init_img"]:

-                if type(send_image) == gr.Gallery:

-                    button.click(

-                        fn=lambda x: image_from_url_text(x),

-                        _js="extract_image_from_gallery",

-                        inputs=[send_image],

-                        outputs=[paste_fields[tab]["init_img"]],

-                    )

+            destination_image_component = paste_fields[tab]["init_img"]

+            fields = paste_fields[tab]["fields"]

+

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

+

+            if source_image_component and destination_image_component:

+                if isinstance(source_image_component, gr.Gallery):

+                    func = send_image_and_dimensions if destination_width_component else image_from_url_text

+                    jsfunc = "extract_image_from_gallery"

                 else:

-                    button.click(

-                        fn=lambda x: x,

-                        inputs=[send_image],

-                        outputs=[paste_fields[tab]["init_img"]],

-                    )

+                    func = send_image_and_dimensions if destination_width_component else lambda x: x

+                    jsfunc = None

 

-            if send_generate_info and paste_fields[tab]["fields"] is not None:

+                button.click(

+                    fn=func,

+                    _js=jsfunc,

+                    inputs=[source_image_component],

+                    outputs=[destination_image_component, destination_width_component, destination_height_component] if destination_width_component else [destination_image_component],

+                )

+

+            if send_generate_info and fields is not None:

                 if send_generate_info in paste_fields:

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

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

                     button.click(

                         fn=lambda *x: x,

                         inputs=[field for field, name in paste_fields[send_generate_info]["fields"] if name in paste_field_names],

-                        outputs=[field for field, name in paste_fields[tab]["fields"] if name in paste_field_names],

+                        outputs=[field for field, name in fields if name in paste_field_names],

                     )

                 else:

-                    connect_paste(button, paste_fields[tab]["fields"], send_generate_info)

+                    connect_paste(button, fields, send_generate_info)

 

             button.click(

                 fn=None,

@@ -164,6 +191,39 @@ def find_hypernetwork_key(hypernet_name, hypernet_hash=None):
     return None

 

 

+def restore_old_hires_fix_params(res):

+    """for infotexts that specify old First pass size parameter, convert it into

+    width, height, and hr scale"""

+

+    firstpass_width = res.get('First pass size-1', None)

+    firstpass_height = res.get('First pass size-2', None)

+

+    if shared.opts.use_old_hires_fix_width_height:

+        hires_width = int(res.get("Hires resize-1", 0))

+        hires_height = int(res.get("Hires resize-2", 0))

+

+        if hires_width and hires_height:

+            res['Size-1'] = hires_width

+            res['Size-2'] = hires_height

+            return

+

+    if firstpass_width is None or firstpass_height is None:

+        return

+

+    firstpass_width, firstpass_height = int(firstpass_width), int(firstpass_height)

+    width = int(res.get("Size-1", 512))

+    height = int(res.get("Size-2", 512))

+

+    if firstpass_width == 0 or firstpass_height == 0:

+        from modules import processing

+        firstpass_width, firstpass_height = processing.old_hires_fix_first_pass_dimensions(width, height)

+

+    res['Size-1'] = firstpass_width

+    res['Size-2'] = firstpass_height

+    res['Hires resize-1'] = width

+    res['Hires resize-2'] = height

+

+

 def parse_generation_parameters(x: str):

     """parses generation parameters string, the one you see in text field under the picture in UI:

 ```

@@ -213,13 +273,15 @@ Steps: 20, Sampler: Euler a, CFG scale: 7, Seed: 965400086, Size: 512x512, Model
     if "Clip skip" not in res:

         res["Clip skip"] = "1"

 

-    if "Hypernet strength" not in res:

-        res["Hypernet strength"] = "1"

+    hypernet = res.get("Hypernet", None)

+    if hypernet is not None:

+        res["Prompt"] += f"""<hypernet:{hypernet}:{res.get("Hypernet strength", "1.0")}>"""

+

+    if "Hires resize-1" not in res:

+        res["Hires resize-1"] = 0

+        res["Hires resize-2"] = 0

 

-    if "Hypernet" in res:

-        hypernet_name = res["Hypernet"]

-        hypernet_hash = res.get("Hypernet hash", None)

-        res["Hypernet"] = find_hypernetwork_key(hypernet_name, hypernet_hash)

+    restore_old_hires_fix_params(res)

 

     return res

 

@@ -227,12 +289,13 @@ Steps: 20, Sampler: Euler a, CFG scale: 7, Seed: 965400086, Size: 512x512, Model
 def connect_paste(button, paste_fields, input_comp, jsfunc=None):

     def paste_func(prompt):

         if not prompt and not shared.cmd_opts.hide_ui_dir_config:

-            filename = os.path.join(script_path, "params.txt")

+            filename = os.path.join(data_path, "params.txt")

             if os.path.exists(filename):

                 with open(filename, "r", encoding="utf8") as file:

                     prompt = file.read()

 

         params = parse_generation_parameters(prompt)

+        script_callbacks.infotext_pasted_callback(prompt, params)

         res = []

 

         for output, key in paste_fields:

diff --git a/modules/gfpgan_model.py b/modules/gfpgan_model.py
index 1e2dbc32..fbe6215a 100644
--- a/modules/gfpgan_model.py
+++ b/modules/gfpgan_model.py
@@ -6,12 +6,11 @@ import facexlib
 import gfpgan

 

 import modules.face_restoration

-from modules import shared, devices, modelloader

-from modules.paths import models_path

+from modules import paths, shared, devices, modelloader

 

 model_dir = "GFPGAN"

 user_path = None

-model_path = os.path.join(models_path, model_dir)

+model_path = os.path.join(paths.models_path, model_dir)

 model_url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"

 have_gfpgan = False

 loaded_gfpgan_model = None

diff --git a/modules/hashes.py b/modules/hashes.py
new file mode 100644
index 00000000..819362a3
--- /dev/null
+++ b/modules/hashes.py
@@ -0,0 +1,87 @@
+import hashlib

+import json

+import os.path

+

+import filelock

+

+from modules.paths import data_path

+

+

+cache_filename = os.path.join(data_path, "cache.json")

+cache_data = None

+

+

+def dump_cache():

+    with filelock.FileLock(cache_filename+".lock"):

+        with open(cache_filename, "w", encoding="utf8") as file:

+            json.dump(cache_data, file, indent=4)

+

+

+def cache(subsection):

+    global cache_data

+

+    if cache_data is None:

+        with filelock.FileLock(cache_filename+".lock"):

+            if not os.path.isfile(cache_filename):

+                cache_data = {}

+            else:

+                with open(cache_filename, "r", encoding="utf8") as file:

+                    cache_data = json.load(file)

+

+    s = cache_data.get(subsection, {})

+    cache_data[subsection] = s

+

+    return s

+

+

+def calculate_sha256(filename):

+    hash_sha256 = hashlib.sha256()

+    blksize = 1024 * 1024

+

+    with open(filename, "rb") as f:

+        for chunk in iter(lambda: f.read(blksize), b""):

+            hash_sha256.update(chunk)

+

+    return hash_sha256.hexdigest()

+

+

+def sha256_from_cache(filename, title):

+    hashes = cache("hashes")

+    ondisk_mtime = os.path.getmtime(filename)

+

+    if title not in hashes:

+        return None

+

+    cached_sha256 = hashes[title].get("sha256", None)

+    cached_mtime = hashes[title].get("mtime", 0)

+

+    if ondisk_mtime > cached_mtime or cached_sha256 is None:

+        return None

+

+    return cached_sha256

+

+

+def sha256(filename, title):

+    hashes = cache("hashes")

+

+    sha256_value = sha256_from_cache(filename, title)

+    if sha256_value is not None:

+        return sha256_value

+

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

+    sha256_value = calculate_sha256(filename)

+    print(f"{sha256_value}")

+

+    hashes[title] = {

+        "mtime": os.path.getmtime(filename),

+        "sha256": sha256_value,

+    }

+

+    dump_cache()

+

+    return sha256_value

+

+

+

+

+

diff --git a/modules/hypernetworks/hypernetwork.py b/modules/hypernetworks/hypernetwork.py
index 109e8078..503534e2 100644
--- a/modules/hypernetworks/hypernetwork.py
+++ b/modules/hypernetworks/hypernetwork.py
@@ -12,8 +12,8 @@ import torch
 import tqdm

 from einops import rearrange, repeat

 from ldm.util import default

-from modules import devices, processing, sd_models, shared, sd_samplers

-from modules.textual_inversion import textual_inversion

+from modules import devices, processing, sd_models, shared, sd_samplers, hashes, sd_hijack_checkpoint

+from modules.textual_inversion import textual_inversion, logging

 from modules.textual_inversion.learn_schedule import LearnRateScheduler

 from torch import einsum

 from torch.nn.init import normal_, xavier_normal_, xavier_uniform_, kaiming_normal_, kaiming_uniform_, zeros_

@@ -25,7 +25,6 @@ from statistics import stdev, mean
 optimizer_dict = {optim_name : cls_obj for optim_name, cls_obj in inspect.getmembers(torch.optim, inspect.isclass) if optim_name != "Optimizer"}

 

 class HypernetworkModule(torch.nn.Module):

-    multiplier = 1.0

     activation_dict = {

         "linear": torch.nn.Identity,

         "relu": torch.nn.ReLU,

@@ -38,9 +37,11 @@ class HypernetworkModule(torch.nn.Module):
     activation_dict.update({cls_name.lower(): cls_obj for cls_name, cls_obj in inspect.getmembers(torch.nn.modules.activation) if inspect.isclass(cls_obj) and cls_obj.__module__ == 'torch.nn.modules.activation'})

 

     def __init__(self, dim, state_dict=None, layer_structure=None, activation_func=None, weight_init='Normal',

-                 add_layer_norm=False, use_dropout=False, activate_output=False, last_layer_dropout=False):

+                 add_layer_norm=False, activate_output=False, dropout_structure=None):

         super().__init__()

 

+        self.multiplier = 1.0

+

         assert layer_structure is not None, "layer_structure must not be None"

         assert layer_structure[0] == 1, "Multiplier Sequence should start with size 1!"

         assert layer_structure[-1] == 1, "Multiplier Sequence should end with size 1!"

@@ -63,9 +64,12 @@ class HypernetworkModule(torch.nn.Module):
             if add_layer_norm:

                 linears.append(torch.nn.LayerNorm(int(dim * layer_structure[i+1])))

 

-            # Add dropout except last layer

-            if use_dropout and (i < len(layer_structure) - 3 or last_layer_dropout and i < len(layer_structure) - 2):

-                linears.append(torch.nn.Dropout(p=0.3))

+            # Everything should be now parsed into dropout structure, and applied here.

+            # Since we only have dropouts after layers, dropout structure should start with 0 and end with 0.

+            if dropout_structure is not None and dropout_structure[i+1] > 0:

+                assert 0 < dropout_structure[i+1] < 1, "Dropout probability should be 0 or float between 0 and 1!"

+                linears.append(torch.nn.Dropout(p=dropout_structure[i+1]))

+            # Code explanation : [1, 2, 1] -> dropout is missing when last_layer_dropout is false. [1, 2, 2, 1] -> [0, 0.3, 0, 0], when its True, [0, 0.3, 0.3, 0].

 

         self.linear = torch.nn.Sequential(*linears)

 

@@ -112,7 +116,7 @@ class HypernetworkModule(torch.nn.Module):
             state_dict[to] = x

 

     def forward(self, x):

-        return x + self.linear(x) * self.multiplier

+        return x + self.linear(x) * (self.multiplier if not self.training else 1)

 

     def trainables(self):

         layer_structure = []

@@ -122,8 +126,20 @@ class HypernetworkModule(torch.nn.Module):
         return layer_structure

 

 

-def apply_strength(value=None):

-    HypernetworkModule.multiplier = value if value is not None else shared.opts.sd_hypernetwork_strength

+#param layer_structure : sequence used for length, use_dropout : controlling boolean, last_layer_dropout : for compatibility check.

+def parse_dropout_structure(layer_structure, use_dropout, last_layer_dropout):

+    if layer_structure is None:

+        layer_structure = [1, 2, 1]

+    if not use_dropout:

+        return [0] * len(layer_structure)

+    dropout_values = [0]

+    dropout_values.extend([0.3] * (len(layer_structure) - 3))

+    if last_layer_dropout:

+        dropout_values.append(0.3)

+    else:

+        dropout_values.append(0)

+    dropout_values.append(0)

+    return dropout_values

 

 

 class Hypernetwork:

@@ -143,18 +159,22 @@ class Hypernetwork:
         self.add_layer_norm = add_layer_norm

         self.use_dropout = use_dropout

         self.activate_output = activate_output

-        self.last_layer_dropout = kwargs['last_layer_dropout'] if 'last_layer_dropout' in kwargs else True

+        self.last_layer_dropout = kwargs.get('last_layer_dropout', True)

+        self.dropout_structure = kwargs.get('dropout_structure', None)

+        if self.dropout_structure is None:

+            self.dropout_structure = parse_dropout_structure(self.layer_structure, self.use_dropout, self.last_layer_dropout)

         self.optimizer_name = None

         self.optimizer_state_dict = None

+        self.optional_info = None

 

         for size in enable_sizes or []:

             self.layers[size] = (

                 HypernetworkModule(size, None, self.layer_structure, self.activation_func, self.weight_init,

-                                   self.add_layer_norm, self.use_dropout, self.activate_output, last_layer_dropout=self.last_layer_dropout),

+                                   self.add_layer_norm, self.activate_output, dropout_structure=self.dropout_structure),

                 HypernetworkModule(size, None, self.layer_structure, self.activation_func, self.weight_init,

-                                   self.add_layer_norm, self.use_dropout, self.activate_output, last_layer_dropout=self.last_layer_dropout),

+                                   self.add_layer_norm, self.activate_output, dropout_structure=self.dropout_structure),

             )

-        self.eval_mode()

+        self.eval()

 

     def weights(self):

         res = []

@@ -163,14 +183,28 @@ class Hypernetwork:
                 res += layer.parameters()

         return res

 

-    def train_mode(self):

+    def train(self, mode=True):

         for k, layers in self.layers.items():

             for layer in layers:

-                layer.train()

+                layer.train(mode=mode)

                 for param in layer.parameters():

-                    param.requires_grad = True

+                    param.requires_grad = mode

+

+    def to(self, device):

+        for k, layers in self.layers.items():

+            for layer in layers:

+                layer.to(device)

+

+        return self

 

-    def eval_mode(self):

+    def set_multiplier(self, multiplier):

+        for k, layers in self.layers.items():

+            for layer in layers:

+                layer.multiplier = multiplier

+

+        return self

+

+    def eval(self):

         for k, layers in self.layers.items():

             for layer in layers:

                 layer.eval()

@@ -190,18 +224,20 @@ class Hypernetwork:
         state_dict['activation_func'] = self.activation_func

         state_dict['is_layer_norm'] = self.add_layer_norm

         state_dict['weight_initialization'] = self.weight_init

-        state_dict['use_dropout'] = self.use_dropout

         state_dict['sd_checkpoint'] = self.sd_checkpoint

         state_dict['sd_checkpoint_name'] = self.sd_checkpoint_name

         state_dict['activate_output'] = self.activate_output

-        state_dict['last_layer_dropout'] = self.last_layer_dropout

+        state_dict['use_dropout'] = self.use_dropout

+        state_dict['dropout_structure'] = self.dropout_structure

+        state_dict['last_layer_dropout'] = (self.dropout_structure[-2] != 0) if self.dropout_structure is not None else self.last_layer_dropout

+        state_dict['optional_info'] = self.optional_info if self.optional_info else None

 

         if self.optimizer_name is not None:

             optimizer_saved_dict['optimizer_name'] = self.optimizer_name

 

         torch.save(state_dict, filename)

         if shared.opts.save_optimizer_state and self.optimizer_state_dict:

-            optimizer_saved_dict['hash'] = sd_models.model_hash(filename)

+            optimizer_saved_dict['hash'] = self.shorthash()

             optimizer_saved_dict['optimizer_state_dict'] = self.optimizer_state_dict

             torch.save(optimizer_saved_dict, filename + '.optim')

 

@@ -213,44 +249,65 @@ class Hypernetwork:
         state_dict = torch.load(filename, map_location='cpu')

 

         self.layer_structure = state_dict.get('layer_structure', [1, 2, 1])

-        print(self.layer_structure)

+        self.optional_info = state_dict.get('optional_info', None)

         self.activation_func = state_dict.get('activation_func', None)

-        print(f"Activation function is {self.activation_func}")

         self.weight_init = state_dict.get('weight_initialization', 'Normal')

-        print(f"Weight initialization is {self.weight_init}")

         self.add_layer_norm = state_dict.get('is_layer_norm', False)

-        print(f"Layer norm is set to {self.add_layer_norm}")

-        self.use_dropout = state_dict.get('use_dropout', False)

-        print(f"Dropout usage is set to {self.use_dropout}" )

+        self.dropout_structure = state_dict.get('dropout_structure', None)

+        self.use_dropout = True if self.dropout_structure is not None and any(self.dropout_structure) else state_dict.get('use_dropout', False)

         self.activate_output = state_dict.get('activate_output', True)

-        print(f"Activate last layer is set to {self.activate_output}")

         self.last_layer_dropout = state_dict.get('last_layer_dropout', False)

+        # Dropout structure should have same length as layer structure, Every digits should be in [0,1), and last digit must be 0.

+        if self.dropout_structure is None:

+            self.dropout_structure = parse_dropout_structure(self.layer_structure, self.use_dropout, self.last_layer_dropout)

+

+        if shared.opts.print_hypernet_extra:

+            if self.optional_info is not None:

+                print(f"  INFO:\n {self.optional_info}\n")

+

+            print(f"  Layer structure: {self.layer_structure}")

+            print(f"  Activation function: {self.activation_func}")

+            print(f"  Weight initialization: {self.weight_init}")

+            print(f"  Layer norm: {self.add_layer_norm}")

+            print(f"  Dropout usage: {self.use_dropout}" )

+            print(f"  Activate last layer: {self.activate_output}")

+            print(f"  Dropout structure: {self.dropout_structure}")

+

+        optimizer_saved_dict = torch.load(self.filename + '.optim', map_location='cpu') if os.path.exists(self.filename + '.optim') else {}

 

-        optimizer_saved_dict = torch.load(self.filename + '.optim', map_location = 'cpu') if os.path.exists(self.filename + '.optim') else {}

-        self.optimizer_name = optimizer_saved_dict.get('optimizer_name', 'AdamW')

-        print(f"Optimizer name is {self.optimizer_name}")

-        if sd_models.model_hash(filename) == optimizer_saved_dict.get('hash', None):

+        if self.shorthash() == optimizer_saved_dict.get('hash', None):

             self.optimizer_state_dict = optimizer_saved_dict.get('optimizer_state_dict', None)

         else:

             self.optimizer_state_dict = None

         if self.optimizer_state_dict:

-            print("Loaded existing optimizer from checkpoint")

+            self.optimizer_name = optimizer_saved_dict.get('optimizer_name', 'AdamW')

+            if shared.opts.print_hypernet_extra:

+                print("Loaded existing optimizer from checkpoint")

+                print(f"Optimizer name is {self.optimizer_name}")

         else:

-            print("No saved optimizer exists in checkpoint")

+            self.optimizer_name = "AdamW"

+            if shared.opts.print_hypernet_extra:

+                print("No saved optimizer exists in checkpoint")

 

         for size, sd in state_dict.items():

             if type(size) == int:

                 self.layers[size] = (

                     HypernetworkModule(size, sd[0], self.layer_structure, self.activation_func, self.weight_init,

-                                       self.add_layer_norm, self.use_dropout, self.activate_output, last_layer_dropout=self.last_layer_dropout),

+                                       self.add_layer_norm, self.activate_output, self.dropout_structure),

                     HypernetworkModule(size, sd[1], self.layer_structure, self.activation_func, self.weight_init,

-                                       self.add_layer_norm, self.use_dropout, self.activate_output, last_layer_dropout=self.last_layer_dropout),

+                                       self.add_layer_norm, self.activate_output, self.dropout_structure),

                 )

 

         self.name = state_dict.get('name', self.name)

         self.step = state_dict.get('step', 0)

         self.sd_checkpoint = state_dict.get('sd_checkpoint', None)

         self.sd_checkpoint_name = state_dict.get('sd_checkpoint_name', None)

+        self.eval()

+

+    def shorthash(self):

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

+

+        return sha256[0:10]

 

 

 def list_hypernetworks(path):

@@ -259,27 +316,47 @@ def list_hypernetworks(path):
         name = os.path.splitext(os.path.basename(filename))[0]

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

         if name != "None":

-            res[name + f"({sd_models.model_hash(filename)})"] = filename

+            res[name] = filename

     return res

 

 

-def load_hypernetwork(filename):

-    path = shared.hypernetworks.get(filename, None)

-    # Prevent any file named "None.pt" from being loaded.

-    if path is not None and filename != "None":

-        print(f"Loading hypernetwork {filename}")

-        try:

-            shared.loaded_hypernetwork = Hypernetwork()

-            shared.loaded_hypernetwork.load(path)

+def load_hypernetwork(name):

+    path = shared.hypernetworks.get(name, None)

 

-        except Exception:

-            print(f"Error loading hypernetwork {path}", file=sys.stderr)

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

-    else:

-        if shared.loaded_hypernetwork is not None:

-            print("Unloading hypernetwork")

+    if path is None:

+        return None

+

+    hypernetwork = Hypernetwork()

+

+    try:

+        hypernetwork.load(path)

+    except Exception:

+        print(f"Error loading hypernetwork {path}", file=sys.stderr)

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

+        return None

+

+    return hypernetwork

+

+

+def load_hypernetworks(names, multipliers=None):

+    already_loaded = {}

 

-        shared.loaded_hypernetwork = None

+    for hypernetwork in shared.loaded_hypernetworks:

+        if hypernetwork.name in names:

+            already_loaded[hypernetwork.name] = hypernetwork

+

+    shared.loaded_hypernetworks.clear()

+

+    for i, name in enumerate(names):

+        hypernetwork = already_loaded.get(name, None)

+        if hypernetwork is None:

+            hypernetwork = load_hypernetwork(name)

+

+        if hypernetwork is None:

+            continue

+

+        hypernetwork.set_multiplier(multipliers[i] if multipliers else 1.0)

+        shared.loaded_hypernetworks.append(hypernetwork)

 

 

 def find_closest_hypernetwork_name(search: str):

@@ -293,18 +370,27 @@ def find_closest_hypernetwork_name(search: str):
     return applicable[0]

 

 

-def apply_hypernetwork(hypernetwork, context, layer=None):

-    hypernetwork_layers = (hypernetwork.layers if hypernetwork is not None else {}).get(context.shape[2], None)

+def apply_single_hypernetwork(hypernetwork, context_k, context_v, layer=None):

+    hypernetwork_layers = (hypernetwork.layers if hypernetwork is not None else {}).get(context_k.shape[2], None)

 

     if hypernetwork_layers is None:

-        return context, context

+        return context_k, context_v

 

     if layer is not None:

         layer.hyper_k = hypernetwork_layers[0]

         layer.hyper_v = hypernetwork_layers[1]

 

-    context_k = hypernetwork_layers[0](context)

-    context_v = hypernetwork_layers[1](context)

+    context_k = hypernetwork_layers[0](context_k)

+    context_v = hypernetwork_layers[1](context_v)

+    return context_k, context_v

+

+

+def apply_hypernetworks(hypernetworks, context, layer=None):

+    context_k = context

+    context_v = context

+    for hypernetwork in hypernetworks:

+        context_k, context_v = apply_single_hypernetwork(hypernetwork, context_k, context_v, layer)

+

     return context_k, context_v

 

 

@@ -314,7 +400,7 @@ def attention_CrossAttention_forward(self, x, context=None, mask=None):
     q = self.to_q(x)

     context = default(context, x)

 

-    context_k, context_v = apply_hypernetwork(shared.loaded_hypernetwork, context, self)

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

     k = self.to_k(context_k)

     v = self.to_v(context_v)

 

@@ -378,9 +464,10 @@ def report_statistics(loss_info:dict):
             print(e)

 

 

-def create_hypernetwork(name, enable_sizes, overwrite_old, layer_structure=None, activation_func=None, weight_init=None, add_layer_norm=False, use_dropout=False):

+def create_hypernetwork(name, enable_sizes, overwrite_old, layer_structure=None, activation_func=None, weight_init=None, add_layer_norm=False, use_dropout=False, dropout_structure=None):

     # Remove illegal characters from name.

     name = "".join( x for x in name if (x.isalnum() or x in "._- "))

+    assert name, "Name cannot be empty!"

 

     fn = os.path.join(shared.cmd_opts.hypernetwork_dir, f"{name}.pt")

     if not overwrite_old:

@@ -389,6 +476,11 @@ def create_hypernetwork(name, enable_sizes, overwrite_old, layer_structure=None,
     if type(layer_structure) == str:

         layer_structure = [float(x.strip()) for x in layer_structure.split(",")]

 

+    if use_dropout and dropout_structure and type(dropout_structure) == str:

+        dropout_structure = [float(x.strip()) for x in dropout_structure.split(",")]

+    else:

+        dropout_structure = [0] * len(layer_structure)

+

     hypernet = modules.hypernetworks.hypernetwork.Hypernetwork(

         name=name,

         enable_sizes=[int(x) for x in enable_sizes],

@@ -397,26 +489,29 @@ def create_hypernetwork(name, enable_sizes, overwrite_old, layer_structure=None,
         weight_init=weight_init,

         add_layer_norm=add_layer_norm,

         use_dropout=use_dropout,

+        dropout_structure=dropout_structure

     )

     hypernet.save(fn)

 

     shared.reload_hypernetworks()

 

-    return fn

-

 

-def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, gradient_step, data_root, log_directory, training_width, training_height, steps, shuffle_tags, tag_drop_out, latent_sampling_method, create_image_every, save_hypernetwork_every, template_file, 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, 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

 

     save_hypernetwork_every = save_hypernetwork_every or 0

     create_image_every = create_image_every or 0

-    textual_inversion.validate_train_inputs(hypernetwork_name, learn_rate, batch_size, gradient_step, data_root, template_file, steps, save_hypernetwork_every, create_image_every, log_directory, name="hypernetwork")

+    template_file = textual_inversion.textual_inversion_templates.get(template_filename, None)

+    textual_inversion.validate_train_inputs(hypernetwork_name, learn_rate, batch_size, gradient_step, data_root, template_file, template_filename, steps, save_hypernetwork_every, create_image_every, log_directory, name="hypernetwork")

+    template_file = template_file.path

 

     path = shared.hypernetworks.get(hypernetwork_name, None)

-    shared.loaded_hypernetwork = Hypernetwork()

-    shared.loaded_hypernetwork.load(path)

+    hypernetwork = Hypernetwork()

+    hypernetwork.load(path)

+    shared.loaded_hypernetworks = [hypernetwork]

 

+    shared.state.job = "train-hypernetwork"

     shared.state.textinfo = "Initializing hypernetwork training..."

     shared.state.job_count = steps

 

@@ -438,7 +533,6 @@ def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, gradient_step,
     else:

         images_dir = None

 

-    hypernetwork = shared.loaded_hypernetwork

     checkpoint = sd_models.select_checkpoint()

 

     initial_step = hypernetwork.step or 0

@@ -447,14 +541,28 @@ def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, gradient_step,
         return hypernetwork, filename

 

     scheduler = LearnRateScheduler(learn_rate, steps, initial_step)

+    

+    clip_grad = torch.nn.utils.clip_grad_value_ if clip_grad_mode == "value" else torch.nn.utils.clip_grad_norm_ if clip_grad_mode == "norm" else None

+    if clip_grad:

+        clip_grad_sched = LearnRateScheduler(clip_grad_value, steps, initial_step, verbose=False)

+

+    if shared.opts.training_enable_tensorboard:

+        tensorboard_writer = textual_inversion.tensorboard_setup(log_directory)

 

     # dataset loading may take a while, so input validations and early returns should be done before this

     shared.state.textinfo = f"Preparing dataset from {html.escape(data_root)}..."

 

     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)

-    

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

+

+    if shared.opts.save_training_settings_to_txt:

+        saved_params = dict(

+            model_name=checkpoint.model_name, model_hash=checkpoint.shorthash, num_of_dataset_images=len(ds),

+            **{field: getattr(hypernetwork, field) for field in ['layer_structure', 'activation_func', 'weight_init', 'add_layer_norm', 'use_dropout', ]}

+        )

+        logging.save_settings_to_file(log_directory, {**saved_params, **locals()})

+

     latent_sampling_method = ds.latent_sampling_method

 

     dl = modules.textual_inversion.dataset.PersonalizedDataLoader(ds, latent_sampling_method=latent_sampling_method, batch_size=ds.batch_size, pin_memory=pin_memory)

@@ -465,9 +573,9 @@ def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, gradient_step,
         shared.parallel_processing_allowed = False

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

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

-    

+

     weights = hypernetwork.weights()

-    hypernetwork.train_mode()

+    hypernetwork.train()

 

     # Here we use optimizer from saved HN, or we can specify as UI option.

     if hypernetwork.optimizer_name in optimizer_dict:

@@ -496,6 +604,7 @@ def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, gradient_step,
     _loss_step = 0 #internal

     # size = len(ds.indexes)

     # loss_dict = defaultdict(lambda : deque(maxlen = 1024))

+    loss_logging = deque(maxlen=len(ds) * 3)  # this should be configurable parameter, this is 3 * epoch(dataset size)

     # losses = torch.zeros((size,))

     # previous_mean_losses = [0]

     # previous_mean_loss = 0

@@ -509,6 +618,8 @@ def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, gradient_step,
 

     pbar = tqdm.tqdm(total=steps - initial_step)

     try:

+        sd_hijack_checkpoint.add()

+

         for i in range((steps-initial_step) * gradient_step):

             if scheduler.finished:

                 break

@@ -524,6 +635,9 @@ def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, gradient_step,
                 if shared.state.interrupted:

                     break

 

+                if clip_grad:

+                    clip_grad_sched.step(hypernetwork.step)

+                

                 with devices.autocast():

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

                     if tag_drop_out != 0 or shuffle_tags:

@@ -538,14 +652,14 @@ def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, gradient_step,
 

                     _loss_step += loss.item()

                 scaler.scale(loss).backward()

+                

                 # go back until we reach gradient accumulation steps

                 if (j + 1) % gradient_step != 0:

                     continue

-                # print(f"grad:{weights[0].grad.detach().cpu().abs().mean().item():.7f}")

-                # scaler.unscale_(optimizer)

-                # print(f"grad:{weights[0].grad.detach().cpu().abs().mean().item():.15f}")

-                # torch.nn.utils.clip_grad_norm_(weights, max_norm=1.0)

-                # print(f"grad:{weights[0].grad.detach().cpu().abs().mean().item():.15f}")

+                loss_logging.append(_loss_step)

+                if clip_grad:

+                    clip_grad(weights, clip_grad_sched.learn_rate)

+                

                 scaler.step(optimizer)

                 scaler.update()

                 hypernetwork.step += 1

@@ -559,7 +673,8 @@ def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, gradient_step,
                 epoch_num = hypernetwork.step // steps_per_epoch

                 epoch_step = hypernetwork.step % steps_per_epoch

 

-                pbar.set_description(f"[Epoch {epoch_num}: {epoch_step+1}/{steps_per_epoch}]loss: {loss_step:.7f}")

+                description = f"Training hypernetwork [Epoch {epoch_num}: {epoch_step+1}/{steps_per_epoch}]loss: {loss_step:.7f}"

+                pbar.set_description(description)

                 if hypernetwork_dir is not None and steps_done % save_hypernetwork_every == 0:

                     # Before saving, change name to match current checkpoint.

                     hypernetwork_name_every = f'{hypernetwork_name}-{steps_done}'

@@ -570,6 +685,14 @@ def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, gradient_step,
                     save_hypernetwork(hypernetwork, checkpoint, hypernetwork_name, last_saved_file)

                     hypernetwork.optimizer_state_dict = None  # dereference it after saving, to save memory.

 

+

+

+                if shared.opts.training_enable_tensorboard:

+                    epoch_num = hypernetwork.step // len(ds)

+                    epoch_step = hypernetwork.step - (epoch_num * len(ds)) + 1

+                    mean_loss = sum(loss_logging) / len(loss_logging)

+                    textual_inversion.tensorboard_add(tensorboard_writer, loss=mean_loss, global_step=hypernetwork.step, step=epoch_step, learn_rate=scheduler.learn_rate, epoch_num=epoch_num)

+

                 textual_inversion.write_loss(log_directory, "hypernetwork_loss.csv", hypernetwork.step, steps_per_epoch, {

                     "loss": f"{loss_step:.7f}",

                     "learn_rate": scheduler.learn_rate

@@ -578,7 +701,11 @@ def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, gradient_step,
                 if images_dir is not None and steps_done % create_image_every == 0:

                     forced_filename = f'{hypernetwork_name}-{steps_done}'

                     last_saved_image = os.path.join(images_dir, forced_filename)

-                    hypernetwork.eval_mode()

+                    hypernetwork.eval()

+                    rng_state = torch.get_rng_state()

+                    cuda_rng_state = None

+                    if torch.cuda.is_available():

+                        cuda_rng_state = torch.cuda.get_rng_state_all()

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

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

 

@@ -588,6 +715,8 @@ def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, gradient_step,
                         do_not_save_samples=True,

                     )

 

+                    p.disable_extra_networks = True

+

                     if preview_from_txt2img:

                         p.prompt = preview_prompt

                         p.negative_prompt = preview_negative_prompt

@@ -611,9 +740,16 @@ def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, gradient_step,
                     if unload:

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

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

-                    hypernetwork.train_mode()

+                    torch.set_rng_state(rng_state)

+                    if torch.cuda.is_available():

+                        torch.cuda.set_rng_state_all(cuda_rng_state)

+                    hypernetwork.train()

                     if image is not None:

-                        shared.state.current_image = image

+                        shared.state.assign_current_image(image)

+                        if shared.opts.training_enable_tensorboard and shared.opts.training_tensorboard_save_images:

+                            textual_inversion.tensorboard_add_image(tensorboard_writer,

+                                                                    f"Validation at epoch {epoch_num}", image,

+                                                                    hypernetwork.step)

                         last_saved_image, last_text_info = images.save_image(image, images_dir, "", p.seed, p.prompt, shared.opts.samples_format, processed.infotexts[0], p=p, forced_filename=forced_filename, save_to_dirs=False)

                         last_saved_image += f", prompt: {preview_text}"

 

@@ -633,8 +769,11 @@ Last saved image: {html.escape(last_saved_image)}<br/>
     finally:

         pbar.leave = False

         pbar.close()

-        hypernetwork.eval_mode()

+        hypernetwork.eval()

         #report_statistics(loss_dict)

+        sd_hijack_checkpoint.remove()

+

+

 

     filename = os.path.join(shared.cmd_opts.hypernetwork_dir, f'{hypernetwork_name}.pt')

     hypernetwork.optimizer_name = optimizer_name

@@ -655,7 +794,7 @@ def save_hypernetwork(hypernetwork, checkpoint, hypernetwork_name, filename):
     old_sd_checkpoint = hypernetwork.sd_checkpoint if hasattr(hypernetwork, "sd_checkpoint") else None

     old_sd_checkpoint_name = hypernetwork.sd_checkpoint_name if hasattr(hypernetwork, "sd_checkpoint_name") else None

     try:

-        hypernetwork.sd_checkpoint = checkpoint.hash

+        hypernetwork.sd_checkpoint = checkpoint.shorthash

         hypernetwork.sd_checkpoint_name = checkpoint.model_name

         hypernetwork.name = hypernetwork_name

         hypernetwork.save(filename)

diff --git a/modules/hypernetworks/ui.py b/modules/hypernetworks/ui.py
index e7f9e593..76599f5a 100644
--- a/modules/hypernetworks/ui.py
+++ b/modules/hypernetworks/ui.py
@@ -9,15 +9,15 @@ from modules import devices, sd_hijack, shared
 not_available = ["hardswish", "multiheadattention"]

 keys = list(x for x in modules.hypernetworks.hypernetwork.HypernetworkModule.activation_dict.keys() if x not in not_available)

 

-def create_hypernetwork(name, enable_sizes, overwrite_old, layer_structure=None, activation_func=None, weight_init=None, add_layer_norm=False, use_dropout=False):

-    filename = modules.hypernetworks.hypernetwork.create_hypernetwork(name, enable_sizes, overwrite_old, layer_structure, activation_func, weight_init, add_layer_norm, use_dropout)

+

+def create_hypernetwork(name, enable_sizes, overwrite_old, layer_structure=None, activation_func=None, weight_init=None, add_layer_norm=False, use_dropout=False, dropout_structure=None):

+    filename = modules.hypernetworks.hypernetwork.create_hypernetwork(name, enable_sizes, overwrite_old, layer_structure, activation_func, weight_init, add_layer_norm, use_dropout, dropout_structure)

 

     return gr.Dropdown.update(choices=sorted([x for x in shared.hypernetworks.keys()])), f"Created: {filename}", ""

 

 

 def train_hypernetwork(*args):

-

-    initial_hypernetwork = shared.loaded_hypernetwork

+    shared.loaded_hypernetworks = []

 

     assert not shared.cmd_opts.lowvram, 'Training models with lowvram is not possible'

 

@@ -34,7 +34,6 @@ Hypernetwork saved to {html.escape(filename)}
     except Exception:

         raise

     finally:

-        shared.loaded_hypernetwork = initial_hypernetwork

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

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

         sd_hijack.apply_optimizations()

diff --git a/modules/images.py b/modules/images.py
index 31d4528d..0bc3d524 100644
--- a/modules/images.py
+++ b/modules/images.py
@@ -39,11 +39,14 @@ def image_grid(imgs, batch_size=1, rows=None):
 

     cols = math.ceil(len(imgs) / rows)

 

+    params = script_callbacks.ImageGridLoopParams(imgs, cols, rows)

+    script_callbacks.image_grid_callback(params)

+

     w, h = imgs[0].size

-    grid = Image.new('RGB', size=(cols * w, rows * h), color='black')

+    grid = Image.new('RGB', size=(params.cols * w, params.rows * h), color='black')

 

-    for i, img in enumerate(imgs):

-        grid.paste(img, box=(i % cols * w, i // cols * h))

+    for i, img in enumerate(params.imgs):

+        grid.paste(img, box=(i % params.cols * w, i // params.cols * h))

 

     return grid

 

@@ -192,7 +195,7 @@ def draw_grid_annotations(im, width, height, hor_texts, 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 = max(hor_text_heights) + line_spacing * 2

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

@@ -227,16 +230,32 @@ def draw_prompt_matrix(im, width, height, all_prompts):
     return draw_grid_annotations(im, width, height, hor_texts, ver_texts)

 

 

-def resize_image(resize_mode, im, width, height):

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

+    """

+    Resizes an image with the specified resize_mode, width, and height.

+

+    Args:

+        resize_mode: The mode to use when resizing the image.

+            0: Resize the image to the specified width and height.

+            1: Resize the image to fill the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, cropping the excess.

+            2: Resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, filling empty with data from image.

+        im: The image to resize.

+        width: The width to resize the image to.

+        height: The height to resize the image to.

+        upscaler_name: The name of the upscaler to use. If not provided, defaults to opts.upscaler_for_img2img.

+    """

+

+    upscaler_name = upscaler_name or opts.upscaler_for_img2img

+

     def resize(im, w, h):

-        if opts.upscaler_for_img2img is None or opts.upscaler_for_img2img == "None" or im.mode == 'L':

+        if upscaler_name is None or upscaler_name == "None" or im.mode == 'L':

             return im.resize((w, h), resample=LANCZOS)

 

         scale = max(w / im.width, h / im.height)

 

         if scale > 1.0:

-            upscalers = [x for x in shared.sd_upscalers if x.name == opts.upscaler_for_img2img]

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

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

 

             upscaler = upscalers[0]

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

@@ -525,6 +544,9 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
             image_to_save.save(temp_file_path, format=image_format, quality=opts.jpeg_quality, pnginfo=pnginfo_data)

 

         elif extension.lower() in (".jpg", ".jpeg", ".webp"):

+            if image_to_save.mode == 'RGBA':

+                image_to_save = image_to_save.convert("RGB")

+

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

 

             if opts.enable_pnginfo and info is not None:

@@ -583,8 +605,9 @@ def read_info_from_image(image):
         except ValueError:

             exif_comment = exif_comment.decode('utf8', errors="ignore")

 

-        items['exif comment'] = exif_comment

-        geninfo = exif_comment

+        if exif_comment:

+            items['exif comment'] = exif_comment

+            geninfo = exif_comment

 

         for field in ['jfif', 'jfif_version', 'jfif_unit', 'jfif_density', 'dpi', 'exif',

                       'loop', 'background', 'timestamp', 'duration']:

diff --git a/modules/img2img.py b/modules/img2img.py
index 81da4b13..fe9447c7 100644
--- a/modules/img2img.py
+++ b/modules/img2img.py
@@ -16,11 +16,16 @@ import modules.images as images
 import modules.scripts

 

 

-def process_batch(p, input_dir, output_dir, args):

+def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args):

     processing.fix_seed(p)

 

     images = shared.listfiles(input_dir)

 

+    inpaint_masks = shared.listfiles(inpaint_mask_dir)

+    is_inpaint_batch = inpaint_mask_dir and len(inpaint_masks) > 0

+    if is_inpaint_batch:

+        print(f"\nInpaint batch is enabled. {len(inpaint_masks)} masks found.")

+

     print(f"Will process {len(images)} images, creating {p.n_iter * p.batch_size} new images for each.")

 

     save_normally = output_dir == ''

@@ -43,6 +48,15 @@ def process_batch(p, input_dir, output_dir, args):
         img = ImageOps.exif_transpose(img)

         p.init_images = [img] * p.batch_size

 

+        if is_inpaint_batch:

+            # try to find corresponding mask for an image using simple filename matching

+            mask_image_path = os.path.join(inpaint_mask_dir, os.path.basename(image))

+            # if not found use first one ("same mask for all images" use-case)

+            if not mask_image_path in inpaint_masks:

+                mask_image_path = inpaint_masks[0]

+            mask_image = Image.open(mask_image_path)

+            p.image_mask = mask_image

+

         proc = modules.scripts.scripts_img2img.run(p, *args)

         if proc is None:

             proc = process_images(p)

@@ -59,38 +73,34 @@ def process_batch(p, input_dir, output_dir, args):
                 processed_image.save(os.path.join(output_dir, filename))

 

 

-def img2img(mode: int, prompt: str, negative_prompt: str, prompt_style: str, prompt_style2: str, init_img, init_img_with_mask, init_img_with_mask_orig, init_img_inpaint, init_mask_inpaint, mask_mode, 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, *args):

-    is_inpaint = mode == 1

-    is_batch = mode == 2

-

-    if is_inpaint:

-        # Drawn mask

-        if mask_mode == 0:

-            is_mask_sketch = isinstance(init_img_with_mask, dict)

-            is_mask_paint = not is_mask_sketch

-            if is_mask_sketch:

-                # Sketch: mask iff. not transparent

-                image, mask = init_img_with_mask["image"], init_img_with_mask["mask"]

-                alpha_mask = ImageOps.invert(image.split()[-1]).convert('L').point(lambda x: 255 if x > 0 else 0, mode='1')

-                mask = ImageChops.lighter(alpha_mask, mask.convert('L')).convert('L')

-            else:

-                # Color-sketch: mask iff. painted over

-                image = init_img_with_mask

-                orig = init_img_with_mask_orig or init_img_with_mask

-                pred = np.any(np.array(image) != np.array(orig), axis=-1)

-                mask = Image.fromarray(pred.astype(np.uint8) * 255, "L")

-                mask = ImageEnhance.Brightness(mask).enhance(1 - mask_alpha / 100)

-                blur = ImageFilter.GaussianBlur(mask_blur)

-                image = Image.composite(image.filter(blur), orig, mask.filter(blur))

-

-            image = image.convert("RGB")

-        # Uploaded mask

-        else:

-            image = init_img_inpaint

-            mask = init_mask_inpaint

-    # No mask

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

+    is_batch = mode == 5

+

+    if mode == 0:  # img2img

+        image = init_img.convert("RGB")

+        mask = None

+    elif mode == 1:  # img2img sketch

+        image = sketch.convert("RGB")

+        mask = None

+    elif mode == 2:  # inpaint

+        image, mask = init_img_with_mask["image"], init_img_with_mask["mask"]

+        alpha_mask = ImageOps.invert(image.split()[-1]).convert('L').point(lambda x: 255 if x > 0 else 0, mode='1')

+        mask = ImageChops.lighter(alpha_mask, mask.convert('L')).convert('L')

+        image = image.convert("RGB")

+    elif mode == 3:  # inpaint sketch

+        image = inpaint_color_sketch

+        orig = inpaint_color_sketch_orig or inpaint_color_sketch

+        pred = np.any(np.array(image) != np.array(orig), axis=-1)

+        mask = Image.fromarray(pred.astype(np.uint8) * 255, "L")

+        mask = ImageEnhance.Brightness(mask).enhance(1 - mask_alpha / 100)

+        blur = ImageFilter.GaussianBlur(mask_blur)

+        image = Image.composite(image.filter(blur), orig, mask.filter(blur))

+        image = image.convert("RGB")

+    elif mode == 4:  # inpaint upload mask

+        image = init_img_inpaint

+        mask = init_mask_inpaint

     else:

-        image = init_img

+        image = None

         mask = None

 

     # Use the EXIF orientation of photos taken by smartphones.

@@ -105,7 +115,7 @@ def img2img(mode: int, prompt: str, negative_prompt: str, prompt_style: str, pro
         outpath_grids=opts.outdir_grids or opts.outdir_img2img_grids,

         prompt=prompt,

         negative_prompt=negative_prompt,

-        styles=[prompt_style, prompt_style2],

+        styles=prompt_styles,

         seed=seed,

         subseed=subseed,

         subseed_strength=subseed_strength,

@@ -143,7 +153,7 @@ def img2img(mode: int, prompt: str, negative_prompt: str, prompt_style: str, pro
     if is_batch:

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

 

-        process_batch(p, img2img_batch_input_dir, img2img_batch_output_dir, args)

+        process_batch(p, img2img_batch_input_dir, img2img_batch_output_dir, img2img_batch_inpaint_mask_dir, args)

 

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

     else:

@@ -162,4 +172,4 @@ def img2img(mode: int, prompt: str, negative_prompt: str, prompt_style: str, pro
     if opts.do_not_show_images:

         processed.images = []

 

-    return processed.images, generation_info_js, plaintext_to_html(processed.info)

+    return processed.images, generation_info_js, plaintext_to_html(processed.info), plaintext_to_html(processed.comments)

diff --git a/modules/interrogate.py b/modules/interrogate.py
index 46935210..cbb80683 100644
--- a/modules/interrogate.py
+++ b/modules/interrogate.py
@@ -2,15 +2,17 @@ import os
 import sys

 import traceback

 from collections import namedtuple

+from pathlib import Path

 import re

 

 import torch

+import torch.hub

 

 from torchvision import transforms

 from torchvision.transforms.functional import InterpolationMode

 

 import modules.shared as shared

-from modules import devices, paths, lowvram, modelloader

+from modules import devices, paths, shared, lowvram, modelloader, errors

 

 blip_image_eval_size = 384

 clip_model_name = 'ViT-L/14'

@@ -19,30 +21,76 @@ Category = namedtuple("Category", ["name", "topn", "items"])
 

 re_topn = re.compile(r"\.top(\d+)\.")

 

+def category_types():

+    return [f.stem for f in Path(shared.interrogator.content_dir).glob('*.txt')]

+

+

+def download_default_clip_interrogate_categories(content_dir):

+    print("Downloading CLIP categories...")

+

+    tmpdir = content_dir + "_tmp"

+    category_types = ["artists", "flavors", "mediums", "movements"]

+

+    try:

+        os.makedirs(tmpdir)

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

+

+    except Exception as e:

+        errors.display(e, "downloading default CLIP interrogate categories")

+    finally:

+        if os.path.exists(tmpdir):

+            os.remove(tmpdir)

+

 

 class InterrogateModels:

     blip_model = None

     clip_model = None

     clip_preprocess = None

-    categories = None

     dtype = None

     running_on_cpu = None

 

     def __init__(self, content_dir):

-        self.categories = []

+        self.loaded_categories = None

+        self.skip_categories = []

+        self.content_dir = content_dir

         self.running_on_cpu = devices.device_interrogate == torch.device("cpu")

 

-        if os.path.exists(content_dir):

-            for filename in os.listdir(content_dir):

-                m = re_topn.search(filename)

-                topn = 1 if m is None else int(m.group(1))

+    def categories(self):

+        if not os.path.exists(self.content_dir):

+            download_default_clip_interrogate_categories(self.content_dir)

+

+        if self.loaded_categories is not None and self.skip_categories == shared.opts.interrogate_clip_skip_categories:

+           return self.loaded_categories

 

-                with open(os.path.join(content_dir, filename), "r", encoding="utf8") as file:

+        self.loaded_categories = []

+

+        if os.path.exists(self.content_dir):

+            self.skip_categories = shared.opts.interrogate_clip_skip_categories

+            category_types = []

+            for filename in Path(self.content_dir).glob('*.txt'):

+                category_types.append(filename.stem)

+                if filename.stem in self.skip_categories:

+                    continue

+                m = re_topn.search(filename.stem)

+                topn = 1 if m is None else int(m.group(1))

+                with open(filename, "r", encoding="utf8") as file:

                     lines = [x.strip() for x in file.readlines()]

 

-                self.categories.append(Category(name=filename, topn=topn, items=lines))

+                self.loaded_categories.append(Category(name=filename.stem, topn=topn, items=lines))

+

+        return self.loaded_categories

+

+    def create_fake_fairscale(self):

+        class FakeFairscale:

+            def checkpoint_wrapper(self):

+                pass

+

+        sys.modules["fairscale.nn.checkpoint.checkpoint_activations"] = FakeFairscale

 

     def load_blip_model(self):

+        self.create_fake_fairscale()

         import models.blip

 

         files = modelloader.load_models(

@@ -106,6 +154,8 @@ class InterrogateModels:
     def rank(self, image_features, text_array, top_count=1):

         import clip

 

+        devices.torch_gc()

+

         if shared.opts.interrogate_clip_dict_limit != 0:

             text_array = text_array[0:int(shared.opts.interrogate_clip_dict_limit)]

 

@@ -135,10 +185,10 @@ class InterrogateModels:
         return caption[0]

 

     def interrogate(self, pil_image):

-        res = None

-

+        res = ""

+        shared.state.begin()

+        shared.state.job = 'interrogate'

         try:

-

             if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:

                 lowvram.send_everything_to_cpu()

                 devices.torch_gc()

@@ -158,12 +208,7 @@ class InterrogateModels:
 

                 image_features /= image_features.norm(dim=-1, keepdim=True)

 

-                if shared.opts.interrogate_use_builtin_artists:

-                    artist = self.rank(image_features, ["by " + artist.name for artist in shared.artist_db.artists])[0]

-

-                    res += ", " + artist[0]

-

-                for name, topn, items in self.categories:

+                for name, topn, items in self.categories():

                     matches = self.rank(image_features, items, top_count=topn)

                     for match, score in matches:

                         if shared.opts.interrogate_return_ranks:

@@ -177,5 +222,6 @@ class InterrogateModels:
             res += "<error>"

 

         self.unload()

+        shared.state.end()

 

         return res

diff --git a/modules/memmon.py b/modules/memmon.py
index 9fb9b687..a7060f58 100644
--- a/modules/memmon.py
+++ b/modules/memmon.py
@@ -71,10 +71,13 @@ class MemUsageMonitor(threading.Thread):
     def read(self):
         if not self.disabled:
             free, total = torch.cuda.mem_get_info()
+            self.data["free"] = free
             self.data["total"] = total
 
             torch_stats = torch.cuda.memory_stats(self.device)
+            self.data["active"] = torch_stats["active.all.current"]
             self.data["active_peak"] = torch_stats["active_bytes.all.peak"]
+            self.data["reserved"] = torch_stats["reserved_bytes.all.current"]
             self.data["reserved_peak"] = torch_stats["reserved_bytes.all.peak"]
             self.data["system_peak"] = total - self.data["min_free"]
 
diff --git a/modules/modelloader.py b/modules/modelloader.py
index e647f6fa..e9aa514e 100644
--- a/modules/modelloader.py
+++ b/modules/modelloader.py
@@ -10,7 +10,7 @@ from modules.upscaler import Upscaler
 from modules.paths import script_path, models_path
 
 
-def load_models(model_path: str, model_url: str = None, command_path: str = None, ext_filter=None, download_name=None) -> list:
+def load_models(model_path: str, model_url: str = None, command_path: str = None, ext_filter=None, download_name=None, ext_blacklist=None) -> list:
     """
     A one-and done loader to try finding the desired models in specified directories.
 
@@ -45,6 +45,8 @@ 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 ext_blacklist is not None and any([full_path.endswith(x) for x in ext_blacklist]):
+                        continue
                     if len(ext_filter) != 0:
                         model_name, extension = os.path.splitext(file)
                         if extension not in ext_filter:
@@ -123,6 +125,23 @@ def move_files(src_path: str, dest_path: str, ext_filter: str = None):
         pass
 
 
+builtin_upscaler_classes = []
+forbidden_upscaler_classes = set()
+
+
+def list_builtin_upscalers():
+    load_upscalers()
+
+    builtin_upscaler_classes.clear()
+    builtin_upscaler_classes.extend(Upscaler.__subclasses__())
+
+
+def forbid_loaded_nonbuiltin_upscalers():
+    for cls in Upscaler.__subclasses__():
+        if cls not in builtin_upscaler_classes:
+            forbidden_upscaler_classes.add(cls)
+
+
 def load_upscalers():
     # We can only do this 'magic' method to dynamically load upscalers if they are referenced,
     # so we'll try to import any _model.py files before looking in __subclasses__
@@ -139,6 +158,9 @@ def load_upscalers():
     datas = []
     commandline_options = vars(shared.cmd_opts)
     for cls in Upscaler.__subclasses__():
+        if cls in forbidden_upscaler_classes:
+            continue
+
         name = cls.__name__
         cmd_name = f"{name.lower().replace('upscaler', '')}_models_path"
         scaler = cls(commandline_options.get(cmd_name, None))
diff --git a/modules/models/diffusion/ddpm_edit.py b/modules/models/diffusion/ddpm_edit.py
new file mode 100644
index 00000000..f3d49c44
--- /dev/null
+++ b/modules/models/diffusion/ddpm_edit.py
@@ -0,0 +1,1459 @@
+"""
+wild mixture of
+https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
+https://github.com/openai/improved-diffusion/blob/e94489283bb876ac1477d5dd7709bbbd2d9902ce/improved_diffusion/gaussian_diffusion.py
+https://github.com/CompVis/taming-transformers
+-- merci
+"""
+
+# File modified by authors of InstructPix2Pix from original (https://github.com/CompVis/stable-diffusion).
+# See more details in LICENSE.
+
+import torch
+import torch.nn as nn
+import numpy as np
+import pytorch_lightning as pl
+from torch.optim.lr_scheduler import LambdaLR
+from einops import rearrange, repeat
+from contextlib import contextmanager
+from functools import partial
+from tqdm import tqdm
+from torchvision.utils import make_grid
+from pytorch_lightning.utilities.distributed import rank_zero_only
+
+from ldm.util import log_txt_as_img, exists, default, ismap, isimage, mean_flat, count_params, instantiate_from_config
+from ldm.modules.ema import LitEma
+from ldm.modules.distributions.distributions import normal_kl, DiagonalGaussianDistribution
+from ldm.models.autoencoder import VQModelInterface, IdentityFirstStage, AutoencoderKL
+from ldm.modules.diffusionmodules.util import make_beta_schedule, extract_into_tensor, noise_like
+from ldm.models.diffusion.ddim import DDIMSampler
+
+
+__conditioning_keys__ = {'concat': 'c_concat',
+                         'crossattn': 'c_crossattn',
+                         'adm': 'y'}
+
+
+def disabled_train(self, mode=True):
+    """Overwrite model.train with this function to make sure train/eval mode
+    does not change anymore."""
+    return self
+
+
+def uniform_on_device(r1, r2, shape, device):
+    return (r1 - r2) * torch.rand(*shape, device=device) + r2
+
+
+class DDPM(pl.LightningModule):
+    # classic DDPM with Gaussian diffusion, in image space
+    def __init__(self,
+                 unet_config,
+                 timesteps=1000,
+                 beta_schedule="linear",
+                 loss_type="l2",
+                 ckpt_path=None,
+                 ignore_keys=[],
+                 load_only_unet=False,
+                 monitor="val/loss",
+                 use_ema=True,
+                 first_stage_key="image",
+                 image_size=256,
+                 channels=3,
+                 log_every_t=100,
+                 clip_denoised=True,
+                 linear_start=1e-4,
+                 linear_end=2e-2,
+                 cosine_s=8e-3,
+                 given_betas=None,
+                 original_elbo_weight=0.,
+                 v_posterior=0.,  # weight for choosing posterior variance as sigma = (1-v) * beta_tilde + v * beta
+                 l_simple_weight=1.,
+                 conditioning_key=None,
+                 parameterization="eps",  # all assuming fixed variance schedules
+                 scheduler_config=None,
+                 use_positional_encodings=False,
+                 learn_logvar=False,
+                 logvar_init=0.,
+                 load_ema=True,
+                 ):
+        super().__init__()
+        assert parameterization in ["eps", "x0"], 'currently only supporting "eps" and "x0"'
+        self.parameterization = parameterization
+        print(f"{self.__class__.__name__}: Running in {self.parameterization}-prediction mode")
+        self.cond_stage_model = None
+        self.clip_denoised = clip_denoised
+        self.log_every_t = log_every_t
+        self.first_stage_key = first_stage_key
+        self.image_size = image_size  # try conv?
+        self.channels = channels
+        self.use_positional_encodings = use_positional_encodings
+        self.model = DiffusionWrapper(unet_config, conditioning_key)
+        count_params(self.model, verbose=True)
+        self.use_ema = use_ema
+
+        self.use_scheduler = scheduler_config is not None
+        if self.use_scheduler:
+            self.scheduler_config = scheduler_config
+
+        self.v_posterior = v_posterior
+        self.original_elbo_weight = original_elbo_weight
+        self.l_simple_weight = l_simple_weight
+
+        if monitor is not None:
+            self.monitor = monitor
+
+        if self.use_ema and load_ema:
+            self.model_ema = LitEma(self.model)
+            print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
+
+        if ckpt_path is not None:
+            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys, only_model=load_only_unet)
+
+            # If initialing from EMA-only checkpoint, create EMA model after loading.
+            if self.use_ema and not load_ema:
+                self.model_ema = LitEma(self.model)
+                print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
+
+        self.register_schedule(given_betas=given_betas, beta_schedule=beta_schedule, timesteps=timesteps,
+                               linear_start=linear_start, linear_end=linear_end, cosine_s=cosine_s)
+
+        self.loss_type = loss_type
+
+        self.learn_logvar = learn_logvar
+        self.logvar = torch.full(fill_value=logvar_init, size=(self.num_timesteps,))
+        if self.learn_logvar:
+            self.logvar = nn.Parameter(self.logvar, requires_grad=True)
+
+
+    def register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000,
+                          linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
+        if exists(given_betas):
+            betas = given_betas
+        else:
+            betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end,
+                                       cosine_s=cosine_s)
+        alphas = 1. - betas
+        alphas_cumprod = np.cumprod(alphas, axis=0)
+        alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
+
+        timesteps, = betas.shape
+        self.num_timesteps = int(timesteps)
+        self.linear_start = linear_start
+        self.linear_end = linear_end
+        assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep'
+
+        to_torch = partial(torch.tensor, dtype=torch.float32)
+
+        self.register_buffer('betas', to_torch(betas))
+        self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
+        self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
+        self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
+        self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
+        self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
+        self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))
+
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+        posterior_variance = (1 - self.v_posterior) * betas * (1. - alphas_cumprod_prev) / (
+                    1. - alphas_cumprod) + self.v_posterior * betas
+        # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
+        self.register_buffer('posterior_variance', to_torch(posterior_variance))
+        # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
+        self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20))))
+        self.register_buffer('posterior_mean_coef1', to_torch(
+            betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)))
+        self.register_buffer('posterior_mean_coef2', to_torch(
+            (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod)))
+
+        if self.parameterization == "eps":
+            lvlb_weights = self.betas ** 2 / (
+                        2 * self.posterior_variance * to_torch(alphas) * (1 - self.alphas_cumprod))
+        elif self.parameterization == "x0":
+            lvlb_weights = 0.5 * np.sqrt(torch.Tensor(alphas_cumprod)) / (2. * 1 - torch.Tensor(alphas_cumprod))
+        else:
+            raise NotImplementedError("mu not supported")
+        # TODO how to choose this term
+        lvlb_weights[0] = lvlb_weights[1]
+        self.register_buffer('lvlb_weights', lvlb_weights, persistent=False)
+        assert not torch.isnan(self.lvlb_weights).all()
+
+    @contextmanager
+    def ema_scope(self, context=None):
+        if self.use_ema:
+            self.model_ema.store(self.model.parameters())
+            self.model_ema.copy_to(self.model)
+            if context is not None:
+                print(f"{context}: Switched to EMA weights")
+        try:
+            yield None
+        finally:
+            if self.use_ema:
+                self.model_ema.restore(self.model.parameters())
+                if context is not None:
+                    print(f"{context}: Restored training weights")
+
+    def init_from_ckpt(self, path, ignore_keys=list(), only_model=False):
+        sd = torch.load(path, map_location="cpu")
+        if "state_dict" in list(sd.keys()):
+            sd = sd["state_dict"]
+        keys = list(sd.keys())
+
+        # Our model adds additional channels to the first layer to condition on an input image.
+        # For the first layer, copy existing channel weights and initialize new channel weights to zero.
+        input_keys = [
+            "model.diffusion_model.input_blocks.0.0.weight",
+            "model_ema.diffusion_modelinput_blocks00weight",
+        ]
+
+        self_sd = self.state_dict()
+        for input_key in input_keys:
+            if input_key not in sd or input_key not in self_sd:
+                continue
+
+            input_weight = self_sd[input_key]
+
+            if input_weight.size() != sd[input_key].size():
+                print(f"Manual init: {input_key}")
+                input_weight.zero_()
+                input_weight[:, :4, :, :].copy_(sd[input_key])
+                ignore_keys.append(input_key)
+
+        for k in keys:
+            for ik in ignore_keys:
+                if k.startswith(ik):
+                    print("Deleting key {} from state_dict.".format(k))
+                    del sd[k]
+        missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict(
+            sd, strict=False)
+        print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
+        if len(missing) > 0:
+            print(f"Missing Keys: {missing}")
+        if len(unexpected) > 0:
+            print(f"Unexpected Keys: {unexpected}")
+
+    def q_mean_variance(self, x_start, t):
+        """
+        Get the distribution q(x_t | x_0).
+        :param x_start: the [N x C x ...] tensor of noiseless inputs.
+        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+        :return: A tuple (mean, variance, log_variance), all of x_start's shape.
+        """
+        mean = (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start)
+        variance = extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)
+        log_variance = extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape)
+        return mean, variance, log_variance
+
+    def predict_start_from_noise(self, x_t, t, noise):
+        return (
+                extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
+                extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
+        )
+
+    def q_posterior(self, x_start, x_t, t):
+        posterior_mean = (
+                extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start +
+                extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        posterior_variance = extract_into_tensor(self.posterior_variance, t, x_t.shape)
+        posterior_log_variance_clipped = extract_into_tensor(self.posterior_log_variance_clipped, t, x_t.shape)
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+    def p_mean_variance(self, x, t, clip_denoised: bool):
+        model_out = self.model(x, t)
+        if self.parameterization == "eps":
+            x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)
+        elif self.parameterization == "x0":
+            x_recon = model_out
+        if clip_denoised:
+            x_recon.clamp_(-1., 1.)
+
+        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
+        return model_mean, posterior_variance, posterior_log_variance
+
+    @torch.no_grad()
+    def p_sample(self, x, t, clip_denoised=True, repeat_noise=False):
+        b, *_, device = *x.shape, x.device
+        model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, clip_denoised=clip_denoised)
+        noise = noise_like(x.shape, device, repeat_noise)
+        # no noise when t == 0
+        nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
+        return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
+
+    @torch.no_grad()
+    def p_sample_loop(self, shape, return_intermediates=False):
+        device = self.betas.device
+        b = shape[0]
+        img = torch.randn(shape, device=device)
+        intermediates = [img]
+        for i in tqdm(reversed(range(0, self.num_timesteps)), desc='Sampling t', total=self.num_timesteps):
+            img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long),
+                                clip_denoised=self.clip_denoised)
+            if i % self.log_every_t == 0 or i == self.num_timesteps - 1:
+                intermediates.append(img)
+        if return_intermediates:
+            return img, intermediates
+        return img
+
+    @torch.no_grad()
+    def sample(self, batch_size=16, return_intermediates=False):
+        image_size = self.image_size
+        channels = self.channels
+        return self.p_sample_loop((batch_size, channels, image_size, image_size),
+                                  return_intermediates=return_intermediates)
+
+    def q_sample(self, x_start, t, noise=None):
+        noise = default(noise, lambda: torch.randn_like(x_start))
+        return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
+                extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise)
+
+    def get_loss(self, pred, target, mean=True):
+        if self.loss_type == 'l1':
+            loss = (target - pred).abs()
+            if mean:
+                loss = loss.mean()
+        elif self.loss_type == 'l2':
+            if mean:
+                loss = torch.nn.functional.mse_loss(target, pred)
+            else:
+                loss = torch.nn.functional.mse_loss(target, pred, reduction='none')
+        else:
+            raise NotImplementedError("unknown loss type '{loss_type}'")
+
+        return loss
+
+    def p_losses(self, x_start, t, noise=None):
+        noise = default(noise, lambda: torch.randn_like(x_start))
+        x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
+        model_out = self.model(x_noisy, t)
+
+        loss_dict = {}
+        if self.parameterization == "eps":
+            target = noise
+        elif self.parameterization == "x0":
+            target = x_start
+        else:
+            raise NotImplementedError(f"Paramterization {self.parameterization} not yet supported")
+
+        loss = self.get_loss(model_out, target, mean=False).mean(dim=[1, 2, 3])
+
+        log_prefix = 'train' if self.training else 'val'
+
+        loss_dict.update({f'{log_prefix}/loss_simple': loss.mean()})
+        loss_simple = loss.mean() * self.l_simple_weight
+
+        loss_vlb = (self.lvlb_weights[t] * loss).mean()
+        loss_dict.update({f'{log_prefix}/loss_vlb': loss_vlb})
+
+        loss = loss_simple + self.original_elbo_weight * loss_vlb
+
+        loss_dict.update({f'{log_prefix}/loss': loss})
+
+        return loss, loss_dict
+
+    def forward(self, x, *args, **kwargs):
+        # b, c, h, w, device, img_size, = *x.shape, x.device, self.image_size
+        # assert h == img_size and w == img_size, f'height and width of image must be {img_size}'
+        t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long()
+        return self.p_losses(x, t, *args, **kwargs)
+
+    def get_input(self, batch, k):
+        return batch[k]
+
+    def shared_step(self, batch):
+        x = self.get_input(batch, self.first_stage_key)
+        loss, loss_dict = self(x)
+        return loss, loss_dict
+
+    def training_step(self, batch, batch_idx):
+        loss, loss_dict = self.shared_step(batch)
+
+        self.log_dict(loss_dict, prog_bar=True,
+                      logger=True, on_step=True, on_epoch=True)
+
+        self.log("global_step", self.global_step,
+                 prog_bar=True, logger=True, on_step=True, on_epoch=False)
+
+        if self.use_scheduler:
+            lr = self.optimizers().param_groups[0]['lr']
+            self.log('lr_abs', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False)
+
+        return loss
+
+    @torch.no_grad()
+    def validation_step(self, batch, batch_idx):
+        _, loss_dict_no_ema = self.shared_step(batch)
+        with self.ema_scope():
+            _, loss_dict_ema = self.shared_step(batch)
+            loss_dict_ema = {key + '_ema': loss_dict_ema[key] for key in loss_dict_ema}
+        self.log_dict(loss_dict_no_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True)
+        self.log_dict(loss_dict_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True)
+
+    def on_train_batch_end(self, *args, **kwargs):
+        if self.use_ema:
+            self.model_ema(self.model)
+
+    def _get_rows_from_list(self, samples):
+        n_imgs_per_row = len(samples)
+        denoise_grid = rearrange(samples, 'n b c h w -> b n c h w')
+        denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w')
+        denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row)
+        return denoise_grid
+
+    @torch.no_grad()
+    def log_images(self, batch, N=8, n_row=2, sample=True, return_keys=None, **kwargs):
+        log = dict()
+        x = self.get_input(batch, self.first_stage_key)
+        N = min(x.shape[0], N)
+        n_row = min(x.shape[0], n_row)
+        x = x.to(self.device)[:N]
+        log["inputs"] = x
+
+        # get diffusion row
+        diffusion_row = list()
+        x_start = x[:n_row]
+
+        for t in range(self.num_timesteps):
+            if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
+                t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
+                t = t.to(self.device).long()
+                noise = torch.randn_like(x_start)
+                x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
+                diffusion_row.append(x_noisy)
+
+        log["diffusion_row"] = self._get_rows_from_list(diffusion_row)
+
+        if sample:
+            # get denoise row
+            with self.ema_scope("Plotting"):
+                samples, denoise_row = self.sample(batch_size=N, return_intermediates=True)
+
+            log["samples"] = samples
+            log["denoise_row"] = self._get_rows_from_list(denoise_row)
+
+        if return_keys:
+            if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0:
+                return log
+            else:
+                return {key: log[key] for key in return_keys}
+        return log
+
+    def configure_optimizers(self):
+        lr = self.learning_rate
+        params = list(self.model.parameters())
+        if self.learn_logvar:
+            params = params + [self.logvar]
+        opt = torch.optim.AdamW(params, lr=lr)
+        return opt
+
+
+class LatentDiffusion(DDPM):
+    """main class"""
+    def __init__(self,
+                 first_stage_config,
+                 cond_stage_config,
+                 num_timesteps_cond=None,
+                 cond_stage_key="image",
+                 cond_stage_trainable=False,
+                 concat_mode=True,
+                 cond_stage_forward=None,
+                 conditioning_key=None,
+                 scale_factor=1.0,
+                 scale_by_std=False,
+                 load_ema=True,
+                 *args, **kwargs):
+        self.num_timesteps_cond = default(num_timesteps_cond, 1)
+        self.scale_by_std = scale_by_std
+        assert self.num_timesteps_cond <= kwargs['timesteps']
+        # for backwards compatibility after implementation of DiffusionWrapper
+        if conditioning_key is None:
+            conditioning_key = 'concat' if concat_mode else 'crossattn'
+        if cond_stage_config == '__is_unconditional__':
+            conditioning_key = None
+        ckpt_path = kwargs.pop("ckpt_path", None)
+        ignore_keys = kwargs.pop("ignore_keys", [])
+        super().__init__(conditioning_key=conditioning_key, *args, load_ema=load_ema, **kwargs)
+        self.concat_mode = concat_mode
+        self.cond_stage_trainable = cond_stage_trainable
+        self.cond_stage_key = cond_stage_key
+        try:
+            self.num_downs = len(first_stage_config.params.ddconfig.ch_mult) - 1
+        except:
+            self.num_downs = 0
+        if not scale_by_std:
+            self.scale_factor = scale_factor
+        else:
+            self.register_buffer('scale_factor', torch.tensor(scale_factor))
+        self.instantiate_first_stage(first_stage_config)
+        self.instantiate_cond_stage(cond_stage_config)
+        self.cond_stage_forward = cond_stage_forward
+        self.clip_denoised = False
+        self.bbox_tokenizer = None
+
+        self.restarted_from_ckpt = False
+        if ckpt_path is not None:
+            self.init_from_ckpt(ckpt_path, ignore_keys)
+            self.restarted_from_ckpt = True
+
+            if self.use_ema and not load_ema:
+                self.model_ema = LitEma(self.model)
+                print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
+
+    def make_cond_schedule(self, ):
+        self.cond_ids = torch.full(size=(self.num_timesteps,), fill_value=self.num_timesteps - 1, dtype=torch.long)
+        ids = torch.round(torch.linspace(0, self.num_timesteps - 1, self.num_timesteps_cond)).long()
+        self.cond_ids[:self.num_timesteps_cond] = ids
+
+    @rank_zero_only
+    @torch.no_grad()
+    def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
+        # only for very first batch
+        if self.scale_by_std and self.current_epoch == 0 and self.global_step == 0 and batch_idx == 0 and not self.restarted_from_ckpt:
+            assert self.scale_factor == 1., 'rather not use custom rescaling and std-rescaling simultaneously'
+            # set rescale weight to 1./std of encodings
+            print("### USING STD-RESCALING ###")
+            x = super().get_input(batch, self.first_stage_key)
+            x = x.to(self.device)
+            encoder_posterior = self.encode_first_stage(x)
+            z = self.get_first_stage_encoding(encoder_posterior).detach()
+            del self.scale_factor
+            self.register_buffer('scale_factor', 1. / z.flatten().std())
+            print(f"setting self.scale_factor to {self.scale_factor}")
+            print("### USING STD-RESCALING ###")
+
+    def register_schedule(self,
+                          given_betas=None, beta_schedule="linear", timesteps=1000,
+                          linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
+        super().register_schedule(given_betas, beta_schedule, timesteps, linear_start, linear_end, cosine_s)
+
+        self.shorten_cond_schedule = self.num_timesteps_cond > 1
+        if self.shorten_cond_schedule:
+            self.make_cond_schedule()
+
+    def instantiate_first_stage(self, config):
+        model = instantiate_from_config(config)
+        self.first_stage_model = model.eval()
+        self.first_stage_model.train = disabled_train
+        for param in self.first_stage_model.parameters():
+            param.requires_grad = False
+
+    def instantiate_cond_stage(self, config):
+        if not self.cond_stage_trainable:
+            if config == "__is_first_stage__":
+                print("Using first stage also as cond stage.")
+                self.cond_stage_model = self.first_stage_model
+            elif config == "__is_unconditional__":
+                print(f"Training {self.__class__.__name__} as an unconditional model.")
+                self.cond_stage_model = None
+                # self.be_unconditional = True
+            else:
+                model = instantiate_from_config(config)
+                self.cond_stage_model = model.eval()
+                self.cond_stage_model.train = disabled_train
+                for param in self.cond_stage_model.parameters():
+                    param.requires_grad = False
+        else:
+            assert config != '__is_first_stage__'
+            assert config != '__is_unconditional__'
+            model = instantiate_from_config(config)
+            self.cond_stage_model = model
+
+    def _get_denoise_row_from_list(self, samples, desc='', force_no_decoder_quantization=False):
+        denoise_row = []
+        for zd in tqdm(samples, desc=desc):
+            denoise_row.append(self.decode_first_stage(zd.to(self.device),
+                                                            force_not_quantize=force_no_decoder_quantization))
+        n_imgs_per_row = len(denoise_row)
+        denoise_row = torch.stack(denoise_row)  # n_log_step, n_row, C, H, W
+        denoise_grid = rearrange(denoise_row, 'n b c h w -> b n c h w')
+        denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w')
+        denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row)
+        return denoise_grid
+
+    def get_first_stage_encoding(self, encoder_posterior):
+        if isinstance(encoder_posterior, DiagonalGaussianDistribution):
+            z = encoder_posterior.sample()
+        elif isinstance(encoder_posterior, torch.Tensor):
+            z = encoder_posterior
+        else:
+            raise NotImplementedError(f"encoder_posterior of type '{type(encoder_posterior)}' not yet implemented")
+        return self.scale_factor * z
+
+    def get_learned_conditioning(self, c):
+        if self.cond_stage_forward is None:
+            if hasattr(self.cond_stage_model, 'encode') and callable(self.cond_stage_model.encode):
+                c = self.cond_stage_model.encode(c)
+                if isinstance(c, DiagonalGaussianDistribution):
+                    c = c.mode()
+            else:
+                c = self.cond_stage_model(c)
+        else:
+            assert hasattr(self.cond_stage_model, self.cond_stage_forward)
+            c = getattr(self.cond_stage_model, self.cond_stage_forward)(c)
+        return c
+
+    def meshgrid(self, h, w):
+        y = torch.arange(0, h).view(h, 1, 1).repeat(1, w, 1)
+        x = torch.arange(0, w).view(1, w, 1).repeat(h, 1, 1)
+
+        arr = torch.cat([y, x], dim=-1)
+        return arr
+
+    def delta_border(self, h, w):
+        """
+        :param h: height
+        :param w: width
+        :return: normalized distance to image border,
+         wtith min distance = 0 at border and max dist = 0.5 at image center
+        """
+        lower_right_corner = torch.tensor([h - 1, w - 1]).view(1, 1, 2)
+        arr = self.meshgrid(h, w) / lower_right_corner
+        dist_left_up = torch.min(arr, dim=-1, keepdims=True)[0]
+        dist_right_down = torch.min(1 - arr, dim=-1, keepdims=True)[0]
+        edge_dist = torch.min(torch.cat([dist_left_up, dist_right_down], dim=-1), dim=-1)[0]
+        return edge_dist
+
+    def get_weighting(self, h, w, Ly, Lx, device):
+        weighting = self.delta_border(h, w)
+        weighting = torch.clip(weighting, self.split_input_params["clip_min_weight"],
+                               self.split_input_params["clip_max_weight"], )
+        weighting = weighting.view(1, h * w, 1).repeat(1, 1, Ly * Lx).to(device)
+
+        if self.split_input_params["tie_braker"]:
+            L_weighting = self.delta_border(Ly, Lx)
+            L_weighting = torch.clip(L_weighting,
+                                     self.split_input_params["clip_min_tie_weight"],
+                                     self.split_input_params["clip_max_tie_weight"])
+
+            L_weighting = L_weighting.view(1, 1, Ly * Lx).to(device)
+            weighting = weighting * L_weighting
+        return weighting
+
+    def get_fold_unfold(self, x, kernel_size, stride, uf=1, df=1):  # todo load once not every time, shorten code
+        """
+        :param x: img of size (bs, c, h, w)
+        :return: n img crops of size (n, bs, c, kernel_size[0], kernel_size[1])
+        """
+        bs, nc, h, w = x.shape
+
+        # number of crops in image
+        Ly = (h - kernel_size[0]) // stride[0] + 1
+        Lx = (w - kernel_size[1]) // stride[1] + 1
+
+        if uf == 1 and df == 1:
+            fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
+            unfold = torch.nn.Unfold(**fold_params)
+
+            fold = torch.nn.Fold(output_size=x.shape[2:], **fold_params)
+
+            weighting = self.get_weighting(kernel_size[0], kernel_size[1], Ly, Lx, x.device).to(x.dtype)
+            normalization = fold(weighting).view(1, 1, h, w)  # normalizes the overlap
+            weighting = weighting.view((1, 1, kernel_size[0], kernel_size[1], Ly * Lx))
+
+        elif uf > 1 and df == 1:
+            fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
+            unfold = torch.nn.Unfold(**fold_params)
+
+            fold_params2 = dict(kernel_size=(kernel_size[0] * uf, kernel_size[0] * uf),
+                                dilation=1, padding=0,
+                                stride=(stride[0] * uf, stride[1] * uf))
+            fold = torch.nn.Fold(output_size=(x.shape[2] * uf, x.shape[3] * uf), **fold_params2)
+
+            weighting = self.get_weighting(kernel_size[0] * uf, kernel_size[1] * uf, Ly, Lx, x.device).to(x.dtype)
+            normalization = fold(weighting).view(1, 1, h * uf, w * uf)  # normalizes the overlap
+            weighting = weighting.view((1, 1, kernel_size[0] * uf, kernel_size[1] * uf, Ly * Lx))
+
+        elif df > 1 and uf == 1:
+            fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
+            unfold = torch.nn.Unfold(**fold_params)
+
+            fold_params2 = dict(kernel_size=(kernel_size[0] // df, kernel_size[0] // df),
+                                dilation=1, padding=0,
+                                stride=(stride[0] // df, stride[1] // df))
+            fold = torch.nn.Fold(output_size=(x.shape[2] // df, x.shape[3] // df), **fold_params2)
+
+            weighting = self.get_weighting(kernel_size[0] // df, kernel_size[1] // df, Ly, Lx, x.device).to(x.dtype)
+            normalization = fold(weighting).view(1, 1, h // df, w // df)  # normalizes the overlap
+            weighting = weighting.view((1, 1, kernel_size[0] // df, kernel_size[1] // df, Ly * Lx))
+
+        else:
+            raise NotImplementedError
+
+        return fold, unfold, normalization, weighting
+
+    @torch.no_grad()
+    def get_input(self, batch, k, return_first_stage_outputs=False, force_c_encode=False,
+                  cond_key=None, return_original_cond=False, bs=None, uncond=0.05):
+        x = super().get_input(batch, k)
+        if bs is not None:
+            x = x[:bs]
+        x = x.to(self.device)
+        encoder_posterior = self.encode_first_stage(x)
+        z = self.get_first_stage_encoding(encoder_posterior).detach()
+        cond_key = cond_key or self.cond_stage_key
+        xc = super().get_input(batch, cond_key)
+        if bs is not None:
+            xc["c_crossattn"] = xc["c_crossattn"][:bs]
+            xc["c_concat"] = xc["c_concat"][:bs]
+        cond = {}
+
+        # To support classifier-free guidance, randomly drop out only text conditioning 5%, only image conditioning 5%, and both 5%.
+        random = torch.rand(x.size(0), device=x.device)
+        prompt_mask = rearrange(random < 2 * uncond, "n -> n 1 1")
+        input_mask = 1 - rearrange((random >= uncond).float() * (random < 3 * uncond).float(), "n -> n 1 1 1")
+
+        null_prompt = self.get_learned_conditioning([""])
+        cond["c_crossattn"] = [torch.where(prompt_mask, null_prompt, self.get_learned_conditioning(xc["c_crossattn"]).detach())]
+        cond["c_concat"] = [input_mask * self.encode_first_stage((xc["c_concat"].to(self.device))).mode().detach()]
+
+        out = [z, cond]
+        if return_first_stage_outputs:
+            xrec = self.decode_first_stage(z)
+            out.extend([x, xrec])
+        if return_original_cond:
+            out.append(xc)
+        return out
+
+    @torch.no_grad()
+    def decode_first_stage(self, z, predict_cids=False, force_not_quantize=False):
+        if predict_cids:
+            if z.dim() == 4:
+                z = torch.argmax(z.exp(), dim=1).long()
+            z = self.first_stage_model.quantize.get_codebook_entry(z, shape=None)
+            z = rearrange(z, 'b h w c -> b c h w').contiguous()
+
+        z = 1. / self.scale_factor * z
+
+        if hasattr(self, "split_input_params"):
+            if self.split_input_params["patch_distributed_vq"]:
+                ks = self.split_input_params["ks"]  # eg. (128, 128)
+                stride = self.split_input_params["stride"]  # eg. (64, 64)
+                uf = self.split_input_params["vqf"]
+                bs, nc, h, w = z.shape
+                if ks[0] > h or ks[1] > w:
+                    ks = (min(ks[0], h), min(ks[1], w))
+                    print("reducing Kernel")
+
+                if stride[0] > h or stride[1] > w:
+                    stride = (min(stride[0], h), min(stride[1], w))
+                    print("reducing stride")
+
+                fold, unfold, normalization, weighting = self.get_fold_unfold(z, ks, stride, uf=uf)
+
+                z = unfold(z)  # (bn, nc * prod(**ks), L)
+                # 1. Reshape to img shape
+                z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1]))  # (bn, nc, ks[0], ks[1], L )
+
+                # 2. apply model loop over last dim
+                if isinstance(self.first_stage_model, VQModelInterface):
+                    output_list = [self.first_stage_model.decode(z[:, :, :, :, i],
+                                                                 force_not_quantize=predict_cids or force_not_quantize)
+                                   for i in range(z.shape[-1])]
+                else:
+
+                    output_list = [self.first_stage_model.decode(z[:, :, :, :, i])
+                                   for i in range(z.shape[-1])]
+
+                o = torch.stack(output_list, axis=-1)  # # (bn, nc, ks[0], ks[1], L)
+                o = o * weighting
+                # Reverse 1. reshape to img shape
+                o = o.view((o.shape[0], -1, o.shape[-1]))  # (bn, nc * ks[0] * ks[1], L)
+                # stitch crops together
+                decoded = fold(o)
+                decoded = decoded / normalization  # norm is shape (1, 1, h, w)
+                return decoded
+            else:
+                if isinstance(self.first_stage_model, VQModelInterface):
+                    return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize)
+                else:
+                    return self.first_stage_model.decode(z)
+
+        else:
+            if isinstance(self.first_stage_model, VQModelInterface):
+                return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize)
+            else:
+                return self.first_stage_model.decode(z)
+
+    # same as above but without decorator
+    def differentiable_decode_first_stage(self, z, predict_cids=False, force_not_quantize=False):
+        if predict_cids:
+            if z.dim() == 4:
+                z = torch.argmax(z.exp(), dim=1).long()
+            z = self.first_stage_model.quantize.get_codebook_entry(z, shape=None)
+            z = rearrange(z, 'b h w c -> b c h w').contiguous()
+
+        z = 1. / self.scale_factor * z
+
+        if hasattr(self, "split_input_params"):
+            if self.split_input_params["patch_distributed_vq"]:
+                ks = self.split_input_params["ks"]  # eg. (128, 128)
+                stride = self.split_input_params["stride"]  # eg. (64, 64)
+                uf = self.split_input_params["vqf"]
+                bs, nc, h, w = z.shape
+                if ks[0] > h or ks[1] > w:
+                    ks = (min(ks[0], h), min(ks[1], w))
+                    print("reducing Kernel")
+
+                if stride[0] > h or stride[1] > w:
+                    stride = (min(stride[0], h), min(stride[1], w))
+                    print("reducing stride")
+
+                fold, unfold, normalization, weighting = self.get_fold_unfold(z, ks, stride, uf=uf)
+
+                z = unfold(z)  # (bn, nc * prod(**ks), L)
+                # 1. Reshape to img shape
+                z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1]))  # (bn, nc, ks[0], ks[1], L )
+
+                # 2. apply model loop over last dim
+                if isinstance(self.first_stage_model, VQModelInterface):
+                    output_list = [self.first_stage_model.decode(z[:, :, :, :, i],
+                                                                 force_not_quantize=predict_cids or force_not_quantize)
+                                   for i in range(z.shape[-1])]
+                else:
+
+                    output_list = [self.first_stage_model.decode(z[:, :, :, :, i])
+                                   for i in range(z.shape[-1])]
+
+                o = torch.stack(output_list, axis=-1)  # # (bn, nc, ks[0], ks[1], L)
+                o = o * weighting
+                # Reverse 1. reshape to img shape
+                o = o.view((o.shape[0], -1, o.shape[-1]))  # (bn, nc * ks[0] * ks[1], L)
+                # stitch crops together
+                decoded = fold(o)
+                decoded = decoded / normalization  # norm is shape (1, 1, h, w)
+                return decoded
+            else:
+                if isinstance(self.first_stage_model, VQModelInterface):
+                    return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize)
+                else:
+                    return self.first_stage_model.decode(z)
+
+        else:
+            if isinstance(self.first_stage_model, VQModelInterface):
+                return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize)
+            else:
+                return self.first_stage_model.decode(z)
+
+    @torch.no_grad()
+    def encode_first_stage(self, x):
+        if hasattr(self, "split_input_params"):
+            if self.split_input_params["patch_distributed_vq"]:
+                ks = self.split_input_params["ks"]  # eg. (128, 128)
+                stride = self.split_input_params["stride"]  # eg. (64, 64)
+                df = self.split_input_params["vqf"]
+                self.split_input_params['original_image_size'] = x.shape[-2:]
+                bs, nc, h, w = x.shape
+                if ks[0] > h or ks[1] > w:
+                    ks = (min(ks[0], h), min(ks[1], w))
+                    print("reducing Kernel")
+
+                if stride[0] > h or stride[1] > w:
+                    stride = (min(stride[0], h), min(stride[1], w))
+                    print("reducing stride")
+
+                fold, unfold, normalization, weighting = self.get_fold_unfold(x, ks, stride, df=df)
+                z = unfold(x)  # (bn, nc * prod(**ks), L)
+                # Reshape to img shape
+                z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1]))  # (bn, nc, ks[0], ks[1], L )
+
+                output_list = [self.first_stage_model.encode(z[:, :, :, :, i])
+                               for i in range(z.shape[-1])]
+
+                o = torch.stack(output_list, axis=-1)
+                o = o * weighting
+
+                # Reverse reshape to img shape
+                o = o.view((o.shape[0], -1, o.shape[-1]))  # (bn, nc * ks[0] * ks[1], L)
+                # stitch crops together
+                decoded = fold(o)
+                decoded = decoded / normalization
+                return decoded
+
+            else:
+                return self.first_stage_model.encode(x)
+        else:
+            return self.first_stage_model.encode(x)
+
+    def shared_step(self, batch, **kwargs):
+        x, c = self.get_input(batch, self.first_stage_key)
+        loss = self(x, c)
+        return loss
+
+    def forward(self, x, c, *args, **kwargs):
+        t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long()
+        if self.model.conditioning_key is not None:
+            assert c is not None
+            if self.cond_stage_trainable:
+                c = self.get_learned_conditioning(c)
+            if self.shorten_cond_schedule:  # TODO: drop this option
+                tc = self.cond_ids[t].to(self.device)
+                c = self.q_sample(x_start=c, t=tc, noise=torch.randn_like(c.float()))
+        return self.p_losses(x, c, t, *args, **kwargs)
+
+    def _rescale_annotations(self, bboxes, crop_coordinates):  # TODO: move to dataset
+        def rescale_bbox(bbox):
+            x0 = clamp((bbox[0] - crop_coordinates[0]) / crop_coordinates[2])
+            y0 = clamp((bbox[1] - crop_coordinates[1]) / crop_coordinates[3])
+            w = min(bbox[2] / crop_coordinates[2], 1 - x0)
+            h = min(bbox[3] / crop_coordinates[3], 1 - y0)
+            return x0, y0, w, h
+
+        return [rescale_bbox(b) for b in bboxes]
+
+    def apply_model(self, x_noisy, t, cond, return_ids=False):
+
+        if isinstance(cond, dict):
+            # hybrid case, cond is exptected to be a dict
+            pass
+        else:
+            if not isinstance(cond, list):
+                cond = [cond]
+            key = 'c_concat' if self.model.conditioning_key == 'concat' else 'c_crossattn'
+            cond = {key: cond}
+
+        if hasattr(self, "split_input_params"):
+            assert len(cond) == 1  # todo can only deal with one conditioning atm
+            assert not return_ids
+            ks = self.split_input_params["ks"]  # eg. (128, 128)
+            stride = self.split_input_params["stride"]  # eg. (64, 64)
+
+            h, w = x_noisy.shape[-2:]
+
+            fold, unfold, normalization, weighting = self.get_fold_unfold(x_noisy, ks, stride)
+
+            z = unfold(x_noisy)  # (bn, nc * prod(**ks), L)
+            # Reshape to img shape
+            z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1]))  # (bn, nc, ks[0], ks[1], L )
+            z_list = [z[:, :, :, :, i] for i in range(z.shape[-1])]
+
+            if self.cond_stage_key in ["image", "LR_image", "segmentation",
+                                       'bbox_img'] and self.model.conditioning_key:  # todo check for completeness
+                c_key = next(iter(cond.keys()))  # get key
+                c = next(iter(cond.values()))  # get value
+                assert (len(c) == 1)  # todo extend to list with more than one elem
+                c = c[0]  # get element
+
+                c = unfold(c)
+                c = c.view((c.shape[0], -1, ks[0], ks[1], c.shape[-1]))  # (bn, nc, ks[0], ks[1], L )
+
+                cond_list = [{c_key: [c[:, :, :, :, i]]} for i in range(c.shape[-1])]
+
+            elif self.cond_stage_key == 'coordinates_bbox':
+                assert 'original_image_size' in self.split_input_params, 'BoudingBoxRescaling is missing original_image_size'
+
+                # assuming padding of unfold is always 0 and its dilation is always 1
+                n_patches_per_row = int((w - ks[0]) / stride[0] + 1)
+                full_img_h, full_img_w = self.split_input_params['original_image_size']
+                # as we are operating on latents, we need the factor from the original image size to the
+                # spatial latent size to properly rescale the crops for regenerating the bbox annotations
+                num_downs = self.first_stage_model.encoder.num_resolutions - 1
+                rescale_latent = 2 ** (num_downs)
+
+                # get top left postions of patches as conforming for the bbbox tokenizer, therefore we
+                # need to rescale the tl patch coordinates to be in between (0,1)
+                tl_patch_coordinates = [(rescale_latent * stride[0] * (patch_nr % n_patches_per_row) / full_img_w,
+                                         rescale_latent * stride[1] * (patch_nr // n_patches_per_row) / full_img_h)
+                                        for patch_nr in range(z.shape[-1])]
+
+                # patch_limits are tl_coord, width and height coordinates as (x_tl, y_tl, h, w)
+                patch_limits = [(x_tl, y_tl,
+                                 rescale_latent * ks[0] / full_img_w,
+                                 rescale_latent * ks[1] / full_img_h) for x_tl, y_tl in tl_patch_coordinates]
+                # patch_values = [(np.arange(x_tl,min(x_tl+ks, 1.)),np.arange(y_tl,min(y_tl+ks, 1.))) for x_tl, y_tl in tl_patch_coordinates]
+
+                # tokenize crop coordinates for the bounding boxes of the respective patches
+                patch_limits_tknzd = [torch.LongTensor(self.bbox_tokenizer._crop_encoder(bbox))[None].to(self.device)
+                                      for bbox in patch_limits]  # list of length l with tensors of shape (1, 2)
+                print(patch_limits_tknzd[0].shape)
+                # cut tknzd crop position from conditioning
+                assert isinstance(cond, dict), 'cond must be dict to be fed into model'
+                cut_cond = cond['c_crossattn'][0][..., :-2].to(self.device)
+                print(cut_cond.shape)
+
+                adapted_cond = torch.stack([torch.cat([cut_cond, p], dim=1) for p in patch_limits_tknzd])
+                adapted_cond = rearrange(adapted_cond, 'l b n -> (l b) n')
+                print(adapted_cond.shape)
+                adapted_cond = self.get_learned_conditioning(adapted_cond)
+                print(adapted_cond.shape)
+                adapted_cond = rearrange(adapted_cond, '(l b) n d -> l b n d', l=z.shape[-1])
+                print(adapted_cond.shape)
+
+                cond_list = [{'c_crossattn': [e]} for e in adapted_cond]
+
+            else:
+                cond_list = [cond for i in range(z.shape[-1])]  # Todo make this more efficient
+
+            # apply model by loop over crops
+            output_list = [self.model(z_list[i], t, **cond_list[i]) for i in range(z.shape[-1])]
+            assert not isinstance(output_list[0],
+                                  tuple)  # todo cant deal with multiple model outputs check this never happens
+
+            o = torch.stack(output_list, axis=-1)
+            o = o * weighting
+            # Reverse reshape to img shape
+            o = o.view((o.shape[0], -1, o.shape[-1]))  # (bn, nc * ks[0] * ks[1], L)
+            # stitch crops together
+            x_recon = fold(o) / normalization
+
+        else:
+            x_recon = self.model(x_noisy, t, **cond)
+
+        if isinstance(x_recon, tuple) and not return_ids:
+            return x_recon[0]
+        else:
+            return x_recon
+
+    def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
+        return (extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart) / \
+               extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+
+    def _prior_bpd(self, x_start):
+        """
+        Get the prior KL term for the variational lower-bound, measured in
+        bits-per-dim.
+        This term can't be optimized, as it only depends on the encoder.
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :return: a batch of [N] KL values (in bits), one per batch element.
+        """
+        batch_size = x_start.shape[0]
+        t = torch.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
+        qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
+        kl_prior = normal_kl(mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0)
+        return mean_flat(kl_prior) / np.log(2.0)
+
+    def p_losses(self, x_start, cond, t, noise=None):
+        noise = default(noise, lambda: torch.randn_like(x_start))
+        x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
+        model_output = self.apply_model(x_noisy, t, cond)
+
+        loss_dict = {}
+        prefix = 'train' if self.training else 'val'
+
+        if self.parameterization == "x0":
+            target = x_start
+        elif self.parameterization == "eps":
+            target = noise
+        else:
+            raise NotImplementedError()
+
+        loss_simple = self.get_loss(model_output, target, mean=False).mean([1, 2, 3])
+        loss_dict.update({f'{prefix}/loss_simple': loss_simple.mean()})
+
+        logvar_t = self.logvar[t].to(self.device)
+        loss = loss_simple / torch.exp(logvar_t) + logvar_t
+        # loss = loss_simple / torch.exp(self.logvar) + self.logvar
+        if self.learn_logvar:
+            loss_dict.update({f'{prefix}/loss_gamma': loss.mean()})
+            loss_dict.update({'logvar': self.logvar.data.mean()})
+
+        loss = self.l_simple_weight * loss.mean()
+
+        loss_vlb = self.get_loss(model_output, target, mean=False).mean(dim=(1, 2, 3))
+        loss_vlb = (self.lvlb_weights[t] * loss_vlb).mean()
+        loss_dict.update({f'{prefix}/loss_vlb': loss_vlb})
+        loss += (self.original_elbo_weight * loss_vlb)
+        loss_dict.update({f'{prefix}/loss': loss})
+
+        return loss, loss_dict
+
+    def p_mean_variance(self, x, c, t, clip_denoised: bool, return_codebook_ids=False, quantize_denoised=False,
+                        return_x0=False, score_corrector=None, corrector_kwargs=None):
+        t_in = t
+        model_out = self.apply_model(x, t_in, c, return_ids=return_codebook_ids)
+
+        if score_corrector is not None:
+            assert self.parameterization == "eps"
+            model_out = score_corrector.modify_score(self, model_out, x, t, c, **corrector_kwargs)
+
+        if return_codebook_ids:
+            model_out, logits = model_out
+
+        if self.parameterization == "eps":
+            x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)
+        elif self.parameterization == "x0":
+            x_recon = model_out
+        else:
+            raise NotImplementedError()
+
+        if clip_denoised:
+            x_recon.clamp_(-1., 1.)
+        if quantize_denoised:
+            x_recon, _, [_, _, indices] = self.first_stage_model.quantize(x_recon)
+        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
+        if return_codebook_ids:
+            return model_mean, posterior_variance, posterior_log_variance, logits
+        elif return_x0:
+            return model_mean, posterior_variance, posterior_log_variance, x_recon
+        else:
+            return model_mean, posterior_variance, posterior_log_variance
+
+    @torch.no_grad()
+    def p_sample(self, x, c, t, clip_denoised=False, repeat_noise=False,
+                 return_codebook_ids=False, quantize_denoised=False, return_x0=False,
+                 temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None):
+        b, *_, device = *x.shape, x.device
+        outputs = self.p_mean_variance(x=x, c=c, t=t, clip_denoised=clip_denoised,
+                                       return_codebook_ids=return_codebook_ids,
+                                       quantize_denoised=quantize_denoised,
+                                       return_x0=return_x0,
+                                       score_corrector=score_corrector, corrector_kwargs=corrector_kwargs)
+        if return_codebook_ids:
+            raise DeprecationWarning("Support dropped.")
+            model_mean, _, model_log_variance, logits = outputs
+        elif return_x0:
+            model_mean, _, model_log_variance, x0 = outputs
+        else:
+            model_mean, _, model_log_variance = outputs
+
+        noise = noise_like(x.shape, device, repeat_noise) * temperature
+        if noise_dropout > 0.:
+            noise = torch.nn.functional.dropout(noise, p=noise_dropout)
+        # no noise when t == 0
+        nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
+
+        if return_codebook_ids:
+            return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, logits.argmax(dim=1)
+        if return_x0:
+            return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, x0
+        else:
+            return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
+
+    @torch.no_grad()
+    def progressive_denoising(self, cond, shape, verbose=True, callback=None, quantize_denoised=False,
+                              img_callback=None, mask=None, x0=None, temperature=1., noise_dropout=0.,
+                              score_corrector=None, corrector_kwargs=None, batch_size=None, x_T=None, start_T=None,
+                              log_every_t=None):
+        if not log_every_t:
+            log_every_t = self.log_every_t
+        timesteps = self.num_timesteps
+        if batch_size is not None:
+            b = batch_size if batch_size is not None else shape[0]
+            shape = [batch_size] + list(shape)
+        else:
+            b = batch_size = shape[0]
+        if x_T is None:
+            img = torch.randn(shape, device=self.device)
+        else:
+            img = x_T
+        intermediates = []
+        if cond is not None:
+            if isinstance(cond, dict):
+                cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else
+                list(map(lambda x: x[:batch_size], cond[key])) for key in cond}
+            else:
+                cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size]
+
+        if start_T is not None:
+            timesteps = min(timesteps, start_T)
+        iterator = tqdm(reversed(range(0, timesteps)), desc='Progressive Generation',
+                        total=timesteps) if verbose else reversed(
+            range(0, timesteps))
+        if type(temperature) == float:
+            temperature = [temperature] * timesteps
+
+        for i in iterator:
+            ts = torch.full((b,), i, device=self.device, dtype=torch.long)
+            if self.shorten_cond_schedule:
+                assert self.model.conditioning_key != 'hybrid'
+                tc = self.cond_ids[ts].to(cond.device)
+                cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond))
+
+            img, x0_partial = self.p_sample(img, cond, ts,
+                                            clip_denoised=self.clip_denoised,
+                                            quantize_denoised=quantize_denoised, return_x0=True,
+                                            temperature=temperature[i], noise_dropout=noise_dropout,
+                                            score_corrector=score_corrector, corrector_kwargs=corrector_kwargs)
+            if mask is not None:
+                assert x0 is not None
+                img_orig = self.q_sample(x0, ts)
+                img = img_orig * mask + (1. - mask) * img
+
+            if i % log_every_t == 0 or i == timesteps - 1:
+                intermediates.append(x0_partial)
+            if callback: callback(i)
+            if img_callback: img_callback(img, i)
+        return img, intermediates
+
+    @torch.no_grad()
+    def p_sample_loop(self, cond, shape, return_intermediates=False,
+                      x_T=None, verbose=True, callback=None, timesteps=None, quantize_denoised=False,
+                      mask=None, x0=None, img_callback=None, start_T=None,
+                      log_every_t=None):
+
+        if not log_every_t:
+            log_every_t = self.log_every_t
+        device = self.betas.device
+        b = shape[0]
+        if x_T is None:
+            img = torch.randn(shape, device=device)
+        else:
+            img = x_T
+
+        intermediates = [img]
+        if timesteps is None:
+            timesteps = self.num_timesteps
+
+        if start_T is not None:
+            timesteps = min(timesteps, start_T)
+        iterator = tqdm(reversed(range(0, timesteps)), desc='Sampling t', total=timesteps) if verbose else reversed(
+            range(0, timesteps))
+
+        if mask is not None:
+            assert x0 is not None
+            assert x0.shape[2:3] == mask.shape[2:3]  # spatial size has to match
+
+        for i in iterator:
+            ts = torch.full((b,), i, device=device, dtype=torch.long)
+            if self.shorten_cond_schedule:
+                assert self.model.conditioning_key != 'hybrid'
+                tc = self.cond_ids[ts].to(cond.device)
+                cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond))
+
+            img = self.p_sample(img, cond, ts,
+                                clip_denoised=self.clip_denoised,
+                                quantize_denoised=quantize_denoised)
+            if mask is not None:
+                img_orig = self.q_sample(x0, ts)
+                img = img_orig * mask + (1. - mask) * img
+
+            if i % log_every_t == 0 or i == timesteps - 1:
+                intermediates.append(img)
+            if callback: callback(i)
+            if img_callback: img_callback(img, i)
+
+        if return_intermediates:
+            return img, intermediates
+        return img
+
+    @torch.no_grad()
+    def sample(self, cond, batch_size=16, return_intermediates=False, x_T=None,
+               verbose=True, timesteps=None, quantize_denoised=False,
+               mask=None, x0=None, shape=None,**kwargs):
+        if shape is None:
+            shape = (batch_size, self.channels, self.image_size, self.image_size)
+        if cond is not None:
+            if isinstance(cond, dict):
+                cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else
+                list(map(lambda x: x[:batch_size], cond[key])) for key in cond}
+            else:
+                cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size]
+        return self.p_sample_loop(cond,
+                                  shape,
+                                  return_intermediates=return_intermediates, x_T=x_T,
+                                  verbose=verbose, timesteps=timesteps, quantize_denoised=quantize_denoised,
+                                  mask=mask, x0=x0)
+
+    @torch.no_grad()
+    def sample_log(self,cond,batch_size,ddim, ddim_steps,**kwargs):
+
+        if ddim:
+            ddim_sampler = DDIMSampler(self)
+            shape = (self.channels, self.image_size, self.image_size)
+            samples, intermediates =ddim_sampler.sample(ddim_steps,batch_size,
+                                                        shape,cond,verbose=False,**kwargs)
+
+        else:
+            samples, intermediates = self.sample(cond=cond, batch_size=batch_size,
+                                                 return_intermediates=True,**kwargs)
+
+        return samples, intermediates
+
+
+    @torch.no_grad()
+    def log_images(self, batch, N=4, n_row=4, sample=True, ddim_steps=200, ddim_eta=1., return_keys=None,
+                   quantize_denoised=True, inpaint=False, plot_denoise_rows=False, plot_progressive_rows=False,
+                   plot_diffusion_rows=False, **kwargs):
+
+        use_ddim = False
+
+        log = dict()
+        z, c, x, xrec, xc = self.get_input(batch, self.first_stage_key,
+                                           return_first_stage_outputs=True,
+                                           force_c_encode=True,
+                                           return_original_cond=True,
+                                           bs=N, uncond=0)
+        N = min(x.shape[0], N)
+        n_row = min(x.shape[0], n_row)
+        log["inputs"] = x
+        log["reals"] = xc["c_concat"]
+        log["reconstruction"] = xrec
+        if self.model.conditioning_key is not None:
+            if hasattr(self.cond_stage_model, "decode"):
+                xc = self.cond_stage_model.decode(c)
+                log["conditioning"] = xc
+            elif self.cond_stage_key in ["caption"]:
+                xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["caption"])
+                log["conditioning"] = xc
+            elif self.cond_stage_key == 'class_label':
+                xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["human_label"])
+                log['conditioning'] = xc
+            elif isimage(xc):
+                log["conditioning"] = xc
+            if ismap(xc):
+                log["original_conditioning"] = self.to_rgb(xc)
+
+        if plot_diffusion_rows:
+            # get diffusion row
+            diffusion_row = list()
+            z_start = z[:n_row]
+            for t in range(self.num_timesteps):
+                if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
+                    t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
+                    t = t.to(self.device).long()
+                    noise = torch.randn_like(z_start)
+                    z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise)
+                    diffusion_row.append(self.decode_first_stage(z_noisy))
+
+            diffusion_row = torch.stack(diffusion_row)  # n_log_step, n_row, C, H, W
+            diffusion_grid = rearrange(diffusion_row, 'n b c h w -> b n c h w')
+            diffusion_grid = rearrange(diffusion_grid, 'b n c h w -> (b n) c h w')
+            diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0])
+            log["diffusion_row"] = diffusion_grid
+
+        if sample:
+            # get denoise row
+            with self.ema_scope("Plotting"):
+                samples, z_denoise_row = self.sample_log(cond=c,batch_size=N,ddim=use_ddim,
+                                                         ddim_steps=ddim_steps,eta=ddim_eta)
+                # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True)
+            x_samples = self.decode_first_stage(samples)
+            log["samples"] = x_samples
+            if plot_denoise_rows:
+                denoise_grid = self._get_denoise_row_from_list(z_denoise_row)
+                log["denoise_row"] = denoise_grid
+
+            if quantize_denoised and not isinstance(self.first_stage_model, AutoencoderKL) and not isinstance(
+                    self.first_stage_model, IdentityFirstStage):
+                # also display when quantizing x0 while sampling
+                with self.ema_scope("Plotting Quantized Denoised"):
+                    samples, z_denoise_row = self.sample_log(cond=c,batch_size=N,ddim=use_ddim,
+                                                             ddim_steps=ddim_steps,eta=ddim_eta,
+                                                             quantize_denoised=True)
+                    # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True,
+                    #                                      quantize_denoised=True)
+                x_samples = self.decode_first_stage(samples.to(self.device))
+                log["samples_x0_quantized"] = x_samples
+
+            if inpaint:
+                # make a simple center square
+                b, h, w = z.shape[0], z.shape[2], z.shape[3]
+                mask = torch.ones(N, h, w).to(self.device)
+                # zeros will be filled in
+                mask[:, h // 4:3 * h // 4, w // 4:3 * w // 4] = 0.
+                mask = mask[:, None, ...]
+                with self.ema_scope("Plotting Inpaint"):
+
+                    samples, _ = self.sample_log(cond=c,batch_size=N,ddim=use_ddim, eta=ddim_eta,
+                                                ddim_steps=ddim_steps, x0=z[:N], mask=mask)
+                x_samples = self.decode_first_stage(samples.to(self.device))
+                log["samples_inpainting"] = x_samples
+                log["mask"] = mask
+
+                # outpaint
+                with self.ema_scope("Plotting Outpaint"):
+                    samples, _ = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,eta=ddim_eta,
+                                                ddim_steps=ddim_steps, x0=z[:N], mask=mask)
+                x_samples = self.decode_first_stage(samples.to(self.device))
+                log["samples_outpainting"] = x_samples
+
+        if plot_progressive_rows:
+            with self.ema_scope("Plotting Progressives"):
+                img, progressives = self.progressive_denoising(c,
+                                                               shape=(self.channels, self.image_size, self.image_size),
+                                                               batch_size=N)
+            prog_row = self._get_denoise_row_from_list(progressives, desc="Progressive Generation")
+            log["progressive_row"] = prog_row
+
+        if return_keys:
+            if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0:
+                return log
+            else:
+                return {key: log[key] for key in return_keys}
+        return log
+
+    def configure_optimizers(self):
+        lr = self.learning_rate
+        params = list(self.model.parameters())
+        if self.cond_stage_trainable:
+            print(f"{self.__class__.__name__}: Also optimizing conditioner params!")
+            params = params + list(self.cond_stage_model.parameters())
+        if self.learn_logvar:
+            print('Diffusion model optimizing logvar')
+            params.append(self.logvar)
+        opt = torch.optim.AdamW(params, lr=lr)
+        if self.use_scheduler:
+            assert 'target' in self.scheduler_config
+            scheduler = instantiate_from_config(self.scheduler_config)
+
+            print("Setting up LambdaLR scheduler...")
+            scheduler = [
+                {
+                    'scheduler': LambdaLR(opt, lr_lambda=scheduler.schedule),
+                    'interval': 'step',
+                    'frequency': 1
+                }]
+            return [opt], scheduler
+        return opt
+
+    @torch.no_grad()
+    def to_rgb(self, x):
+        x = x.float()
+        if not hasattr(self, "colorize"):
+            self.colorize = torch.randn(3, x.shape[1], 1, 1).to(x)
+        x = nn.functional.conv2d(x, weight=self.colorize)
+        x = 2. * (x - x.min()) / (x.max() - x.min()) - 1.
+        return x
+
+
+class DiffusionWrapper(pl.LightningModule):
+    def __init__(self, diff_model_config, conditioning_key):
+        super().__init__()
+        self.diffusion_model = instantiate_from_config(diff_model_config)
+        self.conditioning_key = conditioning_key
+        assert self.conditioning_key in [None, 'concat', 'crossattn', 'hybrid', 'adm']
+
+    def forward(self, x, t, c_concat: list = None, c_crossattn: list = None):
+        if self.conditioning_key is None:
+            out = self.diffusion_model(x, t)
+        elif self.conditioning_key == 'concat':
+            xc = torch.cat([x] + c_concat, dim=1)
+            out = self.diffusion_model(xc, t)
+        elif self.conditioning_key == 'crossattn':
+            cc = torch.cat(c_crossattn, 1)
+            out = self.diffusion_model(x, t, context=cc)
+        elif self.conditioning_key == 'hybrid':
+            xc = torch.cat([x] + c_concat, dim=1)
+            cc = torch.cat(c_crossattn, 1)
+            out = self.diffusion_model(xc, t, context=cc)
+        elif self.conditioning_key == 'adm':
+            cc = c_crossattn[0]
+            out = self.diffusion_model(x, t, y=cc)
+        else:
+            raise NotImplementedError()
+
+        return out
+
+
+class Layout2ImgDiffusion(LatentDiffusion):
+    # TODO: move all layout-specific hacks to this class
+    def __init__(self, cond_stage_key, *args, **kwargs):
+        assert cond_stage_key == 'coordinates_bbox', 'Layout2ImgDiffusion only for cond_stage_key="coordinates_bbox"'
+        super().__init__(cond_stage_key=cond_stage_key, *args, **kwargs)
+
+    def log_images(self, batch, N=8, *args, **kwargs):
+        logs = super().log_images(batch=batch, N=N, *args, **kwargs)
+
+        key = 'train' if self.training else 'validation'
+        dset = self.trainer.datamodule.datasets[key]
+        mapper = dset.conditional_builders[self.cond_stage_key]
+
+        bbox_imgs = []
+        map_fn = lambda catno: dset.get_textual_label(dset.get_category_id(catno))
+        for tknzd_bbox in batch[self.cond_stage_key][:N]:
+            bboximg = mapper.plot(tknzd_bbox.detach().cpu(), map_fn, (256, 256))
+            bbox_imgs.append(bboximg)
+
+        cond_img = torch.stack(bbox_imgs, dim=0)
+        logs['bbox_image'] = cond_img
+        return logs
diff --git a/modules/paths.py b/modules/paths.py
index 4dd03a35..d991cc71 100644
--- a/modules/paths.py
+++ b/modules/paths.py
@@ -4,7 +4,15 @@ import sys
 import modules.safe

 

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

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

+

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

 

 # search for directory of stable diffusion in following places

@@ -38,3 +46,17 @@ for d, must_exist, what, options in path_dirs:
         else:

             sys.path.append(d)

         paths[what] = d

+

+

+class Prioritize:

+    def __init__(self, name):

+        self.name = name

+        self.path = None

+

+    def __enter__(self):

+        self.path = sys.path.copy()

+        sys.path = [paths[self.name]] + sys.path

+

+    def __exit__(self, exc_type, exc_val, exc_tb):

+        sys.path = self.path

+        self.path = None

diff --git a/modules/postprocessing.py b/modules/postprocessing.py
new file mode 100644
index 00000000..09d8e605
--- /dev/null
+++ b/modules/postprocessing.py
@@ -0,0 +1,103 @@
+import os

+

+from PIL import Image

+

+from modules import shared, images, devices, scripts, scripts_postprocessing, ui_common, generation_parameters_copypaste

+from modules.shared import opts

+

+

+def run_postprocessing(extras_mode, image, image_folder, input_dir, output_dir, show_extras_results, *args, save_output: bool = True):

+    devices.torch_gc()

+

+    shared.state.begin()

+    shared.state.job = 'extras'

+

+    image_data = []

+    image_names = []

+    outputs = []

+

+    if extras_mode == 1:

+        for img in image_folder:

+            image = Image.open(img)

+            image_data.append(image)

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

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

+

+        image_list = shared.listfiles(input_dir)

+        for filename in image_list:

+            try:

+                image = Image.open(filename)

+            except Exception:

+                continue

+            image_data.append(image)

+            image_names.append(filename)

+    else:

+        assert image, 'image not selected'

+

+        image_data.append(image)

+        image_names.append(None)

+

+    if extras_mode == 2 and output_dir != '':

+        outpath = output_dir

+    else:

+        outpath = opts.outdir_samples or opts.outdir_extras_samples

+

+    infotext = ''

+

+    for image, name in zip(image_data, image_names):

+        shared.state.textinfo = name

+

+        existing_pnginfo = image.info or {}

+

+        pp = scripts_postprocessing.PostprocessedImage(image.convert("RGB"))

+

+        scripts.scripts_postproc.run(pp, args)

+

+        if opts.use_original_name_batch and name is not None:

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

+        else:

+            basename = ''

+

+        infotext = ", ".join([k if k == v else f'{k}: {generation_parameters_copypaste.quote(v)}' for k, v in pp.info.items() if v is not None])

+

+        if opts.enable_pnginfo:

+            pp.image.info = existing_pnginfo

+            pp.image.info["postprocessing"] = infotext

+

+        if save_output:

+            images.save_image(pp.image, path=outpath, basename=basename, seed=None, prompt=None, extension=opts.samples_format, info=infotext, short_filename=True, no_prompt=True, grid=False, pnginfo_section_name="extras", existing_info=existing_pnginfo, forced_filename=None)

+

+        if extras_mode != 2 or show_extras_results:

+            outputs.append(pp.image)

+

+    devices.torch_gc()

+

+    return outputs, ui_common.plaintext_to_html(infotext), ''

+

+

+def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_dir, show_extras_results, gfpgan_visibility, codeformer_visibility, codeformer_weight, upscaling_resize, upscaling_resize_w, upscaling_resize_h, upscaling_crop, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility, upscale_first: bool, save_output: bool = True):

+    """old handler for API"""

+

+    args = scripts.scripts_postproc.create_args_for_run({

+        "Upscale": {

+            "upscale_mode": resize_mode,

+            "upscale_by": upscaling_resize,

+            "upscale_to_width": upscaling_resize_w,

+            "upscale_to_height": upscaling_resize_h,

+            "upscale_crop": upscaling_crop,

+            "upscaler_1_name": extras_upscaler_1,

+            "upscaler_2_name": extras_upscaler_2,

+            "upscaler_2_visibility": extras_upscaler_2_visibility,

+        },

+        "GFPGAN": {

+            "gfpgan_visibility": gfpgan_visibility,

+        },

+        "CodeFormer": {

+            "codeformer_visibility": codeformer_visibility,

+            "codeformer_weight": codeformer_weight,

+        },

+    })

+

+    return run_postprocessing(extras_mode, image, image_folder, input_dir, output_dir, show_extras_results, *args, save_output=save_output)

diff --git a/modules/processing.py b/modules/processing.py
index 4a406084..5072fc40 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -13,10 +13,11 @@ from skimage import exposure
 from typing import Any, Dict, List, Optional

 

 import modules.sd_hijack

-from modules import devices, prompt_parser, masking, sd_samplers, lowvram, generation_parameters_copypaste, script_callbacks

+from modules import devices, prompt_parser, masking, sd_samplers, lowvram, generation_parameters_copypaste, script_callbacks, extra_networks, sd_vae_approx, scripts

 from modules.sd_hijack import model_hijack

 from modules.shared import opts, cmd_opts, state

 import modules.shared as shared

+import modules.paths as paths

 import modules.face_restoration

 import modules.images as images

 import modules.styles

@@ -50,9 +51,9 @@ def apply_color_correction(correction, original_image):
         correction,

         channel_axis=2

     ), cv2.COLOR_LAB2RGB).astype("uint8"))

-    

+

     image = blendLayers(image, original_image, BlendType.LUMINOSITY)

-    

+

     return image

 

 

@@ -76,15 +77,32 @@ def apply_overlay(image, paste_loc, index, overlays):
     return image

 

 

-class StableDiffusionProcessing():

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

+

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

+

+    return image_conditioning

+

+

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

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

         if sampler_index is not None:

             print("sampler_index argument for StableDiffusionProcessing does not do anything; use sampler_name", file=sys.stderr)

 

-        self.sd_model = sd_model

         self.outpath_samples: str = outpath_samples

         self.outpath_grids: str = outpath_grids

         self.prompt: str = prompt

@@ -123,6 +141,7 @@ class StableDiffusionProcessing():
         self.override_settings = {k: v for k, v in (override_settings or {}).items() if k not in shared.restricted_opts}

         self.override_settings_restore_afterwards = override_settings_restore_afterwards

         self.is_using_inpainting_conditioning = False

+        self.disable_extra_networks = False

 

         if not seed_enable_extras:

             self.subseed = -1

@@ -131,33 +150,21 @@ class StableDiffusionProcessing():
             self.seed_resize_from_w = 0

 

         self.scripts = None

-        self.script_args = None

+        self.script_args = script_args

         self.all_prompts = None

         self.all_negative_prompts = None

         self.all_seeds = None

         self.all_subseeds = None

+        self.iteration = 0

 

-    def txt2img_image_conditioning(self, x, width=None, height=None):

-        if self.sampler.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)

+    @property

+    def sd_model(self):

+        return shared.sd_model

 

-        self.is_using_inpainting_conditioning = True

-

-        height = height or self.height

-        width = width or self.width

-

-        # 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 = self.sd_model.get_first_stage_encoding(self.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)

+    def txt2img_image_conditioning(self, x, width=None, height=None):

+        self.is_using_inpainting_conditioning = self.sd_model.model.conditioning_key in {'hybrid', 'concat'}

 

-        return image_conditioning

+        return txt2img_image_conditioning(self.sd_model, x, width or self.width, height or self.height)

 

     def depth2img_image_conditioning(self, source_image):

         # Use the AddMiDaS helper to Format our source image to suit the MiDaS model

@@ -166,7 +173,8 @@ class StableDiffusionProcessing():
         midas_in = torch.from_numpy(transformed["midas_in"][None, ...]).to(device=shared.device)

         midas_in = repeat(midas_in, "1 ... -> n ...", n=self.batch_size)

 

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

+        conditioning_image = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(source_image.to(devices.dtype_vae) if devices.unet_needs_upcast else source_image))

+        conditioning_image = conditioning_image.float() if devices.unet_needs_upcast else conditioning_image

         conditioning = torch.nn.functional.interpolate(

             self.sd_model.depth_model(midas_in),

             size=conditioning_image.shape[2:],

@@ -178,7 +186,12 @@ class StableDiffusionProcessing():
         conditioning = 2. * (conditioning - depth_min) / (depth_max - depth_min) - 1.

         return conditioning

 

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

+    def edit_image_conditioning(self, source_image):

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

+

+        return conditioning_image

+

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

         self.is_using_inpainting_conditioning = True

 

         # Handle the different mask inputs

@@ -197,7 +210,7 @@ class StableDiffusionProcessing():
 

         # Create another latent image, this time with a masked version of the original input.

         # Smoothly interpolate between the masked and unmasked latent conditioning image using a parameter.

-        conditioning_mask = conditioning_mask.to(source_image.device).to(source_image.dtype)

+        conditioning_mask = conditioning_mask.to(device=source_image.device, dtype=source_image.dtype)

         conditioning_image = torch.lerp(

             source_image,

             source_image * (1.0 - conditioning_mask),

@@ -205,7 +218,7 @@ class StableDiffusionProcessing():
         )

 

         # Encode the new masked image using first stage of network.

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

+        conditioning_image = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(conditioning_image.to(devices.dtype_vae) if devices.unet_needs_upcast else conditioning_image))

 

         # Create the concatenated conditioning tensor to be fed to `c_concat`

         conditioning_mask = torch.nn.functional.interpolate(conditioning_mask, size=latent_image.shape[-2:])

@@ -219,10 +232,13 @@ class StableDiffusionProcessing():
         # HACK: Using introspection as the Depth2Image model doesn't appear to uniquely

         # identify itself with a field common to all models. The conditioning_key is also hybrid.

         if isinstance(self.sd_model, LatentDepth2ImageDiffusion):

-            return self.depth2img_image_conditioning(source_image)

+            return self.depth2img_image_conditioning(source_image.float() if devices.unet_needs_upcast else source_image)

+

+        if self.sd_model.cond_stage_key == "edit":

+            return self.edit_image_conditioning(source_image)

 

         if self.sampler.conditioning_key in {'hybrid', 'concat'}:

-            return self.inpainting_image_conditioning(source_image, latent_image, image_mask=image_mask)

+            return self.inpainting_image_conditioning(source_image.float() if devices.unet_needs_upcast else source_image, latent_image, image_mask=image_mask)

 

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

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

@@ -234,12 +250,11 @@ class StableDiffusionProcessing():
         raise NotImplementedError()

 

     def close(self):

-        self.sd_model = None

         self.sampler = None

 

 

 class Processed:

-    def __init__(self, p: StableDiffusionProcessing, images_list, seed=-1, info="", subseed=None, all_prompts=None, all_negative_prompts=None, all_seeds=None, all_subseeds=None, index_of_first_image=0, infotexts=None):

+    def __init__(self, p: StableDiffusionProcessing, images_list, seed=-1, info="", subseed=None, all_prompts=None, all_negative_prompts=None, all_seeds=None, all_subseeds=None, index_of_first_image=0, infotexts=None, comments=""):

         self.images = images_list

         self.prompt = p.prompt

         self.negative_prompt = p.negative_prompt

@@ -247,6 +262,7 @@ class Processed:
         self.subseed = subseed

         self.subseed_strength = p.subseed_strength

         self.info = info

+        self.comments = comments

         self.width = p.width

         self.height = p.height

         self.sampler_name = p.sampler_name

@@ -338,13 +354,14 @@ def slerp(val, low, high):
 

 

 def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, seed_resize_from_h=0, seed_resize_from_w=0, p=None):

+    eta_noise_seed_delta = opts.eta_noise_seed_delta or 0

     xs = []

 

     # if we have multiple seeds, this means we are working with batch size>1; this then

     # enables the generation of additional tensors with noise that the sampler will use during its processing.

     # Using those pre-generated tensors instead of simple torch.randn allows a batch with seeds [100, 101] to

     # produce the same images as with two batches [100], [101].

-    if p is not None and p.sampler is not None and (len(seeds) > 1 and opts.enable_batch_seeds or opts.eta_noise_seed_delta > 0):

+    if p is not None and p.sampler is not None and (len(seeds) > 1 and opts.enable_batch_seeds or eta_noise_seed_delta > 0):

         sampler_noises = [[] for _ in range(p.sampler.number_of_needed_noises(p))]

     else:

         sampler_noises = None

@@ -384,8 +401,8 @@ def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, see
         if sampler_noises is not None:

             cnt = p.sampler.number_of_needed_noises(p)

 

-            if opts.eta_noise_seed_delta > 0:

-                torch.manual_seed(seed + opts.eta_noise_seed_delta)

+            if eta_noise_seed_delta > 0:

+                torch.manual_seed(seed + eta_noise_seed_delta)

 

             for j in range(cnt):

                 sampler_noises[j].append(devices.randn_without_seed(tuple(noise_shape)))

@@ -401,7 +418,7 @@ def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, see
 

 def decode_first_stage(model, x):

     with devices.autocast(disable=x.dtype == devices.dtype_vae):

-        x = model.decode_first_stage(x)

+        x = model.decode_first_stage(x.to(devices.dtype_vae) if devices.unet_needs_upcast else x)

 

     return x

 

@@ -418,7 +435,7 @@ def fix_seed(p):
     p.subseed = get_fixed_seed(p.subseed)

 

 

-def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments, iteration=0, position_in_batch=0):

+def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iteration=0, position_in_batch=0):

     index = position_in_batch + iteration * p.batch_size

 

     clip_skip = getattr(p, 'clip_skip', opts.CLIP_stop_at_last_layers)

@@ -432,9 +449,6 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments, iteration
         "Size": f"{p.width}x{p.height}",

         "Model hash": getattr(p, 'sd_model_hash', None if not opts.add_model_hash_to_info or not shared.sd_model.sd_model_hash else shared.sd_model.sd_model_hash),

         "Model": (None if not opts.add_model_name_to_info or not shared.sd_model.sd_checkpoint_info.model_name else shared.sd_model.sd_checkpoint_info.model_name.replace(',', '').replace(':', '')),

-        "Hypernet": (None if shared.loaded_hypernetwork is None else shared.loaded_hypernetwork.name),

-        "Hypernet hash": (None if shared.loaded_hypernetwork is None else sd_models.model_hash(shared.loaded_hypernetwork.filename)),

-        "Hypernet strength": (None if shared.loaded_hypernetwork is None or shared.opts.sd_hypernetwork_strength >= 1 else shared.opts.sd_hypernetwork_strength),

         "Batch size": (None if p.batch_size < 2 else p.batch_size),

         "Batch pos": (None if p.batch_size < 2 else position_in_batch),

         "Variation seed": (None if p.subseed_strength == 0 else all_subseeds[index]),

@@ -462,9 +476,12 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
     try:

         for k, v in p.override_settings.items():

             setattr(opts, k, v)

-            if k == 'sd_hypernetwork': shared.reload_hypernetworks()  # make onchange call for changing hypernet

-            if k == 'sd_model_checkpoint': sd_models.reload_model_weights()  # make onchange call for changing SD model

-            if k == 'sd_vae': sd_vae.reload_vae_weights()  # make onchange call for changing VAE

+

+            if k == 'sd_model_checkpoint':

+                sd_models.reload_model_weights()

+

+            if k == 'sd_vae':

+                sd_vae.reload_vae_weights()

 

         res = process_images_inner(p)

 

@@ -473,9 +490,11 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
         if p.override_settings_restore_afterwards:

             for k, v in stored_opts.items():

                 setattr(opts, k, v)

-                if k == 'sd_hypernetwork': shared.reload_hypernetworks()

-                if k == 'sd_model_checkpoint': sd_models.reload_model_weights()

-                if k == 'sd_vae': sd_vae.reload_vae_weights()

+                if k == 'sd_model_checkpoint':

+                    sd_models.reload_model_weights()

+

+                if k == 'sd_vae':

+                    sd_vae.reload_vae_weights()

 

     return res

 

@@ -521,27 +540,61 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
     def infotext(iteration=0, position_in_batch=0):

         return create_infotext(p, p.all_prompts, p.all_seeds, p.all_subseeds, comments, iteration, position_in_batch)

 

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

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

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

-

     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)

 

     infotexts = []

     output_images = []

 

+    cached_uc = [None, None]

+    cached_c = [None, None]

+

+    def get_conds_with_caching(function, required_prompts, steps, cache):

+        """

+        Returns the result of calling function(shared.sd_model, required_prompts, steps)

+        using a cache to store the result if the same arguments have been used before.

+

+        cache is an array containing two elements. The first element is a tuple

+        representing the previously used arguments, or None if no arguments

+        have been used before. The second element is where the previously

+        computed result is stored.

+        """

+

+        if cache[0] is not None and (required_prompts, steps) == cache[0]:

+            return cache[1]

+

+        with devices.autocast():

+            cache[1] = function(shared.sd_model, required_prompts, steps)

+

+        cache[0] = (required_prompts, steps)

+        return cache[1]

+

     with torch.no_grad(), p.sd_model.ema_scope():

         with devices.autocast():

             p.init(p.all_prompts, p.all_seeds, p.all_subseeds)

 

+            # for OSX, loading the model during sampling changes the generated picture, so it is loaded here

+            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

 

         for n in range(p.n_iter):

+            p.iteration = n

+

             if state.skipped:

                 state.skipped = False

 

@@ -556,12 +609,13 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             if len(prompts) == 0:

                 break

 

+            prompts, _ = extra_networks.parse_prompts(prompts)

+

             if p.scripts is not None:

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

 

-            with devices.autocast():

-                uc = prompt_parser.get_learned_conditioning(shared.sd_model, negative_prompts, p.steps)

-                c = prompt_parser.get_multicond_learned_conditioning(shared.sd_model, prompts, p.steps)

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

 

             if len(model_hijack.comments) > 0:

                 for comment in model_hijack.comments:

@@ -570,10 +624,13 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             if p.n_iter > 1:

                 shared.state.job = f"Batch {n+1} out of {p.n_iter}"

 

-            with devices.autocast():

+            with devices.without_autocast() if devices.unet_needs_upcast else devices.autocast():

                 samples_ddim = p.sample(conditioning=c, unconditional_conditioning=uc, seeds=seeds, subseeds=subseeds, subseed_strength=p.subseed_strength, prompts=prompts)

 

             x_samples_ddim = [decode_first_stage(p.sd_model, samples_ddim[i:i+1].to(dtype=devices.dtype_vae))[0].cpu() for i in range(samples_ddim.size(0))]

+            for x in x_samples_ddim:

+                devices.test_for_nans(x, "vae")

+

             x_samples_ddim = torch.stack(x_samples_ddim).float()

             x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)

 

@@ -602,6 +659,11 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
 

                 image = Image.fromarray(x_sample)

 

+                if p.scripts is not None:

+                    pp = scripts.PostprocessImageArgs(image)

+                    p.scripts.postprocess_image(p, pp)

+                    image = pp.image

+

                 if p.color_corrections is not None and i < len(p.color_corrections):

                     if opts.save and not p.do_not_save_samples and opts.save_images_before_color_correction:

                         image_without_cc = apply_overlay(image, p.paste_to, i, p.overlay_images)

@@ -643,9 +705,12 @@ 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:

+        extra_networks.deactivate(p, extra_network_data)

+

     devices.torch_gc()

 

-    res = Processed(p, output_images, p.all_seeds[0], infotext() + "".join(["\n\n" + x for x in comments]), subseed=p.all_subseeds[0], index_of_first_image=index_of_first_image, infotexts=infotexts)

+    res = Processed(p, output_images, p.all_seeds[0], infotext(), comments="".join(["\n\n" + x for x in comments]), subseed=p.all_subseeds[0], index_of_first_image=index_of_first_image, infotexts=infotexts)

 

     if p.scripts is not None:

         p.scripts.postprocess(p, res)

@@ -653,79 +718,138 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
     return res

 

 

+def old_hires_fix_first_pass_dimensions(width, height):

+    """old algorithm for auto-calculating first pass size"""

+

+    desired_pixel_count = 512 * 512

+    actual_pixel_count = width * height

+    scale = math.sqrt(desired_pixel_count / actual_pixel_count)

+    width = math.ceil(scale * width / 64) * 64

+    height = math.ceil(scale * height / 64) * 64

+

+    return width, height

+

+

 class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):

     sampler = None

 

-    def __init__(self, enable_hr: bool=False, denoising_strength: float=0.75, firstphase_width: int=0, firstphase_height: int=0, **kwargs):

+    def __init__(self, enable_hr: bool = False, denoising_strength: float = 0.75, firstphase_width: int = 0, firstphase_height: int = 0, hr_scale: float = 2.0, hr_upscaler: str = None, hr_second_pass_steps: int = 0, hr_resize_x: int = 0, hr_resize_y: int = 0, **kwargs):

         super().__init__(**kwargs)

         self.enable_hr = enable_hr

         self.denoising_strength = denoising_strength

-        self.firstphase_width = firstphase_width

-        self.firstphase_height = firstphase_height

+        self.hr_scale = hr_scale

+        self.hr_upscaler = hr_upscaler

+        self.hr_second_pass_steps = hr_second_pass_steps

+        self.hr_resize_x = hr_resize_x

+        self.hr_resize_y = hr_resize_y

+        self.hr_upscale_to_x = hr_resize_x

+        self.hr_upscale_to_y = hr_resize_y

+

+        if firstphase_width != 0 or firstphase_height != 0:

+            self.hr_upscale_to_x = self.width

+            self.hr_upscale_to_y = self.height

+            self.width = firstphase_width

+            self.height = firstphase_height

+

         self.truncate_x = 0

         self.truncate_y = 0

+        self.applied_old_hires_behavior_to = None

 

     def init(self, all_prompts, all_seeds, all_subseeds):

         if self.enable_hr:

-            if state.job_count == -1:

-                state.job_count = self.n_iter * 2

+            if opts.use_old_hires_fix_width_height and self.applied_old_hires_behavior_to != (self.width, self.height):

+                self.hr_resize_x = self.width

+                self.hr_resize_y = self.height

+                self.hr_upscale_to_x = self.width

+                self.hr_upscale_to_y = self.height

+

+                self.width, self.height = old_hires_fix_first_pass_dimensions(self.width, self.height)

+                self.applied_old_hires_behavior_to = (self.width, self.height)

+

+            if self.hr_resize_x == 0 and self.hr_resize_y == 0:

+                self.extra_generation_params["Hires upscale"] = self.hr_scale

+                self.hr_upscale_to_x = int(self.width * self.hr_scale)

+                self.hr_upscale_to_y = int(self.height * self.hr_scale)

             else:

-                state.job_count = state.job_count * 2

-

-            self.extra_generation_params["First pass size"] = f"{self.firstphase_width}x{self.firstphase_height}"

+                self.extra_generation_params["Hires resize"] = f"{self.hr_resize_x}x{self.hr_resize_y}"

+

+                if self.hr_resize_y == 0:

+                    self.hr_upscale_to_x = self.hr_resize_x

+                    self.hr_upscale_to_y = self.hr_resize_x * self.height // self.width

+                elif self.hr_resize_x == 0:

+                    self.hr_upscale_to_x = self.hr_resize_y * self.width // self.height

+                    self.hr_upscale_to_y = self.hr_resize_y

+                else:

+                    target_w = self.hr_resize_x

+                    target_h = self.hr_resize_y

+                    src_ratio = self.width / self.height

+                    dst_ratio = self.hr_resize_x / self.hr_resize_y

+

+                    if src_ratio < dst_ratio:

+                        self.hr_upscale_to_x = self.hr_resize_x

+                        self.hr_upscale_to_y = self.hr_resize_x * self.height // self.width

+                    else:

+                        self.hr_upscale_to_x = self.hr_resize_y * self.width // self.height

+                        self.hr_upscale_to_y = self.hr_resize_y

+

+                    self.truncate_x = (self.hr_upscale_to_x - target_w) // opt_f

+                    self.truncate_y = (self.hr_upscale_to_y - target_h) // opt_f

+

+            # special case: the user has chosen to do nothing

+            if self.hr_upscale_to_x == self.width and self.hr_upscale_to_y == self.height:

+                self.enable_hr = False

+                self.denoising_strength = None

+                self.extra_generation_params.pop("Hires upscale", None)

+                self.extra_generation_params.pop("Hires resize", None)

+                return

 

-            if self.firstphase_width == 0 or self.firstphase_height == 0:

-                desired_pixel_count = 512 * 512

-                actual_pixel_count = self.width * self.height

-                scale = math.sqrt(desired_pixel_count / actual_pixel_count)

-                self.firstphase_width = math.ceil(scale * self.width / 64) * 64

-                self.firstphase_height = math.ceil(scale * self.height / 64) * 64

-                firstphase_width_truncated = int(scale * self.width)

-                firstphase_height_truncated = int(scale * self.height)

+            if not state.processing_has_refined_job_count:

+                if state.job_count == -1:

+                    state.job_count = self.n_iter

 

-            else:

-

-                width_ratio = self.width / self.firstphase_width

-                height_ratio = self.height / self.firstphase_height

+                shared.total_tqdm.updateTotal((self.steps + (self.hr_second_pass_steps or self.steps)) * state.job_count)

+                state.job_count = state.job_count * 2

+                state.processing_has_refined_job_count = True

 

-                if width_ratio > height_ratio:

-                    firstphase_width_truncated = self.firstphase_width

-                    firstphase_height_truncated = self.firstphase_width * self.height / self.width

-                else:

-                    firstphase_width_truncated = self.firstphase_height * self.width / self.height

-                    firstphase_height_truncated = self.firstphase_height

+            if self.hr_second_pass_steps:

+                self.extra_generation_params["Hires steps"] = self.hr_second_pass_steps

 

-            self.truncate_x = int(self.firstphase_width - firstphase_width_truncated) // opt_f

-            self.truncate_y = int(self.firstphase_height - firstphase_height_truncated) // opt_f

+            if self.hr_upscaler is not None:

+                self.extra_generation_params["Hires upscaler"] = self.hr_upscaler

 

     def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts):

         self.sampler = sd_samplers.create_sampler(self.sampler_name, self.sd_model)

 

+        latent_scale_mode = shared.latent_upscale_modes.get(self.hr_upscaler, None) if self.hr_upscaler is not None else shared.latent_upscale_modes.get(shared.latent_upscale_default_mode, "nearest")

+        if self.enable_hr and latent_scale_mode is None:

+            assert len([x for x in shared.sd_upscalers if x.name == self.hr_upscaler]) > 0, f"could not find upscaler named {self.hr_upscaler}"

+

+        x = create_random_tensors([opt_C, self.height // opt_f, self.width // opt_f], seeds=seeds, subseeds=subseeds, subseed_strength=self.subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w, p=self)

+        samples = self.sampler.sample(self, x, conditioning, unconditional_conditioning, image_conditioning=self.txt2img_image_conditioning(x))

+

         if not self.enable_hr:

-            x = create_random_tensors([opt_C, self.height // opt_f, self.width // opt_f], seeds=seeds, subseeds=subseeds, subseed_strength=self.subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w, p=self)

-            samples = self.sampler.sample(self, x, conditioning, unconditional_conditioning, image_conditioning=self.txt2img_image_conditioning(x))

             return samples

 

-        x = create_random_tensors([opt_C, self.firstphase_height // opt_f, self.firstphase_width // opt_f], seeds=seeds, subseeds=subseeds, subseed_strength=self.subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w, p=self)

-        samples = self.sampler.sample(self, x, conditioning, unconditional_conditioning, image_conditioning=self.txt2img_image_conditioning(x, self.firstphase_width, self.firstphase_height))

-

-        samples = samples[:, :, self.truncate_y//2:samples.shape[2]-self.truncate_y//2, self.truncate_x//2:samples.shape[3]-self.truncate_x//2]

+        target_width = self.hr_upscale_to_x

+        target_height = self.hr_upscale_to_y

 

-        """saves image before applying hires fix, if enabled in options; takes as an argument either an image or batch with latent space images"""

         def save_intermediate(image, index):

+            """saves image before applying hires fix, if enabled in options; takes as an argument either an image or batch with latent space images"""

+

             if not opts.save or self.do_not_save_samples or not opts.save_images_before_highres_fix:

                 return

 

             if not isinstance(image, Image.Image):

-                image = sd_samplers.sample_to_image(image, index)

+                image = sd_samplers.sample_to_image(image, index, approximation=0)

 

-            images.save_image(image, self.outpath_samples, "", seeds[index], prompts[index], opts.samples_format, suffix="-before-highres-fix")

+            info = create_infotext(self, self.all_prompts, self.all_seeds, self.all_subseeds, [], iteration=self.iteration, position_in_batch=index)

+            images.save_image(image, self.outpath_samples, "", seeds[index], prompts[index], opts.samples_format, info=info, suffix="-before-highres-fix")

 

-        if opts.use_scale_latent_for_hires_fix:

+        if latent_scale_mode is not None:

             for i in range(samples.shape[0]):

                 save_intermediate(samples, i)

 

-            samples = torch.nn.functional.interpolate(samples, size=(self.height // opt_f, self.width // opt_f), mode="bilinear")

+            samples = torch.nn.functional.interpolate(samples, size=(target_height // opt_f, target_width // opt_f), mode=latent_scale_mode["mode"], antialias=latent_scale_mode["antialias"])

 

             # Avoid making the inpainting conditioning unless necessary as

             # this does need some extra compute to decode / encode the image again.

@@ -745,7 +869,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 

                 save_intermediate(image, i)

 

-                image = images.resize_image(0, image, self.width, self.height)

+                image = images.resize_image(0, image, target_width, target_height, upscaler_name=self.hr_upscaler)

                 image = np.array(image).astype(np.float32) / 255.0

                 image = np.moveaxis(image, 2, 0)

                 batch_images.append(image)

@@ -760,15 +884,18 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 

         shared.state.nextjob()

 

-        self.sampler = sd_samplers.create_sampler(self.sampler_name, self.sd_model)

+        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

+        self.sampler = sd_samplers.create_sampler(img2img_sampler_name, self.sd_model)

+

+        samples = samples[:, :, self.truncate_y//2:samples.shape[2]-(self.truncate_y+1)//2, self.truncate_x//2:samples.shape[3]-(self.truncate_x+1)//2]

 

-        noise = create_random_tensors(samples.shape[1:], seeds=seeds, subseeds=subseeds, subseed_strength=subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w, p=self)

+        noise = create_random_tensors(samples.shape[1:], seeds=seeds, subseeds=subseeds, subseed_strength=subseed_strength, p=self)

 

         # GC now before running the next img2img to prevent running out of memory

         x = None

         devices.torch_gc()

 

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

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

 

         return samples

 

@@ -880,7 +1007,7 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
 

         image = torch.from_numpy(batch_images)

         image = 2. * image - 1.

-        image = image.to(shared.device)

+        image = image.to(device=shared.device, dtype=devices.dtype_vae if devices.unet_needs_upcast else None)

 

         self.init_latent = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(image))

 

diff --git a/modules/progress.py b/modules/progress.py
new file mode 100644
index 00000000..c69ecf3d
--- /dev/null
+++ b/modules/progress.py
@@ -0,0 +1,99 @@
+import base64

+import io

+import time

+

+import gradio as gr

+from pydantic import BaseModel, Field

+

+from modules.shared import opts

+

+import modules.shared as shared

+

+

+current_task = None

+pending_tasks = {}

+finished_tasks = []

+

+

+def start_task(id_task):

+    global current_task

+

+    current_task = id_task

+    pending_tasks.pop(id_task, None)

+

+

+def finish_task(id_task):

+    global current_task

+

+    if current_task == id_task:

+        current_task = None

+

+    finished_tasks.append(id_task)

+    if len(finished_tasks) > 16:

+        finished_tasks.pop(0)

+

+

+def add_task_to_queue(id_job):

+    pending_tasks[id_job] = time.time()

+

+

+class ProgressRequest(BaseModel):

+    id_task: str = Field(default=None, title="Task ID", description="id of the task to get progress for")

+    id_live_preview: int = Field(default=-1, title="Live preview image ID", description="id of last received last preview image")

+

+

+class ProgressResponse(BaseModel):

+    active: bool = Field(title="Whether the task is being worked on right now")

+    queued: bool = Field(title="Whether the task is in queue")

+    completed: bool = Field(title="Whether the task has already finished")

+    progress: float = Field(default=None, title="Progress", description="The progress with a range of 0 to 1")

+    eta: float = Field(default=None, title="ETA in secs")

+    live_preview: str = Field(default=None, title="Live preview image", description="Current live preview; a data: uri")

+    id_live_preview: int = Field(default=None, title="Live preview image ID", description="Send this together with next request to prevent receiving same image")

+    textinfo: str = Field(default=None, title="Info text", description="Info text used by WebUI.")

+

+

+def setup_progress_api(app):

+    return app.add_api_route("/internal/progress", progressapi, methods=["POST"], response_model=ProgressResponse)

+

+

+def progressapi(req: ProgressRequest):

+    active = req.id_task == current_task

+    queued = req.id_task in pending_tasks

+    completed = req.id_task in finished_tasks

+

+    if not active:

+        return ProgressResponse(active=active, queued=queued, completed=completed, id_live_preview=-1, textinfo="In queue..." if queued else "Waiting...")

+

+    progress = 0

+

+    job_count, job_no = shared.state.job_count, shared.state.job_no

+    sampling_steps, sampling_step = shared.state.sampling_steps, shared.state.sampling_step

+

+    if job_count > 0:

+        progress += job_no / job_count

+    if sampling_steps > 0 and job_count > 0:

+        progress += 1 / job_count * sampling_step / sampling_steps

+

+    progress = min(progress, 1)

+

+    elapsed_since_start = time.time() - shared.state.time_start

+    predicted_duration = elapsed_since_start / progress if progress > 0 else None

+    eta = predicted_duration - elapsed_since_start if predicted_duration is not None else None

+

+    id_live_preview = req.id_live_preview

+    shared.state.set_current_image()

+    if opts.live_previews_enable and shared.state.id_live_preview != req.id_live_preview:

+        image = shared.state.current_image

+        if image is not None:

+            buffered = io.BytesIO()

+            image.save(buffered, format="png")

+            live_preview = 'data:image/png;base64,' + base64.b64encode(buffered.getvalue()).decode("ascii")

+            id_live_preview = shared.state.id_live_preview

+        else:

+            live_preview = None

+    else:

+        live_preview = None

+

+    return ProgressResponse(active=active, queued=queued, completed=completed, progress=progress, eta=eta, live_preview=live_preview, id_live_preview=id_live_preview, textinfo=shared.state.textinfo)

+

diff --git a/modules/prompt_parser.py b/modules/prompt_parser.py
index f70872c4..69665372 100644
--- a/modules/prompt_parser.py
+++ b/modules/prompt_parser.py
@@ -49,6 +49,8 @@ def get_learned_conditioning_prompt_schedules(prompts, steps):
     [[5, 'a  c'], [10, 'a {b|d{ c']]

     >>> g("((a][:b:c [d:3]")

     [[3, '((a][:b:c '], [10, '((a][:b:c d']]

+    >>> g("[a|(b:1.1)]")

+    [[1, 'a'], [2, '(b:1.1)'], [3, 'a'], [4, '(b:1.1)'], [5, 'a'], [6, '(b:1.1)'], [7, 'a'], [8, '(b:1.1)'], [9, 'a'], [10, '(b:1.1)']]

     """

 

     def collect_steps(steps, tree):

@@ -84,7 +86,7 @@ def get_learned_conditioning_prompt_schedules(prompts, steps):
                 yield args[0].value

             def __default__(self, data, children, meta):

                 for child in children:

-                    yield from child

+                    yield child

         return AtStep().transform(tree)

 

     def get_schedule(prompt):

@@ -272,6 +274,7 @@ re_attention = re.compile(r"""
 :

 """, re.X)

 

+re_break = re.compile(r"\s*\bBREAK\b\s*", re.S)

 

 def parse_prompt_attention(text):

     """

@@ -337,7 +340,11 @@ def parse_prompt_attention(text):
         elif text == ']' and len(square_brackets) > 0:

             multiply_range(square_brackets.pop(), square_bracket_multiplier)

         else:

-            res.append([text, 1.0])

+            parts = re.split(re_break, text)

+            for i, part in enumerate(parts):

+                if i > 0:

+                    res.append(["BREAK", -1])

+                res.append([part, 1.0])

 

     for pos in round_brackets:

         multiply_range(pos, round_bracket_multiplier)

diff --git a/modules/realesrgan_model.py b/modules/realesrgan_model.py
index 3ac0b97a..47f70251 100644
--- a/modules/realesrgan_model.py
+++ b/modules/realesrgan_model.py
@@ -38,13 +38,13 @@ class UpscalerRealESRGAN(Upscaler):
             return img

 

         info = self.load_model(path)

-        if not os.path.exists(info.data_path):

+        if not os.path.exists(info.local_data_path):

             print("Unable to load RealESRGAN model: %s" % info.name)

             return img

 

         upsampler = RealESRGANer(

             scale=info.scale,

-            model_path=info.data_path,

+            model_path=info.local_data_path,

             model=info.model(),

             half=not cmd_opts.no_half,

             tile=opts.ESRGAN_tile,

@@ -58,17 +58,13 @@ class UpscalerRealESRGAN(Upscaler):
 

     def load_model(self, path):

         try:

-            info = None

-            for scaler in self.scalers:

-                if scaler.data_path == path:

-                    info = scaler

+            info = next(iter([scaler for scaler in self.scalers if scaler.data_path == path]), None)

 

             if info is None:

                 print(f"Unable to find model info: {path}")

                 return None

 

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

-            info.data_path = model_file

+            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/script_callbacks.py b/modules/script_callbacks.py
index 8e22f875..4bb45ec7 100644
--- a/modules/script_callbacks.py
+++ b/modules/script_callbacks.py
@@ -2,7 +2,7 @@ import sys
 import traceback

 from collections import namedtuple

 import inspect

-from typing import Optional

+from typing import Optional, Dict, Any

 

 from fastapi import FastAPI

 from gradio import Blocks

@@ -51,6 +51,13 @@ class UiTrainTabParams:
         self.txt2img_preview_params = txt2img_preview_params

 

 

+class ImageGridLoopParams:

+    def __init__(self, imgs, cols, rows):

+        self.imgs = imgs

+        self.cols = cols

+        self.rows = rows

+

+

 ScriptCallback = namedtuple("ScriptCallback", ["script", "callback"])

 callback_map = dict(

     callbacks_app_started=[],

@@ -63,6 +70,10 @@ callback_map = dict(
     callbacks_cfg_denoiser=[],

     callbacks_before_component=[],

     callbacks_after_component=[],

+    callbacks_image_grid=[],

+    callbacks_infotext_pasted=[],

+    callbacks_script_unloaded=[],

+    callbacks_before_ui=[],

 )

 

 

@@ -155,6 +166,38 @@ def after_component_callback(component, **kwargs):
             report_exception(c, 'after_component_callback')

 

 

+def image_grid_callback(params: ImageGridLoopParams):

+    for c in callback_map['callbacks_image_grid']:

+        try:

+            c.callback(params)

+        except Exception:

+            report_exception(c, 'image_grid')

+

+

+def infotext_pasted_callback(infotext: str, params: Dict[str, Any]):

+    for c in callback_map['callbacks_infotext_pasted']:

+        try:

+            c.callback(infotext, params)

+        except Exception:

+            report_exception(c, 'infotext_pasted')

+

+

+def script_unloaded_callback():

+    for c in reversed(callback_map['callbacks_script_unloaded']):

+        try:

+            c.callback()

+        except Exception:

+            report_exception(c, 'script_unloaded')

+

+

+def before_ui_callback():

+    for c in reversed(callback_map['callbacks_before_ui']):

+        try:

+            c.callback()

+        except Exception:

+            report_exception(c, 'before_ui')

+

+

 def add_callback(callbacks, fun):

     stack = [x for x in inspect.stack() if x.filename != __file__]

     filename = stack[0].filename if len(stack) > 0 else 'unknown file'

@@ -186,7 +229,7 @@ def on_app_started(callback):
 

 def on_model_loaded(callback):

     """register a function to be called when the stable diffusion model is created; the model is

-    passed as an argument"""

+    passed as an argument; this function is also called when the script is reloaded. """

     add_callback(callback_map['callbacks_model_loaded'], callback)

 

 

@@ -255,3 +298,33 @@ def on_before_component(callback):
 def on_after_component(callback):

     """register a function to be called after a component is created. See on_before_component for more."""

     add_callback(callback_map['callbacks_after_component'], callback)

+

+

+def on_image_grid(callback):

+    """register a function to be called before making an image grid.

+    The callback is called with one argument:

+       - params: ImageGridLoopParams - parameters to be used for grid creation. Can be modified.

+    """

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

+

+

+def on_infotext_pasted(callback):

+    """register a function to be called before applying an infotext.

+    The callback is called with two arguments:

+       - infotext: str - raw infotext.

+       - result: Dict[str, any] - parsed infotext parameters.

+    """

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

+

+

+def on_script_unloaded(callback):

+    """register a function to be called before the script is unloaded. Any hooks/hijacks/monkeying about that

+    the script did should be reverted here"""

+

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

+

+

+def on_before_ui(callback):

+    """register a function to be called before the UI is created."""

+

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

diff --git a/modules/script_loading.py b/modules/script_loading.py
index f93f0951..a7d2203f 100644
--- a/modules/script_loading.py
+++ b/modules/script_loading.py
@@ -1,16 +1,14 @@
 import os

 import sys

 import traceback

+import importlib.util

 from types import ModuleType

 

 

 def load_module(path):

-    with open(path, "r", encoding="utf8") as file:

-        text = file.read()

-

-    compiled = compile(text, path, 'exec')

-    module = ModuleType(os.path.basename(path))

-    exec(compiled, module.__dict__)

+    module_spec = importlib.util.spec_from_file_location(os.path.basename(path), path)

+    module = importlib.util.module_from_spec(module_spec)

+    module_spec.loader.exec_module(module)

 

     return module

 

diff --git a/modules/scripts.py b/modules/scripts.py
index 722f8685..6e9dc0c0 100644
--- a/modules/scripts.py
+++ b/modules/scripts.py
@@ -1,16 +1,21 @@
 import os

+import re

 import sys

 import traceback

 from collections import namedtuple

 

 import gradio as gr

 

-from modules.processing import StableDiffusionProcessing

-from modules import shared, paths, script_callbacks, extensions, script_loading

+from modules import shared, paths, script_callbacks, extensions, script_loading, scripts_postprocessing

 

 AlwaysVisible = object()

 

 

+class PostprocessImageArgs:

+    def __init__(self, image):

+        self.image = image

+

+

 class Script:

     filename = None

     args_from = None

@@ -64,7 +69,7 @@ class Script:
         args contains all values returned by components from ui()

         """

 

-        raise NotImplementedError()

+        pass

 

     def process(self, p, *args):

         """

@@ -99,6 +104,13 @@ class Script:
 

         pass

 

+    def postprocess_image(self, p, pp: PostprocessImageArgs, *args):

+        """

+        Called for every image after it has been generated.

+        """

+

+        pass

+

     def postprocess(self, p, processed, *args):

         """

         This function is called after processing ends for AlwaysVisible scripts.

@@ -128,6 +140,15 @@ class Script:
         """unused"""

         return ""

 

+    def elem_id(self, item_id):

+        """helper function to generate id for a HTML element, constructs final id out of script name, tab and user-supplied item_id"""

+

+        need_tabname = self.show(True) == self.show(False)

+        tabname = ('img2img' if self.is_img2img else 'txt2txt') + "_" if need_tabname else ""

+        title = re.sub(r'[^a-z_0-9]', '', re.sub(r'\s', '_', self.title().lower()))

+

+        return f'script_{tabname}{title}_{item_id}'

+

 

 current_basedir = paths.script_path

 

@@ -140,9 +161,11 @@ def basedir():
     return current_basedir

 

 

-scripts_data = []

 ScriptFile = namedtuple("ScriptFile", ["basedir", "filename", "path"])

-ScriptClassData = namedtuple("ScriptClassData", ["script_class", "path", "basedir"])

+

+scripts_data = []

+postprocessing_scripts_data = []

+ScriptClassData = namedtuple("ScriptClassData", ["script_class", "path", "basedir", "module"])

 

 

 def list_scripts(scriptdirname, extension):

@@ -180,23 +203,31 @@ def list_files_with_name(filename):
 def load_scripts():

     global current_basedir

     scripts_data.clear()

+    postprocessing_scripts_data.clear()

     script_callbacks.clear_callbacks()

 

     scripts_list = list_scripts("scripts", ".py")

 

     syspath = sys.path

 

+    def register_scripts_from_module(module):

+        for key, script_class in module.__dict__.items():

+            if type(script_class) != type:

+                continue

+

+            if issubclass(script_class, Script):

+                scripts_data.append(ScriptClassData(script_class, scriptfile.path, scriptfile.basedir, module))

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

         try:

             if scriptfile.basedir != paths.script_path:

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

             current_basedir = scriptfile.basedir

 

-            module = script_loading.load_module(scriptfile.path)

-

-            for key, script_class in module.__dict__.items():

-                if type(script_class) == type and issubclass(script_class, Script):

-                    scripts_data.append(ScriptClassData(script_class, scriptfile.path, scriptfile.basedir))

+            script_module = script_loading.load_module(scriptfile.path)

+            register_scripts_from_module(script_module)

 

         except Exception:

             print(f"Error loading script: {scriptfile.filename}", file=sys.stderr)

@@ -227,11 +258,15 @@ class ScriptRunner:
         self.infotext_fields = []

 

     def initialize_scripts(self, is_img2img):

+        from modules import scripts_auto_postprocessing

+

         self.scripts.clear()

         self.alwayson_scripts.clear()

         self.selectable_scripts.clear()

 

-        for script_class, path, basedir in scripts_data:

+        auto_processing_scripts = scripts_auto_postprocessing.create_auto_preprocessing_script_data()

+

+        for script_class, path, basedir, script_module in auto_processing_scripts + scripts_data:

             script = script_class()

             script.filename = path

             script.is_txt2img = not is_img2img

@@ -280,7 +315,6 @@ class ScriptRunner:
             script.group = group

 

         dropdown = gr.Dropdown(label="Script", elem_id="script_list", choices=["None"] + self.titles, value="None", type="index")

-        dropdown.save_to_config = True

         inputs[0] = dropdown

 

         for script in self.selectable_scripts:

@@ -313,7 +347,7 @@ class ScriptRunner:
 

         return inputs

 

-    def run(self, p: StableDiffusionProcessing, *args):

+    def run(self, p, *args):

         script_index = args[0]

 

         if script_index == 0:

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

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

 

+    def postprocess_image(self, p, pp: PostprocessImageArgs):

+        for script in self.alwayson_scripts:

+            try:

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

+                script.postprocess_image(p, pp, *script_args)

+            except Exception:

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

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

+

     def before_component(self, component, **kwargs):

         for script in self.scripts:

             try:

@@ -404,6 +447,7 @@ class ScriptRunner:
 

 scripts_txt2img = ScriptRunner()

 scripts_img2img = ScriptRunner()

+scripts_postproc = scripts_postprocessing.ScriptPostprocessingRunner()

 scripts_current: ScriptRunner = None

 

 

@@ -414,12 +458,13 @@ def reload_script_body_only():
 

 

 def reload_scripts():

-    global scripts_txt2img, scripts_img2img

+    global scripts_txt2img, scripts_img2img, scripts_postproc

 

     load_scripts()

 

     scripts_txt2img = ScriptRunner()

     scripts_img2img = ScriptRunner()

+    scripts_postproc = scripts_postprocessing.ScriptPostprocessingRunner()

 

 

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

diff --git a/modules/scripts_auto_postprocessing.py b/modules/scripts_auto_postprocessing.py
new file mode 100644
index 00000000..30d6d658
--- /dev/null
+++ b/modules/scripts_auto_postprocessing.py
@@ -0,0 +1,42 @@
+from modules import scripts, scripts_postprocessing, shared

+

+

+class ScriptPostprocessingForMainUI(scripts.Script):

+    def __init__(self, script_postproc):

+        self.script: scripts_postprocessing.ScriptPostprocessing = script_postproc

+        self.postprocessing_controls = None

+

+    def title(self):

+        return self.script.name

+

+    def show(self, is_img2img):

+        return scripts.AlwaysVisible

+

+    def ui(self, is_img2img):

+        self.postprocessing_controls = self.script.ui()

+        return self.postprocessing_controls.values()

+

+    def postprocess_image(self, p, script_pp, *args):

+        args_dict = {k: v for k, v in zip(self.postprocessing_controls, args)}

+

+        pp = scripts_postprocessing.PostprocessedImage(script_pp.image)

+        pp.info = {}

+        self.script.process(pp, **args_dict)

+        p.extra_generation_params.update(pp.info)

+        script_pp.image = pp.image

+

+

+def create_auto_preprocessing_script_data():

+    from modules import scripts

+

+    res = []

+

+    for name in shared.opts.postprocessing_enable_in_main_ui:

+        script = next(iter([x for x in scripts.postprocessing_scripts_data if x.script_class.name == name]), None)

+        if script is None:

+            continue

+

+        constructor = lambda s=script: ScriptPostprocessingForMainUI(s.script_class())

+        res.append(scripts.ScriptClassData(script_class=constructor, path=script.path, basedir=script.basedir, module=script.module))

+

+    return res

diff --git a/modules/scripts_postprocessing.py b/modules/scripts_postprocessing.py
new file mode 100644
index 00000000..ce0ebb61
--- /dev/null
+++ b/modules/scripts_postprocessing.py
@@ -0,0 +1,152 @@
+import os

+import gradio as gr

+

+from modules import errors, shared

+

+

+class PostprocessedImage:

+    def __init__(self, image):

+        self.image = image

+        self.info = {}

+

+

+class ScriptPostprocessing:

+    filename = None

+    controls = None

+    args_from = None

+    args_to = None

+

+    order = 1000

+    """scripts will be ordred by this value in postprocessing UI"""

+

+    name = None

+    """this function should return the title of the script."""

+

+    group = None

+    """A gr.Group component that has all script's UI inside it"""

+

+    def ui(self):

+        """

+        This function should create gradio UI elements. See https://gradio.app/docs/#components

+        The return value should be a dictionary that maps parameter names to components used in processing.

+        Values of those components will be passed to process() function.

+        """

+

+        pass

+

+    def process(self, pp: PostprocessedImage, **args):

+        """

+        This function is called to postprocess the image.

+        args contains a dictionary with all values returned by components from ui()

+        """

+

+        pass

+

+    def image_changed(self):

+        pass

+

+

+

+

+def wrap_call(func, filename, funcname, *args, default=None, **kwargs):

+    try:

+        res = func(*args, **kwargs)

+        return res

+    except Exception as e:

+        errors.display(e, f"calling {filename}/{funcname}")

+

+    return default

+

+

+class ScriptPostprocessingRunner:

+    def __init__(self):

+        self.scripts = None

+        self.ui_created = False

+

+    def initialize_scripts(self, scripts_data):

+        self.scripts = []

+

+        for script_class, path, basedir, script_module in scripts_data:

+            script: ScriptPostprocessing = script_class()

+            script.filename = path

+

+            if script.name == "Simple Upscale":

+                continue

+

+            self.scripts.append(script)

+

+    def create_script_ui(self, script, inputs):

+        script.args_from = len(inputs)

+        script.args_to = len(inputs)

+

+        script.controls = wrap_call(script.ui, script.filename, "ui")

+

+        for control in script.controls.values():

+            control.custom_script_source = os.path.basename(script.filename)

+

+        inputs += list(script.controls.values())

+        script.args_to = len(inputs)

+

+    def scripts_in_preferred_order(self):

+        if self.scripts is None:

+            import modules.scripts

+            self.initialize_scripts(modules.scripts.postprocessing_scripts_data)

+

+        scripts_order = shared.opts.postprocessing_operation_order

+

+        def script_score(name):

+            for i, possible_match in enumerate(scripts_order):

+                if possible_match == name:

+                    return i

+

+            return len(self.scripts)

+

+        script_scores = {script.name: (script_score(script.name), script.order, script.name, original_index) for original_index, script in enumerate(self.scripts)}

+

+        return sorted(self.scripts, key=lambda x: script_scores[x.name])

+

+    def setup_ui(self):

+        inputs = []

+

+        for script in self.scripts_in_preferred_order():

+            with gr.Box() as group:

+                self.create_script_ui(script, inputs)

+

+            script.group = group

+

+        self.ui_created = True

+        return inputs

+

+    def run(self, pp: PostprocessedImage, args):

+        for script in self.scripts_in_preferred_order():

+            shared.state.job = script.name

+

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

+

+            process_args = {}

+            for (name, component), value in zip(script.controls.items(), script_args):

+                process_args[name] = value

+

+            script.process(pp, **process_args)

+

+    def create_args_for_run(self, scripts_args):

+        if not self.ui_created:

+            with gr.Blocks(analytics_enabled=False):

+                self.setup_ui()

+

+        scripts = self.scripts_in_preferred_order()

+        args = [None] * max([x.args_to for x in scripts])

+

+        for script in scripts:

+            script_args_dict = scripts_args.get(script.name, None)

+            if script_args_dict is not None:

+

+                for i, name in enumerate(script.controls):

+                    args[script.args_from + i] = script_args_dict.get(name, None)

+

+        return args

+

+    def image_changed(self):

+        for script in self.scripts_in_preferred_order():

+            script.image_changed()

+

diff --git a/modules/sd_disable_initialization.py b/modules/sd_disable_initialization.py
new file mode 100644
index 00000000..e90aa9fe
--- /dev/null
+++ b/modules/sd_disable_initialization.py
@@ -0,0 +1,90 @@
+import ldm.modules.encoders.modules

+import open_clip

+import torch

+import transformers.utils.hub

+

+

+class DisableInitialization:

+    """

+    When an object of this class enters a `with` block, it starts:

+    - preventing torch's layer initialization functions from working

+    - changes CLIP and OpenCLIP to not download model weights

+    - changes CLIP to not make requests to check if there is a new version of a file you already have

+

+    When it leaves the block, it reverts everything to how it was before.

+

+    Use it like this:

+    ```

+    with DisableInitialization():

+        do_things()

+    ```

+    """

+

+    def __init__(self):

+        self.replaced = []

+

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

+        original = getattr(obj, field, None)

+        if original is None:

+            return None

+

+        self.replaced.append((obj, field, original))

+        setattr(obj, field, func)

+

+        return original

+

+    def __enter__(self):

+        def do_nothing(*args, **kwargs):

+            pass

+

+        def create_model_and_transforms_without_pretrained(*args, pretrained=None, **kwargs):

+            return self.create_model_and_transforms(*args, pretrained=None, **kwargs)

+

+        def CLIPTextModel_from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs):

+            res = self.CLIPTextModel_from_pretrained(None, *model_args, config=pretrained_model_name_or_path, state_dict={}, **kwargs)

+            res.name_or_path = pretrained_model_name_or_path

+            return res

+

+        def transformers_modeling_utils_load_pretrained_model(*args, **kwargs):

+            args = args[0:3] + ('/', ) + args[4:]  # resolved_archive_file; must set it to something to prevent what seems to be a bug

+            return self.transformers_modeling_utils_load_pretrained_model(*args, **kwargs)

+

+        def transformers_utils_hub_get_file_from_cache(original, url, *args, **kwargs):

+

+            # this file is always 404, prevent making request

+            if url == 'https://huggingface.co/openai/clip-vit-large-patch14/resolve/main/added_tokens.json' or url == 'openai/clip-vit-large-patch14' and args[0] == 'added_tokens.json':

+                return None

+

+            try:

+                res = original(url, *args, local_files_only=True, **kwargs)

+                if res is None:

+                    res = original(url, *args, local_files_only=False, **kwargs)

+                return res

+            except Exception as e:

+                return original(url, *args, local_files_only=False, **kwargs)

+

+        def transformers_utils_hub_get_from_cache(url, *args, local_files_only=False, **kwargs):

+            return transformers_utils_hub_get_file_from_cache(self.transformers_utils_hub_get_from_cache, url, *args, **kwargs)

+

+        def transformers_tokenization_utils_base_cached_file(url, *args, local_files_only=False, **kwargs):

+            return transformers_utils_hub_get_file_from_cache(self.transformers_tokenization_utils_base_cached_file, url, *args, **kwargs)

+

+        def transformers_configuration_utils_cached_file(url, *args, local_files_only=False, **kwargs):

+            return transformers_utils_hub_get_file_from_cache(self.transformers_configuration_utils_cached_file, url, *args, **kwargs)

+

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

+

+    def __exit__(self, exc_type, exc_val, exc_tb):

+        for obj, field, original in self.replaced:

+            setattr(obj, field, original)

+

+        self.replaced.clear()

+

diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index 690a9ec2..f9652d21 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -5,9 +5,7 @@ import modules.textual_inversion.textual_inversion
 from modules import devices, sd_hijack_optimizations, shared, sd_hijack_checkpoint

 from modules.hypernetworks import hypernetwork

 from modules.shared import cmd_opts

-from modules import sd_hijack_clip, sd_hijack_open_clip, sd_hijack_unet

-

-from modules.sd_hijack_optimizations import invokeAI_mps_available

+from modules import sd_hijack_clip, sd_hijack_open_clip, sd_hijack_unet, sd_hijack_xlmr, xlmr

 

 import ldm.modules.attention

 import ldm.modules.diffusionmodules.model

@@ -35,26 +33,34 @@ def apply_optimizations():
 

     ldm.modules.diffusionmodules.model.nonlinearity = silu

     ldm.modules.diffusionmodules.openaimodel.th = sd_hijack_unet.th

+    

+    optimization_method = None

 

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

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

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

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

+        optimization_method = 'sub-quadratic'

     elif cmd_opts.opt_split_attention_v1:

         print("Applying v1 cross attention optimization.")

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

-    elif not cmd_opts.disable_opt_split_attention and (cmd_opts.opt_split_attention_invokeai or not torch.cuda.is_available()):

-        if not invokeAI_mps_available and shared.device.type == 'mps':

-            print("The InvokeAI cross attention optimization for MPS requires the psutil package which is not installed.")

-            print("Applying v1 cross attention optimization.")

-            ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.split_cross_attention_forward_v1

-        else:

-            print("Applying cross attention optimization (InvokeAI).")

-            ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.split_cross_attention_forward_invokeAI

+        optimization_method = 'V1'

+    elif not cmd_opts.disable_opt_split_attention and (cmd_opts.opt_split_attention_invokeai or not cmd_opts.opt_split_attention and not torch.cuda.is_available()):

+        print("Applying cross attention optimization (InvokeAI).")

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

+        optimization_method = 'InvokeAI'

     elif not cmd_opts.disable_opt_split_attention and (cmd_opts.opt_split_attention or torch.cuda.is_available()):

         print("Applying cross attention optimization (Doggettx).")

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

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

+        optimization_method = 'Doggettx'

+

+    return optimization_method

 

 

 def undo_optimizations():

@@ -64,9 +70,11 @@ def undo_optimizations():
 

 

 def fix_checkpoint():

-    ldm.modules.attention.BasicTransformerBlock.forward = sd_hijack_checkpoint.BasicTransformerBlock_forward

-    ldm.modules.diffusionmodules.openaimodel.ResBlock.forward = sd_hijack_checkpoint.ResBlock_forward

-    ldm.modules.diffusionmodules.openaimodel.AttentionBlock.forward = sd_hijack_checkpoint.AttentionBlock_forward

+    """checkpoints are now added and removed in embedding/hypernet code, since torch doesn't want

+    checkpoints to be added when not training (there's a warning)"""

+

+    pass

+

 

 class StableDiffusionModelHijack:

     fixes = None

@@ -74,22 +82,31 @@ class StableDiffusionModelHijack:
     layers = None

     circular_enabled = False

     clip = None

+    optimization_method = None

+

+    embedding_db = modules.textual_inversion.textual_inversion.EmbeddingDatabase()

 

-    embedding_db = modules.textual_inversion.textual_inversion.EmbeddingDatabase(cmd_opts.embeddings_dir)

+    def __init__(self):

+        self.embedding_db.add_embedding_dir(cmd_opts.embeddings_dir)

 

     def hijack(self, m):

-        if type(m.cond_stage_model) == ldm.modules.encoders.modules.FrozenCLIPEmbedder:

+        if type(m.cond_stage_model) == xlmr.BertSeriesModelWithTransformation:

+            model_embeddings = m.cond_stage_model.roberta.embeddings

+            model_embeddings.token_embedding = EmbeddingsWithFixes(model_embeddings.word_embeddings, self)

+            m.cond_stage_model = sd_hijack_xlmr.FrozenXLMREmbedderWithCustomWords(m.cond_stage_model, self)

+

+        elif type(m.cond_stage_model) == ldm.modules.encoders.modules.FrozenCLIPEmbedder:

             model_embeddings = m.cond_stage_model.transformer.text_model.embeddings

             model_embeddings.token_embedding = EmbeddingsWithFixes(model_embeddings.token_embedding, self)

             m.cond_stage_model = sd_hijack_clip.FrozenCLIPEmbedderWithCustomWords(m.cond_stage_model, self)

+

         elif type(m.cond_stage_model) == ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder:

             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)

 

-        self.clip = m.cond_stage_model

+        self.optimization_method = apply_optimizations()

 

-        apply_optimizations()

-        fix_checkpoint()

+        self.clip = m.cond_stage_model

 

         def flatten(el):

             flattened = [flatten(children) for children in el.children()]

@@ -101,7 +118,10 @@ class StableDiffusionModelHijack:
         self.layers = flatten(m)

 

     def undo_hijack(self, m):

-        if type(m.cond_stage_model) == sd_hijack_clip.FrozenCLIPEmbedderWithCustomWords:

+        if type(m.cond_stage_model) == xlmr.BertSeriesModelWithTransformation:

+            m.cond_stage_model = m.cond_stage_model.wrapped 

+

+        elif type(m.cond_stage_model) == sd_hijack_clip.FrozenCLIPEmbedderWithCustomWords:

             m.cond_stage_model = m.cond_stage_model.wrapped

 

             model_embeddings = m.cond_stage_model.transformer.text_model.embeddings

@@ -127,10 +147,10 @@ class StableDiffusionModelHijack:
     def clear_comments(self):

         self.comments = []

 

-    def tokenize(self, text):

-        _, remade_batch_tokens, _, _, _, token_count = self.clip.process_text([text])

-        return remade_batch_tokens[0], token_count, sd_hijack_clip.get_target_prompt_token_count(token_count)

+    def get_prompt_lengths(self, text):

+        _, token_count = self.clip.process_texts([text])

 

+        return token_count, self.clip.get_target_prompt_token_count(token_count)

 

 

 class EmbeddingsWithFixes(torch.nn.Module):

diff --git a/modules/sd_hijack_checkpoint.py b/modules/sd_hijack_checkpoint.py
index 5712972f..2604d969 100644
--- a/modules/sd_hijack_checkpoint.py
+++ b/modules/sd_hijack_checkpoint.py
@@ -1,10 +1,46 @@
 from torch.utils.checkpoint import checkpoint
 
+import ldm.modules.attention
+import ldm.modules.diffusionmodules.openaimodel
+
+
 def BasicTransformerBlock_forward(self, x, context=None):
     return checkpoint(self._forward, x, context)
 
+
 def AttentionBlock_forward(self, x):
     return checkpoint(self._forward, x)
 
+
 def ResBlock_forward(self, x, emb):
-    return checkpoint(self._forward, x, emb)
\ No newline at end of file
+    return checkpoint(self._forward, x, emb)
+
+
+stored = []
+
+
+def add():
+    if len(stored) != 0:
+        return
+
+    stored.extend([
+        ldm.modules.attention.BasicTransformerBlock.forward,
+        ldm.modules.diffusionmodules.openaimodel.ResBlock.forward,
+        ldm.modules.diffusionmodules.openaimodel.AttentionBlock.forward
+    ])
+
+    ldm.modules.attention.BasicTransformerBlock.forward = BasicTransformerBlock_forward
+    ldm.modules.diffusionmodules.openaimodel.ResBlock.forward = ResBlock_forward
+    ldm.modules.diffusionmodules.openaimodel.AttentionBlock.forward = AttentionBlock_forward
+
+
+def remove():
+    if len(stored) == 0:
+        return
+
+    ldm.modules.attention.BasicTransformerBlock.forward = stored[0]
+    ldm.modules.diffusionmodules.openaimodel.ResBlock.forward = stored[1]
+    ldm.modules.diffusionmodules.openaimodel.AttentionBlock.forward = stored[2]
+
+    stored.clear()
+
diff --git a/modules/sd_hijack_clip.py b/modules/sd_hijack_clip.py
index b451d1cf..9fa5c5c5 100644
--- a/modules/sd_hijack_clip.py
+++ b/modules/sd_hijack_clip.py
@@ -1,31 +1,89 @@
 import math

+from collections import namedtuple

 

 import torch

 

-from modules import prompt_parser, devices

+from modules import prompt_parser, devices, sd_hijack

 from modules.shared import opts

 

 

-def get_target_prompt_token_count(token_count):

-    return math.ceil(max(token_count, 1) / 75) * 75

+class PromptChunk:

+    """

+    This object contains token ids, weight (multipliers:1.4) and textual inversion embedding info for a chunk of prompt.

+    If a prompt is short, it is represented by one PromptChunk, otherwise, multiple are necessary.

+    Each PromptChunk contains an exact amount of tokens - 77, which includes one for start and end token,

+    so just 75 tokens from prompt.

+    """

+

+    def __init__(self):

+        self.tokens = []

+        self.multipliers = []

+        self.fixes = []

+

+

+PromptChunkFix = namedtuple('PromptChunkFix', ['offset', 'embedding'])

+"""An object of this type is a marker showing that textual inversion embedding's vectors have to placed at offset in the prompt

+chunk. Thos objects are found in PromptChunk.fixes and, are placed into FrozenCLIPEmbedderWithCustomWordsBase.hijack.fixes, and finally

+are applied by sd_hijack.EmbeddingsWithFixes's forward function."""

 

 

 class FrozenCLIPEmbedderWithCustomWordsBase(torch.nn.Module):

+    """A pytorch module that is a wrapper for FrozenCLIPEmbedder module. it enhances FrozenCLIPEmbedder, making it possible to

+    have unlimited prompt length and assign weights to tokens in prompt.

+    """

+

     def __init__(self, wrapped, hijack):

         super().__init__()

+

         self.wrapped = wrapped

-        self.hijack = hijack

+        """Original FrozenCLIPEmbedder module; can also be FrozenOpenCLIPEmbedder or xlmr.BertSeriesModelWithTransformation,

+        depending on model."""

+

+        self.hijack: sd_hijack.StableDiffusionModelHijack = hijack

+        self.chunk_length = 75

+

+    def empty_chunk(self):

+        """creates an empty PromptChunk and returns it"""

+

+        chunk = PromptChunk()

+        chunk.tokens = [self.id_start] + [self.id_end] * (self.chunk_length + 1)

+        chunk.multipliers = [1.0] * (self.chunk_length + 2)

+        return chunk

+

+    def get_target_prompt_token_count(self, token_count):

+        """returns the maximum number of tokens a prompt of a known length can have before it requires one more PromptChunk to be represented"""

+

+        return math.ceil(max(token_count, 1) / self.chunk_length) * self.chunk_length

 

     def tokenize(self, texts):

+        """Converts a batch of texts into a batch of token ids"""

+

         raise NotImplementedError

 

     def encode_with_transformers(self, tokens):

+        """

+        converts a batch of token ids (in python lists) into a single tensor with numeric respresentation of those tokens;

+        All python lists with tokens are assumed to have same length, usually 77.

+        if input is a list with B elements and each element has T tokens, expected output shape is (B, T, C), where C depends on

+        model - can be 768 and 1024.

+        Among other things, this call will read self.hijack.fixes, apply it to its inputs, and clear it (setting it to None).

+        """

+

         raise NotImplementedError

 

     def encode_embedding_init_text(self, init_text, nvpt):

+        """Converts text into a tensor with this text's tokens' embeddings. Note that those are embeddings before they are passed through

+        transformers. nvpt is used as a maximum length in tokens. If text produces less teokens than nvpt, only this many is returned."""

+

         raise NotImplementedError

 

-    def tokenize_line(self, line, used_custom_terms, hijack_comments):

+    def tokenize_line(self, line):

+        """

+        this transforms a single prompt into a list of PromptChunk objects - as many as needed to

+        represent the prompt.

+        Returns the list and the total number of tokens in the prompt.

+        """

+

         if opts.enable_emphasis:

             parsed = prompt_parser.parse_prompt_attention(line)

         else:

@@ -33,205 +91,161 @@ class FrozenCLIPEmbedderWithCustomWordsBase(torch.nn.Module):
 

         tokenized = self.tokenize([text for text, _ in parsed])

 

-        fixes = []

-        remade_tokens = []

-        multipliers = []

+        chunks = []

+        chunk = PromptChunk()

+        token_count = 0

         last_comma = -1

 

+        def next_chunk(is_last=False):

+            """puts current chunk into the list of results and produces the next one - empty;

+            if is_last is true, tokens <end-of-text> tokens at the end won't add to token_count"""

+            nonlocal token_count

+            nonlocal last_comma

+            nonlocal chunk

+

+            if is_last:

+                token_count += len(chunk.tokens)

+            else:

+                token_count += self.chunk_length

+

+            to_add = self.chunk_length - len(chunk.tokens)

+            if to_add > 0:

+                chunk.tokens += [self.id_end] * to_add

+                chunk.multipliers += [1.0] * to_add

+

+            chunk.tokens = [self.id_start] + chunk.tokens + [self.id_end]

+            chunk.multipliers = [1.0] + chunk.multipliers + [1.0]

+

+            last_comma = -1

+            chunks.append(chunk)

+            chunk = PromptChunk()

+

         for tokens, (text, weight) in zip(tokenized, parsed):

-            i = 0

-            while i < len(tokens):

-                token = tokens[i]

+            if text == 'BREAK' and weight == -1:

+                next_chunk()

+                continue

 

-                embedding, embedding_length_in_tokens = self.hijack.embedding_db.find_embedding_at_position(tokens, i)

+            position = 0

+            while position < len(tokens):

+                token = tokens[position]

 

                 if token == self.comma_token:

-                    last_comma = len(remade_tokens)

-                elif opts.comma_padding_backtrack != 0 and max(len(remade_tokens), 1) % 75 == 0 and last_comma != -1 and len(remade_tokens) - last_comma <= opts.comma_padding_backtrack:

-                    last_comma += 1

-                    reloc_tokens = remade_tokens[last_comma:]

-                    reloc_mults = multipliers[last_comma:]

+                    last_comma = len(chunk.tokens)

+

+                # this is when we are at the end of alloted 75 tokens for the current chunk, and the current token is not a comma. opts.comma_padding_backtrack

+                # is a setting that specifies that if there is a comma nearby, the text after the comma should be moved out of this chunk and into the next.

+                elif opts.comma_padding_backtrack != 0 and len(chunk.tokens) == self.chunk_length and last_comma != -1 and len(chunk.tokens) - last_comma <= opts.comma_padding_backtrack:

+                    break_location = last_comma + 1

+

+                    reloc_tokens = chunk.tokens[break_location:]

+                    reloc_mults = chunk.multipliers[break_location:]

 

-                    remade_tokens = remade_tokens[:last_comma]

-                    length = len(remade_tokens)

+                    chunk.tokens = chunk.tokens[:break_location]

+                    chunk.multipliers = chunk.multipliers[:break_location]

 

-                    rem = int(math.ceil(length / 75)) * 75 - length

-                    remade_tokens += [self.id_end] * rem + reloc_tokens

-                    multipliers = multipliers[:last_comma] + [1.0] * rem + reloc_mults

+                    next_chunk()

+                    chunk.tokens = reloc_tokens

+                    chunk.multipliers = reloc_mults

 

+                if len(chunk.tokens) == self.chunk_length:

+                    next_chunk()

+

+                embedding, embedding_length_in_tokens = self.hijack.embedding_db.find_embedding_at_position(tokens, position)

                 if embedding is None:

-                    remade_tokens.append(token)

-                    multipliers.append(weight)

-                    i += 1

-                else:

-                    emb_len = int(embedding.vec.shape[0])

-                    iteration = len(remade_tokens) // 75

-                    if (len(remade_tokens) + emb_len) // 75 != iteration:

-                        rem = (75 * (iteration + 1) - len(remade_tokens))

-                        remade_tokens += [self.id_end] * rem

-                        multipliers += [1.0] * rem

-                        iteration += 1

-                    fixes.append((iteration, (len(remade_tokens) % 75, embedding)))

-                    remade_tokens += [0] * emb_len

-                    multipliers += [weight] * emb_len

-                    used_custom_terms.append((embedding.name, embedding.checksum()))

-                    i += embedding_length_in_tokens

-

-        token_count = len(remade_tokens)

-        prompt_target_length = get_target_prompt_token_count(token_count)

-        tokens_to_add = prompt_target_length - len(remade_tokens)

-

-        remade_tokens = remade_tokens + [self.id_end] * tokens_to_add

-        multipliers = multipliers + [1.0] * tokens_to_add

-

-        return remade_tokens, fixes, multipliers, token_count

-

-    def process_text(self, texts):

-        used_custom_terms = []

-        remade_batch_tokens = []

-        hijack_comments = []

-        hijack_fixes = []

+                    chunk.tokens.append(token)

+                    chunk.multipliers.append(weight)

+                    position += 1

+                    continue

+

+                emb_len = int(embedding.vec.shape[0])

+                if len(chunk.tokens) + emb_len > self.chunk_length:

+                    next_chunk()

+

+                chunk.fixes.append(PromptChunkFix(len(chunk.tokens), embedding))

+

+                chunk.tokens += [0] * emb_len

+                chunk.multipliers += [weight] * emb_len

+                position += embedding_length_in_tokens

+

+        if len(chunk.tokens) > 0 or len(chunks) == 0:

+            next_chunk(is_last=True)

+

+        return chunks, token_count

+

+    def process_texts(self, texts):

+        """

+        Accepts a list of texts and calls tokenize_line() on each, with cache. Returns the list of results and maximum

+        length, in tokens, of all texts.

+        """

+

         token_count = 0

 

         cache = {}

-        batch_multipliers = []

+        batch_chunks = []

         for line in texts:

             if line in cache:

-                remade_tokens, fixes, multipliers = cache[line]

+                chunks = cache[line]

             else:

-                remade_tokens, fixes, multipliers, current_token_count = self.tokenize_line(line, used_custom_terms, hijack_comments)

+                chunks, current_token_count = self.tokenize_line(line)

                 token_count = max(current_token_count, token_count)

 

-                cache[line] = (remade_tokens, fixes, multipliers)

+                cache[line] = chunks

 

-            remade_batch_tokens.append(remade_tokens)

-            hijack_fixes.append(fixes)

-            batch_multipliers.append(multipliers)

+            batch_chunks.append(chunks)

 

-        return batch_multipliers, remade_batch_tokens, used_custom_terms, hijack_comments, hijack_fixes, token_count

+        return batch_chunks, token_count

 

-    def process_text_old(self, texts):

-        id_start = self.id_start

-        id_end = self.id_end

-        maxlen = self.wrapped.max_length  # you get to stay at 77

-        used_custom_terms = []

-        remade_batch_tokens = []

-        hijack_comments = []

-        hijack_fixes = []

-        token_count = 0

+    def forward(self, texts):

+        """

+        Accepts an array of texts; Passes texts through transformers network to create a tensor with numerical representation of those texts.

+        Returns a tensor with shape of (B, T, C), where B is length of the array; T is length, in tokens, of texts (including padding) - T will

+        be a multiple of 77; and C is dimensionality of each token - for SD1 it's 768, and for SD2 it's 1024.

+        An example shape returned by this function can be: (2, 77, 768).

+        Webui usually sends just one text at a time through this function - the only time when texts is an array with more than one elemenet

+        is when you do prompt editing: "a picture of a [cat:dog:0.4] eating ice cream"

+        """

 

-        cache = {}

-        batch_tokens = self.tokenize(texts)

-        batch_multipliers = []

-        for tokens in batch_tokens:

-            tuple_tokens = tuple(tokens)

+        if opts.use_old_emphasis_implementation:

+            import modules.sd_hijack_clip_old

+            return modules.sd_hijack_clip_old.forward_old(self, texts)

 

-            if tuple_tokens in cache:

-                remade_tokens, fixes, multipliers = cache[tuple_tokens]

-            else:

-                fixes = []

-                remade_tokens = []

-                multipliers = []

-                mult = 1.0

-

-                i = 0

-                while i < len(tokens):

-                    token = tokens[i]

-

-                    embedding, embedding_length_in_tokens = self.hijack.embedding_db.find_embedding_at_position(tokens, i)

-

-                    mult_change = self.token_mults.get(token) if opts.enable_emphasis else None

-                    if mult_change is not None:

-                        mult *= mult_change

-                        i += 1

-                    elif embedding is None:

-                        remade_tokens.append(token)

-                        multipliers.append(mult)

-                        i += 1

-                    else:

-                        emb_len = int(embedding.vec.shape[0])

-                        fixes.append((len(remade_tokens), embedding))

-                        remade_tokens += [0] * emb_len

-                        multipliers += [mult] * emb_len

-                        used_custom_terms.append((embedding.name, embedding.checksum()))

-                        i += embedding_length_in_tokens

-

-                if len(remade_tokens) > maxlen - 2:

-                    vocab = {v: k for k, v in self.wrapped.tokenizer.get_vocab().items()}

-                    ovf = remade_tokens[maxlen - 2:]

-                    overflowing_words = [vocab.get(int(x), "") for x in ovf]

-                    overflowing_text = self.wrapped.tokenizer.convert_tokens_to_string(''.join(overflowing_words))

-                    hijack_comments.append(f"Warning: too many input tokens; some ({len(overflowing_words)}) have been truncated:\n{overflowing_text}\n")

-

-                token_count = len(remade_tokens)

-                remade_tokens = remade_tokens + [id_end] * (maxlen - 2 - len(remade_tokens))

-                remade_tokens = [id_start] + remade_tokens[0:maxlen - 2] + [id_end]

-                cache[tuple_tokens] = (remade_tokens, fixes, multipliers)

-

-            multipliers = multipliers + [1.0] * (maxlen - 2 - len(multipliers))

-            multipliers = [1.0] + multipliers[0:maxlen - 2] + [1.0]

-

-            remade_batch_tokens.append(remade_tokens)

-            hijack_fixes.append(fixes)

-            batch_multipliers.append(multipliers)

-        return batch_multipliers, remade_batch_tokens, used_custom_terms, hijack_comments, hijack_fixes, token_count

-

-    def forward(self, text):

-        use_old = opts.use_old_emphasis_implementation

-        if use_old:

-            batch_multipliers, remade_batch_tokens, used_custom_terms, hijack_comments, hijack_fixes, token_count = self.process_text_old(text)

-        else:

-            batch_multipliers, remade_batch_tokens, used_custom_terms, hijack_comments, hijack_fixes, token_count = self.process_text(text)

-

-        self.hijack.comments += hijack_comments

-

-        if len(used_custom_terms) > 0:

-            self.hijack.comments.append("Used embeddings: " + ", ".join([f'{word} [{checksum}]' for word, checksum in used_custom_terms]))

-

-        if use_old:

-            self.hijack.fixes = hijack_fixes

-            return self.process_tokens(remade_batch_tokens, batch_multipliers)

-

-        z = None

-        i = 0

-        while max(map(len, remade_batch_tokens)) != 0:

-            rem_tokens = [x[75:] for x in remade_batch_tokens]

-            rem_multipliers = [x[75:] for x in batch_multipliers]

-

-            self.hijack.fixes = []

-            for unfiltered in hijack_fixes:

-                fixes = []

-                for fix in unfiltered:

-                    if fix[0] == i:

-                        fixes.append(fix[1])

-                self.hijack.fixes.append(fixes)

-

-            tokens = []

-            multipliers = []

-            for j in range(len(remade_batch_tokens)):

-                if len(remade_batch_tokens[j]) > 0:

-                    tokens.append(remade_batch_tokens[j][:75])

-                    multipliers.append(batch_multipliers[j][:75])

-                else:

-                    tokens.append([self.id_end] * 75)

-                    multipliers.append([1.0] * 75)

-

-            z1 = self.process_tokens(tokens, multipliers)

-            z = z1 if z is None else torch.cat((z, z1), axis=-2)

-

-            remade_batch_tokens = rem_tokens

-            batch_multipliers = rem_multipliers

-            i += 1

+        batch_chunks, token_count = self.process_texts(texts)

 

-        return z

+        used_embeddings = {}

+        chunk_count = max([len(x) for x in batch_chunks])

 

-    def process_tokens(self, remade_batch_tokens, batch_multipliers):

-        if not opts.use_old_emphasis_implementation:

-            remade_batch_tokens = [[self.id_start] + x[:75] + [self.id_end] for x in remade_batch_tokens]

-            batch_multipliers = [[1.0] + x[:75] + [1.0] for x in batch_multipliers]

+        zs = []

+        for i in range(chunk_count):

+            batch_chunk = [chunks[i] if i < len(chunks) else self.empty_chunk() for chunks in batch_chunks]

+

+            tokens = [x.tokens for x in batch_chunk]

+            multipliers = [x.multipliers for x in batch_chunk]

+            self.hijack.fixes = [x.fixes for x in batch_chunk]

+

+            for fixes in self.hijack.fixes:

+                for position, embedding in fixes:

+                    used_embeddings[embedding.name] = embedding

+

+            z = self.process_tokens(tokens, multipliers)

+            zs.append(z)

 

+        if len(used_embeddings) > 0:

+            embeddings_list = ", ".join([f'{name} [{embedding.checksum()}]' for name, embedding in used_embeddings.items()])

+            self.hijack.comments.append(f"Used embeddings: {embeddings_list}")

+

+        return torch.hstack(zs)

+

+    def process_tokens(self, remade_batch_tokens, batch_multipliers):

+        """

+        sends one single prompt chunk to be encoded by transformers neural network.

+        remade_batch_tokens is a batch of tokens - a list, where every element is a list of tokens; usually

+        there are exactly 77 tokens in the list. batch_multipliers is the same but for multipliers instead of tokens.

+        Multipliers are used to give more or less weight to the outputs of transformers network. Each multiplier

+        corresponds to one token.

+        """

         tokens = torch.asarray(remade_batch_tokens).to(devices.device)

 

+        # this is for SD2: SD1 uses the same token for padding and end of text, while SD2 uses different ones.

         if self.id_end != self.id_pad:

             for batch_pos in range(len(remade_batch_tokens)):

                 index = remade_batch_tokens[batch_pos].index(self.id_end)

@@ -240,12 +254,11 @@ class FrozenCLIPEmbedderWithCustomWordsBase(torch.nn.Module):
         z = self.encode_with_transformers(tokens)

 

         # restoring original mean is likely not correct, but it seems to work well to prevent artifacts that happen otherwise

-        batch_multipliers_of_same_length = [x + [1.0] * (75 - len(x)) for x in batch_multipliers]

-        batch_multipliers = torch.asarray(batch_multipliers_of_same_length).to(devices.device)

+        batch_multipliers = torch.asarray(batch_multipliers).to(devices.device)

         original_mean = z.mean()

-        z *= batch_multipliers.reshape(batch_multipliers.shape + (1,)).expand(z.shape)

+        z = z * batch_multipliers.reshape(batch_multipliers.shape + (1,)).expand(z.shape)

         new_mean = z.mean()

-        z *= original_mean / new_mean

+        z = z * (original_mean / new_mean)

 

         return z

 

@@ -254,10 +267,13 @@ class FrozenCLIPEmbedderWithCustomWords(FrozenCLIPEmbedderWithCustomWordsBase):
     def __init__(self, wrapped, hijack):

         super().__init__(wrapped, hijack)

         self.tokenizer = wrapped.tokenizer

-        self.comma_token = [v for k, v in self.tokenizer.get_vocab().items() if k == ',</w>'][0]

+

+        vocab = self.tokenizer.get_vocab()

+

+        self.comma_token = vocab.get(',</w>', None)

 

         self.token_mults = {}

-        tokens_with_parens = [(k, v) for k, v in self.tokenizer.get_vocab().items() if '(' in k or ')' in k or '[' in k or ']' in k]

+        tokens_with_parens = [(k, v) for k, v in vocab.items() if '(' in k or ')' in k or '[' in k or ']' in k]

         for text, ident in tokens_with_parens:

             mult = 1.0

             for c in text:

@@ -296,6 +312,6 @@ class FrozenCLIPEmbedderWithCustomWords(FrozenCLIPEmbedderWithCustomWordsBase):
     def encode_embedding_init_text(self, init_text, nvpt):

         embedding_layer = self.wrapped.transformer.text_model.embeddings

         ids = self.wrapped.tokenizer(init_text, max_length=nvpt, return_tensors="pt", add_special_tokens=False)["input_ids"]

-        embedded = embedding_layer.token_embedding.wrapped(ids.to(devices.device)).squeeze(0)

+        embedded = embedding_layer.token_embedding.wrapped(ids.to(embedding_layer.token_embedding.wrapped.weight.device)).squeeze(0)

 

         return embedded

diff --git a/modules/sd_hijack_clip_old.py b/modules/sd_hijack_clip_old.py
new file mode 100644
index 00000000..6d9fbbe6
--- /dev/null
+++ b/modules/sd_hijack_clip_old.py
@@ -0,0 +1,81 @@
+from modules import sd_hijack_clip

+from modules import shared

+

+

+def process_text_old(self: sd_hijack_clip.FrozenCLIPEmbedderWithCustomWordsBase, texts):

+    id_start = self.id_start

+    id_end = self.id_end

+    maxlen = self.wrapped.max_length  # you get to stay at 77

+    used_custom_terms = []

+    remade_batch_tokens = []

+    hijack_comments = []

+    hijack_fixes = []

+    token_count = 0

+

+    cache = {}

+    batch_tokens = self.tokenize(texts)

+    batch_multipliers = []

+    for tokens in batch_tokens:

+        tuple_tokens = tuple(tokens)

+

+        if tuple_tokens in cache:

+            remade_tokens, fixes, multipliers = cache[tuple_tokens]

+        else:

+            fixes = []

+            remade_tokens = []

+            multipliers = []

+            mult = 1.0

+

+            i = 0

+            while i < len(tokens):

+                token = tokens[i]

+

+                embedding, embedding_length_in_tokens = self.hijack.embedding_db.find_embedding_at_position(tokens, i)

+

+                mult_change = self.token_mults.get(token) if shared.opts.enable_emphasis else None

+                if mult_change is not None:

+                    mult *= mult_change

+                    i += 1

+                elif embedding is None:

+                    remade_tokens.append(token)

+                    multipliers.append(mult)

+                    i += 1

+                else:

+                    emb_len = int(embedding.vec.shape[0])

+                    fixes.append((len(remade_tokens), embedding))

+                    remade_tokens += [0] * emb_len

+                    multipliers += [mult] * emb_len

+                    used_custom_terms.append((embedding.name, embedding.checksum()))

+                    i += embedding_length_in_tokens

+

+            if len(remade_tokens) > maxlen - 2:

+                vocab = {v: k for k, v in self.wrapped.tokenizer.get_vocab().items()}

+                ovf = remade_tokens[maxlen - 2:]

+                overflowing_words = [vocab.get(int(x), "") for x in ovf]

+                overflowing_text = self.wrapped.tokenizer.convert_tokens_to_string(''.join(overflowing_words))

+                hijack_comments.append(f"Warning: too many input tokens; some ({len(overflowing_words)}) have been truncated:\n{overflowing_text}\n")

+

+            token_count = len(remade_tokens)

+            remade_tokens = remade_tokens + [id_end] * (maxlen - 2 - len(remade_tokens))

+            remade_tokens = [id_start] + remade_tokens[0:maxlen - 2] + [id_end]

+            cache[tuple_tokens] = (remade_tokens, fixes, multipliers)

+

+        multipliers = multipliers + [1.0] * (maxlen - 2 - len(multipliers))

+        multipliers = [1.0] + multipliers[0:maxlen - 2] + [1.0]

+

+        remade_batch_tokens.append(remade_tokens)

+        hijack_fixes.append(fixes)

+        batch_multipliers.append(multipliers)

+    return batch_multipliers, remade_batch_tokens, used_custom_terms, hijack_comments, hijack_fixes, token_count

+

+

+def forward_old(self: sd_hijack_clip.FrozenCLIPEmbedderWithCustomWordsBase, texts):

+    batch_multipliers, remade_batch_tokens, used_custom_terms, hijack_comments, hijack_fixes, token_count = process_text_old(self, texts)

+

+    self.hijack.comments += hijack_comments

+

+    if len(used_custom_terms) > 0:

+        self.hijack.comments.append("Used embeddings: " + ", ".join([f'{word} [{checksum}]' for word, checksum in used_custom_terms]))

+

+    self.hijack.fixes = hijack_fixes

+    return self.process_tokens(remade_batch_tokens, batch_multipliers)

diff --git a/modules/sd_hijack_inpainting.py b/modules/sd_hijack_inpainting.py
index bb5499b3..478cd499 100644
--- a/modules/sd_hijack_inpainting.py
+++ b/modules/sd_hijack_inpainting.py
@@ -12,191 +12,6 @@ from ldm.models.diffusion.ddpm import LatentDiffusion
 from ldm.models.diffusion.plms import PLMSSampler
 from ldm.models.diffusion.ddim import DDIMSampler, noise_like
 
-# =================================================================================================
-# Monkey patch DDIMSampler methods from RunwayML repo directly.
-# Adapted from:
-# https://github.com/runwayml/stable-diffusion/blob/main/ldm/models/diffusion/ddim.py
-# =================================================================================================
-@torch.no_grad()
-def sample_ddim(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}")
-        else:
-            if conditioning.shape[0] != batch_size:
-                print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
-
-    self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
-    # sampling
-    C, H, W = shape
-    size = (batch_size, C, H, W)
-    print(f'Data shape for DDIM sampling is {size}, eta {eta}')
-
-    samples, intermediates = self.ddim_sampling(conditioning, size,
-                                                callback=callback,
-                                                img_callback=img_callback,
-                                                quantize_denoised=quantize_x0,
-                                                mask=mask, x0=x0,
-                                                ddim_use_original_steps=False,
-                                                noise_dropout=noise_dropout,
-                                                temperature=temperature,
-                                                score_corrector=score_corrector,
-                                                corrector_kwargs=corrector_kwargs,
-                                                x_T=x_T,
-                                                log_every_t=log_every_t,
-                                                unconditional_guidance_scale=unconditional_guidance_scale,
-                                                unconditional_conditioning=unconditional_conditioning,
-                                                )
-    return samples, intermediates
-
-@torch.no_grad()
-def p_sample_ddim(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
-                    temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
-                    unconditional_guidance_scale=1., unconditional_conditioning=None):
-    b, *_, device = *x.shape, x.device
-
-    if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
-        e_t = self.model.apply_model(x, t, c)
-    else:
-        x_in = torch.cat([x] * 2)
-        t_in = torch.cat([t] * 2)
-        if isinstance(c, dict):
-            assert isinstance(unconditional_conditioning, dict)
-            c_in = dict()
-            for k in c:
-                if isinstance(c[k], list):
-                    c_in[k] = [
-                        torch.cat([unconditional_conditioning[k][i], c[k][i]])
-                        for i in range(len(c[k]))
-                    ]
-                else:
-                    c_in[k] = torch.cat([unconditional_conditioning[k], c[k]])
-        else:
-            c_in = torch.cat([unconditional_conditioning, c])
-        e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
-        e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
-
-    if score_corrector is not None:
-        assert self.model.parameterization == "eps"
-        e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
-
-    alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
-    alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
-    sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
-    sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
-    # select parameters corresponding to the currently considered timestep
-    a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
-    a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
-    sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
-    sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device)
-
-    # current prediction for x_0
-    pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
-    if quantize_denoised:
-        pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
-    # direction pointing to x_t
-    dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
-    noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
-    if noise_dropout > 0.:
-        noise = torch.nn.functional.dropout(noise, p=noise_dropout)
-    x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
-    return x_prev, pred_x0
-
-
-# =================================================================================================
-# Monkey patch PLMSSampler methods.
-# This one was not actually patched correctly in the RunwayML repo, but we can replicate the changes.
-# Adapted from:
-# https://github.com/CompVis/stable-diffusion/blob/main/ldm/models/diffusion/plms.py
-# =================================================================================================
-@torch.no_grad()
-def sample_plms(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}")
-        else:
-            if conditioning.shape[0] != batch_size:
-                print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
-
-    self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
-    # sampling
-    C, H, W = shape
-    size = (batch_size, C, H, W)
-    print(f'Data shape for PLMS sampling is {size}')
-
-    samples, intermediates = self.plms_sampling(conditioning, size,
-                                                callback=callback,
-                                                img_callback=img_callback,
-                                                quantize_denoised=quantize_x0,
-                                                mask=mask, x0=x0,
-                                                ddim_use_original_steps=False,
-                                                noise_dropout=noise_dropout,
-                                                temperature=temperature,
-                                                score_corrector=score_corrector,
-                                                corrector_kwargs=corrector_kwargs,
-                                                x_T=x_T,
-                                                log_every_t=log_every_t,
-                                                unconditional_guidance_scale=unconditional_guidance_scale,
-                                                unconditional_conditioning=unconditional_conditioning,
-                                                )
-    return samples, intermediates
-
 
 @torch.no_grad()
 def p_sample_plms(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
@@ -280,61 +95,8 @@ def p_sample_plms(self, x, c, t, index, repeat_noise=False, use_original_steps=F
 
     return x_prev, pred_x0, e_t
 
-# =================================================================================================
-# Monkey patch LatentInpaintDiffusion to load the checkpoint with a proper config.
-# Adapted from:
-# https://github.com/runwayml/stable-diffusion/blob/main/ldm/models/diffusion/ddpm.py
-# =================================================================================================
-
-@torch.no_grad()
-def get_unconditional_conditioning(self, batch_size, null_label=None):
-    if null_label is not None:
-        xc = null_label
-        if isinstance(xc, ListConfig):
-            xc = list(xc)
-        if isinstance(xc, dict) or isinstance(xc, list):
-            c = self.get_learned_conditioning(xc)
-        else:
-            if hasattr(xc, "to"):
-                xc = xc.to(self.device)
-            c = self.get_learned_conditioning(xc)
-    else:
-        # todo: get null label from cond_stage_model
-        raise NotImplementedError()
-    c = repeat(c, "1 ... -> b ...", b=batch_size).to(self.device)
-    return c
-
-
-class LatentInpaintDiffusion(LatentDiffusion):
-    def __init__(
-        self,
-        concat_keys=("mask", "masked_image"),
-        masked_image_key="masked_image",
-        *args,
-        **kwargs,
-    ):
-        super().__init__(*args, **kwargs)
-        self.masked_image_key = masked_image_key
-        assert self.masked_image_key in concat_keys
-        self.concat_keys = concat_keys
-
-
-def should_hijack_inpainting(checkpoint_info):
-    ckpt_basename = os.path.basename(checkpoint_info.filename).lower()
-    cfg_basename = os.path.basename(checkpoint_info.config).lower()
-    return "inpainting" in ckpt_basename and not "inpainting" in cfg_basename
-
 
 def do_inpainting_hijack():
-    # most of this stuff seems to no longer be needed because it is already included into SD2.0
     # p_sample_plms is needed because PLMS can't work with dicts as conditionings
-    # this file should be cleaned up later if everything turns out to work fine
-
-    # ldm.models.diffusion.ddpm.get_unconditional_conditioning = get_unconditional_conditioning
-    # ldm.models.diffusion.ddpm.LatentInpaintDiffusion = LatentInpaintDiffusion
-
-    # ldm.models.diffusion.ddim.DDIMSampler.p_sample_ddim = p_sample_ddim
-    # ldm.models.diffusion.ddim.DDIMSampler.sample = sample_ddim
 
     ldm.models.diffusion.plms.PLMSSampler.p_sample_plms = p_sample_plms
-    # ldm.models.diffusion.plms.PLMSSampler.sample = sample_plms
diff --git a/modules/sd_hijack_ip2p.py b/modules/sd_hijack_ip2p.py
new file mode 100644
index 00000000..3c727d3b
--- /dev/null
+++ b/modules/sd_hijack_ip2p.py
@@ -0,0 +1,13 @@
+import collections
+import os.path
+import sys
+import gc
+import time
+
+def should_hijack_ip2p(checkpoint_info):
+    from modules import sd_models_config
+
+    ckpt_basename = os.path.basename(checkpoint_info.filename).lower()
+    cfg_basename = os.path.basename(sd_models_config.find_checkpoint_config_near_filename(checkpoint_info)).lower()
+
+    return "pix2pix" in ckpt_basename and not "pix2pix" in cfg_basename
diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 02c87f40..c02d954c 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -1,7 +1,7 @@
 import math

 import sys

 import traceback

-import importlib

+import psutil

 

 import torch

 from torch import einsum

@@ -9,9 +9,11 @@ from torch import einsum
 from ldm.util import default

 from einops import rearrange

 

-from modules import shared

+from modules import shared, errors, devices

 from modules.hypernetworks import hypernetwork

 

+from .sub_quadratic_attention import efficient_dot_product_attention

+

 

 if shared.cmd_opts.xformers or shared.cmd_opts.force_enable_xformers:

     try:

@@ -22,6 +24,19 @@ if shared.cmd_opts.xformers or shared.cmd_opts.force_enable_xformers:
         print(traceback.format_exc(), file=sys.stderr)

 

 

+def get_available_vram():

+    if shared.device.type == 'cuda':

+        stats = torch.cuda.memory_stats(shared.device)

+        mem_active = stats['active_bytes.all.current']

+        mem_reserved = stats['reserved_bytes.all.current']

+        mem_free_cuda, _ = torch.cuda.mem_get_info(torch.cuda.current_device())

+        mem_free_torch = mem_reserved - mem_active

+        mem_free_total = mem_free_cuda + mem_free_torch

+        return mem_free_total

+    else:

+        return psutil.virtual_memory().available

+

+

 # see https://github.com/basujindal/stable-diffusion/pull/117 for discussion

 def split_cross_attention_forward_v1(self, x, context=None, mask=None):

     h = self.heads

@@ -29,7 +44,7 @@ def split_cross_attention_forward_v1(self, x, context=None, mask=None):
     q_in = self.to_q(x)

     context = default(context, x)

 

-    context_k, context_v = hypernetwork.apply_hypernetwork(shared.loaded_hypernetwork, context)

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

     del context, context_k, context_v, x

@@ -37,18 +52,25 @@ def split_cross_attention_forward_v1(self, x, context=None, mask=None):
     q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))

     del q_in, k_in, v_in

 

-    r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device)

-    for i in range(0, q.shape[0], 2):

-        end = i + 2

-        s1 = einsum('b i d, b j d -> b i j', q[i:end], k[i:end])

-        s1 *= self.scale

+    dtype = q.dtype

+    if shared.opts.upcast_attn:

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

 

-        s2 = s1.softmax(dim=-1)

-        del s1

+    with devices.without_autocast(disable=not shared.opts.upcast_attn):

+        r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)

+        for i in range(0, q.shape[0], 2):

+            end = i + 2

+            s1 = einsum('b i d, b j d -> b i j', q[i:end], k[i:end])

+            s1 *= self.scale

+    

+            s2 = s1.softmax(dim=-1)

+            del s1

+    

+            r1[i:end] = einsum('b i j, b j d -> b i d', s2, v[i:end])

+            del s2

+        del q, k, v

 

-        r1[i:end] = einsum('b i j, b j d -> b i d', s2, v[i:end])

-        del s2

-    del q, k, v

+    r1 = r1.to(dtype)

 

     r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h)

     del r1

@@ -63,54 +85,56 @@ def split_cross_attention_forward(self, x, context=None, mask=None):
     q_in = self.to_q(x)

     context = default(context, x)

 

-    context_k, context_v = hypernetwork.apply_hypernetwork(shared.loaded_hypernetwork, context)

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

 

-    k_in *= self.scale

-

-    del context, x

-

-    q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))

-    del q_in, k_in, v_in

-

-    r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)

-

-    stats = torch.cuda.memory_stats(q.device)

-    mem_active = stats['active_bytes.all.current']

-    mem_reserved = stats['reserved_bytes.all.current']

-    mem_free_cuda, _ = torch.cuda.mem_get_info(torch.cuda.current_device())

-    mem_free_torch = mem_reserved - mem_active

-    mem_free_total = mem_free_cuda + mem_free_torch

-

-    gb = 1024 ** 3

-    tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * q.element_size()

-    modifier = 3 if q.element_size() == 2 else 2.5

-    mem_required = tensor_size * modifier

-    steps = 1

-

-    if mem_required > mem_free_total:

-        steps = 2 ** (math.ceil(math.log(mem_required / mem_free_total, 2)))

-        # print(f"Expected tensor size:{tensor_size/gb:0.1f}GB, cuda free:{mem_free_cuda/gb:0.1f}GB "

-        #       f"torch free:{mem_free_torch/gb:0.1f} total:{mem_free_total/gb:0.1f} steps:{steps}")

+    dtype = q_in.dtype

+    if shared.opts.upcast_attn:

+        q_in, k_in, v_in = q_in.float(), k_in.float(), v_in if v_in.device.type == 'mps' else v_in.float()

 

-    if steps > 64:

-        max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64

-        raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). '

-                           f'Need: {mem_required / 64 / gb:0.1f}GB free, Have:{mem_free_total / gb:0.1f}GB free')

-

-    slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]

-    for i in range(0, q.shape[1], slice_size):

-        end = i + slice_size

-        s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k)

-

-        s2 = s1.softmax(dim=-1, dtype=q.dtype)

-        del s1

-

-        r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v)

-        del s2

+    with devices.without_autocast(disable=not shared.opts.upcast_attn):

+        k_in = k_in * self.scale

+    

+        del context, x

+    

+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))

+        del q_in, k_in, v_in

+    

+        r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)

+    

+        mem_free_total = get_available_vram()

+    

+        gb = 1024 ** 3

+        tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * q.element_size()

+        modifier = 3 if q.element_size() == 2 else 2.5

+        mem_required = tensor_size * modifier

+        steps = 1

+    

+        if mem_required > mem_free_total:

+            steps = 2 ** (math.ceil(math.log(mem_required / mem_free_total, 2)))

+            # print(f"Expected tensor size:{tensor_size/gb:0.1f}GB, cuda free:{mem_free_cuda/gb:0.1f}GB "

+            #       f"torch free:{mem_free_torch/gb:0.1f} total:{mem_free_total/gb:0.1f} steps:{steps}")

+    

+        if steps > 64:

+            max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64

+            raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). '

+                               f'Need: {mem_required / 64 / gb:0.1f}GB free, Have:{mem_free_total / gb:0.1f}GB free')

+    

+        slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]

+        for i in range(0, q.shape[1], slice_size):

+            end = i + slice_size

+            s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k)

+    

+            s2 = s1.softmax(dim=-1, dtype=q.dtype)

+            del s1

+    

+            r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v)

+            del s2

+    

+        del q, k, v

 

-    del q, k, v

+    r1 = r1.to(dtype)

 

     r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h)

     del r1

@@ -118,19 +142,8 @@ def split_cross_attention_forward(self, x, context=None, mask=None):
     return self.to_out(r2)

 

 

-def check_for_psutil():

-    try:

-        spec = importlib.util.find_spec('psutil')

-        return spec is not None

-    except ModuleNotFoundError:

-        return False

-

-invokeAI_mps_available = check_for_psutil()

-

 # -- Taken from https://github.com/invoke-ai/InvokeAI and modified --

-if invokeAI_mps_available:

-    import psutil

-    mem_total_gb = psutil.virtual_memory().total // (1 << 30)

+mem_total_gb = psutil.virtual_memory().total // (1 << 30)

 

 def einsum_op_compvis(q, k, v):

     s = einsum('b i d, b j d -> b i j', q, k)

@@ -204,29 +217,131 @@ def split_cross_attention_forward_invokeAI(self, x, context=None, mask=None):
     q = self.to_q(x)

     context = default(context, x)

 

-    context_k, context_v = hypernetwork.apply_hypernetwork(shared.loaded_hypernetwork, context)

-    k = self.to_k(context_k) * self.scale

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

+    k = self.to_k(context_k)

     v = self.to_v(context_v)

     del context, context_k, context_v, x

 

-    q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))

-    r = einsum_op(q, k, v)

+    dtype = q.dtype

+    if shared.opts.upcast_attn:

+        q, k, v = q.float(), k.float(), v if v.device.type == 'mps' else v.float()

+

+    with devices.without_autocast(disable=not shared.opts.upcast_attn):

+        k = k * self.scale

+    

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

+        r = einsum_op(q, k, v)

+    r = r.to(dtype)

     return self.to_out(rearrange(r, '(b h) n d -> b n (h d)', h=h))

 

 # -- End of code from https://github.com/invoke-ai/InvokeAI --

 

+

+# Based on Birch-san's modified implementation of sub-quadratic attention from https://github.com/Birch-san/diffusers/pull/1

+# The sub_quad_attention_forward function is under the MIT License listed under Memory Efficient Attention in the Licenses section of the web UI interface

+def sub_quad_attention_forward(self, x, context=None, mask=None):

+    assert mask is None, "attention-mask not currently implemented for SubQuadraticCrossAttnProcessor."

+

+    h = self.heads

+

+    q = self.to_q(x)

+    context = default(context, x)

+

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

+    k = self.to_k(context_k)

+    v = self.to_v(context_v)

+    del context, context_k, context_v, x

+

+    q = q.unflatten(-1, (h, -1)).transpose(1,2).flatten(end_dim=1)

+    k = k.unflatten(-1, (h, -1)).transpose(1,2).flatten(end_dim=1)

+    v = v.unflatten(-1, (h, -1)).transpose(1,2).flatten(end_dim=1)

+

+    dtype = q.dtype

+    if shared.opts.upcast_attn:

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

+

+    x = sub_quad_attention(q, k, v, q_chunk_size=shared.cmd_opts.sub_quad_q_chunk_size, kv_chunk_size=shared.cmd_opts.sub_quad_kv_chunk_size, chunk_threshold=shared.cmd_opts.sub_quad_chunk_threshold, use_checkpoint=self.training)

+

+    x = x.to(dtype)

+

+    x = x.unflatten(0, (-1, h)).transpose(1,2).flatten(start_dim=2)

+

+    out_proj, dropout = self.to_out

+    x = out_proj(x)

+    x = dropout(x)

+

+    return x

+

+def sub_quad_attention(q, k, v, q_chunk_size=1024, kv_chunk_size=None, kv_chunk_size_min=None, chunk_threshold=None, use_checkpoint=True):

+    bytes_per_token = torch.finfo(q.dtype).bits//8

+    batch_x_heads, q_tokens, _ = q.shape

+    _, k_tokens, _ = k.shape

+    qk_matmul_size_bytes = batch_x_heads * bytes_per_token * q_tokens * k_tokens

+

+    if chunk_threshold is None:

+        chunk_threshold_bytes = int(get_available_vram() * 0.9) if q.device.type == 'mps' else int(get_available_vram() * 0.7)

+    elif chunk_threshold == 0:

+        chunk_threshold_bytes = None

+    else:

+        chunk_threshold_bytes = int(0.01 * chunk_threshold * get_available_vram())

+

+    if kv_chunk_size_min is None and chunk_threshold_bytes is not None:

+        kv_chunk_size_min = chunk_threshold_bytes // (batch_x_heads * bytes_per_token * (k.shape[2] + v.shape[2]))

+    elif kv_chunk_size_min == 0:

+        kv_chunk_size_min = None

+

+    if chunk_threshold_bytes is not None and qk_matmul_size_bytes <= chunk_threshold_bytes:

+        # the big matmul fits into our memory limit; do everything in 1 chunk,

+        # i.e. send it down the unchunked fast-path

+        query_chunk_size = q_tokens

+        kv_chunk_size = k_tokens

+

+    with devices.without_autocast(disable=q.dtype == v.dtype):

+        return efficient_dot_product_attention(

+            q,

+            k,

+            v,

+            query_chunk_size=q_chunk_size,

+            kv_chunk_size=kv_chunk_size,

+            kv_chunk_size_min = kv_chunk_size_min,

+            use_checkpoint=use_checkpoint,

+        )

+

+

+def get_xformers_flash_attention_op(q, k, v):

+    if not shared.cmd_opts.xformers_flash_attention:

+        return None

+

+    try:

+        flash_attention_op = xformers.ops.MemoryEfficientAttentionFlashAttentionOp

+        fw, bw = flash_attention_op

+        if fw.supports(xformers.ops.fmha.Inputs(query=q, key=k, value=v, attn_bias=None)):

+            return flash_attention_op

+    except Exception as e:

+        errors.display_once(e, "enabling flash attention")

+

+    return None

+

+

 def xformers_attention_forward(self, x, context=None, mask=None):

     h = self.heads

     q_in = self.to_q(x)

     context = default(context, x)

 

-    context_k, context_v = hypernetwork.apply_hypernetwork(shared.loaded_hypernetwork, context)

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

 

     q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b n h d', h=h), (q_in, k_in, v_in))

     del q_in, k_in, v_in

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

+

+    dtype = q.dtype

+    if shared.opts.upcast_attn:

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

+

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

+

+    out = out.to(dtype)

 

     out = rearrange(out, 'b n h d -> b n (h d)', h=h)

     return self.to_out(out)

@@ -252,12 +367,7 @@ def cross_attention_attnblock_forward(self, x):
 

         h_ = torch.zeros_like(k, device=q.device)

 

-        stats = torch.cuda.memory_stats(q.device)

-        mem_active = stats['active_bytes.all.current']

-        mem_reserved = stats['reserved_bytes.all.current']

-        mem_free_cuda, _ = torch.cuda.mem_get_info(torch.cuda.current_device())

-        mem_free_torch = mem_reserved - mem_active

-        mem_free_total = mem_free_cuda + mem_free_torch

+        mem_free_total = get_available_vram()

 

         tensor_size = q.shape[0] * q.shape[1] * k.shape[2] * q.element_size()

         mem_required = tensor_size * 2.5

@@ -303,12 +413,32 @@ def xformers_attnblock_forward(self, x):
         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 = xformers.ops.memory_efficient_attention(q, k, v)

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

+        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

     except NotImplementedError:

         return cross_attention_attnblock_forward(self, x)

+

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

+    q = q.contiguous()

+    k = k.contiguous()

+    v = v.contiguous()

+    out = sub_quad_attention(q, k, v, q_chunk_size=shared.cmd_opts.sub_quad_q_chunk_size, kv_chunk_size=shared.cmd_opts.sub_quad_kv_chunk_size, chunk_threshold=shared.cmd_opts.sub_quad_chunk_threshold, use_checkpoint=self.training)

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

+    out = self.proj_out(out)

+    return x + out

diff --git a/modules/sd_hijack_unet.py b/modules/sd_hijack_unet.py
index 18daf8c1..a6ee577c 100644
--- a/modules/sd_hijack_unet.py
+++ b/modules/sd_hijack_unet.py
@@ -1,4 +1,8 @@
 import torch

+from packaging import version

+

+from modules import devices

+from modules.sd_hijack_utils import CondFunc

 

 

 class TorchHijackForUnet:

@@ -28,3 +32,31 @@ class TorchHijackForUnet:
 

 

 th = TorchHijackForUnet()

+

+

+# Below are monkey patches to enable upcasting a float16 UNet for float32 sampling

+def apply_model(orig_func, self, x_noisy, t, cond, **kwargs):

+

+    if isinstance(cond, dict):

+        for y in cond.keys():

+            cond[y] = [x.to(devices.dtype_unet) if isinstance(x, torch.Tensor) else x for x in cond[y]]

+

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

+    def forward(self, x):

+        if devices.unet_needs_upcast:

+            return torch.nn.GELU.forward(self.float(), x.float()).to(devices.dtype_unet)

+        else:

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

+

+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, *args, **kwargs: orig_func(*args, **kwargs).to(devices.dtype_unet), unet_needs_upcast)

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

+    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_hijack_utils.py b/modules/sd_hijack_utils.py
new file mode 100644
index 00000000..f8684475
--- /dev/null
+++ b/modules/sd_hijack_utils.py
@@ -0,0 +1,28 @@
+import importlib

+

+class CondFunc:

+    def __new__(cls, orig_func, sub_func, cond_func):

+        self = super(CondFunc, cls).__new__(cls)

+        if isinstance(orig_func, str):

+            func_path = orig_func.split('.')

+            for i in range(len(func_path)-1, -1, -1):

+                try:

+                    resolved_obj = importlib.import_module('.'.join(func_path[:i]))

+                    break

+                except ImportError:

+                    pass

+            for attr_name in func_path[i:-1]:

+                resolved_obj = getattr(resolved_obj, attr_name)

+            orig_func = getattr(resolved_obj, func_path[-1])

+            setattr(resolved_obj, func_path[-1], lambda *args, **kwargs: self(*args, **kwargs))

+        self.__init__(orig_func, sub_func, cond_func)

+        return lambda *args, **kwargs: self(*args, **kwargs)
+    def __init__(self, orig_func, sub_func, cond_func):
+        self.__orig_func = orig_func
+        self.__sub_func = sub_func
+        self.__cond_func = cond_func
+    def __call__(self, *args, **kwargs):
+        if not self.__cond_func or self.__cond_func(self.__orig_func, *args, **kwargs):
+            return self.__sub_func(self.__orig_func, *args, **kwargs)
+        else:
+            return self.__orig_func(*args, **kwargs)
diff --git a/modules/sd_hijack_xlmr.py b/modules/sd_hijack_xlmr.py
new file mode 100644
index 00000000..4ac51c38
--- /dev/null
+++ b/modules/sd_hijack_xlmr.py
@@ -0,0 +1,34 @@
+import open_clip.tokenizer

+import torch

+

+from modules import sd_hijack_clip, devices

+from modules.shared import opts

+

+

+class FrozenXLMREmbedderWithCustomWords(sd_hijack_clip.FrozenCLIPEmbedderWithCustomWords):

+    def __init__(self, wrapped, hijack):

+        super().__init__(wrapped, hijack)

+

+        self.id_start = wrapped.config.bos_token_id

+        self.id_end = wrapped.config.eos_token_id

+        self.id_pad = wrapped.config.pad_token_id

+

+        self.comma_token = self.tokenizer.get_vocab().get(',', None)  # alt diffusion doesn't have </w> bits for comma

+

+    def encode_with_transformers(self, tokens):

+        # there's no CLIP Skip here because all hidden layers have size of 1024 and the last one uses a

+        # trained layer to transform those 1024 into 768 for unet; so you can't choose which transformer

+        # layer to work with - you have to use the last

+

+        attention_mask = (tokens != self.id_pad).to(device=tokens.device, dtype=torch.int64)

+        features = self.wrapped(input_ids=tokens, attention_mask=attention_mask)

+        z = features['projection_state']

+

+        return z

+

+    def encode_embedding_init_text(self, init_text, nvpt):

+        embedding_layer = self.wrapped.roberta.embeddings

+        ids = self.wrapped.tokenizer(init_text, max_length=nvpt, return_tensors="pt", add_special_tokens=False)["input_ids"]

+        embedded = embedding_layer.token_embedding.wrapped(ids.to(devices.device)).squeeze(0)

+

+        return embedded

diff --git a/modules/sd_models.py b/modules/sd_models.py
index ecdd91c5..b2d48a51 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -2,7 +2,6 @@ import collections
 import os.path

 import sys

 import gc

-from collections import namedtuple

 import torch

 import re

 import safetensors.torch

@@ -13,17 +12,63 @@ import ldm.modules.midas as midas
 

 from ldm.util import instantiate_from_config

 

-from modules import shared, modelloader, devices, script_callbacks, sd_vae

+from modules import paths, shared, modelloader, devices, script_callbacks, sd_vae, sd_disable_initialization, errors, hashes, sd_models_config

 from modules.paths import models_path

-from modules.sd_hijack_inpainting import do_inpainting_hijack, should_hijack_inpainting

+from modules.sd_hijack_inpainting import do_inpainting_hijack

+from modules.timer import Timer

 

 model_dir = "Stable-diffusion"

-model_path = os.path.abspath(os.path.join(models_path, model_dir))

+model_path = os.path.abspath(os.path.join(paths.models_path, model_dir))

 

-CheckpointInfo = namedtuple("CheckpointInfo", ['filename', 'title', 'hash', 'model_name', 'config'])

 checkpoints_list = {}

+checkpoint_alisases = {}

 checkpoints_loaded = collections.OrderedDict()

 

+

+class CheckpointInfo:

+    def __init__(self, filename):

+        self.filename = filename

+        abspath = os.path.abspath(filename)

+

+        if shared.cmd_opts.ckpt_dir is not None and abspath.startswith(shared.cmd_opts.ckpt_dir):

+            name = abspath.replace(shared.cmd_opts.ckpt_dir, '')

+        elif abspath.startswith(model_path):

+            name = abspath.replace(model_path, '')

+        else:

+            name = os.path.basename(filename)

+

+        if name.startswith("\\") or name.startswith("/"):

+            name = name[1:]

+

+        self.name = name

+        self.model_name = os.path.splitext(name.replace("/", "_").replace("\\", "_"))[0]

+        self.hash = model_hash(filename)

+

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

+        self.shorthash = self.sha256[0:10] if self.sha256 else None

+

+        self.title = name if self.shorthash is None else f'{name} [{self.shorthash}]'

+

+        self.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 [])

+

+    def register(self):

+        checkpoints_list[self.title] = self

+        for id in self.ids:

+            checkpoint_alisases[id] = self

+

+    def calculate_shorthash(self):

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

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

+

+        if self.shorthash not in self.ids:

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

+            self.register()

+

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

+

+        return self.shorthash

+

+

 try:

     # this silences the annoying "Some weights of the model checkpoint were not used when initializing..." message at start.

 

@@ -42,61 +87,50 @@ def setup_model():
     enable_midas_autodownload()

 

 

-def checkpoint_tiles(): 

-    convert = lambda name: int(name) if name.isdigit() else name.lower() 

-    alphanumeric_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)] 

-    return sorted([x.title for x in checkpoints_list.values()], key = alphanumeric_key)

-

-

-def list_models():

-    checkpoints_list.clear()

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

-

-    def modeltitle(path, shorthash):

-        abspath = os.path.abspath(path)

+def checkpoint_tiles():

+    def convert(name):

+        return int(name) if name.isdigit() else name.lower()

 

-        if shared.cmd_opts.ckpt_dir is not None and abspath.startswith(shared.cmd_opts.ckpt_dir):

-            name = abspath.replace(shared.cmd_opts.ckpt_dir, '')

-        elif abspath.startswith(model_path):

-            name = abspath.replace(model_path, '')

-        else:

-            name = os.path.basename(path)

+    def alphanumeric_key(key):

+        return [convert(c) for c in re.split('([0-9]+)', key)]

 

-        if name.startswith("\\") or name.startswith("/"):

-            name = name[1:]

+    return sorted([x.title for x in checkpoints_list.values()], key=alphanumeric_key)

 

-        shortname = os.path.splitext(name.replace("/", "_").replace("\\", "_"))[0]

 

-        return f'{name} [{shorthash}]', shortname

+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 os.path.exists(cmd_ckpt):

-        h = model_hash(cmd_ckpt)

-        title, short_model_name = modeltitle(cmd_ckpt, h)

-        checkpoints_list[title] = CheckpointInfo(cmd_ckpt, title, h, short_model_name, shared.cmd_opts.config)

-        shared.opts.data['sd_model_checkpoint'] = title

+        checkpoint_info = CheckpointInfo(cmd_ckpt)

+        checkpoint_info.register()

+

+        shared.opts.data['sd_model_checkpoint'] = checkpoint_info.title

     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:

-        h = model_hash(filename)

-        title, short_model_name = modeltitle(filename, h)

+        checkpoint_info = CheckpointInfo(filename)

+        checkpoint_info.register()

 

-        basename, _ = os.path.splitext(filename)

-        config = basename + ".yaml"

-        if not os.path.exists(config):

-            config = shared.cmd_opts.config

 

-        checkpoints_list[title] = CheckpointInfo(filename, title, h, short_model_name, config)

+def get_closet_checkpoint_match(search_string):

+    checkpoint_info = checkpoint_alisases.get(search_string, None)

+    if checkpoint_info is not None:

+        return checkpoint_info

 

+    found = sorted([info for info in checkpoints_list.values() if search_string in info.title], key=lambda x: len(x.title))

+    if found:

+        return found[0]

 

-def get_closet_checkpoint_match(searchString):

-    applicable = sorted([info for info in checkpoints_list.values() if searchString in info.title], key = lambda x:len(x.title))

-    if len(applicable) > 0:

-        return applicable[0]

     return None

 

 

 def model_hash(filename):

+    """old hash that only looks at a small part of the file and is prone to collisions"""

+

     try:

         with open(filename, "rb") as file:

             import hashlib

@@ -112,7 +146,7 @@ def model_hash(filename):
 def select_checkpoint():

     model_checkpoint = shared.opts.sd_model_checkpoint

         

-    checkpoint_info = checkpoints_list.get(model_checkpoint, None)

+    checkpoint_info = checkpoint_alisases.get(model_checkpoint, None)

     if checkpoint_info is not None:

         return checkpoint_info

 

@@ -168,7 +202,8 @@ def get_state_dict_from_checkpoint(pl_sd):
 def read_state_dict(checkpoint_file, print_global_state=False, map_location=None):

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

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

-        pl_sd = safetensors.torch.load_file(checkpoint_file, device=map_location or shared.weight_load_location)

+        device = map_location or shared.weight_load_location or devices.get_optimal_device_name()

+        pl_sd = safetensors.torch.load_file(checkpoint_file, device=device)

     else:

         pl_sd = torch.load(checkpoint_file, map_location=map_location or shared.weight_load_location)

 

@@ -179,59 +214,87 @@ def read_state_dict(checkpoint_file, print_global_state=False, map_location=None
     return sd

 

 

-def load_model_weights(model, checkpoint_info, vae_file="auto"):

-    checkpoint_file = checkpoint_info.filename

-    sd_model_hash = checkpoint_info.hash

-

-    cache_enabled = shared.opts.sd_checkpoint_cache > 0

+def get_checkpoint_state_dict(checkpoint_info: CheckpointInfo, timer):

+    sd_model_hash = checkpoint_info.calculate_shorthash()

+    timer.record("calculate hash")

 

-    if cache_enabled and checkpoint_info in checkpoints_loaded:

+    if checkpoint_info in checkpoints_loaded:

         # use checkpoint cache

         print(f"Loading weights [{sd_model_hash}] from cache")

-        model.load_state_dict(checkpoints_loaded[checkpoint_info])

-    else:

-        # load from file

-        print(f"Loading weights [{sd_model_hash}] from {checkpoint_file}")

+        return checkpoints_loaded[checkpoint_info]

 

-        sd = read_state_dict(checkpoint_file)

-        model.load_state_dict(sd, strict=False)

-        del sd

-        

-        if cache_enabled:

-            # cache newly loaded model

-            checkpoints_loaded[checkpoint_info] = model.state_dict().copy()

+    print(f"Loading weights [{sd_model_hash}] from {checkpoint_info.filename}")

+    res = read_state_dict(checkpoint_info.filename)

+    timer.record("load weights from disk")

+

+    return res

+

+

+def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer):

+    title = checkpoint_info.title

+    sd_model_hash = checkpoint_info.calculate_shorthash()

+    timer.record("calculate hash")

+

+    if checkpoint_info.title != title:

+        shared.opts.data["sd_model_checkpoint"] = checkpoint_info.title

+

+    if state_dict is None:

+        state_dict = get_checkpoint_state_dict(checkpoint_info, timer)

+

+    model.load_state_dict(state_dict, strict=False)

+    del state_dict

+    timer.record("apply weights to model")

+

+    if shared.opts.sd_checkpoint_cache > 0:

+        # cache newly loaded model

+        checkpoints_loaded[checkpoint_info] = model.state_dict().copy()

 

-        if shared.cmd_opts.opt_channelslast:

-            model.to(memory_format=torch.channels_last)

+    if shared.cmd_opts.opt_channelslast:

+        model.to(memory_format=torch.channels_last)

+        timer.record("apply channels_last")

 

-        if not shared.cmd_opts.no_half:

-            vae = model.first_stage_model

+    if not shared.cmd_opts.no_half:

+        vae = model.first_stage_model

+        depth_model = getattr(model, 'depth_model', None)

 

-            # with --no-half-vae, remove VAE from model when doing half() to prevent its weights from being converted to float16

-            if shared.cmd_opts.no_half_vae:

-                model.first_stage_model = None

+        # with --no-half-vae, remove VAE from model when doing half() to prevent its weights from being converted to float16

+        if shared.cmd_opts.no_half_vae:

+            model.first_stage_model = None

+        # with --upcast-sampling, don't convert the depth model weights to float16

+        if shared.cmd_opts.upcast_sampling and depth_model:

+            model.depth_model = None

 

-            model.half()

-            model.first_stage_model = vae

+        model.half()

+        model.first_stage_model = vae

+        if depth_model:

+            model.depth_model = depth_model

 

-        devices.dtype = torch.float32 if shared.cmd_opts.no_half else torch.float16

-        devices.dtype_vae = torch.float32 if shared.cmd_opts.no_half or shared.cmd_opts.no_half_vae else torch.float16

+        timer.record("apply half()")

 

-        model.first_stage_model.to(devices.dtype_vae)

+    devices.dtype = torch.float32 if shared.cmd_opts.no_half else torch.float16

+    devices.dtype_vae = torch.float32 if shared.cmd_opts.no_half or shared.cmd_opts.no_half_vae else torch.float16

+    devices.dtype_unet = model.model.diffusion_model.dtype

+    devices.unet_needs_upcast = shared.cmd_opts.upcast_sampling and devices.dtype == torch.float16 and devices.dtype_unet == torch.float16

+

+    model.first_stage_model.to(devices.dtype_vae)

+    timer.record("apply dtype to VAE")

 

     # clean up cache if limit is reached

-    if cache_enabled:

-        while len(checkpoints_loaded) > shared.opts.sd_checkpoint_cache + 1: # we need to count the current model

-            checkpoints_loaded.popitem(last=False)  # LRU

+    while len(checkpoints_loaded) > shared.opts.sd_checkpoint_cache:

+        checkpoints_loaded.popitem(last=False)

 

     model.sd_model_hash = sd_model_hash

-    model.sd_model_checkpoint = checkpoint_file

+    model.sd_model_checkpoint = checkpoint_info.filename

     model.sd_checkpoint_info = checkpoint_info

+    shared.opts.data["sd_checkpoint_hash"] = checkpoint_info.sha256

+

+    model.logvar = model.logvar.to(devices.device)  # fix for training

 

     sd_vae.delete_base_vae()

     sd_vae.clear_loaded_vae()

-    vae_file = sd_vae.resolve_vae(checkpoint_file, vae_file=vae_file)

-    sd_vae.load_vae(model, vae_file)

+    vae_file, vae_source = sd_vae.resolve_vae(checkpoint_info.filename)

+    sd_vae.load_vae(model, vae_file, vae_source)

+    timer.record("load VAE")

 

 

 def enable_midas_autodownload():

@@ -244,7 +307,7 @@ def enable_midas_autodownload():
     location automatically.

     """

 

-    midas_path = os.path.join(models_path, 'midas')

+    midas_path = os.path.join(paths.models_path, 'midas')

 

     # stable-diffusion-stability-ai hard-codes the midas model path to

     # a location that differs from where other scripts using this model look.

@@ -276,88 +339,144 @@ def enable_midas_autodownload():
 

     midas.api.load_model = load_model_wrapper

 

-def load_model(checkpoint_info=None):

+

+def repair_config(sd_config):

+

+    if not hasattr(sd_config.model.params, "use_ema"):

+        sd_config.model.params.use_ema = False

+

+    if shared.cmd_opts.no_half:

+        sd_config.model.params.unet_config.params.use_fp16 = False

+    elif shared.cmd_opts.upcast_sampling:

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

+

+

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

     from modules import lowvram, sd_hijack

     checkpoint_info = checkpoint_info or select_checkpoint()

 

-    if checkpoint_info.config != shared.cmd_opts.config:

-        print(f"Loading config from: {checkpoint_info.config}")

-

     if shared.sd_model:

         sd_hijack.model_hijack.undo_hijack(shared.sd_model)

         shared.sd_model = None

         gc.collect()

         devices.torch_gc()

 

-    sd_config = OmegaConf.load(checkpoint_info.config)

-    

-    if should_hijack_inpainting(checkpoint_info):

-        # Hardcoded config for now...

-        sd_config.model.target = "ldm.models.diffusion.ddpm.LatentInpaintDiffusion"

-        sd_config.model.params.conditioning_key = "hybrid"

-        sd_config.model.params.unet_config.params.in_channels = 9

-        sd_config.model.params.finetune_keys = None

+    do_inpainting_hijack()

 

-        # Create a "fake" config with a different name so that we know to unload it when switching models.

-        checkpoint_info = checkpoint_info._replace(config=checkpoint_info.config.replace(".yaml", "-inpainting.yaml"))

+    timer = Timer()

 

-    if not hasattr(sd_config.model.params, "use_ema"):

-        sd_config.model.params.use_ema = False

+    if already_loaded_state_dict is not None:

+        state_dict = already_loaded_state_dict

+    else:

+        state_dict = get_checkpoint_state_dict(checkpoint_info, timer)

 

-    do_inpainting_hijack()

+    checkpoint_config = sd_models_config.find_checkpoint_config(state_dict, checkpoint_info)

 

-    if shared.cmd_opts.no_half:

-        sd_config.model.params.unet_config.params.use_fp16 = False

+    timer.record("find config")

+

+    sd_config = OmegaConf.load(checkpoint_config)

+    repair_config(sd_config)

+

+    timer.record("load config")

+

+    print(f"Creating model from config: {checkpoint_config}")

+

+    sd_model = None

+    try:

+        with sd_disable_initialization.DisableInitialization():

+            sd_model = instantiate_from_config(sd_config.model)

+    except Exception as e:

+        pass

 

-    sd_model = instantiate_from_config(sd_config.model)

-    load_model_weights(sd_model, checkpoint_info)

+    if sd_model is None:

+        print('Failed to create model quickly; will retry using slow method.', file=sys.stderr)

+        sd_model = instantiate_from_config(sd_config.model)

+

+    sd_model.used_config = checkpoint_config

+

+    timer.record("create model")

+

+    load_model_weights(sd_model, checkpoint_info, state_dict, timer)

 

     if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:

         lowvram.setup_for_low_vram(sd_model, shared.cmd_opts.medvram)

     else:

         sd_model.to(shared.device)

 

+    timer.record("move model to device")

+

     sd_hijack.model_hijack.hijack(sd_model)

 

+    timer.record("hijack")

+

     sd_model.eval()

     shared.sd_model = sd_model

 

+    sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings(force_reload=True)  # Reload embeddings after model load as they may or may not fit the model

+

+    timer.record("load textual inversion embeddings")

+

     script_callbacks.model_loaded_callback(sd_model)

 

-    print("Model loaded.")

+    timer.record("scripts callbacks")

+

+    print(f"Model loaded in {timer.summary()}.")

+

     return sd_model

 

 

 def reload_model_weights(sd_model=None, info=None):

     from modules import lowvram, devices, sd_hijack

     checkpoint_info = info or select_checkpoint()

- 

+

     if not sd_model:

         sd_model = shared.sd_model

 

-    if sd_model.sd_model_checkpoint == checkpoint_info.filename:

-        return

+    if sd_model is None:  # previous model load failed

+        current_checkpoint_info = None

+    else:

+        current_checkpoint_info = sd_model.sd_checkpoint_info

+        if sd_model.sd_model_checkpoint == checkpoint_info.filename:

+            return

+

+        if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:

+            lowvram.send_everything_to_cpu()

+        else:

+            sd_model.to(devices.cpu)

+

+        sd_hijack.model_hijack.undo_hijack(sd_model)

+

+    timer = Timer()

+

+    state_dict = get_checkpoint_state_dict(checkpoint_info, timer)

+

+    checkpoint_config = sd_models_config.find_checkpoint_config(state_dict, checkpoint_info)

+

+    timer.record("find config")

 

-    if sd_model.sd_checkpoint_info.config != checkpoint_info.config or should_hijack_inpainting(checkpoint_info) != should_hijack_inpainting(sd_model.sd_checkpoint_info):

+    if sd_model is None or checkpoint_config != sd_model.used_config:

         del sd_model

         checkpoints_loaded.clear()

-        load_model(checkpoint_info)

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

         return shared.sd_model

 

-    if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:

-        lowvram.send_everything_to_cpu()

-    else:

-        sd_model.to(devices.cpu)

-

-    sd_hijack.model_hijack.undo_hijack(sd_model)

+    try:

+        load_model_weights(sd_model, checkpoint_info, state_dict, timer)

+    except Exception as e:

+        print("Failed to load checkpoint, restoring previous")

+        load_model_weights(sd_model, current_checkpoint_info, None, timer)

+        raise

+    finally:

+        sd_hijack.model_hijack.hijack(sd_model)

+        timer.record("hijack")

 

-    load_model_weights(sd_model, checkpoint_info)

+        script_callbacks.model_loaded_callback(sd_model)

+        timer.record("script callbacks")

 

-    sd_hijack.model_hijack.hijack(sd_model)

-    script_callbacks.model_loaded_callback(sd_model)

+        if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram:

+            sd_model.to(devices.device)

+            timer.record("move model to device")

 

-    if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram:

-        sd_model.to(devices.device)

+    print(f"Weights loaded in {timer.summary()}.")

 

-    print("Weights loaded.")

     return sd_model

diff --git a/modules/sd_models_config.py b/modules/sd_models_config.py
new file mode 100644
index 00000000..00217990
--- /dev/null
+++ b/modules/sd_models_config.py
@@ -0,0 +1,71 @@
+import re

+import os

+

+from modules import shared, paths

+

+sd_configs_path = shared.sd_configs_path

+sd_repo_configs_path = os.path.join(paths.paths['Stable Diffusion'], "configs", "stable-diffusion")

+

+

+config_default = shared.sd_default_config

+config_sd2 = os.path.join(sd_repo_configs_path, "v2-inference.yaml")

+config_sd2v = os.path.join(sd_repo_configs_path, "v2-inference-v.yaml")

+config_sd2_inpainting = os.path.join(sd_repo_configs_path, "v2-inpainting-inference.yaml")

+config_depth_model = os.path.join(sd_repo_configs_path, "v2-midas-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")

+

+re_parametrization_v = re.compile(r'-v\b')

+

+

+def guess_model_config_from_state_dict(sd, filename):

+    fn = os.path.basename(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)

+

+    if sd.get('depth_model.model.pretrained.act_postprocess3.0.project.0.bias', None) is not None:

+        return config_depth_model

+

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

+        if diffusion_model_input.shape[1] == 9:

+            return config_sd2_inpainting

+        elif re.search(re_parametrization_v, fn):

+            return config_sd2v

+        else:

+            return config_sd2

+

+    if diffusion_model_input is not None:

+        if diffusion_model_input.shape[1] == 9:

+            return config_inpainting

+        if diffusion_model_input.shape[1] == 8:

+            return config_instruct_pix2pix

+

+    if sd.get('cond_stage_model.roberta.embeddings.word_embeddings.weight', None) is not None:

+        return config_alt_diffusion

+

+    return config_default

+

+

+def find_checkpoint_config(state_dict, info):

+    if info is None:

+        return guess_model_config_from_state_dict(state_dict, "")

+

+    config = find_checkpoint_config_near_filename(info)

+    if config is not None:

+        return config

+

+    return guess_model_config_from_state_dict(state_dict, info.filename)

+

+

+def find_checkpoint_config_near_filename(info):

+    if info is None:

+        return None

+

+    config = os.path.splitext(info.filename)[0] + ".yaml"

+    if os.path.exists(config):

+        return config

+

+    return None

+

diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index 177b5338..a7910b56 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -97,8 +97,9 @@ sampler_extra_params = {
 

 def setup_img2img_steps(p, steps=None):

     if opts.img2img_fix_steps or steps is not None:

-        steps = int((steps or p.steps) / min(p.denoising_strength, 0.999)) if p.denoising_strength > 0 else 0

-        t_enc = p.steps - 1

+        requested_steps = (steps or p.steps)

+        steps = int(requested_steps / min(p.denoising_strength, 0.999)) if p.denoising_strength > 0 else 0

+        t_enc = requested_steps - 1

     else:

         steps = p.steps

         t_enc = int(min(p.denoising_strength, 0.999) * steps)

@@ -137,9 +138,9 @@ def samples_to_image_grid(samples, approximation=None):
 def store_latent(decoded):

     state.current_latent = decoded

 

-    if opts.show_progress_every_n_steps > 0 and shared.state.sampling_step % opts.show_progress_every_n_steps == 0:

+    if opts.live_previews_enable and opts.show_progress_every_n_steps > 0 and shared.state.sampling_step % opts.show_progress_every_n_steps == 0:

         if not shared.parallel_processing_allowed:

-            shared.state.current_image = sample_to_image(decoded)

+            shared.state.assign_current_image(sample_to_image(decoded))

 

 

 class InterruptedException(BaseException):

@@ -242,7 +243,7 @@ class VanillaStableDiffusionSampler:
         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'):

             valid_step = 999 / (1000 // num_steps)

             if valid_step == floor(valid_step):

                 return int(valid_step) + 1

@@ -265,8 +266,7 @@ class VanillaStableDiffusionSampler:
         if image_conditioning is not None:

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

 

         return samples

@@ -351,6 +351,13 @@ class CFGDenoiser(torch.nn.Module):
 

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

 

+        devices.test_for_nans(x_out, "unet")

+

+        if opts.live_preview_content == "Prompt":

+            store_latent(x_out[0:uncond.shape[0]])

+        elif opts.live_preview_content == "Negative prompt":

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

+

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

 

         if self.mask is not None:

@@ -422,7 +429,8 @@ class KDiffusionSampler:
     def callback_state(self, d):

         step = d['i']

         latent = d["denoised"]

-        store_latent(latent)

+        if opts.live_preview_content == "Combined":

+            store_latent(latent)

         self.last_latent = latent

 

         if self.stop_at is not None and step > self.stop_at:

@@ -446,7 +454,7 @@ class KDiffusionSampler:
     def initialize(self, p):

         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.step = 0

+        self.model_wrap_cfg.step = 0

         self.eta = p.eta or opts.eta_ancestral

 

         k_diffusion.sampling.torch = TorchHijack(self.sampler_noises if self.sampler_noises is not None else [])

@@ -462,14 +470,23 @@ class KDiffusionSampler:
         return extra_params_kwargs

 

     def get_sigmas(self, p, steps):

+        discard_next_to_last_sigma = self.config is not None and self.config.options.get('discard_next_to_last_sigma', False)

+        if opts.always_discard_next_to_last_sigma and not discard_next_to_last_sigma:

+            discard_next_to_last_sigma = True

+            p.extra_generation_params["Discard penultimate sigma"] = True

+

+        steps += 1 if discard_next_to_last_sigma else 0

+

         if p.sampler_noise_scheduler_override:

             sigmas = p.sampler_noise_scheduler_override(steps)

         elif self.config is not None and self.config.options.get('scheduler', None) == 'karras':

-            sigmas = k_diffusion.sampling.get_sigmas_karras(n=steps, sigma_min=0.1, sigma_max=10, device=shared.device)

+            sigma_min, sigma_max = (0.1, 10) if opts.use_old_karras_scheduler_sigmas else (self.model_wrap.sigmas[0].item(), self.model_wrap.sigmas[-1].item())

+

+            sigmas = k_diffusion.sampling.get_sigmas_karras(n=steps, sigma_min=sigma_min, sigma_max=sigma_max, device=shared.device)

         else:

             sigmas = self.model_wrap.get_sigmas(steps)

 

-        if self.config is not None and self.config.options.get('discard_next_to_last_sigma', False):

+        if discard_next_to_last_sigma:

             sigmas = torch.cat([sigmas[:-2], sigmas[-1:]])

 

         return sigmas

diff --git a/modules/sd_vae.py b/modules/sd_vae.py
index 3856418e..9b00f76e 100644
--- a/modules/sd_vae.py
+++ b/modules/sd_vae.py
@@ -1,35 +1,23 @@
 import torch
+import safetensors.torch
 import os
+import collections
 from collections import namedtuple
-from modules import shared, devices, script_callbacks
-from modules.paths import models_path
+from modules import paths, shared, devices, script_callbacks, sd_models
 import glob
 from copy import deepcopy
 
 
-model_dir = "Stable-diffusion"
-model_path = os.path.abspath(os.path.join(models_path, model_dir))
-vae_dir = "VAE"
-vae_path = os.path.abspath(os.path.join(models_path, vae_dir))
-
-
+vae_path = os.path.abspath(os.path.join(paths.models_path, "VAE"))
 vae_ignore_keys = {"model_ema.decay", "model_ema.num_updates"}
-
-
-default_vae_dict = {"auto": "auto", "None": None, None: None}
-default_vae_list = ["auto", "None"]
-
-
-default_vae_values = [default_vae_dict[x] for x in default_vae_list]
-vae_dict = dict(default_vae_dict)
-vae_list = list(default_vae_list)
-first_load = True
+vae_dict = {}
 
 
 base_vae = None
 loaded_vae_file = None
 checkpoint_info = None
 
+checkpoints_loaded = collections.OrderedDict()
 
 def get_base_vae(model):
     if base_vae is not None and checkpoint_info == model.sd_checkpoint_info and model:
@@ -61,126 +49,138 @@ def restore_base_vae(model):
 
 
 def get_filename(filepath):
-    return os.path.splitext(os.path.basename(filepath))[0]
+    return os.path.basename(filepath)
 
 
-def refresh_vae_list(vae_path=vae_path, model_path=model_path):
-    global vae_dict, vae_list
-    res = {}
-    candidates = [
-        *glob.iglob(os.path.join(model_path, '**/*.vae.ckpt'), recursive=True),
-        *glob.iglob(os.path.join(model_path, '**/*.vae.pt'), recursive=True),
-        *glob.iglob(os.path.join(vae_path, '**/*.ckpt'), recursive=True),
-        *glob.iglob(os.path.join(vae_path, '**/*.pt'), recursive=True)
+def refresh_vae_list():
+    vae_dict.clear()
+
+    paths = [
+        os.path.join(sd_models.model_path, '**/*.vae.ckpt'),
+        os.path.join(sd_models.model_path, '**/*.vae.pt'),
+        os.path.join(sd_models.model_path, '**/*.vae.safetensors'),
+        os.path.join(vae_path, '**/*.ckpt'),
+        os.path.join(vae_path, '**/*.pt'),
+        os.path.join(vae_path, '**/*.safetensors'),
     ]
-    if shared.cmd_opts.vae_path is not None and os.path.isfile(shared.cmd_opts.vae_path):
-        candidates.append(shared.cmd_opts.vae_path)
+
+    if shared.cmd_opts.ckpt_dir is not None and os.path.isdir(shared.cmd_opts.ckpt_dir):
+        paths += [
+            os.path.join(shared.cmd_opts.ckpt_dir, '**/*.vae.ckpt'),
+            os.path.join(shared.cmd_opts.ckpt_dir, '**/*.vae.pt'),
+            os.path.join(shared.cmd_opts.ckpt_dir, '**/*.vae.safetensors'),
+        ]
+
+    if shared.cmd_opts.vae_dir is not None and os.path.isdir(shared.cmd_opts.vae_dir):
+        paths += [
+            os.path.join(shared.cmd_opts.vae_dir, '**/*.ckpt'),
+            os.path.join(shared.cmd_opts.vae_dir, '**/*.pt'),
+            os.path.join(shared.cmd_opts.vae_dir, '**/*.safetensors'),
+        ]
+
+    candidates = []
+    for path in paths:
+        candidates += glob.iglob(path, recursive=True)
+
     for filepath in candidates:
         name = get_filename(filepath)
-        res[name] = filepath
-    vae_list.clear()
-    vae_list.extend(default_vae_list)
-    vae_list.extend(list(res.keys()))
-    vae_dict.clear()
-    vae_dict.update(res)
-    vae_dict.update(default_vae_dict)
-    return vae_list
-
-
-def get_vae_from_settings(vae_file="auto"):
-    # else, we load from settings, if not set to be default
-    if vae_file == "auto" and shared.opts.sd_vae is not None:
-        # if saved VAE settings isn't recognized, fallback to auto
-        vae_file = vae_dict.get(shared.opts.sd_vae, "auto")
-        # if VAE selected but not found, fallback to auto
-        if vae_file not in default_vae_values and not os.path.isfile(vae_file):
-            vae_file = "auto"
-            print(f"Selected VAE doesn't exist: {vae_file}")
-    return vae_file
-
-
-def resolve_vae(checkpoint_file=None, vae_file="auto"):
-    global first_load, vae_dict, vae_list
-
-    # if vae_file argument is provided, it takes priority, but not saved
-    if vae_file and vae_file not in default_vae_list:
-        if not os.path.isfile(vae_file):
-            print(f"VAE provided as function argument doesn't exist: {vae_file}")
-            vae_file = "auto"
-    # for the first load, if vae-path is provided, it takes priority, saved, and failure is reported
-    if first_load and shared.cmd_opts.vae_path is not None:
-        if os.path.isfile(shared.cmd_opts.vae_path):
-            vae_file = shared.cmd_opts.vae_path
-            shared.opts.data['sd_vae'] = get_filename(vae_file)
-        else:
-            print(f"VAE provided as command line argument doesn't exist: {vae_file}")
-    # fallback to selector in settings, if vae selector not set to act as default fallback
-    if not shared.opts.sd_vae_as_default:
-        vae_file = get_vae_from_settings(vae_file)
-    # vae-path cmd arg takes priority for auto
-    if vae_file == "auto" and shared.cmd_opts.vae_path is not None:
-        if os.path.isfile(shared.cmd_opts.vae_path):
-            vae_file = shared.cmd_opts.vae_path
-            print(f"Using VAE provided as command line argument: {vae_file}")
-    # if still not found, try look for ".vae.pt" beside model
-    model_path = os.path.splitext(checkpoint_file)[0]
-    if vae_file == "auto":
-        vae_file_try = model_path + ".vae.pt"
-        if os.path.isfile(vae_file_try):
-            vae_file = vae_file_try
-            print(f"Using VAE found similar to selected model: {vae_file}")
-    # if still not found, try look for ".vae.ckpt" beside model
-    if vae_file == "auto":
-        vae_file_try = model_path + ".vae.ckpt"
-        if os.path.isfile(vae_file_try):
-            vae_file = vae_file_try
-            print(f"Using VAE found similar to selected model: {vae_file}")
-    # No more fallbacks for auto
-    if vae_file == "auto":
-        vae_file = None
-    # Last check, just because
-    if vae_file and not os.path.exists(vae_file):
-        vae_file = None
-
-    return vae_file
-
-
-def load_vae(model, vae_file=None):
-    global first_load, vae_dict, vae_list, loaded_vae_file
+        vae_dict[name] = filepath
+
+
+def find_vae_near_checkpoint(checkpoint_file):
+    checkpoint_path = os.path.splitext(checkpoint_file)[0]
+    for vae_location in [checkpoint_path + ".vae.pt", checkpoint_path + ".vae.ckpt", checkpoint_path + ".vae.safetensors"]:
+        if os.path.isfile(vae_location):
+            return vae_location
+
+    return None
+
+
+def resolve_vae(checkpoint_file):
+    if shared.cmd_opts.vae_path is not None:
+        return shared.cmd_opts.vae_path, 'from commandline argument'
+
+    is_automatic = shared.opts.sd_vae in {"Automatic", "auto"}  # "auto" for people with old config
+
+    vae_near_checkpoint = find_vae_near_checkpoint(checkpoint_file)
+    if vae_near_checkpoint is not None and (shared.opts.sd_vae_as_default or is_automatic):
+        return vae_near_checkpoint, 'found near the checkpoint'
+
+    if shared.opts.sd_vae == "None":
+        return None, None
+
+    vae_from_options = vae_dict.get(shared.opts.sd_vae, None)
+    if vae_from_options is not None:
+        return vae_from_options, 'specified in settings'
+
+    if not is_automatic:
+        print(f"Couldn't find VAE named {shared.opts.sd_vae}; using None instead")
+
+    return None, None
+
+
+def load_vae_dict(filename, map_location):
+    vae_ckpt = sd_models.read_state_dict(filename, map_location=map_location)
+    vae_dict_1 = {k: v for k, v in vae_ckpt.items() if k[0:4] != "loss" and k not in vae_ignore_keys}
+    return vae_dict_1
+
+
+def load_vae(model, vae_file=None, vae_source="from unknown source"):
+    global vae_dict, loaded_vae_file
     # save_settings = False
 
+    cache_enabled = shared.opts.sd_vae_checkpoint_cache > 0
+
     if vae_file:
-        assert os.path.isfile(vae_file), f"VAE file doesn't exist: {vae_file}"
-        print(f"Loading VAE weights from: {vae_file}")
-        store_base_vae(model)
-        vae_ckpt = torch.load(vae_file, map_location=shared.weight_load_location)
-        vae_dict_1 = {k: v for k, v in vae_ckpt["state_dict"].items() if k[0:4] != "loss" and k not in vae_ignore_keys}
-        _load_vae_dict(model, vae_dict_1)
+        if cache_enabled and vae_file in checkpoints_loaded:
+            # use vae checkpoint cache
+            print(f"Loading VAE weights {vae_source}: cached {get_filename(vae_file)}")
+            store_base_vae(model)
+            _load_vae_dict(model, checkpoints_loaded[vae_file])
+        else:
+            assert os.path.isfile(vae_file), f"VAE {vae_source} doesn't exist: {vae_file}"
+            print(f"Loading VAE weights {vae_source}: {vae_file}")
+            store_base_vae(model)
+
+            vae_dict_1 = load_vae_dict(vae_file, map_location=shared.weight_load_location)
+            _load_vae_dict(model, vae_dict_1)
+
+            if cache_enabled:
+                # cache newly loaded vae
+                checkpoints_loaded[vae_file] = vae_dict_1.copy()
+
+        # clean up cache if limit is reached
+        if cache_enabled:
+            while len(checkpoints_loaded) > shared.opts.sd_vae_checkpoint_cache + 1: # we need to count the current model
+                checkpoints_loaded.popitem(last=False)  # LRU
 
         # If vae used is not in dict, update it
         # It will be removed on refresh though
         vae_opt = get_filename(vae_file)
         if vae_opt not in vae_dict:
             vae_dict[vae_opt] = vae_file
-            vae_list.append(vae_opt)
+
     elif loaded_vae_file:
         restore_base_vae(model)
 
     loaded_vae_file = vae_file
 
-    first_load = False
-
 
 # don't call this from outside
 def _load_vae_dict(model, vae_dict_1):
     model.first_stage_model.load_state_dict(vae_dict_1)
     model.first_stage_model.to(devices.dtype_vae)
 
+
 def clear_loaded_vae():
     global loaded_vae_file
     loaded_vae_file = None
 
-def reload_vae_weights(sd_model=None, vae_file="auto"):
+
+unspecified = object()
+
+
+def reload_vae_weights(sd_model=None, vae_file=unspecified):
     from modules import lowvram, devices, sd_hijack
 
     if not sd_model:
@@ -188,7 +188,11 @@ def reload_vae_weights(sd_model=None, vae_file="auto"):
 
     checkpoint_info = sd_model.sd_checkpoint_info
     checkpoint_file = checkpoint_info.filename
-    vae_file = resolve_vae(checkpoint_file, vae_file=vae_file)
+
+    if vae_file == unspecified:
+        vae_file, vae_source = resolve_vae(checkpoint_file)
+    else:
+        vae_source = "from function argument"
 
     if loaded_vae_file == vae_file:
         return
@@ -200,7 +204,7 @@ def reload_vae_weights(sd_model=None, vae_file="auto"):
 
     sd_hijack.model_hijack.undo_hijack(sd_model)
 
-    load_vae(sd_model, vae_file)
+    load_vae(sd_model, vae_file, vae_source)
 
     sd_hijack.model_hijack.hijack(sd_model)
     script_callbacks.model_loaded_callback(sd_model)
@@ -208,5 +212,5 @@ def reload_vae_weights(sd_model=None, vae_file="auto"):
     if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram:
         sd_model.to(devices.device)
 
-    print("VAE Weights loaded.")
+    print("VAE weights loaded.")
     return sd_model
diff --git a/modules/sd_vae_approx.py b/modules/sd_vae_approx.py
index 0a58542d..0027343a 100644
--- a/modules/sd_vae_approx.py
+++ b/modules/sd_vae_approx.py
@@ -36,7 +36,7 @@ def model():
 

     if sd_vae_approx_model is None:

         sd_vae_approx_model = VAEApprox()

-        sd_vae_approx_model.load_state_dict(torch.load(os.path.join(paths.models_path, "VAE-approx", "model.pt")))

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

         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 d4ddeea0..474fcc42 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -9,30 +9,35 @@ from PIL import Image
 import gradio as gr

 import tqdm

 

-import modules.artists

 import modules.interrogate

 import modules.memmon

 import modules.styles

 import modules.devices as devices

-from modules import localization, sd_vae, extensions, script_loading

-from modules.paths import models_path, script_path, sd_path

+from modules import localization, extensions, script_loading, errors, ui_components, shared_items

+from modules.paths import models_path, script_path, data_path

 

 

 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("--config", type=str, default=os.path.join(script_path, "v1-inference.yaml"), help="path to config which constructs model",)

+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(script_path, 'embeddings'), help="embeddings directory for textual inversion (default: embeddings)")

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

@@ -42,6 +47,7 @@ parser.add_argument("--lowram", action='store_true', help="load stable diffusion
 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")

@@ -54,34 +60,41 @@ parser.add_argument("--realesrgan-models-path", type=str, help="Path to director
 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(script_path, 'ui-config.json'))

+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(script_path, 'config.json'))

+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='gradio image uploader tool: can be either editor for ctopping, or color-sketch for drawing', choices=["color-sketch", "editor"], default="editor")

-parser.add_argument("--gradio-inpaint-tool", type=str, choices=["sketch", "color-sketch"], default="sketch", help="gradio inpainting editor: can be either sketch to only blur/noise the input, or color-sketch to paint over it")

+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(script_path, 'styles.csv'))

+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='Path to Variational Autoencoders model', default=None)

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

@@ -91,6 +104,8 @@ parser.add_argument("--cors-allow-origins-regex", type=str, help="Allowed CORS o
 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")

+

 

 script_loading.preload_extensions(extensions.extensions_dir, parser)

 script_loading.preload_extensions(extensions.extensions_builtin_dir, parser)

@@ -109,6 +124,18 @@ restricted_opts = {
     "outdir_save",

 }

 

+ui_reorder_categories = [

+    "inpaint",

+    "sampler",

+    "dimensions",

+    "cfg",

+    "seed",

+    "checkboxes",

+    "hires_fix",

+    "batch",

+    "scripts",

+]

+

 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 = \

@@ -124,7 +151,7 @@ config_filename = cmd_opts.ui_settings_file
 

 os.makedirs(cmd_opts.hypernetwork_dir, exist_ok=True)

 hypernetworks = {}

-loaded_hypernetwork = None

+loaded_hypernetworks = []

 

 

 def reload_hypernetworks():

@@ -132,7 +159,6 @@ def reload_hypernetworks():
     global hypernetworks

 

     hypernetworks = hypernetwork.list_hypernetworks(cmd_opts.hypernetwork_dir)

-    hypernetwork.load_hypernetwork(opts.sd_hypernetwork)

 

 

 class State:

@@ -141,15 +167,18 @@ class State:
     job = ""

     job_no = 0

     job_count = 0

+    processing_has_refined_job_count = False

     job_timestamp = '0'

     sampling_step = 0

     sampling_steps = 0

     current_latent = None

     current_image = None

     current_image_sampling_step = 0

+    id_live_preview = 0

     textinfo = None

     time_start = None

     need_restart = False

+    server_start = None

 

     def skip(self):

         self.skipped = True

@@ -158,7 +187,7 @@ class State:
         self.interrupted = True

 

     def nextjob(self):

-        if opts.show_progress_every_n_steps == -1:

+        if opts.live_previews_enable and opts.show_progress_every_n_steps == -1:

             self.do_set_current_image()

 

         self.job_no += 1

@@ -168,9 +197,10 @@ class State:
     def dict(self):

         obj = {

             "skipped": self.skipped,

-            "interrupted": self.skipped,

+            "interrupted": self.interrupted,

             "job": self.job,

             "job_count": self.job_count,

+            "job_timestamp": self.job_timestamp,

             "job_no": self.job_no,

             "sampling_step": self.sampling_step,

             "sampling_steps": self.sampling_steps,

@@ -181,11 +211,13 @@ class State:
     def begin(self):

         self.sampling_step = 0

         self.job_count = -1

+        self.processing_has_refined_job_count = False

         self.job_no = 0

         self.job_timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")

         self.current_latent = None

         self.current_image = None

         self.current_image_sampling_step = 0

+        self.id_live_preview = 0

         self.skipped = False

         self.interrupted = False

         self.textinfo = None

@@ -199,28 +231,33 @@ class State:
 

         devices.torch_gc()

 

-    """sets self.current_image from self.current_latent if enough sampling steps have been made after the last call to this"""

     def set_current_image(self):

-        if self.sampling_step - self.current_image_sampling_step >= opts.show_progress_every_n_steps and opts.show_progress_every_n_steps > 0:

+        """sets self.current_image from self.current_latent if enough sampling steps have been made after the last call to this"""

+        if not parallel_processing_allowed:

+            return

+

+        if self.sampling_step - self.current_image_sampling_step >= opts.show_progress_every_n_steps and opts.live_previews_enable and opts.show_progress_every_n_steps != -1:

             self.do_set_current_image()

 

     def do_set_current_image(self):

-        if not parallel_processing_allowed:

-            return

         if self.current_latent is None:

             return

 

         import modules.sd_samplers

         if opts.show_progress_grid:

-            self.current_image = modules.sd_samplers.samples_to_image_grid(self.current_latent)

+            self.assign_current_image(modules.sd_samplers.samples_to_image_grid(self.current_latent))

         else:

-            self.current_image = modules.sd_samplers.sample_to_image(self.current_latent)

+            self.assign_current_image(modules.sd_samplers.sample_to_image(self.current_latent))

 

         self.current_image_sampling_step = self.sampling_step

 

-state = State()

+    def assign_current_image(self, image):

+        self.current_image = image

+        self.id_live_preview += 1

 

-artist_db = modules.artists.ArtistsDatabase(os.path.join(script_path, 'artists.csv'))

+

+state = State()

+state.server_start = time.time()

 

 styles_filename = cmd_opts.styles_file

 prompt_styles = modules.styles.StyleDatabase(styles_filename)

@@ -229,12 +266,6 @@ interrogator = modules.interrogate.InterrogateModels("interrogate")
 

 face_restorers = []

 

-

-def realesrgan_models_names():

-    import modules.realesrgan_model

-    return [x.name for x in modules.realesrgan_model.get_realesrgan_models(None)]

-

-

 class OptionInfo:

     def __init__(self, default=None, label="", component=None, component_args=None, onchange=None, section=None, refresh=None):

         self.default = default

@@ -325,9 +356,8 @@ options_templates.update(options_section(('saving-to-dirs', "Saving to a directo
 options_templates.update(options_section(('upscaling', "Upscaling"), {

     "ESRGAN_tile": OptionInfo(192, "Tile size for ESRGAN upscalers. 0 = no tiling.", gr.Slider, {"minimum": 0, "maximum": 512, "step": 16}),

     "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": realesrgan_models_names()}),

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

-    "use_scale_latent_for_hires_fix": OptionInfo(False, "Upscale latent space image when doing hires. fix"),

 }))

 

 options_templates.update(options_section(('face-restoration', "Face restoration"), {

@@ -337,51 +367,62 @@ options_templates.update(options_section(('face-restoration', "Face restoration"
 }))

 

 options_templates.update(options_section(('system', "System"), {

+    "show_warnings": OptionInfo(False, "Show warnings in console."),

     "memmon_poll_rate": OptionInfo(8, "VRAM usage polls per second during generation. Set to 0 to disable.", gr.Slider, {"minimum": 0, "maximum": 40, "step": 1}),

     "samples_log_stdout": OptionInfo(False, "Always print all generation info to standard output"),

     "multiple_tqdm": OptionInfo(True, "Add a second progress bar to the console that shows progress for an entire job."),

+    "print_hypernet_extra": OptionInfo(False, "Print extra hypernetwork information to console."),

 }))

 

 options_templates.update(options_section(('training', "Training"), {

     "unload_models_when_training": OptionInfo(False, "Move VAE and CLIP to RAM when training if possible. Saves VRAM."),

     "pin_memory": OptionInfo(False, "Turn on pin_memory for DataLoader. Makes training slightly faster but can increase memory usage."),

-    "save_optimizer_state": OptionInfo(False, "Saves Optimizer state as separate *.optim file. Training can be resumed with HN itself and matching optim file."),

+    "save_optimizer_state": OptionInfo(False, "Saves Optimizer state as separate *.optim file. Training of embedding or HN can be resumed with the matching optim file."),

+    "save_training_settings_to_txt": OptionInfo(True, "Save textual inversion and hypernet settings to a text file whenever training starts."),

     "dataset_filename_word_regex": OptionInfo("", "Filename word regex"),

     "dataset_filename_join_string": OptionInfo(" ", "Filename join string"),

     "training_image_repeats_per_epoch": OptionInfo(1, "Number of repeats for a single input image per epoch; used only for displaying epoch number", gr.Number, {"precision": 0}),

     "training_write_csv_every": OptionInfo(500, "Save an csv containing the loss to log directory every N steps, 0 to disable"),

     "training_xattention_optimizations": OptionInfo(False, "Use cross attention optimizations while training"),

+    "training_enable_tensorboard": OptionInfo(False, "Enable tensorboard logging."),

+    "training_tensorboard_save_images": OptionInfo(False, "Save generated images within tensorboard."),

+    "training_tensorboard_flush_every": OptionInfo(120, "How often, in seconds, to flush the pending tensorboard events and summaries to disk."),

 }))

 

 options_templates.update(options_section(('sd', "Stable Diffusion"), {

     "sd_model_checkpoint": OptionInfo(None, "Stable Diffusion checkpoint", gr.Dropdown, lambda: {"choices": list_checkpoint_tiles()}, refresh=refresh_checkpoints),

     "sd_checkpoint_cache": OptionInfo(0, "Checkpoints to cache in RAM", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}),

-    "sd_vae": OptionInfo("auto", "SD VAE", gr.Dropdown, lambda: {"choices": sd_vae.vae_list}, refresh=sd_vae.refresh_vae_list),

-    "sd_vae_as_default": OptionInfo(False, "Ignore selected VAE for stable diffusion checkpoints that have their own .vae.pt next to them"),

-    "sd_hypernetwork": OptionInfo("None", "Hypernetwork", gr.Dropdown, lambda: {"choices": ["None"] + [x for x in hypernetworks.keys()]}, refresh=reload_hypernetworks),

-    "sd_hypernetwork_strength": OptionInfo(1.0, "Hypernetwork strength", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.001}),

+    "sd_vae_checkpoint_cache": OptionInfo(0, "VAE Checkpoints to cache in RAM", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}),

+    "sd_vae": OptionInfo("Automatic", "SD VAE", gr.Dropdown, lambda: {"choices": shared_items.sd_vae_items()}, refresh=shared_items.refresh_vae_list),

+    "sd_vae_as_default": OptionInfo(True, "Ignore selected VAE for stable diffusion checkpoints that have their own .vae.pt next to them"),

     "inpainting_mask_weight": OptionInfo(1.0, "Inpainting conditioning mask strength", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),

-    "initial_noise_multiplier": OptionInfo(1.0, "Noise multiplier for img2img", gr.Slider, {"minimum": 0.5, "maximum": 1.5, "step": 0.01 }),

+    "initial_noise_multiplier": OptionInfo(1.0, "Noise multiplier for img2img", gr.Slider, {"minimum": 0.5, "maximum": 1.5, "step": 0.01}),

     "img2img_color_correction": OptionInfo(False, "Apply color correction to img2img results to match original colors."),

     "img2img_fix_steps": OptionInfo(False, "With img2img, do exactly the amount of steps the slider specifies (normally you'd do less with less denoising)."),

-    "img2img_background_color": OptionInfo("#ffffff", "With img2img, fill image's transparent parts with this color.", gr.ColorPicker, {}),

+    "img2img_background_color": OptionInfo("#ffffff", "With img2img, fill image's transparent parts with this color.", ui_components.FormColorPicker, {}),

     "enable_quantization": OptionInfo(False, "Enable quantization in K samplers for sharper and cleaner results. This may change existing seeds. Requires restart to apply."),

     "enable_emphasis": OptionInfo(True, "Emphasis: use (text) to make model pay more attention to text and [text] to make it pay less attention"),

-    "use_old_emphasis_implementation": OptionInfo(False, "Use old emphasis implementation. Can be useful to reproduce old seeds."),

     "enable_batch_seeds": OptionInfo(True, "Make K-diffusion samplers produce same images in a batch as when making a single image"),

     "comma_padding_backtrack": OptionInfo(20, "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}),

-    "random_artist_categories": OptionInfo([], "Allowed categories for random artists selection when using the Roll button", gr.CheckboxGroup, {"choices": artist_db.categories()}),

+    "CLIP_stop_at_last_layers": OptionInfo(1, "Clip skip", gr.Slider, {"minimum": 1, "maximum": 12, "step": 1}),

+    "extra_networks_default_multiplier": OptionInfo(1.0, "Multiplier for extra networks", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),

+    "upcast_attn": OptionInfo(False, "Upcast cross attention layer to float32"),

+}))

+

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

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

 }))

 

 options_templates.update(options_section(('interrogate', "Interrogate Options"), {

     "interrogate_keep_models_in_memory": OptionInfo(False, "Interrogate: keep models in VRAM"),

-    "interrogate_use_builtin_artists": OptionInfo(True, "Interrogate: use artists from artists.csv"),

     "interrogate_return_ranks": OptionInfo(False, "Interrogate: include ranks of model tags matches in results (Has no effect on caption-based interrogators)."),

     "interrogate_clip_num_beams": OptionInfo(1, "Interrogate: num_beams for BLIP", gr.Slider, {"minimum": 1, "maximum": 16, "step": 1}),

     "interrogate_clip_min_length": OptionInfo(24, "Interrogate: minimum description length (excluding artists, etc..)", gr.Slider, {"minimum": 1, "maximum": 128, "step": 1}),

     "interrogate_clip_max_length": OptionInfo(48, "Interrogate: maximum description length", gr.Slider, {"minimum": 1, "maximum": 256, "step": 1}),

     "interrogate_clip_dict_limit": OptionInfo(1500, "CLIP: maximum number of lines in text file (0 = No limit)"),

+    "interrogate_clip_skip_categories": OptionInfo([], "CLIP: skip inquire categories", gr.CheckboxGroup, lambda: {"choices": modules.interrogate.category_types()}, refresh=modules.interrogate.category_types),

     "interrogate_deepbooru_score_threshold": OptionInfo(0.5, "Interrogate: deepbooru score threshold", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01}),

     "deepbooru_sort_alpha": OptionInfo(True, "Interrogate: deepbooru sort alphabetically"),

     "deepbooru_use_spaces": OptionInfo(False, "use spaces for tags in deepbooru"),

@@ -389,15 +430,15 @@ options_templates.update(options_section(('interrogate', "Interrogate Options"),
     "deepbooru_filter_tags": OptionInfo("", "filter out those tags from deepbooru output (separated by comma)"),

 }))

 

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

+}))

+

 options_templates.update(options_section(('ui', "User interface"), {

-    "show_progressbar": OptionInfo(True, "Show progressbar"),

-    "show_progress_every_n_steps": OptionInfo(0, "Show image creation progress every N sampling steps. Set to 0 to disable. Set to -1 to show after completion of batch.", gr.Slider, {"minimum": -1, "maximum": 32, "step": 1}),

-    "show_progress_type": OptionInfo("Full", "Image creation progress preview mode", gr.Radio, {"choices": ["Full", "Approx NN", "Approx cheap"]}),

-    "show_progress_grid": OptionInfo(True, "Show previews of all images generated in a batch as a grid"),

     "return_grid": OptionInfo(True, "Show grid 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(False, "Add model name to generation information"),

+    "add_model_name_to_info": OptionInfo(True, "Add model name to generation information"),

     "disable_weights_auto_swap": OptionInfo(False, "When reading generation parameters from text into UI (from PNG info or pasted text), do not change the selected model/checkpoint."),

     "send_seed": OptionInfo(True, "Send seed when sending prompt or image to other interface"),

     "send_size": OptionInfo(True, "Send size when sending prompt or image to another interface"),

@@ -405,8 +446,24 @@ options_templates.update(options_section(('ui', "User interface"), {
     "js_modal_lightbox": OptionInfo(True, "Enable full page image viewer"),

     "js_modal_lightbox_initially_zoomed": OptionInfo(True, "Show images zoomed in by default in full page image viewer"),

     "show_progress_in_title": OptionInfo(True, "Show generation progress in window title."),

-    'quicksettings': OptionInfo("sd_model_checkpoint", "Quicksettings list"),

-    'localization': OptionInfo("None", "Localization (requires restart)", gr.Dropdown, lambda: {"choices": ["None"] + list(localization.localizations.keys())}, refresh=lambda: localization.list_localizations(cmd_opts.localizations_dir)),

+    "samplers_in_dropdown": OptionInfo(True, "Use dropdown for sampler selection instead of radio group"),

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

+    "quicksettings": OptionInfo("sd_model_checkpoint", "Quicksettings list"),

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

+}))

+

+options_templates.update(options_section(('ui', "Live previews"), {

+    "show_progressbar": OptionInfo(True, "Show progressbar"),

+    "live_previews_enable": OptionInfo(True, "Show live previews of the created image"),

+    "show_progress_grid": OptionInfo(True, "Show previews of all images generated in a batch as a grid"),

+    "show_progress_every_n_steps": OptionInfo(10, "Show new live preview image every N sampling steps. Set to -1 to show after completion of batch.", gr.Slider, {"minimum": -1, "maximum": 32, "step": 1}),

+    "show_progress_type": OptionInfo("Approx NN", "Image creation progress preview mode", gr.Radio, {"choices": ["Full", "Approx NN", "Approx cheap"]}),

+    "live_preview_content": OptionInfo("Prompt", "Live preview subject", gr.Radio, {"choices": ["Combined", "Prompt", "Negative prompt"]}),

+    "live_preview_refresh_period": OptionInfo(1000, "Progressbar/preview update period, in milliseconds")

 }))

 

 options_templates.update(options_section(('sampler-params', "Sampler parameters"), {

@@ -418,10 +475,18 @@ options_templates.update(options_section(('sampler-params', "Sampler parameters"
     '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"),

+}))

+

+options_templates.update(options_section(('postprocessing', "Postprocessing"), {

+    'postprocessing_enable_in_main_ui': OptionInfo([], "Enable postprocessing operations in txt2img and img2img tabs", ui_components.DropdownMulti, lambda: {"choices": [x.name for x in shared_items.postprocessing_scripts()]}),

+    'postprocessing_operation_order': OptionInfo([], "Postprocessing operation order", ui_components.DropdownMulti, lambda: {"choices": [x.name for x in shared_items.postprocessing_scripts()]}),

+    'upscaling_max_images_in_cache': OptionInfo(5, "Maximum number of images in upscaling cache", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}),

 }))

 

 options_templates.update(options_section((None, "Hidden options"), {

     "disabled_extensions": OptionInfo([], "Disable those extensions"),

+    "sd_checkpoint_hash": OptionInfo("", "SHA256 hash of the current checkpoint"),

 }))

 

 options_templates.update()

@@ -476,7 +541,12 @@ class Options:
             return False

 

         if self.data_labels[key].onchange is not None:

-            self.data_labels[key].onchange()

+            try:

+                self.data_labels[key].onchange()

+            except Exception as e:

+                errors.display(e, f"changing setting {key} to {value}")

+                setattr(self, key, oldval)

+                return False

 

         return True

 

@@ -539,6 +609,16 @@ opts = Options()
 if os.path.exists(config_filename):

     opts.load(config_filename)

 

+latent_upscale_default_mode = "Latent"

+latent_upscale_modes = {

+    "Latent": {"mode": "bilinear", "antialias": False},

+    "Latent (antialiased)": {"mode": "bilinear", "antialias": True},

+    "Latent (bicubic)": {"mode": "bicubic", "antialias": False},

+    "Latent (bicubic antialiased)": {"mode": "bicubic", "antialias": True},

+    "Latent (nearest)": {"mode": "nearest", "antialias": False},

+    "Latent (nearest-exact)": {"mode": "nearest-exact", "antialias": False},

+}

+

 sd_upscalers = []

 

 sd_model = None

@@ -572,7 +652,7 @@ class TotalTQDM:
             return

         if self._tqdm is None:

             self.reset()

-        self._tqdm.total=new_total

+        self._tqdm.total = new_total

 

     def clear(self):

         if self._tqdm is not None:

@@ -589,3 +669,17 @@ mem_mon.start()
 def listfiles(dirname):

     filenames = [os.path.join(dirname, x) for x in sorted(os.listdir(dirname)) if not x.startswith(".")]

     return [file for file in filenames if os.path.isfile(file)]

+

+

+def html_path(filename):

+    return os.path.join(script_path, "html", filename)

+

+

+def html(filename):

+    path = html_path(filename)

+

+    if os.path.exists(path):

+        with open(path, encoding="utf8") as file:

+            return file.read()

+

+    return ""

diff --git a/modules/shared_items.py b/modules/shared_items.py
new file mode 100644
index 00000000..8b5ec96d
--- /dev/null
+++ b/modules/shared_items.py
@@ -0,0 +1,23 @@
+

+

+def realesrgan_models_names():

+    import modules.realesrgan_model

+    return [x.name for x in modules.realesrgan_model.get_realesrgan_models(None)]

+

+

+def postprocessing_scripts():

+    import modules.scripts

+

+    return modules.scripts.scripts_postproc.scripts

+

+

+def sd_vae_items():

+    import modules.sd_vae

+

+    return ["Automatic", "None"] + list(modules.sd_vae.vae_dict)

+

+

+def refresh_vae_list():

+    import modules.sd_vae

+

+    return modules.sd_vae.refresh_vae_list

diff --git a/modules/styles.py b/modules/styles.py
index ce6e71ca..990d5623 100644
--- a/modules/styles.py
+++ b/modules/styles.py
@@ -40,12 +40,18 @@ def apply_styles_to_prompt(prompt, styles):
 class StyleDatabase:

     def __init__(self, path: str):

         self.no_style = PromptStyle("None", "", "")

-        self.styles = {"None": self.no_style}

+        self.styles = {}

+        self.path = path

 

-        if not os.path.exists(path):

+        self.reload()

+

+    def reload(self):

+        self.styles.clear()

+

+        if not os.path.exists(self.path):

             return

 

-        with open(path, "r", encoding="utf-8-sig", newline='') as file:

+        with open(self.path, "r", encoding="utf-8-sig", newline='') as file:

             reader = csv.DictReader(file)

             for row in reader:

                 # Support loading old CSV format with "name, text"-columns

diff --git a/modules/sub_quadratic_attention.py b/modules/sub_quadratic_attention.py
new file mode 100644
index 00000000..05595323
--- /dev/null
+++ b/modules/sub_quadratic_attention.py
@@ -0,0 +1,214 @@
+# original source:
+#   https://github.com/AminRezaei0x443/memory-efficient-attention/blob/1bc0d9e6ac5f82ea43a375135c4e1d3896ee1694/memory_efficient_attention/attention_torch.py
+# license:
+#   MIT License (see Memory Efficient Attention under the Licenses section in the web UI interface for the full license)
+# credit:
+#   Amin Rezaei (original author)
+#   Alex Birch (optimized algorithm for 3D tensors, at the expense of removing bias, masking and callbacks)
+#   brkirch (modified to use torch.narrow instead of dynamic_slice implementation)
+# implementation of:
+#   Self-attention Does Not Need O(n2) Memory":
+#   https://arxiv.org/abs/2112.05682v2
+
+from functools import partial
+import torch
+from torch import Tensor
+from torch.utils.checkpoint import checkpoint
+import math
+from typing import Optional, NamedTuple, List
+
+
+def narrow_trunc(
+    input: Tensor,
+    dim: int,
+    start: int,
+    length: int
+) -> Tensor:
+    return torch.narrow(input, dim, start, length if input.shape[dim] >= start + length else input.shape[dim] - start)
+
+
+class AttnChunk(NamedTuple):
+    exp_values: Tensor
+    exp_weights_sum: Tensor
+    max_score: Tensor
+
+
+class SummarizeChunk:
+    @staticmethod
+    def __call__(
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+    ) -> AttnChunk: ...
+
+
+class ComputeQueryChunkAttn:
+    @staticmethod
+    def __call__(
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+    ) -> Tensor: ...
+
+
+def _summarize_chunk(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    scale: float,
+) -> AttnChunk:
+    attn_weights = torch.baddbmm(
+        torch.empty(1, 1, 1, device=query.device, dtype=query.dtype),
+        query,
+        key.transpose(1,2),
+        alpha=scale,
+        beta=0,
+    )
+    max_score, _ = torch.max(attn_weights, -1, keepdim=True)
+    max_score = max_score.detach()
+    exp_weights = torch.exp(attn_weights - max_score)
+    exp_values = torch.bmm(exp_weights, value) if query.device.type == 'mps' else torch.bmm(exp_weights, value.to(exp_weights.dtype)).to(value.dtype)
+    max_score = max_score.squeeze(-1)
+    return AttnChunk(exp_values, exp_weights.sum(dim=-1), max_score)
+
+
+def _query_chunk_attention(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    summarize_chunk: SummarizeChunk,
+    kv_chunk_size: int,
+) -> Tensor:
+    batch_x_heads, k_tokens, k_channels_per_head = key.shape
+    _, _, v_channels_per_head = value.shape
+
+    def chunk_scanner(chunk_idx: int) -> AttnChunk:
+        key_chunk = narrow_trunc(
+            key,
+            1,
+            chunk_idx,
+            kv_chunk_size
+        )
+        value_chunk = narrow_trunc(
+            value,
+            1,
+            chunk_idx,
+            kv_chunk_size
+        )
+        return summarize_chunk(query, key_chunk, value_chunk)
+
+    chunks: List[AttnChunk] = [
+        chunk_scanner(chunk) for chunk in torch.arange(0, k_tokens, kv_chunk_size)
+    ]
+    acc_chunk = AttnChunk(*map(torch.stack, zip(*chunks)))
+    chunk_values, chunk_weights, chunk_max = acc_chunk
+
+    global_max, _ = torch.max(chunk_max, 0, keepdim=True)
+    max_diffs = torch.exp(chunk_max - global_max)
+    chunk_values *= torch.unsqueeze(max_diffs, -1)
+    chunk_weights *= max_diffs
+
+    all_values = chunk_values.sum(dim=0)
+    all_weights = torch.unsqueeze(chunk_weights, -1).sum(dim=0)
+    return all_values / all_weights
+
+
+# TODO: refactor CrossAttention#get_attention_scores to share code with this
+def _get_attention_scores_no_kv_chunking(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    scale: float,
+) -> Tensor:
+    attn_scores = torch.baddbmm(
+        torch.empty(1, 1, 1, device=query.device, dtype=query.dtype),
+        query,
+        key.transpose(1,2),
+        alpha=scale,
+        beta=0,
+    )
+    attn_probs = attn_scores.softmax(dim=-1)
+    del attn_scores
+    hidden_states_slice = torch.bmm(attn_probs, value) if query.device.type == 'mps' else torch.bmm(attn_probs, value.to(attn_probs.dtype)).to(value.dtype)
+    return hidden_states_slice
+
+
+class ScannedChunk(NamedTuple):
+    chunk_idx: int
+    attn_chunk: AttnChunk
+
+
+def efficient_dot_product_attention(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    query_chunk_size=1024,
+    kv_chunk_size: Optional[int] = None,
+    kv_chunk_size_min: Optional[int] = None,
+    use_checkpoint=True,
+):
+    """Computes efficient dot-product attention given query, key, and value.
+      This is efficient version of attention presented in
+      https://arxiv.org/abs/2112.05682v2 which comes with O(sqrt(n)) memory requirements.
+      Args:
+        query: queries for calculating attention with shape of
+          `[batch * num_heads, tokens, channels_per_head]`.
+        key: keys for calculating attention with shape of
+          `[batch * num_heads, tokens, channels_per_head]`.
+        value: values to be used in attention with shape of
+          `[batch * num_heads, tokens, channels_per_head]`.
+        query_chunk_size: int: query chunks size
+        kv_chunk_size: Optional[int]: key/value chunks size. if None: defaults to sqrt(key_tokens)
+        kv_chunk_size_min: Optional[int]: key/value minimum chunk size. only considered when kv_chunk_size is None. changes `sqrt(key_tokens)` into `max(sqrt(key_tokens), kv_chunk_size_min)`, to ensure our chunk sizes don't get too small (smaller chunks = more chunks = less concurrent work done).
+        use_checkpoint: bool: whether to use checkpointing (recommended True for training, False for inference)
+      Returns:
+        Output of shape `[batch * num_heads, query_tokens, channels_per_head]`.
+      """
+    batch_x_heads, q_tokens, q_channels_per_head = query.shape
+    _, k_tokens, _ = key.shape
+    scale = q_channels_per_head ** -0.5
+
+    kv_chunk_size = min(kv_chunk_size or int(math.sqrt(k_tokens)), k_tokens)
+    if kv_chunk_size_min is not None:
+        kv_chunk_size = max(kv_chunk_size, kv_chunk_size_min)
+
+    def get_query_chunk(chunk_idx: int) -> Tensor:
+        return narrow_trunc(
+            query,
+            1,
+            chunk_idx,
+            min(query_chunk_size, q_tokens)
+        )
+    
+    summarize_chunk: SummarizeChunk = partial(_summarize_chunk, scale=scale)
+    summarize_chunk: SummarizeChunk = partial(checkpoint, summarize_chunk) if use_checkpoint else summarize_chunk
+    compute_query_chunk_attn: ComputeQueryChunkAttn = partial(
+        _get_attention_scores_no_kv_chunking,
+        scale=scale
+    ) if k_tokens <= kv_chunk_size else (
+        # fast-path for when there's just 1 key-value chunk per query chunk (this is just sliced attention btw)
+        partial(
+            _query_chunk_attention,
+            kv_chunk_size=kv_chunk_size,
+            summarize_chunk=summarize_chunk,
+        )
+    )
+
+    if q_tokens <= query_chunk_size:
+        # fast-path for when there's just 1 query chunk
+        return compute_query_chunk_attn(
+            query=query,
+            key=key,
+            value=value,
+        )
+    
+    # TODO: maybe we should use torch.empty_like(query) to allocate storage in-advance,
+    # and pass slices to be mutated, instead of torch.cat()ing the returned slices
+    res = torch.cat([
+        compute_query_chunk_attn(
+            query=get_query_chunk(i * query_chunk_size),
+            key=key,
+            value=value,
+        ) for i in range(math.ceil(q_tokens / query_chunk_size))
+    ], dim=1)
+    return res
diff --git a/modules/textual_inversion/dataset.py b/modules/textual_inversion/dataset.py
index 88d68c76..d31963d4 100644
--- a/modules/textual_inversion/dataset.py
+++ b/modules/textual_inversion/dataset.py
@@ -3,8 +3,10 @@ import numpy as np
 import PIL

 import torch

 from PIL import Image

-from torch.utils.data import Dataset, DataLoader

+from torch.utils.data import Dataset, DataLoader, Sampler

 from torchvision import transforms

+from collections import defaultdict

+from random import shuffle, choices

 

 import random

 import tqdm

@@ -28,13 +30,11 @@ 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'):

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

         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

 

-        self.width = width

-        self.height = height

         self.flip = transforms.RandomHorizontalFlip(p=flip_p)

 

         self.dataset = []

@@ -50,16 +50,18 @@ class PersonalizedBase(Dataset):
 

         self.image_paths = [os.path.join(data_root, file_path) for file_path in os.listdir(data_root)]

 

-

         self.shuffle_tags = shuffle_tags

         self.tag_drop_out = tag_drop_out

+        groups = defaultdict(list)

 

         print("Preparing dataset...")

         for path in tqdm.tqdm(self.image_paths):

             if shared.state.interrupted:

                 raise Exception("interrupted")

             try:

-                image = Image.open(path).convert('RGB').resize((self.width, self.height), PIL.Image.BICUBIC)

+                image = Image.open(path).convert('RGB')

+                if not varsize:

+                    image = image.resize((width, height), PIL.Image.BICUBIC)

             except Exception:

                 continue

 

@@ -103,18 +105,25 @@ class PersonalizedBase(Dataset):
             if include_cond and not (self.tag_drop_out != 0 or self.shuffle_tags):

                 with devices.autocast():

                     entry.cond = cond_model([entry.cond_text]).to(devices.cpu).squeeze(0)

-

+            groups[image.size].append(len(self.dataset))

             self.dataset.append(entry)

             del torchdata

             del latent_dist

             del latent_sample

 

         self.length = len(self.dataset)

+        self.groups = list(groups.values())

         assert self.length > 0, "No images have been found in the dataset."

         self.batch_size = min(batch_size, self.length)

         self.gradient_step = min(gradient_step, self.length // self.batch_size)

         self.latent_sampling_method = latent_sampling_method

 

+        if len(groups) > 1:

+            print("Buckets:")

+            for (w, h), ids in sorted(groups.items(), key=lambda x: x[0]):

+                print(f"  {w}x{h}: {len(ids)}")

+            print()

+

     def create_text(self, filename_text):

         text = random.choice(self.lines)

         tags = filename_text.split(',')

@@ -137,9 +146,44 @@ class PersonalizedBase(Dataset):
             entry.latent_sample = shared.sd_model.get_first_stage_encoding(entry.latent_dist).to(devices.cpu)

         return entry

 

+

+class GroupedBatchSampler(Sampler):

+    def __init__(self, data_source: PersonalizedBase, batch_size: int):

+        super().__init__(data_source)

+

+        n = len(data_source)

+        self.groups = data_source.groups

+        self.len = n_batch = n // batch_size

+        expected = [len(g) / n * n_batch * batch_size for g in data_source.groups]

+        self.base = [int(e) // batch_size for e in expected]

+        self.n_rand_batches = nrb = n_batch - sum(self.base)

+        self.probs = [e%batch_size/nrb/batch_size if nrb>0 else 0 for e in expected]

+        self.batch_size = batch_size

+

+    def __len__(self):

+        return self.len

+

+    def __iter__(self):

+        b = self.batch_size

+

+        for g in self.groups:

+            shuffle(g)

+

+        batches = []

+        for g in self.groups:

+            batches.extend(g[i*b:(i+1)*b] for i in range(len(g) // b))

+        for _ in range(self.n_rand_batches):

+            rand_group = choices(self.groups, self.probs)[0]

+            batches.append(choices(rand_group, k=b))

+

+        shuffle(batches)

+

+        yield from batches

+

+

 class PersonalizedDataLoader(DataLoader):

     def __init__(self, dataset, latent_sampling_method="once", batch_size=1, pin_memory=False):

-        super(PersonalizedDataLoader, self).__init__(dataset, shuffle=True, drop_last=True, batch_size=batch_size, pin_memory=pin_memory)

+        super(PersonalizedDataLoader, self).__init__(dataset, batch_sampler=GroupedBatchSampler(dataset, batch_size), pin_memory=pin_memory)

         if latent_sampling_method == "random":

             self.collate_fn = collate_wrapper_random

         else:

diff --git a/modules/textual_inversion/image_embedding.py b/modules/textual_inversion/image_embedding.py
index ea653806..5593f88c 100644
--- a/modules/textual_inversion/image_embedding.py
+++ b/modules/textual_inversion/image_embedding.py
@@ -76,10 +76,10 @@ def insert_image_data_embed(image, data):
     next_size = data_np_low.shape[0] + (h-(data_np_low.shape[0] % h))

     next_size = next_size + ((h*d)-(next_size % (h*d)))

 

-    data_np_low.resize(next_size)

+    data_np_low = np.resize(data_np_low, next_size)

     data_np_low = data_np_low.reshape((h, -1, d))

 

-    data_np_high.resize(next_size)

+    data_np_high = np.resize(data_np_high, next_size)

     data_np_high = data_np_high.reshape((h, -1, d))

 

     edge_style = list(data['string_to_param'].values())[0].cpu().detach().numpy().tolist()[0][:1024]

diff --git a/modules/textual_inversion/learn_schedule.py b/modules/textual_inversion/learn_schedule.py
index dd0c0ad1..f63fc72f 100644
--- a/modules/textual_inversion/learn_schedule.py
+++ b/modules/textual_inversion/learn_schedule.py
@@ -58,14 +58,19 @@ class LearnRateScheduler:
 

         self.finished = False

 

-    def apply(self, optimizer, step_number):

+    def step(self, step_number):

         if step_number < self.end_step:

-            return

+            return False

 

         try:

             (self.learn_rate, self.end_step) = next(self.schedules)

-        except Exception:

+        except StopIteration:

             self.finished = True

+            return False

+        return True

+

+    def apply(self, optimizer, step_number):

+        if not self.step(step_number):

             return

 

         if self.verbose:

diff --git a/modules/textual_inversion/logging.py b/modules/textual_inversion/logging.py
new file mode 100644
index 00000000..734a4b6f
--- /dev/null
+++ b/modules/textual_inversion/logging.py
@@ -0,0 +1,24 @@
+import datetime

+import json

+import os

+

+saved_params_shared = {"model_name", "model_hash", "initial_step", "num_of_dataset_images", "learn_rate", "batch_size", "clip_grad_mode", "clip_grad_value", "gradient_step", "data_root", "log_directory", "training_width", "training_height", "steps", "create_image_every", "template_file", "gradient_step", "latent_sampling_method"}

+saved_params_ti = {"embedding_name", "num_vectors_per_token", "save_embedding_every", "save_image_with_stored_embedding"}

+saved_params_hypernet = {"hypernetwork_name", "layer_structure", "activation_func", "weight_init", "add_layer_norm", "use_dropout", "save_hypernetwork_every"}

+saved_params_all = saved_params_shared | saved_params_ti | saved_params_hypernet

+saved_params_previews = {"preview_prompt", "preview_negative_prompt", "preview_steps", "preview_sampler_index", "preview_cfg_scale", "preview_seed", "preview_width", "preview_height"}

+

+

+def save_settings_to_file(log_directory, all_params):

+    now = datetime.datetime.now()

+    params = {"datetime": now.strftime("%Y-%m-%d %H:%M:%S")}

+

+    keys = saved_params_all

+    if all_params.get('preview_from_txt2img'):

+        keys = keys | saved_params_previews

+

+    params.update({k: v for k, v in all_params.items() if k in keys})

+

+    filename = f'settings-{now.strftime("%Y-%m-%d-%H-%M-%S")}.json'

+    with open(os.path.join(log_directory, filename), "w") as file:

+        json.dump(params, file, indent=4)

diff --git a/modules/textual_inversion/preprocess.py b/modules/textual_inversion/preprocess.py
index 56b9b2eb..2239cb84 100644
--- a/modules/textual_inversion/preprocess.py
+++ b/modules/textual_inversion/preprocess.py
@@ -6,13 +6,12 @@ import sys
 import tqdm

 import time

 

-from modules import shared, images, deepbooru

-from modules.paths import models_path

+from modules import paths, shared, images, deepbooru

 from modules.shared import opts, cmd_opts

 from modules.textual_inversion import autocrop

 

 

-def preprocess(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):

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

     try:

         if process_caption:

             shared.interrogator.load()

@@ -20,7 +19,7 @@ def preprocess(process_src, process_dst, process_width, process_height, preproce
         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)

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

 

     finally:

 

@@ -109,8 +108,30 @@ def split_pic(image, inverse_xy, width, height, overlap_ratio):
             splitted = image.crop((0, y, to_w, y + to_h))

         yield splitted

 

-

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

+# not using torchvision.transforms.CenterCrop because it doesn't allow float regions

+def center_crop(image: Image, w: int, h: int):

+    iw, ih = image.size

+    if ih / h < iw / w:

+        sw = w * ih / h

+        box = (iw - sw) / 2, 0, iw - (iw - sw) / 2, ih

+    else:

+        sh = h * iw / w

+        box = 0, (ih - sh) / 2, iw, ih - (ih - sh) / 2

+    return image.resize((w, h), Image.Resampling.LANCZOS, box)

+

+

+def multicrop_pic(image: Image, mindim, maxdim, minarea, maxarea, objective, threshold):

+    iw, ih = image.size

+    err = lambda w, h: 1-(lambda x: x if x < 1 else 1/x)(iw/ih/(w/h))

+    wh = max(((w, h) for w in range(mindim, maxdim+1, 64) for h in range(mindim, maxdim+1, 64)

+        if minarea <= w * h <= maxarea and err(w, h) <= threshold),

+        key= lambda wh: (wh[0]*wh[1], -err(*wh))[::1 if objective=='Maximize area' else -1],

+        default=None

+    )

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

     width = process_width

     height = process_height

     src = os.path.abspath(process_src)

@@ -124,6 +145,7 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
 

     files = listfiles(src)

 

+    shared.state.job = "preprocess"

     shared.state.textinfo = "Preprocessing..."

     shared.state.job_count = len(files)

 

@@ -134,7 +156,8 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
     params.process_caption_deepbooru = process_caption_deepbooru

     params.preprocess_txt_action = preprocess_txt_action

 

-    for index, imagefile in enumerate(tqdm.tqdm(files)):

+    pbar = tqdm.tqdm(files)

+    for index, imagefile in enumerate(pbar):

         params.subindex = 0

         filename = os.path.join(src, imagefile)

         try:

@@ -142,6 +165,10 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
         except Exception:

             continue

 

+        description = f"Preprocessing [Image {index}/{len(files)}]"

+        pbar.set_description(description)

+        shared.state.textinfo = description

+

         params.src = filename

 

         existing_caption = None

@@ -171,7 +198,7 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
 

             dnn_model_path = None

             try:

-                dnn_model_path = autocrop.download_and_cache_models(os.path.join(models_path, "opencv"))

+                dnn_model_path = autocrop.download_and_cache_models(os.path.join(paths.models_path, "opencv"))

             except Exception as e:

                 print("Unable to load face detection model for auto crop selection. Falling back to lower quality haar method.", e)

 

@@ -188,6 +215,14 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
                 save_pic(focal, index, params, existing_caption=existing_caption)

             process_default_resize = False

 

+        if process_multicrop:

+            cropped = multicrop_pic(img, process_multicrop_mindim, process_multicrop_maxdim, process_multicrop_minarea, process_multicrop_maxarea, process_multicrop_objective, process_multicrop_threshold)

+            if cropped is not None:

+                save_pic(cropped, index, params, existing_caption=existing_caption)

+            else:

+                print(f"skipped {img.width}x{img.height} image {filename} (can't find suitable size within error threshold)")

+            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 f6112578..6cf00e65 100644
--- a/modules/textual_inversion/textual_inversion.py
+++ b/modules/textual_inversion/textual_inversion.py
@@ -1,31 +1,56 @@
 import os

 import sys

 import traceback

+import inspect

+from collections import namedtuple

 

 import torch

 import tqdm

 import html

 import datetime

 import csv

+import safetensors.torch

 

+import numpy as np

 from PIL import Image, PngImagePlugin

+from torch.utils.tensorboard import SummaryWriter

 

-from modules import shared, devices, sd_hijack, processing, sd_models, images, sd_samplers

+from modules import shared, devices, sd_hijack, processing, sd_models, images, sd_samplers, sd_hijack_checkpoint

 import modules.textual_inversion.dataset

 from modules.textual_inversion.learn_schedule import LearnRateScheduler

 

-from modules.textual_inversion.image_embedding import (embedding_to_b64, embedding_from_b64,

-                                                       insert_image_data_embed, extract_image_data_embed,

-                                                       caption_image_overlay)

+from modules.textual_inversion.image_embedding import embedding_to_b64, embedding_from_b64, insert_image_data_embed, extract_image_data_embed, caption_image_overlay

+from modules.textual_inversion.logging import save_settings_to_file

+

+

+TextualInversionTemplate = namedtuple("TextualInversionTemplate", ["name", "path"])

+textual_inversion_templates = {}

+

+

+def list_textual_inversion_templates():

+    textual_inversion_templates.clear()

+

+    for root, dirs, fns in os.walk(shared.cmd_opts.textual_inversion_templates_dir):

+        for fn in fns:

+            path = os.path.join(root, fn)

+

+            textual_inversion_templates[fn] = TextualInversionTemplate(fn, path)

+

+    return textual_inversion_templates

+

 

 class Embedding:

     def __init__(self, vec, name, step=None):

         self.vec = vec

         self.name = name

         self.step = step

+        self.shape = None

+        self.vectors = 0

         self.cached_checksum = None

         self.sd_checkpoint = None

         self.sd_checkpoint_name = None

+        self.optimizer_state_dict = None

+        self.filename = None

 

     def save(self, filename):

         embedding_data = {

@@ -39,6 +64,13 @@ class Embedding:
 

         torch.save(embedding_data, filename)

 

+        if shared.opts.save_optimizer_state and self.optimizer_state_dict is not None:

+            optimizer_saved_dict = {

+                'hash': self.checksum(),

+                'optimizer_state_dict': self.optimizer_state_dict,

+            }

+            torch.save(optimizer_saved_dict, filename + '.optim')

+

     def checksum(self):

         if self.cached_checksum is not None:

             return self.cached_checksum

@@ -53,18 +85,43 @@ class Embedding:
         return self.cached_checksum

 

 

+class DirWithTextualInversionEmbeddings:

+    def __init__(self, path):

+        self.path = path

+        self.mtime = None

+

+    def has_changed(self):

+        if not os.path.isdir(self.path):

+            return False

+

+        mt = os.path.getmtime(self.path)

+        if self.mtime is None or mt > self.mtime:

+            return True

+

+    def update(self):

+        if not os.path.isdir(self.path):

+            return

+

+        self.mtime = os.path.getmtime(self.path)

+

+

 class EmbeddingDatabase:

-    def __init__(self, embeddings_dir):

+    def __init__(self):

         self.ids_lookup = {}

         self.word_embeddings = {}

-        self.dir_mtime = None

-        self.embeddings_dir = embeddings_dir

+        self.skipped_embeddings = {}

+        self.expected_shape = -1

+        self.embedding_dirs = {}

 

-    def register_embedding(self, embedding, model):

+    def add_embedding_dir(self, path):

+        self.embedding_dirs[path] = DirWithTextualInversionEmbeddings(path)

+

+    def clear_embedding_dirs(self):

+        self.embedding_dirs.clear()

 

+    def register_embedding(self, embedding, model):

         self.word_embeddings[embedding.name] = embedding

 

-        # TODO changing between clip and open clip changes tokenization, which will cause embeddings to stop working

         ids = model.cond_stage_model.tokenize([embedding.name])[0]

 

         first_id = ids[0]

@@ -75,70 +132,105 @@ class EmbeddingDatabase:
 

         return embedding

 

-    def load_textual_inversion_embeddings(self):

-        mt = os.path.getmtime(self.embeddings_dir)

-        if self.dir_mtime is not None and mt <= self.dir_mtime:

-            return

-

-        self.dir_mtime = mt

-        self.ids_lookup.clear()

-        self.word_embeddings.clear()

+    def get_expected_shape(self):

+        vec = shared.sd_model.cond_stage_model.encode_embedding_init_text(",", 1)

+        return vec.shape[1]

 

-        def process_file(path, filename):

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

+    def load_from_file(self, path, filename):

+        name, ext = os.path.splitext(filename)

+        ext = ext.upper()

 

-            data = []

+        if ext in ['.PNG', '.WEBP', '.JXL', '.AVIF']:

+            _, second_ext = os.path.splitext(name)

+            if second_ext.upper() == '.PREVIEW':

+                return

 

-            if os.path.splitext(filename.upper())[-1] in ['.PNG', '.WEBP', '.JXL', '.AVIF']:

-                embed_image = Image.open(path)

-                if hasattr(embed_image, 'text') and 'sd-ti-embedding' in embed_image.text:

-                    data = embedding_from_b64(embed_image.text['sd-ti-embedding'])

-                    name = data.get('name', name)

-                else:

-                    data = extract_image_data_embed(embed_image)

-                    name = data.get('name', name)

-            else:

-                data = torch.load(path, map_location="cpu")

-

-            # textual inversion embeddings

-            if 'string_to_param' in data:

-                param_dict = data['string_to_param']

-                if hasattr(param_dict, '_parameters'):

-                    param_dict = getattr(param_dict, '_parameters')  # fix for torch 1.12.1 loading saved file from torch 1.11

-                assert len(param_dict) == 1, 'embedding file has multiple terms in it'

-                emb = next(iter(param_dict.items()))[1]

-            # diffuser concepts

-            elif type(data) == dict and type(next(iter(data.values()))) == torch.Tensor:

-                assert len(data.keys()) == 1, 'embedding file has multiple terms in it'

-

-                emb = next(iter(data.values()))

-                if len(emb.shape) == 1:

-                    emb = emb.unsqueeze(0)

+            embed_image = Image.open(path)

+            if hasattr(embed_image, 'text') and 'sd-ti-embedding' in embed_image.text:

+                data = embedding_from_b64(embed_image.text['sd-ti-embedding'])

+                name = data.get('name', name)

             else:

-                raise Exception(f"Couldn't identify {filename} as neither textual inversion embedding nor diffuser concept.")

+                data = extract_image_data_embed(embed_image)

+                name = data.get('name', name)

+        elif ext in ['.BIN', '.PT']:

+            data = torch.load(path, map_location="cpu")

+        elif ext in ['.SAFETENSORS']:

+            data = safetensors.torch.load_file(path, device="cpu")

+        else:

+            return

 

-            vec = emb.detach().to(devices.device, dtype=torch.float32)

-            embedding = Embedding(vec, name)

-            embedding.step = data.get('step', None)

-            embedding.sd_checkpoint = data.get('sd_checkpoint', None)

-            embedding.sd_checkpoint_name = data.get('sd_checkpoint_name', None)

+        # textual inversion embeddings

+        if 'string_to_param' in data:

+            param_dict = data['string_to_param']

+            if hasattr(param_dict, '_parameters'):

+                param_dict = getattr(param_dict, '_parameters')  # fix for torch 1.12.1 loading saved file from torch 1.11

+            assert len(param_dict) == 1, 'embedding file has multiple terms in it'

+            emb = next(iter(param_dict.items()))[1]

+        # diffuser concepts

+        elif type(data) == dict and type(next(iter(data.values()))) == torch.Tensor:

+            assert len(data.keys()) == 1, 'embedding file has multiple terms in it'

+

+            emb = next(iter(data.values()))

+            if len(emb.shape) == 1:

+                emb = emb.unsqueeze(0)

+        else:

+            raise Exception(f"Couldn't identify {filename} as neither textual inversion embedding nor diffuser concept.")

+

+        vec = emb.detach().to(devices.device, dtype=torch.float32)

+        embedding = Embedding(vec, name)

+        embedding.step = data.get('step', None)

+        embedding.sd_checkpoint = data.get('sd_checkpoint', None)

+        embedding.sd_checkpoint_name = data.get('sd_checkpoint_name', None)

+        embedding.vectors = vec.shape[0]

+        embedding.shape = vec.shape[-1]

+        embedding.filename = path

+

+        if self.expected_shape == -1 or self.expected_shape == embedding.shape:

             self.register_embedding(embedding, shared.sd_model)

+        else:

+            self.skipped_embeddings[name] = embedding

 

-        for fn in os.listdir(self.embeddings_dir):

-            try:

-                fullfn = os.path.join(self.embeddings_dir, fn)

+    def load_from_dir(self, embdir):

+        if not os.path.isdir(embdir.path):

+            return

+

+        for root, dirs, fns in os.walk(embdir.path, followlinks=True):

+            for fn in fns:

+                try:

+                    fullfn = os.path.join(root, fn)

+

+                    if os.stat(fullfn).st_size == 0:

+                        continue

 

-                if os.stat(fullfn).st_size == 0:

+                    self.load_from_file(fullfn, fn)

+                except Exception:

+                    print(f"Error loading embedding {fn}:", file=sys.stderr)

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

                     continue

 

-                process_file(fullfn, fn)

-            except Exception:

-                print(f"Error loading embedding {fn}:", file=sys.stderr)

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

-                continue

+    def load_textual_inversion_embeddings(self, force_reload=False):

+        if not force_reload:

+            need_reload = False

+            for path, embdir in self.embedding_dirs.items():

+                if embdir.has_changed():

+                    need_reload = True

+                    break

+

+            if not need_reload:

+                return

+

+        self.ids_lookup.clear()

+        self.word_embeddings.clear()

+        self.skipped_embeddings.clear()

+        self.expected_shape = self.get_expected_shape()

+

+        for path, embdir in self.embedding_dirs.items():

+            self.load_from_dir(embdir)

+            embdir.update()

 

-        print(f"Loaded a total of {len(self.word_embeddings)} textual inversion embeddings.")

-        print("Embeddings:", ', '.join(self.word_embeddings.keys()))

+        print(f"Textual inversion embeddings loaded({len(self.word_embeddings)}): {', '.join(self.word_embeddings.keys())}")

+        if len(self.skipped_embeddings) > 0:

+            print(f"Textual inversion embeddings skipped({len(self.skipped_embeddings)}): {', '.join(self.skipped_embeddings.keys())}")

 

     def find_embedding_at_position(self, tokens, offset):

         token = tokens[offset]

@@ -160,11 +252,14 @@ def create_embedding(name, num_vectors_per_token, overwrite_old, init_text='*'):
     with devices.autocast():

         cond_model([""])  # will send cond model to GPU if lowvram/medvram is active

 

-    embedded = cond_model.encode_embedding_init_text(init_text, num_vectors_per_token)

+    #cond_model expects at least some text, so we provide '*' as backup.

+    embedded = cond_model.encode_embedding_init_text(init_text or '*', num_vectors_per_token)

     vec = torch.zeros((num_vectors_per_token, embedded.shape[1]), device=devices.device)

 

-    for i in range(num_vectors_per_token):

-        vec[i] = embedded[i * int(embedded.shape[0]) // num_vectors_per_token]

+    #Only copy if we provided an init_text, otherwise keep vectors as zeros

+    if init_text:

+        for i in range(num_vectors_per_token):

+            vec[i] = embedded[i * int(embedded.shape[0]) // num_vectors_per_token]

 

     # Remove illegal characters from name.

     name = "".join( x for x in name if (x.isalnum() or x in "._- "))

@@ -203,7 +298,32 @@ def write_loss(log_directory, filename, step, epoch_len, values):
             **values,

         })

 

-def validate_train_inputs(model_name, learn_rate, batch_size, gradient_step, data_root, template_file, steps, save_model_every, create_image_every, log_directory, name="embedding"):

+def tensorboard_setup(log_directory):

+    os.makedirs(os.path.join(log_directory, "tensorboard"), exist_ok=True)

+    return SummaryWriter(

+            log_dir=os.path.join(log_directory, "tensorboard"),

+            flush_secs=shared.opts.training_tensorboard_flush_every)

+

+def tensorboard_add(tensorboard_writer, loss, global_step, step, learn_rate, epoch_num):

+    tensorboard_add_scaler(tensorboard_writer, "Loss/train", loss, global_step)

+    tensorboard_add_scaler(tensorboard_writer, f"Loss/train/epoch-{epoch_num}", loss, step)

+    tensorboard_add_scaler(tensorboard_writer, "Learn rate/train", learn_rate, global_step)

+    tensorboard_add_scaler(tensorboard_writer, f"Learn rate/train/epoch-{epoch_num}", learn_rate, step)

+

+def tensorboard_add_scaler(tensorboard_writer, tag, value, step):

+    tensorboard_writer.add_scalar(tag=tag, 

+        scalar_value=value, global_step=step)

+

+def tensorboard_add_image(tensorboard_writer, tag, pil_image, step):

+    # Convert a pil image to a torch tensor

+    img_tensor = torch.as_tensor(np.array(pil_image, copy=True))

+    img_tensor = img_tensor.view(pil_image.size[1], pil_image.size[0], 

+        len(pil_image.getbands()))

+    img_tensor = img_tensor.permute((2, 0, 1))

+                

+    tensorboard_writer.add_image(tag, img_tensor, global_step=step)

+

+def validate_train_inputs(model_name, learn_rate, batch_size, gradient_step, data_root, template_file, template_filename, steps, save_model_every, create_image_every, log_directory, name="embedding"):

     assert model_name, f"{name} not selected"

     assert learn_rate, "Learning rate is empty or 0"

     assert isinstance(batch_size, int), "Batch size must be integer"

@@ -213,23 +333,28 @@ def validate_train_inputs(model_name, learn_rate, batch_size, gradient_step, dat
     assert data_root, "Dataset directory is empty"

     assert os.path.isdir(data_root), "Dataset directory doesn't exist"

     assert os.listdir(data_root), "Dataset directory is empty"

-    assert template_file, "Prompt template file is empty"

-    assert os.path.isfile(template_file), "Prompt template file doesn't exist"

+    assert template_filename, "Prompt template file not selected"

+    assert template_file, f"Prompt template file {template_filename} not found"

+    assert os.path.isfile(template_file.path), f"Prompt template file {template_filename} doesn't exist"

     assert steps, "Max steps is empty or 0"

     assert isinstance(steps, int), "Max steps must be integer"

-    assert steps > 0 , "Max steps must be positive"

+    assert steps > 0, "Max steps must be positive"

     assert isinstance(save_model_every, int), "Save {name} must be integer"

-    assert save_model_every >= 0 , "Save {name} must be positive or 0"

+    assert save_model_every >= 0, "Save {name} must be positive or 0"

     assert isinstance(create_image_every, int), "Create image must be integer"

-    assert create_image_every >= 0 , "Create image must be positive or 0"

+    assert create_image_every >= 0, "Create image must be positive or 0"

     if save_model_every or create_image_every:

         assert log_directory, "Log directory is empty"

 

-def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_root, log_directory, training_width, training_height, steps, shuffle_tags, tag_drop_out, latent_sampling_method, create_image_every, save_embedding_every, template_file, 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, 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

-    validate_train_inputs(embedding_name, learn_rate, batch_size, gradient_step, data_root, template_file, steps, save_embedding_every, create_image_every, log_directory, name="embedding")

+    template_file = textual_inversion_templates.get(template_filename, None)

+    validate_train_inputs(embedding_name, learn_rate, batch_size, gradient_step, data_root, template_file, template_filename, steps, save_embedding_every, create_image_every, log_directory, name="embedding")

+    template_file = template_file.path

 

+    shared.state.job = "train-embedding"

     shared.state.textinfo = "Initializing textual inversion training..."

     shared.state.job_count = steps

 

@@ -265,15 +390,26 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
     if initial_step >= steps:

         shared.state.textinfo = "Model has already been trained beyond specified max steps"

         return embedding, filename

+    

     scheduler = LearnRateScheduler(learn_rate, steps, initial_step)

-

-   # dataset loading may take a while, so input validations and early returns should be done before this

+    clip_grad = torch.nn.utils.clip_grad_value_ if clip_grad_mode == "value" else \

+        torch.nn.utils.clip_grad_norm_ if clip_grad_mode == "norm" else \

+        None

+    if clip_grad:

+        clip_grad_sched = LearnRateScheduler(clip_grad_value, steps, initial_step, verbose=False)

+    # dataset loading may take a while, so input validations and early returns should be done before this

     shared.state.textinfo = f"Preparing dataset from {html.escape(data_root)}..."

     old_parallel_processing_allowed = shared.parallel_processing_allowed

+    

+    if shared.opts.training_enable_tensorboard:

+        tensorboard_writer = tensorboard_setup(log_directory)

 

     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)

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

+

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

 

     latent_sampling_method = ds.latent_sampling_method

 

@@ -285,6 +421,19 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
 

     embedding.vec.requires_grad = True

     optimizer = torch.optim.AdamW([embedding.vec], lr=scheduler.learn_rate, weight_decay=0.0)

+    if shared.opts.save_optimizer_state:

+        optimizer_state_dict = None

+        if os.path.exists(filename + '.optim'):

+            optimizer_saved_dict = torch.load(filename + '.optim', map_location='cpu')

+            if embedding.checksum() == optimizer_saved_dict.get('hash', None):

+                optimizer_state_dict = optimizer_saved_dict.get('optimizer_state_dict', None)

+    

+        if optimizer_state_dict is not None:

+            optimizer.load_state_dict(optimizer_state_dict)

+            print("Loaded existing optimizer from checkpoint")

+        else:

+            print("No saved optimizer exists in checkpoint")

+

     scaler = torch.cuda.amp.GradScaler()

 

     batch_size = ds.batch_size

@@ -295,14 +444,18 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
     loss_step = 0

     _loss_step = 0 #internal

 

-

     last_saved_file = "<none>"

     last_saved_image = "<none>"

     forced_filename = "<none>"

     embedding_yet_to_be_embedded = False

 

+    is_training_inpainting_model = shared.sd_model.model.conditioning_key in {'hybrid', 'concat'}

+    img_c = None

+

     pbar = tqdm.tqdm(total=steps - initial_step)

     try:

+        sd_hijack_checkpoint.add()

+

         for i in range((steps-initial_step) * gradient_step):

             if scheduler.finished:

                 break

@@ -318,14 +471,22 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
                 if shared.state.interrupted:

                     break

 

+                if clip_grad:

+                    clip_grad_sched.step(embedding.step)

+            

                 with devices.autocast():

-                    # c = stack_conds(batch.cond).to(devices.device)

-                    # mask = torch.tensor(batch.emb_index).to(devices.device, non_blocking=pin_memory)

-                    # print(mask)

-                    # c[:, 1:1+embedding.vec.shape[0]] = embedding.vec.to(devices.device, non_blocking=pin_memory)

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

                     c = shared.sd_model.cond_stage_model(batch.cond_text)

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

+

+                    if is_training_inpainting_model:

+                        if img_c is None:

+                            img_c = processing.txt2img_image_conditioning(shared.sd_model, c, training_width, training_height)

+

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

+                    else:

+                        cond = c

+

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

                     del x

 

                     _loss_step += loss.item()

@@ -334,6 +495,10 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
                 # go back until we reach gradient accumulation steps

                 if (j + 1) % gradient_step != 0:

                     continue

+                

+                if clip_grad:

+                    clip_grad(embedding.vec, clip_grad_sched.learn_rate)

+

                 scaler.step(optimizer)

                 scaler.update()

                 embedding.step += 1

@@ -347,14 +512,13 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
                 epoch_num = embedding.step // steps_per_epoch

                 epoch_step = embedding.step % steps_per_epoch

 

-                pbar.set_description(f"[Epoch {epoch_num}: {epoch_step+1}/{steps_per_epoch}]loss: {loss_step:.7f}")

+                description = f"Training textual inversion [Epoch {epoch_num}: {epoch_step+1}/{steps_per_epoch}] loss: {loss_step:.7f}"

+                pbar.set_description(description)

                 if embedding_dir is not None and steps_done % save_embedding_every == 0:

                     # Before saving, change name to match current checkpoint.

                     embedding_name_every = f'{embedding_name}-{steps_done}'

                     last_saved_file = os.path.join(embedding_dir, f'{embedding_name_every}.pt')

-                    #if shared.opts.save_optimizer_state:

-                        #embedding.optimizer_state_dict = optimizer.state_dict()

-                    save_embedding(embedding, checkpoint, embedding_name_every, last_saved_file, remove_cached_checksum=True)

+                    save_embedding(embedding, optimizer, checkpoint, embedding_name_every, last_saved_file, remove_cached_checksum=True)

                     embedding_yet_to_be_embedded = True

 

                 write_loss(log_directory, "textual_inversion_loss.csv", embedding.step, steps_per_epoch, {

@@ -399,10 +563,14 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
                         shared.sd_model.first_stage_model.to(devices.cpu)

 

                     if image is not None:

-                        shared.state.current_image = image

+                        shared.state.assign_current_image(image)

+

                         last_saved_image, last_text_info = images.save_image(image, images_dir, "", p.seed, p.prompt, shared.opts.samples_format, processed.infotexts[0], p=p, forced_filename=forced_filename, save_to_dirs=False)

                         last_saved_image += f", prompt: {preview_text}"

 

+                        if shared.opts.training_enable_tensorboard and shared.opts.training_tensorboard_save_images:

+                            tensorboard_add_image(tensorboard_writer, f"Validation at epoch {epoch_num}", image, embedding.step)

+

                     if save_image_with_stored_embedding and os.path.exists(last_saved_file) and embedding_yet_to_be_embedded:

 

                         last_saved_image_chunks = os.path.join(images_embeds_dir, f'{embedding_name}-{steps_done}.png')

@@ -420,7 +588,7 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
 

                         checkpoint = sd_models.select_checkpoint()

                         footer_left = checkpoint.model_name

-                        footer_mid = '[{}]'.format(checkpoint.hash)

+                        footer_mid = '[{}]'.format(checkpoint.shorthash)

                         footer_right = '{}v {}s'.format(vectorSize, steps_done)

 

                         captioned_image = caption_image_overlay(image, title, footer_left, footer_mid, footer_right)

@@ -444,7 +612,7 @@ Last saved image: {html.escape(last_saved_image)}<br/>
 </p>

 """

         filename = os.path.join(shared.cmd_opts.embeddings_dir, f'{embedding_name}.pt')

-        save_embedding(embedding, checkpoint, embedding_name, filename, remove_cached_checksum=True)

+        save_embedding(embedding, optimizer, checkpoint, embedding_name, filename, remove_cached_checksum=True)

     except Exception:

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

         pass

@@ -453,20 +621,23 @@ Last saved image: {html.escape(last_saved_image)}<br/>
         pbar.close()

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

         shared.parallel_processing_allowed = old_parallel_processing_allowed

+        sd_hijack_checkpoint.remove()

 

     return embedding, filename

 

-def save_embedding(embedding, checkpoint, embedding_name, filename, remove_cached_checksum=True):

+

+def save_embedding(embedding, optimizer, checkpoint, embedding_name, filename, remove_cached_checksum=True):

     old_embedding_name = embedding.name

     old_sd_checkpoint = embedding.sd_checkpoint if hasattr(embedding, "sd_checkpoint") else None

     old_sd_checkpoint_name = embedding.sd_checkpoint_name if hasattr(embedding, "sd_checkpoint_name") else None

     old_cached_checksum = embedding.cached_checksum if hasattr(embedding, "cached_checksum") else None

     try:

-        embedding.sd_checkpoint = checkpoint.hash

+        embedding.sd_checkpoint = checkpoint.shorthash

         embedding.sd_checkpoint_name = checkpoint.model_name

         if remove_cached_checksum:

             embedding.cached_checksum = None

         embedding.name = embedding_name

+        embedding.optimizer_state_dict = optimizer.state_dict()

         embedding.save(filename)

     except:

         embedding.sd_checkpoint = old_sd_checkpoint

diff --git a/modules/timer.py b/modules/timer.py
new file mode 100644
index 00000000..57a4f17a
--- /dev/null
+++ b/modules/timer.py
@@ -0,0 +1,35 @@
+import time

+

+

+class Timer:

+    def __init__(self):

+        self.start = time.time()

+        self.records = {}

+        self.total = 0

+

+    def elapsed(self):

+        end = time.time()

+        res = end - self.start

+        self.start = end

+        return res

+

+    def record(self, category, extra_time=0):

+        e = self.elapsed()

+        if category not in self.records:

+            self.records[category] = 0

+

+        self.records[category] += e + extra_time

+        self.total += e + extra_time

+

+    def summary(self):

+        res = f"{self.total:.1f}s"

+

+        additions = [x for x in self.records.items() if x[1] >= 0.1]

+        if not additions:

+            return res

+

+        res += " ("

+        res += ", ".join([f"{category}: {time_taken:.1f}s" for category, time_taken in additions])

+        res += ")"

+

+        return res

diff --git a/modules/txt2img.py b/modules/txt2img.py
index c8f81176..e945fd69 100644
--- a/modules/txt2img.py
+++ b/modules/txt2img.py
@@ -8,13 +8,13 @@ import modules.processing as processing
 from modules.ui import plaintext_to_html

 

 

-def txt2img(prompt: str, negative_prompt: str, prompt_style: str, prompt_style2: str, steps: int, sampler_index: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, enable_hr: bool, denoising_strength: float, firstphase_width: int, firstphase_height: int, *args):

+def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, steps: int, sampler_index: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, enable_hr: bool, denoising_strength: float, hr_scale: float, hr_upscaler: str, hr_second_pass_steps: int, hr_resize_x: int, hr_resize_y: int, *args):

     p = StableDiffusionProcessingTxt2Img(

         sd_model=shared.sd_model,

         outpath_samples=opts.outdir_samples or opts.outdir_txt2img_samples,

         outpath_grids=opts.outdir_grids or opts.outdir_txt2img_grids,

         prompt=prompt,

-        styles=[prompt_style, prompt_style2],

+        styles=prompt_styles,

         negative_prompt=negative_prompt,

         seed=seed,

         subseed=subseed,

@@ -33,8 +33,11 @@ def txt2img(prompt: str, negative_prompt: str, prompt_style: str, prompt_style2:
         tiling=tiling,

         enable_hr=enable_hr,

         denoising_strength=denoising_strength if enable_hr else None,

-        firstphase_width=firstphase_width if enable_hr else None,

-        firstphase_height=firstphase_height if enable_hr else None,

+        hr_scale=hr_scale,

+        hr_upscaler=hr_upscaler,

+        hr_second_pass_steps=hr_second_pass_steps,

+        hr_resize_x=hr_resize_x,

+        hr_resize_y=hr_resize_y,

     )

 

     p.scripts = modules.scripts.scripts_txt2img

@@ -59,4 +62,4 @@ def txt2img(prompt: str, negative_prompt: str, prompt_style: str, prompt_style2:
     if opts.do_not_show_images:

         processed.images = []

 

-    return processed.images, generation_info_js, plaintext_to_html(processed.info)

+    return processed.images, generation_info_js, plaintext_to_html(processed.info), plaintext_to_html(processed.comments)

diff --git a/modules/ui.py b/modules/ui.py
index 57ee0465..9f4cfda1 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -5,12 +5,12 @@ import mimetypes
 import os

 import platform

 import random

-import subprocess as sp

 import sys

 import tempfile

 import time

 import traceback

 from functools import partial, reduce

+import warnings

 

 import gradio as gr

 import gradio.routes

@@ -19,8 +19,9 @@ import numpy as np
 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

-from modules.paths import script_path

+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.paths import script_path, data_path

 

 from modules.shared import opts, cmd_opts, restricted_opts

 

@@ -36,9 +37,12 @@ from modules import prompt_parser
 from modules.images import save_image

 from modules.sd_hijack import model_hijack

 from modules.sd_samplers import samplers, samplers_for_img2img

-import modules.textual_inversion.ui

+from modules.textual_inversion import textual_inversion

 import modules.hypernetworks.ui

 from modules.generation_parameters_copypaste import image_from_url_text

+import modules.extras

+

+warnings.filterwarnings("default" if opts.show_warnings else "ignore", category=UserWarning)

 

 # this is a fix for Windows users. Without it, javascript files will be served with text/html content-type and the browser will not show any UI

 mimetypes.init()

@@ -71,6 +75,7 @@ css_hide_progressbar = """
 .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; }

@@ -80,167 +85,24 @@ css_hide_progressbar = """
 # Important that they exactly match script.js for tooltip to work.

 random_symbol = '\U0001f3b2\ufe0f'  # 🎲️

 reuse_symbol = '\u267b\ufe0f'  # ♻️

-art_symbol = '\U0001f3a8'  # 🎨

 paste_symbol = '\u2199\ufe0f'  # ↙

-folder_symbol = '\U0001f4c2'  # 📂

 refresh_symbol = '\U0001f504'  # 🔄

 save_style_symbol = '\U0001f4be'  # 💾

 apply_style_symbol = '\U0001f4cb'  # 📋

 clear_prompt_symbol = '\U0001F5D1'  # 🗑️

+extra_networks_symbol = '\U0001F3B4'  # 🎴

+switch_values_symbol = '\U000021C5' # ⇅

 

 

 def plaintext_to_html(text):

-    text = "<p>" + "<br>\n".join([f"{html.escape(x)}" for x in text.split('\n')]) + "</p>"

-    return text

+    return ui_common.plaintext_to_html(text)

+

 

 def send_gradio_gallery_to_image(x):

     if len(x) == 0:

         return None

     return image_from_url_text(x[0])

 

-def save_files(js_data, images, do_make_zip, index):

-    import csv

-    filenames = []

-    fullfns = []

-

-    #quick dictionary to class object conversion. Its necessary due apply_filename_pattern requiring it

-    class MyObject:

-        def __init__(self, d=None):

-            if d is not None:

-                for key, value in d.items():

-                    setattr(self, key, value)

-

-    data = json.loads(js_data)

-

-    p = MyObject(data)

-    path = opts.outdir_save

-    save_to_dirs = opts.use_save_to_dirs_for_ui

-    extension: str = opts.samples_format

-    start_index = 0

-

-    if index > -1 and opts.save_selected_only and (index >= data["index_of_first_image"]):  # ensures we are looking at a specific non-grid picture, and we have save_selected_only

-

-        images = [images[index]]

-        start_index = index

-

-    os.makedirs(opts.outdir_save, exist_ok=True)

-

-    with open(os.path.join(opts.outdir_save, "log.csv"), "a", encoding="utf8", newline='') as file:

-        at_start = file.tell() == 0

-        writer = csv.writer(file)

-        if at_start:

-            writer.writerow(["prompt", "seed", "width", "height", "sampler", "cfgs", "steps", "filename", "negative_prompt"])

-

-        for image_index, filedata in enumerate(images, start_index):

-            image = image_from_url_text(filedata)

-

-            is_grid = image_index < p.index_of_first_image

-            i = 0 if is_grid else (image_index - p.index_of_first_image)

-

-            fullfn, txt_fullfn = save_image(image, path, "", seed=p.all_seeds[i], prompt=p.all_prompts[i], extension=extension, info=p.infotexts[image_index], grid=is_grid, p=p, save_to_dirs=save_to_dirs)

-

-            filename = os.path.relpath(fullfn, path)

-            filenames.append(filename)

-            fullfns.append(fullfn)

-            if txt_fullfn:

-                filenames.append(os.path.basename(txt_fullfn))

-                fullfns.append(txt_fullfn)

-

-        writer.writerow([data["prompt"], data["seed"], data["width"], data["height"], data["sampler_name"], data["cfg_scale"], data["steps"], filenames[0], data["negative_prompt"]])

-

-    # Make Zip

-    if do_make_zip:

-        zip_filepath = os.path.join(path, "images.zip")

-

-        from zipfile import ZipFile

-        with ZipFile(zip_filepath, "w") as zip_file:

-            for i in range(len(fullfns)):

-                with open(fullfns[i], mode="rb") as f:

-                    zip_file.writestr(filenames[i], f.read())

-        fullfns.insert(0, zip_filepath)

-

-    return gr.File.update(value=fullfns, visible=True), '', '', plaintext_to_html(f"Saved: {filenames[0]}")

-

-

-

-

-def calc_time_left(progress, threshold, label, force_display):

-    if progress == 0:

-        return ""

-    else:

-        time_since_start = time.time() - shared.state.time_start

-        eta = (time_since_start/progress)

-        eta_relative = eta-time_since_start

-        if (eta_relative > threshold and progress > 0.02) or force_display:

-            if eta_relative > 3600:

-                return label + time.strftime('%H:%M:%S', time.gmtime(eta_relative))

-            elif eta_relative > 60:

-                return label + time.strftime('%M:%S',  time.gmtime(eta_relative))

-            else:

-                return label + time.strftime('%Ss',  time.gmtime(eta_relative))

-        else:

-            return ""

-

-

-def check_progress_call(id_part):

-    if shared.state.job_count == 0:

-        return "", gr_show(False), gr_show(False), gr_show(False)

-

-    progress = 0

-

-    if shared.state.job_count > 0:

-        progress += shared.state.job_no / shared.state.job_count

-    if shared.state.sampling_steps > 0:

-        progress += 1 / shared.state.job_count * shared.state.sampling_step / shared.state.sampling_steps

-

-    time_left = calc_time_left( progress, 1, " ETA: ", shared.state.time_left_force_display )

-    if time_left != "":

-        shared.state.time_left_force_display = True

-

-    progress = min(progress, 1)

-

-    progressbar = ""

-    if opts.show_progressbar:

-        progressbar = f"""<div class='progressDiv'><div class='progress' style="overflow:visible;width:{progress * 100}%;white-space:nowrap;">{"&nbsp;" * 2 + str(int(progress*100))+"%" + time_left if progress > 0.01 else ""}</div></div>"""

-

-    image = gr_show(False)

-    preview_visibility = gr_show(False)

-

-    if opts.show_progress_every_n_steps != 0:

-        shared.state.set_current_image()

-        image = shared.state.current_image

-

-        if image is None:

-            image = gr.update(value=None)

-        else:

-            preview_visibility = gr_show(True)

-

-    if shared.state.textinfo is not None:

-        textinfo_result = gr.HTML.update(value=shared.state.textinfo, visible=True)

-    else:

-        textinfo_result = gr_show(False)

-

-    return f"<span id='{id_part}_progress_span' style='display: none'>{time.time()}</span><p>{progressbar}</p>", preview_visibility, image, textinfo_result

-

-

-def check_progress_call_initial(id_part):

-    shared.state.job_count = -1

-    shared.state.current_latent = None

-    shared.state.current_image = None

-    shared.state.textinfo = None

-    shared.state.time_start = time.time()

-    shared.state.time_left_force_display = False

-

-    return check_progress_call(id_part)

-

-

-def roll_artist(prompt):

-    allowed_cats = set([x for x in shared.artist_db.categories() if len(opts.random_artist_categories)==0 or x in opts.random_artist_categories])

-    artist = random.choice([x for x in shared.artist_db.artists if x.category in allowed_cats])

-

-    return prompt + ", " + artist.name if prompt != '' else artist.name

-

-

 def visit(x, func, path=""):

     if hasattr(x, 'children'):

         for c in x.children:

@@ -259,56 +121,88 @@ def add_style(name: str, prompt: str, negative_prompt: str):
     # reserialize all styles every time we save them

     shared.prompt_styles.save_styles(shared.styles_filename)

 

-    return [gr.Dropdown.update(visible=True, choices=list(shared.prompt_styles.styles)) for _ in range(4)]

+    return [gr.Dropdown.update(visible=True, choices=list(shared.prompt_styles.styles)) for _ in range(2)]

 

 

-def apply_styles(prompt, prompt_neg, style1_name, style2_name):

-    prompt = shared.prompt_styles.apply_styles_to_prompt(prompt, [style1_name, style2_name])

-    prompt_neg = shared.prompt_styles.apply_negative_styles_to_prompt(prompt_neg, [style1_name, style2_name])

+def calc_resolution_hires(enable, width, height, hr_scale, hr_resize_x, hr_resize_y):

+    from modules import processing, devices

+

+    if not enable:

+        return ""

 

-    return [gr.Textbox.update(value=prompt), gr.Textbox.update(value=prompt_neg), gr.Dropdown.update(value="None"), gr.Dropdown.update(value="None")]

+    p = processing.StableDiffusionProcessingTxt2Img(width=width, height=height, enable_hr=True, hr_scale=hr_scale, hr_resize_x=hr_resize_x, hr_resize_y=hr_resize_y)

+

+    with devices.autocast():

+        p.init([""], [0], [0])

+

+    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 apply_styles(prompt, prompt_neg, styles):

+    prompt = shared.prompt_styles.apply_styles_to_prompt(prompt, styles)

+    prompt_neg = shared.prompt_styles.apply_negative_styles_to_prompt(prompt_neg, styles)

+

+    return [gr.Textbox.update(value=prompt), gr.Textbox.update(value=prompt_neg), gr.Dropdown.update(value=[])]

+

+

+def process_interrogate(interrogation_function, mode, ii_input_dir, ii_output_dir, *ii_singles):

+    if mode in {0, 1, 3, 4}:

+        return [interrogation_function(ii_singles[mode]), None]

+    elif mode == 2:

+        return [interrogation_function(ii_singles[mode]["image"]), None]

+    elif mode == 5:

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

+        images = shared.listfiles(ii_input_dir)

+        print(f"Will process {len(images)} images.")

+        if ii_output_dir != "":

+            os.makedirs(ii_output_dir, exist_ok=True)

+        else:

+            ii_output_dir = ii_input_dir

+

+        for image in images:

+            img = Image.open(image)

+            filename = os.path.basename(image)

+            left, _ = os.path.splitext(filename)

+            print(interrogation_function(img), file=open(os.path.join(ii_output_dir, left + ".txt"), 'a'))

+

+        return [gr.update(), None]

 

 

 def interrogate(image):

     prompt = shared.interrogator.interrogate(image.convert("RGB"))

-

-    return gr_show(True) if prompt is None else prompt

+    return gr.update() if prompt is None else prompt

 

 

 def interrogate_deepbooru(image):

     prompt = deepbooru.model.tag(image)

-    return gr_show(True) if prompt is None else prompt

+    return gr.update() if prompt is None else prompt

 

 

-def create_seed_inputs():

-    with gr.Row():

-        with gr.Box():

-            with gr.Row(elem_id='seed_row'):

-                seed = (gr.Textbox if cmd_opts.use_textbox_seed else gr.Number)(label='Seed', value=-1)

-                seed.style(container=False)

-                random_seed = gr.Button(random_symbol, elem_id='random_seed')

-                reuse_seed = gr.Button(reuse_symbol, elem_id='reuse_seed')

+def create_seed_inputs(target_interface):

+    with FormRow(elem_id=target_interface + '_seed_row'):

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

 

-        with gr.Box(elem_id='subseed_show_box'):

-            seed_checkbox = gr.Checkbox(label='Extra', elem_id='subseed_show', value=False)

+        with gr.Group(elem_id=target_interface + '_subseed_show_box'):

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

 

-    with gr.Row(visible=False) as seed_extra_row_1:

+    with FormRow(visible=False, elem_id=target_interface + '_subseed_row') as seed_extra_row_1:

         seed_extras.append(seed_extra_row_1)

-        with gr.Box():

-            with gr.Row(elem_id='subseed_row'):

-                subseed = gr.Number(label='Variation seed', value=-1)

-                subseed.style(container=False)

-                random_subseed = gr.Button(random_symbol, elem_id='random_subseed')

-                reuse_subseed = gr.Button(reuse_symbol, elem_id='reuse_subseed')

-        subseed_strength = gr.Slider(label='Variation strength', value=0.0, minimum=0, maximum=1, step=0.01)

-

-    with gr.Row(visible=False) as seed_extra_row_2:

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

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

         seed_extras.append(seed_extra_row_2)

-        seed_resize_from_w = gr.Slider(minimum=0, maximum=2048, step=8, label="Resize seed from width", value=0)

-        seed_resize_from_h = gr.Slider(minimum=0, maximum=2048, step=8, label="Resize seed from height", value=0)

+        seed_resize_from_w = gr.Slider(minimum=0, maximum=2048, step=8, label="Resize seed from width", value=0, elem_id=target_interface + '_seed_resize_from_w')

+        seed_resize_from_h = gr.Slider(minimum=0, maximum=2048, step=8, label="Resize seed from height", value=0, elem_id=target_interface + '_seed_resize_from_h')

 

     random_seed.click(fn=lambda: -1, show_progress=False, inputs=[], outputs=[seed])

     random_subseed.click(fn=lambda: -1, show_progress=False, inputs=[], outputs=[subseed])

@@ -368,6 +262,8 @@ def connect_reuse_seed(seed: gr.Number, reuse_seed: gr.Button, generation_info:
 

 def update_token_counter(text, steps):

     try:

+        text, _ = extra_networks.parse_prompt(text)

+

         _, prompt_flat_list, _ = prompt_parser.get_multicond_prompt_list([text])

         prompt_schedules = prompt_parser.get_learned_conditioning_prompt_schedules(prompt_flat_list, steps)

 

@@ -378,45 +274,24 @@ def update_token_counter(text, steps):
 

     flat_prompts = reduce(lambda list1, list2: list1+list2, prompt_schedules)

     prompts = [prompt_text for step, prompt_text in flat_prompts]

-    tokens, token_count, max_length = max([model_hijack.tokenize(prompt) for prompt in prompts], key=lambda args: args[1])

-    style_class = ' class="red"' if (token_count > max_length) else ""

-    return f"<span {style_class}>{token_count}/{max_length}</span>"

+    token_count, max_length = max([model_hijack.get_prompt_lengths(prompt) for prompt in prompts], key=lambda args: args[0])

+    return f"<span class='gr-box gr-text-input'>{token_count}/{max_length}</span>"

 

 

 def create_toprow(is_img2img):

     id_part = "img2img" if is_img2img else "txt2img"

 

-    with gr.Row(elem_id="toprow"):

-        with gr.Column(scale=6):

+    with gr.Row(elem_id=f"{id_part}_toprow", variant="compact"):

+        with gr.Column(elem_id=f"{id_part}_prompt_container", scale=6):

             with gr.Row():

                 with gr.Column(scale=80):

                     with gr.Row():

-                        prompt = gr.Textbox(label="Prompt", elem_id=f"{id_part}_prompt", show_label=False, lines=2,

-                            placeholder="Prompt (press Ctrl+Enter or Alt+Enter to generate)"

-                        )

+                        prompt = gr.Textbox(label="Prompt", elem_id=f"{id_part}_prompt", show_label=False, lines=3, placeholder="Prompt (press Ctrl+Enter or Alt+Enter to generate)")

 

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

-                        )

-

-        with gr.Column(scale=1, elem_id="roll_col"):

-            roll = gr.Button(value=art_symbol, elem_id="roll", visible=len(shared.artist_db.artists) > 0)

-            paste = gr.Button(value=paste_symbol, elem_id="paste")

-            save_style = gr.Button(value=save_style_symbol, elem_id="style_create")

-            prompt_style_apply = gr.Button(value=apply_style_symbol, elem_id="style_apply")

-            clear_prompt_button = gr.Button(value=clear_prompt_symbol, elem_id=f"{id_part}_clear_prompt")

-            token_counter = gr.HTML(value="<span></span>", elem_id=f"{id_part}_token_counter")

-            token_button = gr.Button(visible=False, elem_id=f"{id_part}_token_button")

-

-            clear_prompt_button.click(

-                fn=lambda *x: x,

-                _js="confirm_clear_prompt",

-                inputs=[prompt, negative_prompt],

-                outputs=[prompt, negative_prompt],

-            )

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

 

         button_interrogate = None

         button_deepbooru = None

@@ -425,10 +300,10 @@ def create_toprow(is_img2img):
                 button_interrogate = gr.Button('Interrogate\nCLIP', elem_id="interrogate")

                 button_deepbooru = gr.Button('Interrogate\nDeepBooru', elem_id="deepbooru")

 

-        with gr.Column(scale=1):

-            with gr.Row():

-                skip = gr.Button('Skip', elem_id=f"{id_part}_skip")

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

                 submit = gr.Button('Generate', elem_id=f"{id_part}_generate", variant='primary')

 

                 skip.click(

@@ -443,37 +318,34 @@ def create_toprow(is_img2img):
                     outputs=[],

                 )

 

-            with gr.Row():

-                with gr.Column(scale=1, elem_id="style_pos_col"):

-                    prompt_style = gr.Dropdown(label="Style 1", elem_id=f"{id_part}_style_index", choices=[k for k, v in shared.prompt_styles.styles.items()], value=next(iter(shared.prompt_styles.styles.keys())))

-                    prompt_style.save_to_config = True

+            with gr.Row(elem_id=f"{id_part}_tools"):

+                paste = ToolButton(value=paste_symbol, elem_id="paste")

+                clear_prompt_button = ToolButton(value=clear_prompt_symbol, elem_id=f"{id_part}_clear_prompt")

+                extra_networks_button = ToolButton(value=extra_networks_symbol, elem_id=f"{id_part}_extra_networks")

+                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_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_button = gr.Button(visible=False, elem_id=f"{id_part}_negative_token_button")

+

+                clear_prompt_button.click(

+                    fn=lambda *x: x,

+                    _js="confirm_clear_prompt",

+                    inputs=[prompt, negative_prompt],

+                    outputs=[prompt, negative_prompt],

+                )

 

-                with gr.Column(scale=1, elem_id="style_neg_col"):

-                    prompt_style2 = gr.Dropdown(label="Style 2", elem_id=f"{id_part}_style2_index", choices=[k for k, v in shared.prompt_styles.styles.items()], value=next(iter(shared.prompt_styles.styles.keys())))

-                    prompt_style2.save_to_config = True

+            with gr.Row(elem_id=f"{id_part}_styles_row"):

+                prompt_styles = gr.Dropdown(label="Styles", elem_id=f"{id_part}_styles", choices=[k for k, v in shared.prompt_styles.styles.items()], value=[], multiselect=True)

+                create_refresh_button(prompt_styles, shared.prompt_styles.reload, lambda: {"choices": [k for k, v in shared.prompt_styles.styles.items()]}, f"refresh_{id_part}_styles")

 

-    return prompt, roll, prompt_style, negative_prompt, prompt_style2, submit, button_interrogate, button_deepbooru, prompt_style_apply, save_style, paste, token_counter, token_button

+    return prompt, prompt_styles, negative_prompt, submit, button_interrogate, button_deepbooru, prompt_style_apply, save_style, paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button

 

 

-def setup_progressbar(progressbar, preview, id_part, textinfo=None):

-    if textinfo is None:

-        textinfo = gr.HTML(visible=False)

-

-    check_progress = gr.Button('Check progress', elem_id=f"{id_part}_check_progress", visible=False)

-    check_progress.click(

-        fn=lambda: check_progress_call(id_part),

-        show_progress=False,

-        inputs=[],

-        outputs=[progressbar, preview, preview, textinfo],

-    )

-

-    check_progress_initial = gr.Button('Check progress (first)', elem_id=f"{id_part}_check_progress_initial", visible=False)

-    check_progress_initial.click(

-        fn=lambda: check_progress_call_initial(id_part),

-        show_progress=False,

-        inputs=[],

-        outputs=[progressbar, preview, preview, textinfo],

-    )

+def setup_progressbar(*args, **kwargs):

+    pass

 

 

 def apply_setting(key, value):

@@ -500,27 +372,13 @@ def apply_setting(key, value):
         return

 

     valtype = type(opts.data_labels[key].default)

-    oldval = opts.data[key]

+    oldval = opts.data.get(key, None)

     opts.data[key] = valtype(value) if valtype != type(None) else value

     if oldval != value and opts.data_labels[key].onchange is not None:

         opts.data_labels[key].onchange()

 

     opts.save(shared.config_filename)

-    return value

-

-

-def update_generation_info(args):

-    generation_info, html_info, img_index = args

-    try:

-        generation_info = json.loads(generation_info)

-        if img_index < 0 or img_index >= len(generation_info["infotexts"]):

-            return html_info

-        return plaintext_to_html(generation_info["infotexts"][img_index])

-    except Exception:

-        pass

-    # if the json parse or anything else fails, just return the old html_info

-    return html_info

-

+    return getattr(opts, key)

 

 def create_refresh_button(refresh_component, refresh_method, refreshed_args, elem_id):

     def refresh():

@@ -532,7 +390,7 @@ def create_refresh_button(refresh_component, refresh_method, refreshed_args, ele
 

         return gr.update(**(args or {}))

 

-    refresh_button = gr.Button(value=refresh_symbol, elem_id=elem_id)

+    refresh_button = ToolButton(value=refresh_symbol, elem_id=elem_id)

     refresh_button.click(

         fn=refresh,

         inputs=[],

@@ -542,89 +400,37 @@ def create_refresh_button(refresh_component, refresh_method, refreshed_args, ele
 

 

 def create_output_panel(tabname, outdir):

-    def open_folder(f):

-        if not os.path.exists(f):

-            print(f'Folder "{f}" does not exist. After you create an image, the folder will be created.')

-            return

-        elif not os.path.isdir(f):

-            print(f"""

-WARNING

-An open_folder request was made with an argument that is not a folder.

-This could be an error or a malicious attempt to run code on your computer.

-Requested path was: {f}

-""", file=sys.stderr)

-            return

-

-        if not shared.cmd_opts.hide_ui_dir_config:

-            path = os.path.normpath(f)

-            if platform.system() == "Windows":

-                os.startfile(path)

-            elif platform.system() == "Darwin":

-                sp.Popen(["open", path])

-            else:

-                sp.Popen(["xdg-open", path])

+    return ui_common.create_output_panel(tabname, outdir)

 

-    with gr.Column(variant='panel'):

-            with gr.Group():

-                result_gallery = gr.Gallery(label='Output', show_label=False, elem_id=f"{tabname}_gallery").style(grid=4)

 

-            generation_info = None

-            with gr.Column():

-                with gr.Row():

-                    if tabname != "extras":

-                        save = gr.Button('Save', elem_id=f'save_{tabname}')

+def create_sampler_and_steps_selection(choices, tabname):

+    if opts.samplers_in_dropdown:

+        with FormRow(elem_id=f"sampler_selection_{tabname}"):

+            sampler_index = gr.Dropdown(label='Sampling method', elem_id=f"{tabname}_sampling", choices=[x.name for x in choices], value=choices[0].name, type="index")

+            steps = gr.Slider(minimum=1, maximum=150, step=1, elem_id=f"{tabname}_steps", label="Sampling steps", value=20)

+    else:

+        with FormGroup(elem_id=f"sampler_selection_{tabname}"):

+            steps = gr.Slider(minimum=1, maximum=150, step=1, elem_id=f"{tabname}_steps", label="Sampling steps", value=20)

+            sampler_index = gr.Radio(label='Sampling method', elem_id=f"{tabname}_sampling", choices=[x.name for x in choices], value=choices[0].name, type="index")

 

-                    buttons = parameters_copypaste.create_buttons(["img2img", "inpaint", "extras"])

-                    button_id = "hidden_element" if shared.cmd_opts.hide_ui_dir_config else 'open_folder'

-                    open_folder_button = gr.Button(folder_symbol, elem_id=button_id)

+    return steps, sampler_index

 

-                open_folder_button.click(

-                    fn=lambda: open_folder(opts.outdir_samples or outdir),

-                    inputs=[],

-                    outputs=[],

-                )

 

-                if tabname != "extras":

-                    with gr.Row():

-                        do_make_zip = gr.Checkbox(label="Make Zip when Save?", value=False)

+def ordered_ui_categories():

+    user_order = {x.strip(): i * 2 + 1 for i, x in enumerate(shared.opts.ui_reorder.split(","))}

 

-                    with gr.Row():

-                        download_files = gr.File(None, file_count="multiple", interactive=False, show_label=False, visible=False)

-

-                    with gr.Group():

-                        html_info = gr.HTML()

-                        generation_info = gr.Textbox(visible=False)

-                        if tabname == 'txt2img' or tabname == 'img2img':

-                            generation_info_button = gr.Button(visible=False, elem_id=f"{tabname}_generation_info_button")

-                            generation_info_button.click(

-                                fn=update_generation_info,

-                                _js="(x, y) => [x, y, selected_gallery_index()]",

-                                inputs=[generation_info, html_info],

-                                outputs=[html_info],

-                                preprocess=False

-                            )

-

-                        save.click(

-                            fn=wrap_gradio_call(save_files),

-                            _js="(x, y, z, w) => [x, y, z, selected_gallery_index()]",

-                            inputs=[

-                                generation_info,

-                                result_gallery,

-                                do_make_zip,

-                                html_info,

-                            ],

-                            outputs=[

-                                download_files,

-                                html_info,

-                                html_info,

-                                html_info,

-                            ]

-                        )

-                else:

-                    html_info_x = gr.HTML()

-                    html_info = gr.HTML()

-                parameters_copypaste.bind_buttons(buttons, result_gallery, "txt2img" if tabname == "txt2img" else None)

-                return result_gallery, generation_info if tabname != "extras" else html_info_x, html_info

+    for i, category in sorted(enumerate(shared.ui_reorder_categories), key=lambda x: user_order.get(x[1], x[0] * 2 + 0)):

+        yield category

+

+

+def get_value_for_setting(key):

+    value = getattr(opts, key)

+

+    info = opts.data_labels[key]

+    args = info.component_args() if callable(info.component_args) else info.component_args or {}

+    args = {k: v for k, v in args.items() if k not in {'precision'}}

+

+    return gr.update(value=value, **args)

 

 

 def create_ui():

@@ -639,67 +445,98 @@ def create_ui():
     modules.scripts.scripts_txt2img.initialize_scripts(is_img2img=False)

 

     with gr.Blocks(analytics_enabled=False) as txt2img_interface:

-        txt2img_prompt, roll, txt2img_prompt_style, txt2img_negative_prompt, txt2img_prompt_style2, submit, _, _,txt2img_prompt_style_apply, txt2img_save_style, txt2img_paste, token_counter, token_button = create_toprow(is_img2img=False)

+        txt2img_prompt, txt2img_prompt_styles, txt2img_negative_prompt, submit, _, _, txt2img_prompt_style_apply, txt2img_save_style, txt2img_paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button = create_toprow(is_img2img=False)

 

         dummy_component = gr.Label(visible=False)

-        txt_prompt_img = gr.File(label="", elem_id="txt2img_prompt_image", file_count="single", type="bytes", visible=False)

-

-

-

+        txt_prompt_img = gr.File(label="", elem_id="txt2img_prompt_image", file_count="single", type="binary", visible=False)

 

-        with gr.Row(elem_id='txt2img_progress_row'):

-            with gr.Column(scale=1):

-                pass

-

-            with gr.Column(scale=1):

-                progressbar = gr.HTML(elem_id="txt2img_progressbar")

-                txt2img_preview = gr.Image(elem_id='txt2img_preview', visible=False)

-                setup_progressbar(progressbar, txt2img_preview, 'txt2img')

+        with FormRow(variant='compact', elem_id="txt2img_extra_networks", visible=False) as extra_networks:

+            from modules import ui_extra_networks

+            extra_networks_ui = ui_extra_networks.create_ui(extra_networks, extra_networks_button, 'txt2img')

 

         with gr.Row().style(equal_height=False):

-            with gr.Column(variant='panel', elem_id="txt2img_settings"):

-                steps = gr.Slider(minimum=1, maximum=150, step=1, label="Sampling Steps", value=20)

-                sampler_index = gr.Radio(label='Sampling method', elem_id="txt2img_sampling", choices=[x.name for x in samplers], value=samplers[0].name, type="index")

-

-                with gr.Group():

-                    width = gr.Slider(minimum=64, maximum=2048, step=8, label="Width", value=512)

-                    height = gr.Slider(minimum=64, maximum=2048, step=8, label="Height", value=512)

-

-                with gr.Row():

-                    restore_faces = gr.Checkbox(label='Restore faces', value=False, visible=len(shared.face_restorers) > 1)

-                    tiling = gr.Checkbox(label='Tiling', value=False)

-                    enable_hr = gr.Checkbox(label='Highres. fix', value=False)

-

-                with gr.Row(visible=False) as hr_options:

-                    firstphase_width = gr.Slider(minimum=0, maximum=1024, step=8, label="Firstpass width", value=0)

-                    firstphase_height = gr.Slider(minimum=0, maximum=1024, step=8, label="Firstpass height", value=0)

-                    denoising_strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Denoising strength', value=0.7)

-

-                with gr.Row(equal_height=True):

-                    batch_count = gr.Slider(minimum=1, step=1, label='Batch count', value=1)

-                    batch_size = gr.Slider(minimum=1, maximum=8, step=1, label='Batch size', value=1)

-

-                cfg_scale = gr.Slider(minimum=1.0, maximum=30.0, step=0.5, label='CFG Scale', value=7.0)

-

-                seed, reuse_seed, subseed, reuse_subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w, seed_checkbox = create_seed_inputs()

-

-                with gr.Group():

-                    custom_inputs = modules.scripts.scripts_txt2img.setup_ui()

+            with gr.Column(variant='compact', elem_id="txt2img_settings"):

+                for category in ordered_ui_categories():

+                    if category == "sampler":

+                        steps, sampler_index = create_sampler_and_steps_selection(samplers, "txt2img")

+

+                    elif category == "dimensions":

+                        with FormRow():

+                            with gr.Column(elem_id="txt2img_column_size", scale=4):

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

+

+                            if opts.dimensions_and_batch_together:

+                                res_switch_btn = ToolButton(value=switch_values_symbol, elem_id="txt2img_res_switch_btn")

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

+                                    batch_size = gr.Slider(minimum=1, maximum=8, step=1, label='Batch size', value=1, elem_id="txt2img_batch_size")

+

+                    elif category == "cfg":

+                        cfg_scale = gr.Slider(minimum=1.0, maximum=30.0, step=0.5, label='CFG Scale', value=7.0, elem_id="txt2img_cfg_scale")

+

+                    elif category == "seed":

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

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

+                            hr_final_resolution = FormHTML(value="", elem_id="txtimg_hr_finalres", label="Upscaled resolution", interactive=False)

+

+                    elif category == "hires_fix":

+                        with FormGroup(visible=False, elem_id="txt2img_hires_fix") as hr_options:

+                            with FormRow(elem_id="txt2img_hires_fix_row1", variant="compact"):

+                                hr_upscaler = gr.Dropdown(label="Upscaler", elem_id="txt2img_hr_upscaler", choices=[*shared.latent_upscale_modes, *[x.name for x in shared.sd_upscalers]], value=shared.latent_upscale_default_mode)

+                                hr_second_pass_steps = gr.Slider(minimum=0, maximum=150, step=1, label='Hires steps', value=0, elem_id="txt2img_hires_steps")

+                                denoising_strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Denoising strength', value=0.7, elem_id="txt2img_denoising_strength")

+

+                            with FormRow(elem_id="txt2img_hires_fix_row2", variant="compact"):

+                                hr_scale = gr.Slider(minimum=1.0, maximum=4.0, step=0.05, label="Upscale by", value=2.0, elem_id="txt2img_hr_scale")

+                                hr_resize_x = gr.Slider(minimum=0, maximum=2048, step=8, label="Resize width to", value=0, elem_id="txt2img_hr_resize_x")

+                                hr_resize_y = gr.Slider(minimum=0, maximum=2048, step=8, label="Resize height to", value=0, elem_id="txt2img_hr_resize_y")

+

+                    elif category == "batch":

+                        if not opts.dimensions_and_batch_together:

+                            with FormRow(elem_id="txt2img_column_batch"):

+                                batch_count = gr.Slider(minimum=1, step=1, label='Batch count', value=1, elem_id="txt2img_batch_count")

+                                batch_size = gr.Slider(minimum=1, maximum=8, step=1, label='Batch size', value=1, elem_id="txt2img_batch_size")

+

+                    elif category == "scripts":

+                        with FormGroup(elem_id="txt2img_script_container"):

+                            custom_inputs = modules.scripts.scripts_txt2img.setup_ui()

+

+            hr_resolution_preview_inputs = [enable_hr, width, height, hr_scale, hr_resize_x, hr_resize_y]

+            for input in hr_resolution_preview_inputs:

+                input.change(

+                    fn=calc_resolution_hires,

+                    inputs=hr_resolution_preview_inputs,

+                    outputs=[hr_final_resolution],

+                    show_progress=False,

+                )

+                input.change(

+                    None,

+                    _js="onCalcResolutionHires",

+                    inputs=hr_resolution_preview_inputs,

+                    outputs=[],

+                    show_progress=False,

+                )

 

-            txt2img_gallery, generation_info, html_info = create_output_panel("txt2img", opts.outdir_txt2img_samples)

+            txt2img_gallery, generation_info, html_info, html_log = create_output_panel("txt2img", opts.outdir_txt2img_samples)

             parameters_copypaste.bind_buttons({"txt2img": txt2img_paste}, None, txt2img_prompt)

 

             connect_reuse_seed(seed, reuse_seed, generation_info, dummy_component, is_subseed=False)

             connect_reuse_seed(subseed, reuse_subseed, generation_info, dummy_component, is_subseed=True)

 

             txt2img_args = dict(

-                fn=wrap_gradio_gpu_call(modules.txt2img.txt2img),

+                fn=wrap_gradio_gpu_call(modules.txt2img.txt2img, extra_outputs=[None, '', '']),

                 _js="submit",

                 inputs=[

+                    dummy_component,

                     txt2img_prompt,

                     txt2img_negative_prompt,

-                    txt2img_prompt_style,

-                    txt2img_prompt_style2,

+                    txt2img_prompt_styles,

                     steps,

                     sampler_index,

                     restore_faces,

@@ -713,14 +550,18 @@ def create_ui():
                     width,

                     enable_hr,

                     denoising_strength,

-                    firstphase_width,

-                    firstphase_height,

+                    hr_scale,

+                    hr_upscaler,

+                    hr_second_pass_steps,

+                    hr_resize_x,

+                    hr_resize_y,

                 ] + custom_inputs,

 

                 outputs=[

                     txt2img_gallery,

                     generation_info,

-                    html_info

+                    html_info,

+                    html_log,

                 ],

                 show_progress=False,

             )

@@ -728,6 +569,8 @@ 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])

+

             txt_prompt_img.change(

                 fn=modules.images.image_data,

                 inputs=[

@@ -743,17 +586,7 @@ def create_ui():
                 fn=lambda x: gr_show(x),

                 inputs=[enable_hr],

                 outputs=[hr_options],

-            )

-

-            roll.click(

-                fn=roll_artist,

-                _js="update_txt2img_tokens",

-                inputs=[

-                    txt2img_prompt,

-                ],

-                outputs=[

-                    txt2img_prompt,

-                ]

+                show_progress = False,

             )

 

             txt2img_paste_fields = [

@@ -774,8 +607,11 @@ def create_ui():
                 (denoising_strength, "Denoising strength"),

                 (enable_hr, lambda d: "Denoising strength" in d),

                 (hr_options, lambda d: gr.Row.update(visible="Denoising strength" in d)),

-                (firstphase_width, "First pass size-1"),

-                (firstphase_height, "First pass size-2"),

+                (hr_scale, "Hires upscale"),

+                (hr_upscaler, "Hires upscaler"),

+                (hr_second_pass_steps, "Hires steps"),

+                (hr_resize_x, "Hires resize-1"),

+                (hr_resize_y, "Hires resize-2"),

                 *modules.scripts.scripts_txt2img.infotext_fields

             ]

             parameters_copypaste.add_paste_fields("txt2img", None, txt2img_paste_fields)

@@ -792,98 +628,174 @@ def create_ui():
             ]

 

             token_button.click(fn=wrap_queued_call(update_token_counter), inputs=[txt2img_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])

+

+            ui_extra_networks.setup_ui(extra_networks_ui, txt2img_gallery)

 

     modules.scripts.scripts_current = modules.scripts.scripts_img2img

     modules.scripts.scripts_img2img.initialize_scripts(is_img2img=True)

 

     with gr.Blocks(analytics_enabled=False) as img2img_interface:

-        img2img_prompt, roll, img2img_prompt_style, img2img_negative_prompt, img2img_prompt_style2, submit, img2img_interrogate, img2img_deepbooru, img2img_prompt_style_apply, img2img_save_style, img2img_paste,token_counter, token_button = create_toprow(is_img2img=True)

+        img2img_prompt, img2img_prompt_styles, img2img_negative_prompt, submit, img2img_interrogate, img2img_deepbooru, img2img_prompt_style_apply, img2img_save_style, img2img_paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button = create_toprow(is_img2img=True)

 

+        img2img_prompt_img = gr.File(label="", elem_id="img2img_prompt_image", file_count="single", type="binary", visible=False)

 

-        with gr.Row(elem_id='img2img_progress_row'):

-            img2img_prompt_img = gr.File(label="", elem_id="img2img_prompt_image", file_count="single", type="bytes", visible=False)

+        with FormRow(variant='compact', elem_id="img2img_extra_networks", visible=False) as extra_networks:

+            from modules import ui_extra_networks

+            extra_networks_ui_img2img = ui_extra_networks.create_ui(extra_networks, extra_networks_button, 'img2img')

 

-            with gr.Column(scale=1):

-                pass

-

-            with gr.Column(scale=1):

-                progressbar = gr.HTML(elem_id="img2img_progressbar")

-                img2img_preview = gr.Image(elem_id='img2img_preview', visible=False)

-                setup_progressbar(progressbar, img2img_preview, 'img2img')

+        with FormRow().style(equal_height=False):

+            with gr.Column(variant='compact', elem_id="img2img_settings"):

+                copy_image_buttons = []

+                copy_image_destinations = {}

 

-        with gr.Row().style(equal_height=False):

-            with gr.Column(variant='panel', elem_id="img2img_settings"):

+                def add_copy_image_controls(tab_name, elem):

+                    with gr.Row(variant="compact", elem_id=f"img2img_copy_to_{tab_name}"):

+                        gr.HTML("Copy image to: ", elem_id=f"img2img_label_copy_to_{tab_name}")

 

-                with gr.Tabs(elem_id="mode_img2img") as tabs_img2img_mode:

-                    with gr.TabItem('img2img', id='img2img'):

-                        init_img = gr.Image(label="Image for img2img", elem_id="img2img_image", show_label=False, source="upload", interactive=True, type="pil", tool=cmd_opts.gradio_img2img_tool, image_mode="RGBA").style(height=480)

+                        for title, name in zip(['img2img', 'sketch', 'inpaint', 'inpaint sketch'], ['img2img', 'sketch', 'inpaint', 'inpaint_sketch']):

+                            if name == tab_name:

+                                gr.Button(title, interactive=False)

+                                copy_image_destinations[name] = elem

+                                continue

 

-                    with gr.TabItem('Inpaint', id='inpaint'):

-                        init_img_with_mask = gr.Image(label="Image for inpainting with mask", show_label=False, elem_id="img2maskimg", source="upload", interactive=True, type="pil", tool=cmd_opts.gradio_inpaint_tool, image_mode="RGBA").style(height=480)

-                        init_img_with_mask_orig = gr.State(None)

+                            button = gr.Button(title)

+                            copy_image_buttons.append((button, name, elem))

 

-                        use_color_sketch = cmd_opts.gradio_inpaint_tool == "color-sketch"

-                        if use_color_sketch:

-                            def update_orig(image, state):

-                                if image is not None:

-                                    same_size = state is not None and state.size == image.size

-                                    has_exact_match = np.any(np.all(np.array(image) == np.array(state), axis=-1))

-                                    edited = same_size and has_exact_match

-                                    return image if not edited or state is None else state

+                with gr.Tabs(elem_id="mode_img2img"):

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

 

-                            init_img_with_mask.change(update_orig, [init_img_with_mask, init_img_with_mask_orig], init_img_with_mask_orig)

+                    with gr.TabItem('Sketch', id='img2img_sketch', elem_id="img2img_img2img_sketch_tab") as tab_sketch:

+                        sketch = gr.Image(label="Image for img2img", elem_id="img2img_sketch", show_label=False, source="upload", interactive=True, type="pil", tool="color-sketch", image_mode="RGBA").style(height=480)

+                        add_copy_image_controls('sketch', sketch)

 

-                        init_img_inpaint = gr.Image(label="Image for img2img", show_label=False, source="upload", interactive=True, type="pil", visible=False, elem_id="img_inpaint_base")

-                        init_mask_inpaint = gr.Image(label="Mask", source="upload", interactive=True, type="pil", visible=False, elem_id="img_inpaint_mask")

+                    with gr.TabItem('Inpaint', id='inpaint', elem_id="img2img_inpaint_tab") as tab_inpaint:

+                        init_img_with_mask = gr.Image(label="Image for inpainting with mask", show_label=False, elem_id="img2maskimg", source="upload", interactive=True, type="pil", tool="sketch", image_mode="RGBA").style(height=480)

+                        add_copy_image_controls('inpaint', init_img_with_mask)

 

-                        with gr.Row():

-                            mask_blur = gr.Slider(label='Mask blur', minimum=0, maximum=64, step=1, value=4)

-                            mask_alpha = gr.Slider(label="Mask transparency", interactive=use_color_sketch, visible=use_color_sketch)

+                    with gr.TabItem('Inpaint sketch', id='inpaint_sketch', elem_id="img2img_inpaint_sketch_tab") as tab_inpaint_color:

+                        inpaint_color_sketch = gr.Image(label="Color sketch inpainting", show_label=False, elem_id="inpaint_sketch", source="upload", interactive=True, type="pil", tool="color-sketch", image_mode="RGBA").style(height=480)

+                        inpaint_color_sketch_orig = gr.State(None)

+                        add_copy_image_controls('inpaint_sketch', inpaint_color_sketch)

 

-                        with gr.Row():

-                            mask_mode = gr.Radio(label="Mask mode", show_label=False, choices=["Draw mask", "Upload mask"], type="index", value="Draw mask", elem_id="mask_mode")

-                            inpainting_mask_invert = gr.Radio(label='Masking mode', show_label=False, choices=['Inpaint masked', 'Inpaint not masked'], value='Inpaint masked', type="index")

+                        def update_orig(image, state):

+                            if image is not None:

+                                same_size = state is not None and state.size == image.size

+                                has_exact_match = np.any(np.all(np.array(image) == np.array(state), axis=-1))

+                                edited = same_size and has_exact_match

+                                return image if not edited or state is None else state

 

-                        inpainting_fill = gr.Radio(label='Masked content', choices=['fill', 'original', 'latent noise', 'latent nothing'], value='original', type="index")

+                        inpaint_color_sketch.change(update_orig, [inpaint_color_sketch, inpaint_color_sketch_orig], inpaint_color_sketch_orig)

 

-                        with gr.Row():

-                            inpaint_full_res = gr.Checkbox(label='Inpaint at full resolution', value=False)

-                            inpaint_full_res_padding = gr.Slider(label='Inpaint at full resolution padding, pixels', minimum=0, maximum=256, step=4, value=32)

+                    with gr.TabItem('Inpaint upload', id='inpaint_upload', elem_id="img2img_inpaint_upload_tab") as tab_inpaint_upload:

+                        init_img_inpaint = gr.Image(label="Image for img2img", show_label=False, source="upload", interactive=True, type="pil", elem_id="img_inpaint_base")

+                        init_mask_inpaint = gr.Image(label="Mask", source="upload", interactive=True, type="pil", elem_id="img_inpaint_mask")

 

-                    with gr.TabItem('Batch img2img', id='batch'):

+                    with gr.TabItem('Batch', id='batch', elem_id="img2img_batch_tab") as tab_batch:

                         hidden = '<br>Disabled when launched with --hide-ui-dir-config.' if shared.cmd_opts.hide_ui_dir_config else ''

-                        gr.HTML(f"<p class=\"text-gray-500\">Process images in a directory on the same machine where the server is running.<br>Use an empty output directory to save pictures normally instead of writing to the output directory.{hidden}</p>")

-                        img2img_batch_input_dir = gr.Textbox(label="Input directory", **shared.hide_dirs)

-                        img2img_batch_output_dir = gr.Textbox(label="Output directory", **shared.hide_dirs)

-

-                with gr.Row():

-                    resize_mode = gr.Radio(label="Resize mode", elem_id="resize_mode", show_label=False, choices=["Just resize", "Crop and resize", "Resize and fill", "Just resize (latent upscale)"], type="index", value="Just resize")

-

-                steps = gr.Slider(minimum=1, maximum=150, step=1, label="Sampling Steps", value=20)

-                sampler_index = gr.Radio(label='Sampling method', choices=[x.name for x in samplers_for_img2img], value=samplers_for_img2img[0].name, type="index")

-

-                with gr.Group():

-                    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.Row():

-                    restore_faces = gr.Checkbox(label='Restore faces', value=False, visible=len(shared.face_restorers) > 1)

-                    tiling = gr.Checkbox(label='Tiling', value=False)

-

-                with gr.Row():

-                    batch_count = gr.Slider(minimum=1, step=1, label='Batch count', value=1)

-                    batch_size = gr.Slider(minimum=1, maximum=8, step=1, label='Batch size', value=1)

+                        gr.HTML(

+                            f"<p style='padding-bottom: 1em;' class=\"text-gray-500\">Process images in a directory on the same machine where the server is running." +

+                            f"<br>Use an empty output directory to save pictures normally instead of writing to the output directory." +

+                            f"<br>Add inpaint batch mask directory to enable inpaint batch processing."

+                            f"{hidden}</p>"

+                        )

+                        img2img_batch_input_dir = gr.Textbox(label="Input directory", **shared.hide_dirs, elem_id="img2img_batch_input_dir")

+                        img2img_batch_output_dir = gr.Textbox(label="Output directory", **shared.hide_dirs, elem_id="img2img_batch_output_dir")

+                        img2img_batch_inpaint_mask_dir = gr.Textbox(label="Inpaint batch mask directory (required for inpaint batch processing only)", **shared.hide_dirs, elem_id="img2img_batch_inpaint_mask_dir")

 

-                with gr.Group():

-                    cfg_scale = gr.Slider(minimum=1.0, maximum=30.0, step=0.5, label='CFG Scale', value=7.0)

-                    denoising_strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Denoising strength', value=0.75)

+                def copy_image(img):

+                    if isinstance(img, dict) and 'image' in img:

+                        return img['image']

 

-                seed, reuse_seed, subseed, reuse_subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w, seed_checkbox = create_seed_inputs()

+                    return img

 

-                with gr.Group():

-                    custom_inputs = modules.scripts.scripts_img2img.setup_ui()

+                for button, name, elem in copy_image_buttons:

+                    button.click(

+                        fn=copy_image,

+                        inputs=[elem],

+                        outputs=[copy_image_destinations[name]],

+                    )

+                    button.click(

+                        fn=lambda: None,

+                        _js="switch_to_"+name.replace(" ", "_"),

+                        inputs=[],

+                        outputs=[],

+                    )

 

-            img2img_gallery, generation_info, html_info = create_output_panel("img2img", opts.outdir_img2img_samples)

+                with FormRow():

+                    resize_mode = gr.Radio(label="Resize mode", elem_id="resize_mode", choices=["Just resize", "Crop and resize", "Resize and fill", "Just resize (latent upscale)"], type="index", value="Just resize")

+

+                for category in ordered_ui_categories():

+                    if category == "sampler":

+                        steps, sampler_index = create_sampler_and_steps_selection(samplers_for_img2img, "img2img")

+

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

+

+                            if opts.dimensions_and_batch_together:

+                                res_switch_btn = ToolButton(value=switch_values_symbol, elem_id="img2img_res_switch_btn")

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

+                                    batch_size = gr.Slider(minimum=1, maximum=8, step=1, label='Batch size', value=1, elem_id="img2img_batch_size")

+

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

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

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

+

+                    elif category == "batch":

+                        if not opts.dimensions_and_batch_together:

+                            with FormRow(elem_id="img2img_column_batch"):

+                                batch_count = gr.Slider(minimum=1, step=1, label='Batch count', value=1, elem_id="img2img_batch_count")

+                                batch_size = gr.Slider(minimum=1, maximum=8, step=1, label='Batch size', value=1, elem_id="img2img_batch_size")

+

+                    elif category == "scripts":

+                        with FormGroup(elem_id="img2img_script_container"):

+                            custom_inputs = modules.scripts.scripts_img2img.setup_ui()

+

+                    elif category == "inpaint":

+                        with FormGroup(elem_id="inpaint_controls", visible=False) as inpaint_controls:

+                            with FormRow():

+                                mask_blur = gr.Slider(label='Mask blur', minimum=0, maximum=64, step=1, value=4, elem_id="img2img_mask_blur")

+                                mask_alpha = gr.Slider(label="Mask transparency", visible=False, elem_id="img2img_mask_alpha")

+

+                            with FormRow():

+                                inpainting_mask_invert = gr.Radio(label='Mask mode', choices=['Inpaint masked', 'Inpaint not masked'], value='Inpaint masked', type="index", elem_id="img2img_mask_mode")

+

+                            with FormRow():

+                                inpainting_fill = gr.Radio(label='Masked content', choices=['fill', 'original', 'latent noise', 'latent nothing'], value='original', type="index", elem_id="img2img_inpainting_fill")

+

+                            with FormRow():

+                                with gr.Column():

+                                    inpaint_full_res = gr.Radio(label="Inpaint area", choices=["Whole picture", "Only masked"], type="index", value="Whole picture", elem_id="img2img_inpaint_full_res")

+

+                                with gr.Column(scale=4):

+                                    inpaint_full_res_padding = gr.Slider(label='Only masked padding, pixels', minimum=0, maximum=256, step=4, value=32, elem_id="img2img_inpaint_full_res_padding")

+

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

+                                elem.select(

+                                    fn=lambda tab=i: select_img2img_tab(tab),

+                                    inputs=[],

+                                    outputs=[inpaint_controls, mask_alpha],

+                                )

+

+            img2img_gallery, generation_info, html_info, html_log = create_output_panel("img2img", opts.outdir_img2img_samples)

             parameters_copypaste.bind_buttons({"img2img": img2img_paste}, None, img2img_prompt)

 

             connect_reuse_seed(seed, reuse_seed, generation_info, dummy_component, is_subseed=False)

@@ -900,35 +812,22 @@ def create_ui():
                 ]

             )

 

-            mask_mode.change(

-                lambda mode, img: {

-                    init_img_with_mask: gr_show(mode == 0),

-                    init_img_inpaint: gr_show(mode == 1),

-                    init_mask_inpaint: gr_show(mode == 1),

-                },

-                inputs=[mask_mode, init_img_with_mask],

-                outputs=[

-                    init_img_with_mask,

-                    init_img_inpaint,

-                    init_mask_inpaint,

-                ],

-            )

-

             img2img_args = dict(

-                fn=wrap_gradio_gpu_call(modules.img2img.img2img),

+                fn=wrap_gradio_gpu_call(modules.img2img.img2img, extra_outputs=[None, '', '']),

                 _js="submit_img2img",

                 inputs=[

                     dummy_component,

+                    dummy_component,

                     img2img_prompt,

                     img2img_negative_prompt,

-                    img2img_prompt_style,

-                    img2img_prompt_style2,

+                    img2img_prompt_styles,

                     init_img,

+                    sketch,

                     init_img_with_mask,

-                    init_img_with_mask_orig,

+                    inpaint_color_sketch,

+                    inpaint_color_sketch_orig,

                     init_img_inpaint,

                     init_mask_inpaint,

-                    mask_mode,

                     steps,

                     sampler_index,

                     mask_blur,

@@ -950,44 +849,48 @@ def create_ui():
                     inpainting_mask_invert,

                     img2img_batch_input_dir,

                     img2img_batch_output_dir,

+                    img2img_batch_inpaint_mask_dir

                 ] + custom_inputs,

                 outputs=[

                     img2img_gallery,

                     generation_info,

-                    html_info

+                    html_info,

+                    html_log,

                 ],

                 show_progress=False,

             )

 

+            interrogate_args = dict(

+                _js="get_img2img_tab_index",

+                inputs=[

+                    dummy_component,

+                    img2img_batch_input_dir,

+                    img2img_batch_output_dir,

+                    init_img,

+                    sketch,

+                    init_img_with_mask,

+                    inpaint_color_sketch,

+                    init_img_inpaint,

+                ],

+                outputs=[img2img_prompt, dummy_component],

+            )

+

             img2img_prompt.submit(**img2img_args)

             submit.click(**img2img_args)

+            res_switch_btn.click(lambda w, h: (h, w), inputs=[width, height], outputs=[width, height])

 

             img2img_interrogate.click(

-                fn=interrogate,

-                inputs=[init_img],

-                outputs=[img2img_prompt],

+                fn=lambda *args: process_interrogate(interrogate, *args),

+                **interrogate_args,

             )

 

             img2img_deepbooru.click(

-                fn=interrogate_deepbooru,

-                inputs=[init_img],

-                outputs=[img2img_prompt],

-            )

-

-

-            roll.click(

-                fn=roll_artist,

-                _js="update_img2img_tokens",

-                inputs=[

-                    img2img_prompt,

-                ],

-                outputs=[

-                    img2img_prompt,

-                ]

+                fn=lambda *args: process_interrogate(interrogate_deepbooru, *args),

+                **interrogate_args,

             )

 

             prompts = [(txt2img_prompt, txt2img_negative_prompt), (img2img_prompt, img2img_negative_prompt)]

-            style_dropdowns = [(txt2img_prompt_style, txt2img_prompt_style2), (img2img_prompt_style, img2img_prompt_style2)]

+            style_dropdowns = [txt2img_prompt_styles, img2img_prompt_styles]

             style_js_funcs = ["update_txt2img_tokens", "update_img2img_tokens"]

 

             for button, (prompt, negative_prompt) in zip([txt2img_save_style, img2img_save_style], prompts):

@@ -997,18 +900,21 @@ def create_ui():
                     # Have to pass empty dummy component here, because the JavaScript and Python function have to accept

                     # the same number of parameters, but we only know the style-name after the JavaScript prompt

                     inputs=[dummy_component, prompt, negative_prompt],

-                    outputs=[txt2img_prompt_style, img2img_prompt_style, txt2img_prompt_style2, img2img_prompt_style2],

+                    outputs=[txt2img_prompt_styles, img2img_prompt_styles],

                 )

 

-            for button, (prompt, negative_prompt), (style1, style2), js_func in zip([txt2img_prompt_style_apply, img2img_prompt_style_apply], prompts, style_dropdowns, style_js_funcs):

+            for button, (prompt, negative_prompt), styles, js_func in zip([txt2img_prompt_style_apply, img2img_prompt_style_apply], prompts, style_dropdowns, style_js_funcs):

                 button.click(

                     fn=apply_styles,

                     _js=js_func,

-                    inputs=[prompt, negative_prompt, style1, style2],

-                    outputs=[prompt, negative_prompt, style1, style2],

+                    inputs=[prompt, negative_prompt, styles],

+                    outputs=[prompt, negative_prompt, styles],

                 )

 

             token_button.click(fn=update_token_counter, inputs=[img2img_prompt, steps], outputs=[token_counter])

+            negative_token_button.click(fn=wrap_queued_call(update_token_counter), inputs=[txt2img_negative_prompt, steps], outputs=[negative_token_counter])

+

+            ui_extra_networks.setup_ui(extra_networks_ui_img2img, img2img_gallery)

 

             img2img_paste_fields = [

                 (img2img_prompt, "Prompt"),

@@ -1035,86 +941,7 @@ def create_ui():
     modules.scripts.scripts_current = None

 

     with gr.Blocks(analytics_enabled=False) as extras_interface:

-        with gr.Row().style(equal_height=False):

-            with gr.Column(variant='panel'):

-                with gr.Tabs(elem_id="mode_extras"):

-                    with gr.TabItem('Single Image'):

-                        extras_image = gr.Image(label="Source", source="upload", interactive=True, type="pil")

-

-                    with gr.TabItem('Batch Process'):

-                        image_batch = gr.File(label="Batch Process", file_count="multiple", interactive=True, type="file")

-

-                    with gr.TabItem('Batch from Directory'):

-                        extras_batch_input_dir = gr.Textbox(label="Input directory", **shared.hide_dirs, placeholder="A directory on the same machine where the server is running.")

-                        extras_batch_output_dir = gr.Textbox(label="Output directory", **shared.hide_dirs, placeholder="Leave blank to save images to the default path.")

-                        show_extras_results = gr.Checkbox(label='Show result images', value=True)

-

-                submit = gr.Button('Generate', elem_id="extras_generate", variant='primary')

-

-                with gr.Tabs(elem_id="extras_resize_mode"):

-                    with gr.TabItem('Scale by'):

-                        upscaling_resize = gr.Slider(minimum=1.0, maximum=8.0, step=0.05, label="Resize", value=4)

-                    with gr.TabItem('Scale to'):

-                        with gr.Group():

-                            with gr.Row():

-                                upscaling_resize_w = gr.Number(label="Width", value=512, precision=0)

-                                upscaling_resize_h = gr.Number(label="Height", value=512, precision=0)

-                            upscaling_crop = gr.Checkbox(label='Crop to fit', value=True)

-

-                with gr.Group():

-                    extras_upscaler_1 = gr.Radio(label='Upscaler 1', elem_id="extras_upscaler_1", choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name, type="index")

-

-                with gr.Group():

-                    extras_upscaler_2 = gr.Radio(label='Upscaler 2', elem_id="extras_upscaler_2", choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name, type="index")

-                    extras_upscaler_2_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Upscaler 2 visibility", value=1)

-

-                with gr.Group():

-                    gfpgan_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="GFPGAN visibility", value=0, interactive=modules.gfpgan_model.have_gfpgan)

-

-                with gr.Group():

-                    codeformer_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="CodeFormer visibility", value=0, interactive=modules.codeformer_model.have_codeformer)

-                    codeformer_weight = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="CodeFormer weight (0 = maximum effect, 1 = minimum effect)", value=0, interactive=modules.codeformer_model.have_codeformer)

-

-                with gr.Group():

-                    upscale_before_face_fix = gr.Checkbox(label='Upscale Before Restoring Faces', value=False)

-

-            result_images, html_info_x, html_info = create_output_panel("extras", opts.outdir_extras_samples)

-

-        submit.click(

-            fn=wrap_gradio_gpu_call(modules.extras.run_extras),

-            _js="get_extras_tab_index",

-            inputs=[

-                dummy_component,

-                dummy_component,

-                extras_image,

-                image_batch,

-                extras_batch_input_dir,

-                extras_batch_output_dir,

-                show_extras_results,

-                gfpgan_visibility,

-                codeformer_visibility,

-                codeformer_weight,

-                upscaling_resize,

-                upscaling_resize_w,

-                upscaling_resize_h,

-                upscaling_crop,

-                extras_upscaler_1,

-                extras_upscaler_2,

-                extras_upscaler_2_visibility,

-                upscale_before_face_fix,

-            ],

-            outputs=[

-                result_images,

-                html_info_x,

-                html_info,

-            ]

-        )

-        parameters_copypaste.add_paste_fields("extras", extras_image, None)

-

-        extras_image.change(

-            fn=modules.extras.clear_cache,

-            inputs=[], outputs=[]

-        )

+        ui_postprocessing.create_ui()

 

     with gr.Blocks(analytics_enabled=False) as pnginfo_interface:

         with gr.Row().style(equal_height=False):

@@ -1123,7 +950,7 @@ def create_ui():
 

             with gr.Column(variant='panel'):

                 html = gr.HTML()

-                generation_info = gr.Textbox(visible=False)

+                generation_info = gr.Textbox(visible=False, elem_id="pnginfo_generation_info")

                 html2 = gr.HTML()

                 with gr.Row():

                     buttons = parameters_copypaste.create_buttons(["txt2img", "img2img", "inpaint", "extras"])

@@ -1135,99 +962,141 @@ def create_ui():
             outputs=[html, generation_info, html2],

         )

 

+    def update_interp_description(value):

+        interp_description_css = "<p style='margin-bottom: 2.5em'>{}</p>"

+        interp_descriptions = {

+            "No interpolation": interp_description_css.format("No interpolation will be used. Requires one model; A. Allows for format conversion and VAE baking."),

+            "Weighted sum": interp_description_css.format("A weighted sum will be used for interpolation. Requires two models; A and B. The result is calculated as A * (1 - M) + B * M"),

+            "Add difference": interp_description_css.format("The difference between the last two models will be added to the first. Requires three models; A, B and C. The result is calculated as A + (B - C) * M")

+        }

+        return interp_descriptions[value]

+

     with gr.Blocks(analytics_enabled=False) as modelmerger_interface:

         with gr.Row().style(equal_height=False):

-            with gr.Column(variant='panel'):

-                gr.HTML(value="<p>A merger of the two checkpoints will be generated in your <b>checkpoint</b> directory.</p>")

+            with gr.Column(variant='compact'):

+                interp_description = gr.HTML(value=update_interp_description("Weighted sum"), elem_id="modelmerger_interp_description")

 

-                with gr.Row():

+                with FormRow(elem_id="modelmerger_models"):

                     primary_model_name = gr.Dropdown(modules.sd_models.checkpoint_tiles(), elem_id="modelmerger_primary_model_name", label="Primary model (A)")

+                    create_refresh_button(primary_model_name, modules.sd_models.list_models, lambda: {"choices": modules.sd_models.checkpoint_tiles()}, "refresh_checkpoint_A")

+

                     secondary_model_name = gr.Dropdown(modules.sd_models.checkpoint_tiles(), elem_id="modelmerger_secondary_model_name", label="Secondary model (B)")

+                    create_refresh_button(secondary_model_name, modules.sd_models.list_models, lambda: {"choices": modules.sd_models.checkpoint_tiles()}, "refresh_checkpoint_B")

+

                     tertiary_model_name = gr.Dropdown(modules.sd_models.checkpoint_tiles(), elem_id="modelmerger_tertiary_model_name", label="Tertiary model (C)")

-                custom_name = gr.Textbox(label="Custom Name (Optional)")

-                interp_amount = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Multiplier (M) - set to 0 to get model A', value=0.3)

-                interp_method = gr.Radio(choices=["Weighted sum", "Add difference"], value="Weighted sum", label="Interpolation Method")

+                    create_refresh_button(tertiary_model_name, modules.sd_models.list_models, lambda: {"choices": modules.sd_models.checkpoint_tiles()}, "refresh_checkpoint_C")

 

-                with gr.Row():

-                    checkpoint_format = gr.Radio(choices=["ckpt", "safetensors"], value="ckpt", label="Checkpoint format")

-                    save_as_half = gr.Checkbox(value=False, label="Save as float16")

+                custom_name = gr.Textbox(label="Custom Name (Optional)", elem_id="modelmerger_custom_name")

+                interp_amount = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Multiplier (M) - set to 0 to get model A', value=0.3, elem_id="modelmerger_interp_amount")

+                interp_method = gr.Radio(choices=["No interpolation", "Weighted sum", "Add difference"], value="Weighted sum", label="Interpolation Method", elem_id="modelmerger_interp_method")

+                interp_method.change(fn=update_interp_description, inputs=[interp_method], outputs=[interp_description])

 

-                modelmerger_merge = gr.Button(elem_id="modelmerger_merge", label="Merge", variant='primary')

+                with FormRow():

+                    checkpoint_format = gr.Radio(choices=["ckpt", "safetensors"], value="ckpt", 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")

 

-            with gr.Column(variant='panel'):

-                submit_result = gr.Textbox(elem_id="modelmerger_result", show_label=False)

+                with FormRow():

+                    with gr.Column():

+                        config_source = gr.Radio(choices=["A, B or C", "B", "C", "Don't"], value="A, B or C", label="Copy config from", type="index", elem_id="modelmerger_config_method")

+

+                    with gr.Column():

+                        with FormRow():

+                            bake_in_vae = gr.Dropdown(choices=["None"] + list(sd_vae.vae_dict), value="None", label="Bake in VAE", elem_id="modelmerger_bake_in_vae")

+                            create_refresh_button(bake_in_vae, sd_vae.refresh_vae_list, lambda: {"choices": ["None"] + list(sd_vae.vae_dict)}, "modelmerger_refresh_bake_in_vae")

+

+                with FormRow():

+                    discard_weights = gr.Textbox(value="", label="Discard weights with matching name", elem_id="modelmerger_discard_weights")

+

+                with gr.Row():

+                    modelmerger_merge = gr.Button(elem_id="modelmerger_merge", value="Merge", variant='primary')

 

-    sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings()

+            with gr.Column(variant='compact', elem_id="modelmerger_results_container"):

+                with gr.Group(elem_id="modelmerger_results_panel"):

+                    modelmerger_result = gr.HTML(elem_id="modelmerger_result", show_label=False)

 

     with gr.Blocks(analytics_enabled=False) as train_interface:

         with gr.Row().style(equal_height=False):

             gr.HTML(value="<p style='margin-bottom: 0.7em'>See <b><a href=\"https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Textual-Inversion\">wiki</a></b> for detailed explanation.</p>")

 

-        with gr.Row().style(equal_height=False):

+        with gr.Row(variant="compact").style(equal_height=False):

             with gr.Tabs(elem_id="train_tabs"):

 

                 with gr.Tab(label="Create embedding"):

-                    new_embedding_name = gr.Textbox(label="Name")

-                    initialization_text = gr.Textbox(label="Initialization text", value="*")

-                    nvpt = gr.Slider(label="Number of vectors per token", minimum=1, maximum=75, step=1, value=1)

-                    overwrite_old_embedding = gr.Checkbox(value=False, label="Overwrite Old 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")

+                    overwrite_old_embedding = gr.Checkbox(value=False, label="Overwrite Old Embedding", elem_id="train_overwrite_old_embedding")

 

                     with gr.Row():

                         with gr.Column(scale=3):

                             gr.HTML(value="")

 

                         with gr.Column():

-                            create_embedding = gr.Button(value="Create embedding", variant='primary')

+                            create_embedding = gr.Button(value="Create embedding", variant='primary', elem_id="train_create_embedding")

 

                 with gr.Tab(label="Create hypernetwork"):

-                    new_hypernetwork_name = gr.Textbox(label="Name")

-                    new_hypernetwork_sizes = gr.CheckboxGroup(label="Modules", value=["768", "320", "640", "1280"], choices=["768", "1024", "320", "640", "1280"])

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

-                    new_hypernetwork_activation_func = gr.Dropdown(value="linear", label="Select activation function of hypernetwork. Recommended : Swish / Linear(none)", choices=modules.hypernetworks.ui.keys)

-                    new_hypernetwork_initialization_option = gr.Dropdown(value = "Normal", label="Select Layer weights initialization. Recommended: Kaiming for relu-like, Xavier for sigmoid-like, Normal otherwise", choices=["Normal", "KaimingUniform", "KaimingNormal", "XavierUniform", "XavierNormal"])

-                    new_hypernetwork_add_layer_norm = gr.Checkbox(label="Add layer normalization")

-                    new_hypernetwork_use_dropout = gr.Checkbox(label="Use dropout")

-                    overwrite_old_hypernetwork = gr.Checkbox(value=False, label="Overwrite Old 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")

+                    new_hypernetwork_activation_func = gr.Dropdown(value="linear", label="Select activation function of hypernetwork. Recommended : Swish / Linear(none)", choices=modules.hypernetworks.ui.keys, elem_id="train_new_hypernetwork_activation_func")

+                    new_hypernetwork_initialization_option = gr.Dropdown(value = "Normal", label="Select Layer weights initialization. Recommended: Kaiming for relu-like, Xavier for sigmoid-like, Normal otherwise", choices=["Normal", "KaimingUniform", "KaimingNormal", "XavierUniform", "XavierNormal"], elem_id="train_new_hypernetwork_initialization_option")

+                    new_hypernetwork_add_layer_norm = gr.Checkbox(label="Add layer normalization", elem_id="train_new_hypernetwork_add_layer_norm")

+                    new_hypernetwork_use_dropout = gr.Checkbox(label="Use dropout", elem_id="train_new_hypernetwork_use_dropout")

+                    new_hypernetwork_dropout_structure = gr.Textbox("0, 0, 0", label="Enter hypernetwork Dropout structure (or empty). Recommended : 0~0.35 incrementing sequence: 0, 0.05, 0.15", placeholder="1st and last digit must be 0 and values should be between 0 and 1. ex:'0, 0.01, 0'")

+                    overwrite_old_hypernetwork = gr.Checkbox(value=False, label="Overwrite Old Hypernetwork", elem_id="train_overwrite_old_hypernetwork")

 

                     with gr.Row():

                         with gr.Column(scale=3):

                             gr.HTML(value="")

 

                         with gr.Column():

-                            create_hypernetwork = gr.Button(value="Create hypernetwork", variant='primary')

+                            create_hypernetwork = gr.Button(value="Create hypernetwork", variant='primary', elem_id="train_create_hypernetwork")

 

                 with gr.Tab(label="Preprocess images"):

-                    process_src = gr.Textbox(label='Source directory')

-                    process_dst = gr.Textbox(label='Destination directory')

-                    process_width = gr.Slider(minimum=64, maximum=2048, step=8, label="Width", value=512)

-                    process_height = gr.Slider(minimum=64, maximum=2048, step=8, label="Height", value=512)

-                    preprocess_txt_action = gr.Dropdown(label='Existing Caption txt Action', value="ignore", choices=["ignore", "copy", "prepend", "append"])

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

+                    process_height = gr.Slider(minimum=64, maximum=2048, step=8, label="Height", value=512, elem_id="train_process_height")

+                    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_flip = gr.Checkbox(label='Create flipped copies')

-                        process_split = gr.Checkbox(label='Split oversized images')

-                        process_focal_crop = gr.Checkbox(label='Auto focal point crop')

-                        process_caption = gr.Checkbox(label='Use BLIP for caption')

-                        process_caption_deepbooru = gr.Checkbox(label='Use deepbooru for caption', visible=True)

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

+                        process_multicrop = gr.Checkbox(label='Auto-sized crop', elem_id="train_process_multicrop")

+                        process_caption = gr.Checkbox(label='Use BLIP for caption', elem_id="train_process_caption")

+                        process_caption_deepbooru = gr.Checkbox(label='Use deepbooru for caption', visible=True, elem_id="train_process_caption_deepbooru")

 

                     with gr.Row(visible=False) as process_split_extra_row:

-                        process_split_threshold = gr.Slider(label='Split image threshold', value=0.5, minimum=0.0, maximum=1.0, step=0.05)

-                        process_overlap_ratio = gr.Slider(label='Split image overlap ratio', value=0.2, minimum=0.0, maximum=0.9, step=0.05)

+                        process_split_threshold = gr.Slider(label='Split image threshold', value=0.5, minimum=0.0, maximum=1.0, step=0.05, elem_id="train_process_split_threshold")

+                        process_overlap_ratio = gr.Slider(label='Split image overlap ratio', value=0.2, minimum=0.0, maximum=0.9, step=0.05, elem_id="train_process_overlap_ratio")

 

                     with gr.Row(visible=False) as process_focal_crop_row:

-                        process_focal_crop_face_weight = gr.Slider(label='Focal point face weight', value=0.9, minimum=0.0, maximum=1.0, step=0.05)

-                        process_focal_crop_entropy_weight = gr.Slider(label='Focal point entropy weight', value=0.15, minimum=0.0, maximum=1.0, step=0.05)

-                        process_focal_crop_edges_weight = gr.Slider(label='Focal point edges weight', value=0.5, minimum=0.0, maximum=1.0, step=0.05)

-                        process_focal_crop_debug = gr.Checkbox(label='Create debug image')

-

+                        process_focal_crop_face_weight = gr.Slider(label='Focal point face weight', value=0.9, minimum=0.0, maximum=1.0, step=0.05, elem_id="train_process_focal_crop_face_weight")

+                        process_focal_crop_entropy_weight = gr.Slider(label='Focal point entropy weight', value=0.15, minimum=0.0, maximum=1.0, step=0.05, elem_id="train_process_focal_crop_entropy_weight")

+                        process_focal_crop_edges_weight = gr.Slider(label='Focal point edges weight', value=0.5, minimum=0.0, maximum=1.0, step=0.05, elem_id="train_process_focal_crop_edges_weight")

+                        process_focal_crop_debug = gr.Checkbox(label='Create debug image', elem_id="train_process_focal_crop_debug")

+                    

+                    with gr.Column(visible=False) as process_multicrop_col:

+                        gr.Markdown('Each image is center-cropped with an automatically chosen width and height.')

+                        with gr.Row():

+                            process_multicrop_mindim = gr.Slider(minimum=64, maximum=2048, step=8, label="Dimension lower bound", value=384, elem_id="train_process_multicrop_mindim")

+                            process_multicrop_maxdim = gr.Slider(minimum=64, maximum=2048, step=8, label="Dimension upper bound", value=768, elem_id="train_process_multicrop_maxdim")

+                        with gr.Row():

+                            process_multicrop_minarea = gr.Slider(minimum=64*64, maximum=2048*2048, step=1, label="Area lower bound", value=64*64, elem_id="train_process_multicrop_minarea")

+                            process_multicrop_maxarea = gr.Slider(minimum=64*64, maximum=2048*2048, step=1, label="Area upper bound", value=640*640, elem_id="train_process_multicrop_maxarea")

+                        with gr.Row():

+                            process_multicrop_objective = gr.Radio(["Maximize area", "Minimize error"], value="Maximize area", label="Resizing objective", elem_id="train_process_multicrop_objective")

+                            process_multicrop_threshold = gr.Slider(minimum=0, maximum=1, step=0.01, label="Error threshold", value=0.1, elem_id="train_process_multicrop_threshold")

+   

                     with gr.Row():

                         with gr.Column(scale=3):

                             gr.HTML(value="")

 

                         with gr.Column():

                             with gr.Row():

-                                interrupt_preprocessing = gr.Button("Interrupt")

-                            run_preprocess = gr.Button(value="Preprocess", variant='primary')

+                                interrupt_preprocessing = gr.Button("Interrupt", elem_id="train_interrupt_preprocessing")

+                            run_preprocess = gr.Button(value="Preprocess", variant='primary', elem_id="train_run_preprocess")

 

                     process_split.change(

                         fn=lambda show: gr_show(show),

@@ -1241,54 +1110,74 @@ def create_ui():
                         outputs=[process_focal_crop_row],

                     )

 

+                    process_multicrop.change(

+                        fn=lambda show: gr_show(show),

+                        inputs=[process_multicrop],

+                        outputs=[process_multicrop_col],

+                    )

+

+                def get_textual_inversion_template_names():

+                    return sorted([x for x in textual_inversion.textual_inversion_templates])

+

                 with gr.Tab(label="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 gr.Row():

+                    with FormRow():

                         train_embedding_name = gr.Dropdown(label='Embedding', elem_id="train_embedding", choices=sorted(sd_hijack.model_hijack.embedding_db.word_embeddings.keys()))

                         create_refresh_button(train_embedding_name, sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings, lambda: {"choices": sorted(sd_hijack.model_hijack.embedding_db.word_embeddings.keys())}, "refresh_train_embedding_name")

-                    with gr.Row():

+

                         train_hypernetwork_name = gr.Dropdown(label='Hypernetwork', elem_id="train_hypernetwork", choices=[x for x in shared.hypernetworks.keys()])

                         create_refresh_button(train_hypernetwork_name, shared.reload_hypernetworks, lambda: {"choices": sorted([x for x in shared.hypernetworks.keys()])}, "refresh_train_hypernetwork_name")

-                    with gr.Row():

-                        embedding_learn_rate = gr.Textbox(label='Embedding Learning rate', placeholder="Embedding Learning rate", value="0.005")

-                        hypernetwork_learn_rate = gr.Textbox(label='Hypernetwork Learning rate', placeholder="Hypernetwork Learning rate", value="0.00001")

-

-                    batch_size = gr.Number(label='Batch size', value=1, precision=0)

-                    gradient_step = gr.Number(label='Gradient accumulation steps', value=1, precision=0)

-                    dataset_directory = gr.Textbox(label='Dataset directory', placeholder="Path to directory with input images")

-                    log_directory = gr.Textbox(label='Log directory', placeholder="Path to directory where to write outputs", value="textual_inversion")

-                    template_file = gr.Textbox(label='Prompt template file', value=os.path.join(script_path, "textual_inversion_templates", "style_filewords.txt"))

-                    training_width = gr.Slider(minimum=64, maximum=2048, step=8, label="Width", value=512)

-                    training_height = gr.Slider(minimum=64, maximum=2048, step=8, label="Height", value=512)

-                    steps = gr.Number(label='Max steps', value=100000, precision=0)

-                    create_image_every = gr.Number(label='Save an image to log directory every N steps, 0 to disable', value=500, precision=0)

-                    save_embedding_every = gr.Number(label='Save a copy of embedding to log directory every N steps, 0 to disable', value=500, precision=0)

-                    save_image_with_stored_embedding = gr.Checkbox(label='Save images with embedding in PNG chunks', value=True)

-                    preview_from_txt2img = gr.Checkbox(label='Read parameters (prompt, etc...) from txt2img tab when making previews', value=False)

-                    with gr.Row():

-                        shuffle_tags = gr.Checkbox(label="Shuffle tags by ',' when creating prompts.", value=False)

-                        tag_drop_out = gr.Slider(minimum=0, maximum=1, step=0.1, label="Drop out tags when creating prompts.", value=0)

-                    with gr.Row():

-                        latent_sampling_method = gr.Radio(label='Choose latent sampling method', value="once", choices=['once', 'deterministic', 'random'])

+

+                    with FormRow():

+                        embedding_learn_rate = gr.Textbox(label='Embedding Learning rate', placeholder="Embedding Learning rate", value="0.005", elem_id="train_embedding_learn_rate")

+                        hypernetwork_learn_rate = gr.Textbox(label='Hypernetwork Learning rate', placeholder="Hypernetwork Learning rate", value="0.00001", elem_id="train_hypernetwork_learn_rate")

+                    

+                    with FormRow():

+                        clip_grad_mode = gr.Dropdown(value="disabled", label="Gradient Clipping", choices=["disabled", "value", "norm"])

+                        clip_grad_value = gr.Textbox(placeholder="Gradient clip value", value="0.1", show_label=False)

+

+                    with FormRow():

+                        batch_size = gr.Number(label='Batch size', value=1, precision=0, elem_id="train_batch_size")

+                        gradient_step = gr.Number(label='Gradient accumulation steps', value=1, precision=0, elem_id="train_gradient_step")

+

+                    dataset_directory = gr.Textbox(label='Dataset directory', placeholder="Path to directory with input images", elem_id="train_dataset_directory")

+                    log_directory = gr.Textbox(label='Log directory', placeholder="Path to directory where to write outputs", value="textual_inversion", elem_id="train_log_directory")

+

+                    with FormRow():

+                        template_file = gr.Dropdown(label='Prompt template', value="style_filewords.txt", elem_id="train_template_file", choices=get_textual_inversion_template_names())

+                        create_refresh_button(template_file, textual_inversion.list_textual_inversion_templates, lambda: {"choices": get_textual_inversion_template_names()}, "refrsh_train_template_file")

+

+                    training_width = gr.Slider(minimum=64, maximum=2048, step=8, label="Width", value=512, elem_id="train_training_width")

+                    training_height = gr.Slider(minimum=64, maximum=2048, step=8, label="Height", value=512, elem_id="train_training_height")

+                    varsize = gr.Checkbox(label="Do not resize images", value=False, elem_id="train_varsize")

+                    steps = gr.Number(label='Max steps', value=100000, precision=0, elem_id="train_steps")

+

+                    with FormRow():

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

+

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

+

+                    shuffle_tags = gr.Checkbox(label="Shuffle tags by ',' when creating prompts.", value=False, elem_id="train_shuffle_tags")

+                    tag_drop_out = gr.Slider(minimum=0, maximum=1, step=0.1, label="Drop out tags when creating prompts.", value=0, elem_id="train_tag_drop_out")

+

+                    latent_sampling_method = gr.Radio(label='Choose latent sampling method', value="once", choices=['once', 'deterministic', 'random'], elem_id="train_latent_sampling_method")

 

                     with gr.Row():

-                        interrupt_training = gr.Button(value="Interrupt")

-                        train_hypernetwork = gr.Button(value="Train Hypernetwork", variant='primary')

-                        train_embedding = gr.Button(value="Train Embedding", variant='primary')

+                        train_embedding = gr.Button(value="Train Embedding", variant='primary', elem_id="train_train_embedding")

+                        interrupt_training = gr.Button(value="Interrupt", elem_id="train_interrupt_training")

+                        train_hypernetwork = gr.Button(value="Train Hypernetwork", variant='primary', elem_id="train_train_hypernetwork")

 

                 params = script_callbacks.UiTrainTabParams(txt2img_preview_params)

 

                 script_callbacks.ui_train_tabs_callback(params)

 

-            with gr.Column():

-                progressbar = gr.HTML(elem_id="ti_progressbar")

+            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_preview = gr.Image(elem_id='ti_preview', visible=False)

                 ti_progress = gr.HTML(elem_id="ti_progress", value="")

                 ti_outcome = gr.HTML(elem_id="ti_error", value="")

-                setup_progressbar(progressbar, ti_preview, 'ti', textinfo=ti_progress)

 

         create_embedding.click(

             fn=modules.textual_inversion.ui.create_embedding,

@@ -1315,7 +1204,8 @@ def create_ui():
                 new_hypernetwork_activation_func,

                 new_hypernetwork_initialization_option,

                 new_hypernetwork_add_layer_norm,

-                new_hypernetwork_use_dropout

+                new_hypernetwork_use_dropout,

+                new_hypernetwork_dropout_structure

             ],

             outputs=[

                 train_hypernetwork_name,

@@ -1328,6 +1218,7 @@ def create_ui():
             fn=wrap_gradio_gpu_call(modules.textual_inversion.ui.preprocess, extra_outputs=[gr.update()]),

             _js="start_training_textual_inversion",

             inputs=[

+                dummy_component,

                 process_src,

                 process_dst,

                 process_width,

@@ -1344,6 +1235,13 @@ def create_ui():
                 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,

             ],

             outputs=[

                 ti_output,

@@ -1355,6 +1253,7 @@ def create_ui():
             fn=wrap_gradio_gpu_call(modules.textual_inversion.ui.train_embedding, extra_outputs=[gr.update()]),

             _js="start_training_textual_inversion",

             inputs=[

+                dummy_component,

                 train_embedding_name,

                 embedding_learn_rate,

                 batch_size,

@@ -1363,7 +1262,10 @@ def create_ui():
                 log_directory,

                 training_width,

                 training_height,

+                varsize,

                 steps,

+                clip_grad_mode,

+                clip_grad_value,

                 shuffle_tags,

                 tag_drop_out,

                 latent_sampling_method,

@@ -1384,6 +1286,7 @@ def create_ui():
             fn=wrap_gradio_gpu_call(modules.hypernetworks.ui.train_hypernetwork, extra_outputs=[gr.update()]),

             _js="start_training_textual_inversion",

             inputs=[

+                dummy_component,

                 train_hypernetwork_name,

                 hypernetwork_learn_rate,

                 batch_size,

@@ -1392,7 +1295,10 @@ def create_ui():
                 log_directory,

                 training_width,

                 training_height,

+                varsize,

                 steps,

+                clip_grad_mode,

+                clip_grad_value,

                 shuffle_tags,

                 tag_drop_out,

                 latent_sampling_method,

@@ -1447,7 +1353,7 @@ def create_ui():
                 res = comp(label=info.label, value=fun(), elem_id=elem_id, **(args or {}))

                 create_refresh_button(res, info.refresh, info.component_args, "refresh_" + key)

             else:

-                with gr.Row(variant="compact"):

+                with FormRow():

                     res = comp(label=info.label, value=fun(), elem_id=elem_id, **(args or {}))

                     create_refresh_button(res, info.refresh, info.component_args, "refresh_" + key)

         else:

@@ -1489,44 +1395,44 @@ def create_ui():
 

         opts.save(shared.config_filename)

 

-        return gr.update(value=value), opts.dumpjson()

+        return get_value_for_setting(key), opts.dumpjson()

 

     with gr.Blocks(analytics_enabled=False) as settings_interface:

-        settings_submit = gr.Button(value="Apply settings", variant='primary')

-        result = gr.HTML()

+        with gr.Row():

+            with gr.Column(scale=6):

+                settings_submit = gr.Button(value="Apply settings", variant='primary', elem_id="settings_submit")

+            with gr.Column():

+                restart_gradio = gr.Button(value='Reload UI', variant='primary', elem_id="settings_restart_gradio")

 

-        settings_cols = 3

-        items_per_col = int(len(opts.data_labels) * 0.9 / settings_cols)

+        result = gr.HTML(elem_id="settings_result")

 

         quicksettings_names = [x.strip() for x in opts.quicksettings.split(",")]

-        quicksettings_names = set(x for x in quicksettings_names if x != 'quicksettings')

+        quicksettings_names = {x: i for i, x in enumerate(quicksettings_names) if x != 'quicksettings'}

 

         quicksettings_list = []

 

-        cols_displayed = 0

-        items_displayed = 0

         previous_section = None

-        column = None

-        with gr.Row(elem_id="settings").style(equal_height=False):

+        current_tab = None

+        current_row = None

+        with gr.Tabs(elem_id="settings"):

             for i, (k, item) in enumerate(opts.data_labels.items()):

                 section_must_be_skipped = item.section[0] is None

 

                 if previous_section != item.section and not section_must_be_skipped:

-                    if cols_displayed < settings_cols and (items_displayed >= items_per_col or previous_section is None):

-                        if column is not None:

-                            column.__exit__()

+                    elem_id, text = item.section

 

-                        column = gr.Column(variant='panel')

-                        column.__enter__()

+                    if current_tab is not None:

+                        current_row.__exit__()

+                        current_tab.__exit__()

 

-                        items_displayed = 0

-                        cols_displayed += 1

+                    gr.Group()

+                    current_tab = gr.TabItem(elem_id="settings_{}".format(elem_id), label=text)

+                    current_tab.__enter__()

+                    current_row = gr.Column(variant='compact')

+                    current_row.__enter__()

 

                     previous_section = item.section

 

-                    elem_id, text = item.section

-                    gr.HTML(elem_id="settings_header_text_{}".format(elem_id), value='<h1 class="gr-button-lg">{}</h1>'.format(text))

-

                 if k in quicksettings_names and not shared.cmd_opts.freeze_settings:

                     quicksettings_list.append((i, k, item))

                     components.append(dummy_component)

@@ -1536,15 +1442,20 @@ def create_ui():
                     component = create_setting_component(k)

                     component_dict[k] = component

                     components.append(component)

-                    items_displayed += 1

 

-        with gr.Row():

-            request_notifications = gr.Button(value='Request browser notifications', elem_id="request_notifications")

-            download_localization = gr.Button(value='Download localization template', elem_id="download_localization")

+            if current_tab is not None:

+                current_row.__exit__()

+                current_tab.__exit__()

 

-        with gr.Row():

-            reload_script_bodies = gr.Button(value='Reload custom script bodies (No ui updates, No restart)', variant='secondary')

-            restart_gradio = gr.Button(value='Restart Gradio and Refresh components (Custom Scripts, ui.py, js and css only)', variant='primary')

+            with gr.TabItem("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.TabItem("Licenses"):

+                gr.HTML(shared.html("licenses.html"), elem_id="licenses")

+

+            gr.Button(value="Show all pages", elem_id="settings_show_all_pages")

 

         request_notifications.click(

             fn=lambda: None,

@@ -1581,9 +1492,6 @@ def create_ui():
             outputs=[],

         )

 

-        if column is not None:

-            column.__exit__()

-

     interfaces = [

         (txt2img_interface, "txt2img", "txt2img"),

         (img2img_interface, "img2img", "img2img"),

@@ -1602,8 +1510,8 @@ def create_ui():
         with open(cssfile, "r", encoding="utf8") as file:

             css += file.read() + "\n"

 

-    if os.path.exists(os.path.join(script_path, "user.css")):

-        with open(os.path.join(script_path, "user.css"), "r", encoding="utf8") as file:

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

@@ -1616,8 +1524,8 @@ def create_ui():
     interfaces += [(extensions_interface, "Extensions", "extensions")]

 

     with gr.Blocks(css=css, analytics_enabled=False, title="Stable Diffusion") as demo:

-        with gr.Row(elem_id="quicksettings"):

-            for i, k, item in quicksettings_list:

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

                 component_dict[k] = component

 

@@ -1632,6 +1540,10 @@ def create_ui():
         if os.path.exists(os.path.join(script_path, "notification.mp3")):

             audio_notification = gr.Audio(interactive=False, value=os.path.join(script_path, "notification.mp3"), elem_id="audio_notification", visible=False)

 

+        footer = shared.html("footer.html")

+        footer = footer.format(versions=versions_html())

+        gr.HTML(footer, elem_id="footer")

+

         text_settings = gr.Textbox(elem_id="settings_json", value=lambda: opts.dumpjson(), visible=False)

         settings_submit.click(

             fn=wrap_gradio_call(run_settings, extra_outputs=[gr.update()]),

@@ -1651,7 +1563,7 @@ def create_ui():
         component_keys = [k for k in opts.data_labels.keys() if k in component_dict]

 

         def get_settings_values():

-            return [getattr(opts, key) for key in component_keys]

+            return [get_value_for_setting(key) for key in component_keys]

 

         demo.load(

             fn=get_settings_values,

@@ -1666,12 +1578,15 @@ def create_ui():
                 print("Error loading/saving model file:", file=sys.stderr)

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

                 modules.sd_models.list_models()  # to remove the potentially missing models from the list

-                return ["Error loading/saving model file. It doesn't exist or the name contains illegal characters"] + [gr.Dropdown.update(choices=modules.sd_models.checkpoint_tiles()) for _ in range(3)]

+                return [*[gr.Dropdown.update(choices=modules.sd_models.checkpoint_tiles()) for _ in range(4)], f"Error merging checkpoints: {e}"]

             return results

 

+        modelmerger_merge.click(fn=lambda: '', inputs=[], outputs=[modelmerger_result])

         modelmerger_merge.click(

-            fn=modelmerger,

+            fn=wrap_gradio_gpu_call(modelmerger, extra_outputs=lambda: [gr.update() for _ in range(4)]),

+            _js='modelmerger',

             inputs=[

+                dummy_component,

                 primary_model_name,

                 secondary_model_name,

                 tertiary_model_name,

@@ -1680,13 +1595,16 @@ def create_ui():
                 save_as_half,

                 custom_name,

                 checkpoint_format,

+                config_source,

+                bake_in_vae,

+                discard_weights,

             ],

             outputs=[

-                submit_result,

                 primary_model_name,

                 secondary_model_name,

                 tertiary_model_name,

                 component_dict['sd_model_checkpoint'],

+                modelmerger_result,

             ]

         )

 

@@ -1718,13 +1636,16 @@ def create_ui():
             if saved_value is None:

                 ui_settings[key] = getattr(obj, field)

             elif condition and not condition(saved_value):

-                print(f'Warning: Bad ui setting value: {key}: {saved_value}; Default value "{getattr(obj, field)}" will be used instead.')

+                pass

+

+                # this warning is generally not useful;

+                # print(f'Warning: Bad ui setting value: {key}: {saved_value}; Default value "{getattr(obj, field)}" will be used instead.')

             else:

                 setattr(obj, field, saved_value)

                 if init_field is not None:

                     init_field(saved_value)

 

-        if type(x) in [gr.Slider, gr.Radio, gr.Checkbox, gr.Textbox, gr.Number] and x.visible:

+        if type(x) in [gr.Slider, gr.Radio, gr.Checkbox, gr.Textbox, gr.Number, gr.Dropdown] and x.visible:

             apply_field(x, 'visible')

 

         if type(x) == gr.Slider:

@@ -1745,43 +1666,46 @@ def create_ui():
         if type(x) == gr.Number:

             apply_field(x, 'value')

 

-        # Since there are many dropdowns that shouldn't be saved,

-        # we only mark dropdowns that should be saved.

-        if type(x) == gr.Dropdown and getattr(x, 'save_to_config', False):

-            apply_field(x, 'value', lambda val: val in x.choices, getattr(x, 'init_field', None))

-            apply_field(x, 'visible')

+        if type(x) == gr.Dropdown:

+            def check_dropdown(val):

+                if getattr(x, 'multiselect', False):

+                    return all([value in x.choices for value in val])

+                else:

+                    return val in x.choices

+

+            apply_field(x, 'value', check_dropdown, getattr(x, 'init_field', None))

 

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

 

     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)

 

+    # Required as a workaround for change() event not triggering when loading values from ui-config.json

+    interp_description.value = update_interp_description(interp_method.value)

+

     return demo

 

 

 def reload_javascript():

-    with open(os.path.join(script_path, "script.js"), "r", encoding="utf8") as jsfile:

-        javascript = f'<script>{jsfile.read()}</script>'

-

-    scripts_list = modules.scripts.list_scripts("javascript", ".js")

-

-    for basedir, filename, path in scripts_list:

-        with open(path, "r", encoding="utf8") as jsfile:

-            javascript += f"\n<!-- {filename} --><script>{jsfile.read()}</script>"

+    head = f'<script type="text/javascript" src="file={os.path.abspath("script.js")}"></script>\n'

 

+    inline = f"{localization.localization_js(shared.opts.localization)};"

     if cmd_opts.theme is not None:

-        javascript += f"\n<script>set_theme('{cmd_opts.theme}');</script>\n"

+        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}"></script>\n'

 

-    javascript += f"\n<script>{localization.localization_js(shared.opts.localization)}</script>"

+    head += f'<script type="text/javascript">{inline}</script>\n'

 

     def template_response(*args, **kwargs):

         res = shared.GradioTemplateResponseOriginal(*args, **kwargs)

-        res.body = res.body.replace(

-            b'</head>', f'{javascript}</head>'.encode("utf8"))

+        res.body = res.body.replace(b'</head>', f'{head}</head>'.encode("utf8"))

         res.init_headers()

         return res

 

@@ -1790,3 +1714,32 @@ def reload_javascript():
 

 if not hasattr(shared, 'GradioTemplateResponseOriginal'):

     shared.GradioTemplateResponseOriginal = gradio.routes.templates.TemplateResponse

+

+

+def versions_html():

+    import torch

+    import launch

+

+    python_version = ".".join([str(x) for x in sys.version_info[0:3]])

+    commit = launch.commit_hash()

+    short_commit = commit[0:8]

+

+    if shared.xformers_available:

+        import xformers

+        xformers_version = xformers.__version__

+    else:

+        xformers_version = "N/A"

+

+    return f"""

+python: <span title="{sys.version}">{python_version}</span>

+ • 

+torch: {torch.__version__}

+ • 

+xformers: {xformers_version}

+ • 

+gradio: {gr.__version__}

+ • 

+commit: <a href="https://github.com/AUTOMATIC1111/stable-diffusion-webui/commit/{commit}">{short_commit}</a>

+ • 

+checkpoint: <a id="sd_checkpoint_hash">N/A</a>

+"""

diff --git a/modules/ui_common.py b/modules/ui_common.py
new file mode 100644
index 00000000..9405ac1f
--- /dev/null
+++ b/modules/ui_common.py
@@ -0,0 +1,202 @@
+import json

+import html

+import os

+import platform

+import sys

+

+import gradio as gr

+import subprocess as sp

+

+from modules import call_queue, shared

+from modules.generation_parameters_copypaste import image_from_url_text

+import modules.images

+

+folder_symbol = '\U0001f4c2'  # 📂

+

+

+def update_generation_info(generation_info, html_info, img_index):

+    try:

+        generation_info = json.loads(generation_info)

+        if img_index < 0 or img_index >= len(generation_info["infotexts"]):

+            return html_info, gr.update()

+        return plaintext_to_html(generation_info["infotexts"][img_index]), gr.update()

+    except Exception:

+        pass

+    # if the json parse or anything else fails, just return the old html_info

+    return html_info, gr.update()

+

+

+def plaintext_to_html(text):

+    text = "<p>" + "<br>\n".join([f"{html.escape(x)}" for x in text.split('\n')]) + "</p>"

+    return text

+

+

+def save_files(js_data, images, do_make_zip, index):

+    import csv

+    filenames = []

+    fullfns = []

+

+    #quick dictionary to class object conversion. Its necessary due apply_filename_pattern requiring it

+    class MyObject:

+        def __init__(self, d=None):

+            if d is not None:

+                for key, value in d.items():

+                    setattr(self, key, value)

+

+    data = json.loads(js_data)

+

+    p = MyObject(data)

+    path = shared.opts.outdir_save

+    save_to_dirs = shared.opts.use_save_to_dirs_for_ui

+    extension: str = shared.opts.samples_format

+    start_index = 0

+

+    if index > -1 and shared.opts.save_selected_only and (index >= data["index_of_first_image"]):  # ensures we are looking at a specific non-grid picture, and we have save_selected_only

+

+        images = [images[index]]

+        start_index = index

+

+    os.makedirs(shared.opts.outdir_save, exist_ok=True)

+

+    with open(os.path.join(shared.opts.outdir_save, "log.csv"), "a", encoding="utf8", newline='') as file:

+        at_start = file.tell() == 0

+        writer = csv.writer(file)

+        if at_start:

+            writer.writerow(["prompt", "seed", "width", "height", "sampler", "cfgs", "steps", "filename", "negative_prompt"])

+

+        for image_index, filedata in enumerate(images, start_index):

+            image = image_from_url_text(filedata)

+

+            is_grid = image_index < p.index_of_first_image

+            i = 0 if is_grid else (image_index - p.index_of_first_image)

+

+            fullfn, txt_fullfn = modules.images.save_image(image, path, "", seed=p.all_seeds[i], prompt=p.all_prompts[i], extension=extension, info=p.infotexts[image_index], grid=is_grid, p=p, save_to_dirs=save_to_dirs)

+

+            filename = os.path.relpath(fullfn, path)

+            filenames.append(filename)

+            fullfns.append(fullfn)

+            if txt_fullfn:

+                filenames.append(os.path.basename(txt_fullfn))

+                fullfns.append(txt_fullfn)

+

+        writer.writerow([data["prompt"], data["seed"], data["width"], data["height"], data["sampler_name"], data["cfg_scale"], data["steps"], filenames[0], data["negative_prompt"]])

+

+    # Make Zip

+    if do_make_zip:

+        zip_filepath = os.path.join(path, "images.zip")

+

+        from zipfile import ZipFile

+        with ZipFile(zip_filepath, "w") as zip_file:

+            for i in range(len(fullfns)):

+                with open(fullfns[i], mode="rb") as f:

+                    zip_file.writestr(filenames[i], f.read())

+        fullfns.insert(0, zip_filepath)

+

+    return gr.File.update(value=fullfns, visible=True), plaintext_to_html(f"Saved: {filenames[0]}")

+

+

+def create_output_panel(tabname, outdir):

+    from modules import shared

+    import modules.generation_parameters_copypaste as parameters_copypaste

+

+    def open_folder(f):

+        if not os.path.exists(f):

+            print(f'Folder "{f}" does not exist. After you create an image, the folder will be created.')

+            return

+        elif not os.path.isdir(f):

+            print(f"""

+WARNING

+An open_folder request was made with an argument that is not a folder.

+This could be an error or a malicious attempt to run code on your computer.

+Requested path was: {f}

+""", file=sys.stderr)

+            return

+

+        if not shared.cmd_opts.hide_ui_dir_config:

+            path = os.path.normpath(f)

+            if platform.system() == "Windows":

+                os.startfile(path)

+            elif platform.system() == "Darwin":

+                sp.Popen(["open", path])

+            elif "microsoft-standard-WSL2" in platform.uname().release:

+                sp.Popen(["wsl-open", path])

+            else:

+                sp.Popen(["xdg-open", path])

+

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

+

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

+

+                if tabname != "extras":

+                    save = gr.Button('Save', elem_id=f'save_{tabname}')

+                    save_zip = gr.Button('Zip', elem_id=f'save_zip_{tabname}')

+

+                buttons = parameters_copypaste.create_buttons(["img2img", "inpaint", "extras"])

+

+            open_folder_button.click(

+                fn=lambda: open_folder(shared.opts.outdir_samples or outdir),

+                inputs=[],

+                outputs=[],

+            )

+

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

+

+                with gr.Group():

+                    html_info = gr.HTML(elem_id=f'html_info_{tabname}')

+                    html_log = gr.HTML(elem_id=f'html_log_{tabname}')

+

+                    generation_info = gr.Textbox(visible=False, elem_id=f'generation_info_{tabname}')

+                    if tabname == 'txt2img' or tabname == 'img2img':

+                        generation_info_button = gr.Button(visible=False, elem_id=f"{tabname}_generation_info_button")

+                        generation_info_button.click(

+                            fn=update_generation_info,

+                            _js="function(x, y, z){ return [x, y, selected_gallery_index()] }",

+                            inputs=[generation_info, html_info, html_info],

+                            outputs=[html_info, html_info],

+                        )

+

+                    save.click(

+                        fn=call_queue.wrap_gradio_call(save_files),

+                        _js="(x, y, z, w) => [x, y, false, selected_gallery_index()]",

+                        inputs=[

+                            generation_info,

+                            result_gallery,

+                            html_info,

+                            html_info,

+                        ],

+                        outputs=[

+                            download_files,

+                            html_log,

+                        ],

+                        show_progress=False,

+                    )

+

+                    save_zip.click(

+                        fn=call_queue.wrap_gradio_call(save_files),

+                        _js="(x, y, z, w) => [x, y, true, selected_gallery_index()]",

+                        inputs=[

+                            generation_info,

+                            result_gallery,

+                            html_info,

+                            html_info,

+                        ],

+                        outputs=[

+                            download_files,

+                            html_log,

+                        ]

+                    )

+

+            else:

+                html_info_x = gr.HTML(elem_id=f'html_info_x_{tabname}')

+                html_info = gr.HTML(elem_id=f'html_info_{tabname}')

+                html_log = gr.HTML(elem_id=f'html_log_{tabname}')

+

+            parameters_copypaste.bind_buttons(buttons, result_gallery, "txt2img" if tabname == "txt2img" else None)

+            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
new file mode 100644
index 00000000..284ca0cf
--- /dev/null
+++ b/modules/ui_components.py
@@ -0,0 +1,58 @@
+import gradio as gr

+

+

+class ToolButton(gr.Button, gr.components.FormComponent):

+    """Small button with single emoji as text, fits inside gradio forms"""

+

+    def __init__(self, **kwargs):

+        super().__init__(variant="tool", **kwargs)

+

+    def get_block_name(self):

+        return "button"

+

+

+class ToolButtonTop(gr.Button, gr.components.FormComponent):

+    """Small button with single emoji as text, with extra margin at top, fits inside gradio forms"""

+

+    def __init__(self, **kwargs):

+        super().__init__(variant="tool-top", **kwargs)

+

+    def get_block_name(self):

+        return "button"

+

+

+class FormRow(gr.Row, gr.components.FormComponent):

+    """Same as gr.Row but fits inside gradio forms"""

+

+    def get_block_name(self):

+        return "row"

+

+

+class FormGroup(gr.Group, gr.components.FormComponent):

+    """Same as gr.Row but fits inside gradio forms"""

+

+    def get_block_name(self):

+        return "group"

+

+

+class FormHTML(gr.HTML, gr.components.FormComponent):

+    """Same as gr.HTML but fits inside gradio forms"""

+

+    def get_block_name(self):

+        return "html"

+

+

+class FormColorPicker(gr.ColorPicker, gr.components.FormComponent):

+    """Same as gr.ColorPicker but fits inside gradio forms"""

+

+    def get_block_name(self):

+        return "colorpicker"

+

+

+class DropdownMulti(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"

diff --git a/modules/ui_extensions.py b/modules/ui_extensions.py
index eec9586f..66a41865 100644
--- a/modules/ui_extensions.py
+++ b/modules/ui_extensions.py
@@ -132,7 +132,7 @@ def install_extension_from_url(dirname, url):
     normalized_url = normalize_git_url(url)

     assert len([x for x in extensions.extensions if normalize_git_url(x.remote) == normalized_url]) == 0, 'Extension with this URL is already installed'

 

-    tmpdir = os.path.join(paths.script_path, "tmp", dirname)

+    tmpdir = os.path.join(paths.data_path, "tmp", dirname)

 

     try:

         shutil.rmtree(tmpdir, True)

@@ -162,15 +162,15 @@ def install_extension_from_url(dirname, url):
         shutil.rmtree(tmpdir, True)

 

 

-def install_extension_from_index(url, hide_tags):

+def install_extension_from_index(url, hide_tags, sort_column):

     ext_table, message = install_extension_from_url(None, url)

 

-    code, _ = refresh_available_extensions_from_data(hide_tags)

+    code, _ = refresh_available_extensions_from_data(hide_tags, sort_column)

 

     return code, ext_table, message

 

 

-def refresh_available_extensions(url, hide_tags):

+def refresh_available_extensions(url, hide_tags, sort_column):

     global available_extensions

 

     import urllib.request

@@ -179,18 +179,28 @@ def refresh_available_extensions(url, hide_tags):
 

     available_extensions = json.loads(text)

 

-    code, tags = refresh_available_extensions_from_data(hide_tags)

+    code, tags = refresh_available_extensions_from_data(hide_tags, sort_column)

 

     return url, code, gr.CheckboxGroup.update(choices=tags), ''

 

 

-def refresh_available_extensions_for_tags(hide_tags):

-    code, _ = refresh_available_extensions_from_data(hide_tags)

+def refresh_available_extensions_for_tags(hide_tags, sort_column):

+    code, _ = refresh_available_extensions_from_data(hide_tags, sort_column)

 

     return code, ''

 

 

-def refresh_available_extensions_from_data(hide_tags):

+sort_ordering = [

+    # (reverse, order_by_function)

+    (True, lambda x: x.get('added', 'z')),

+    (False, lambda x: x.get('added', 'z')),

+    (False, lambda x: x.get('name', 'z')),

+    (True, lambda x: x.get('name', 'z')),

+    (False, lambda x: 'z'),

+]

+

+

+def refresh_available_extensions_from_data(hide_tags, sort_column):

     extlist = available_extensions["extensions"]

     installed_extension_urls = {normalize_git_url(extension.remote): extension.name for extension in extensions.extensions}

 

@@ -210,8 +220,11 @@ def refresh_available_extensions_from_data(hide_tags):
         <tbody>

     """

 

-    for ext in extlist:

+    sort_reverse, sort_function = sort_ordering[sort_column if 0 <= sort_column < len(sort_ordering) else 0]

+

+    for ext in sorted(extlist, key=sort_function, reverse=sort_reverse):

         name = ext.get("name", "noname")

+        added = ext.get('added', 'unknown')

         url = ext.get("url", None)

         description = ext.get("description", "")

         extension_tags = ext.get("tags", [])

@@ -233,7 +246,7 @@ def refresh_available_extensions_from_data(hide_tags):
         code += f"""

             <tr>

                 <td><a href="{html.escape(url)}" target="_blank">{html.escape(name)}</a><br />{tags_text}</td>

-                <td>{html.escape(description)}</td>

+                <td>{html.escape(description)}<p class="info"><span class="date_added">Added: {html.escape(added)}</span></p></td>

                 <td>{install_code}</td>

             </tr>

         

@@ -291,25 +304,32 @@ def create_ui():
 

                 with gr.Row():

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

 

                 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],

+                    inputs=[available_extensions_index, hide_tags, sort_column],

                     outputs=[available_extensions_index, available_extensions_table, hide_tags, install_result],

                 )

 

                 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],

+                    inputs=[extension_to_install, hide_tags, sort_column],

                     outputs=[available_extensions_table, 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],

+                    inputs=[hide_tags, sort_column],

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

                     outputs=[available_extensions_table, install_result]

                 )

 

diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py
new file mode 100644
index 00000000..c6ff889a
--- /dev/null
+++ b/modules/ui_extra_networks.py
@@ -0,0 +1,179 @@
+import os.path

+

+from modules import shared

+import gradio as gr

+import json

+import html

+

+from modules.generation_parameters_copypaste import image_from_url_text

+

+extra_pages = []

+

+

+def register_page(page):

+    """registers extra networks page for the UI; recommend doing it in on_before_ui() callback for extensions"""

+

+    extra_pages.append(page)

+

+

+class ExtraNetworksPage:

+    def __init__(self, title):

+        self.title = title

+        self.name = title.lower()

+        self.card_page = shared.html("extra-networks-card.html")

+        self.allow_negative_prompt = False

+

+    def refresh(self):

+        pass

+

+    def create_html(self, tabname):

+        view = shared.opts.extra_networks_default_view

+        items_html = ''

+

+        for item in self.list_items():

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

+

+        res = f"""

+<div id='{tabname}_{self.name}_cards' class='extra-network-{view}'>

+{items_html}

+</div>

+"""

+

+        return res

+

+    def list_items(self):

+        raise NotImplementedError()

+

+    def allowed_directories_for_previews(self):

+        return []

+

+    def create_html_for_item(self, item, tabname):

+        preview = item.get("preview", None)

+

+        args = {

+            "preview_html": "style='background-image: url(\"" + html.escape(preview) + "\")'" if preview else '',

+            "prompt": item["prompt"],

+            "tabname": json.dumps(tabname),

+            "local_preview": json.dumps(item["local_preview"]),

+            "name": item["name"],

+            "card_clicked": '"' + html.escape(f"""return cardClicked({json.dumps(tabname)}, {item["prompt"]}, {"true" if self.allow_negative_prompt else "false"})""") + '"',

+            "save_card_preview": '"' + html.escape(f"""return saveCardPreview(event, {json.dumps(tabname)}, {json.dumps(item["local_preview"])})""") + '"',

+        }

+

+        return self.card_page.format(**args)

+

+

+def intialize():

+    extra_pages.clear()

+

+

+class ExtraNetworksUi:

+    def __init__(self):

+        self.pages = None

+        self.stored_extra_pages = None

+

+        self.button_save_preview = None

+        self.preview_target_filename = None

+

+        self.tabname = None

+

+

+def pages_in_preferred_order(pages):

+    tab_order = [x.lower().strip() for x in shared.opts.ui_extra_networks_tab_reorder.split(",")]

+

+    def tab_name_score(name):

+        name = name.lower()

+        for i, possible_match in enumerate(tab_order):

+            if possible_match in name:

+                return i

+

+        return len(pages)

+

+    tab_scores = {page.name: (tab_name_score(page.name), original_index) for original_index, page in enumerate(pages)}

+

+    return sorted(pages, key=lambda x: tab_scores[x.name])

+

+

+def create_ui(container, button, tabname):

+    ui = ExtraNetworksUi()

+    ui.pages = []

+    ui.stored_extra_pages = pages_in_preferred_order(extra_pages.copy())

+    ui.tabname = tabname

+

+    with gr.Tabs(elem_id=tabname+"_extra_tabs") as tabs:

+        for page in ui.stored_extra_pages:

+            with gr.Tab(page.title):

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

+

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

+

+    def refresh():

+        res = []

+

+        for pg in ui.stored_extra_pages:

+            pg.refresh()

+            res.append(pg.create_html(ui.tabname))

+

+        return res

+

+    button_refresh.click(fn=refresh, inputs=[], outputs=ui.pages)

+

+    return ui

+

+

+def path_is_parent(parent_path, child_path):

+    parent_path = os.path.abspath(parent_path)

+    child_path = os.path.abspath(child_path)

+

+    return os.path.commonpath([parent_path]) == os.path.commonpath([parent_path, child_path])

+

+

+def setup_ui(ui, gallery):

+    def save_preview(index, images, filename):

+        if len(images) == 0:

+            print("There is no image in gallery to save as a preview.")

+            return [page.create_html(ui.tabname) for page in ui.stored_extra_pages]

+

+        index = int(index)

+        index = 0 if index < 0 else index

+        index = len(images) - 1 if index >= len(images) else index

+

+        img_info = images[index if index >= 0 else 0]

+        image = image_from_url_text(img_info)

+

+        is_allowed = False

+        for extra_page in ui.stored_extra_pages:

+            if any([path_is_parent(x, filename) for x in extra_page.allowed_directories_for_previews()]):

+                is_allowed = True

+                break

+

+        assert is_allowed, f'writing to {filename} is not allowed'

+

+        image.save(filename)

+

+        return [page.create_html(ui.tabname) for page in ui.stored_extra_pages]

+

+    ui.button_save_preview.click(

+        fn=save_preview,

+        _js="function(x, y, z){console.log(x, y, z); return [selected_gallery_index(), y, z]}",

+        inputs=[ui.preview_target_filename, gallery, ui.preview_target_filename],

+        outputs=[*ui.pages]

+    )

diff --git a/modules/ui_extra_networks_hypernets.py b/modules/ui_extra_networks_hypernets.py
new file mode 100644
index 00000000..65d000cf
--- /dev/null
+++ b/modules/ui_extra_networks_hypernets.py
@@ -0,0 +1,35 @@
+import json

+import os

+

+from modules import shared, ui_extra_networks

+

+

+class ExtraNetworksPageHypernetworks(ui_extra_networks.ExtraNetworksPage):

+    def __init__(self):

+        super().__init__('Hypernetworks')

+

+    def refresh(self):

+        shared.reload_hypernetworks()

+

+    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 = "./file=" + file.replace('\\', '/') + "?mtime=" + str(os.path.getmtime(file))

+                    break

+

+            yield {

+                "name": name,

+                "filename": path,

+                "preview": preview,

+                "prompt": json.dumps(f"<hypernet:{name}:") + " + opts.extra_networks_default_multiplier + " + json.dumps(">"),

+                "local_preview": path + ".png",

+            }

+

+    def allowed_directories_for_previews(self):

+        return [shared.cmd_opts.hypernetwork_dir]

+

diff --git a/modules/ui_extra_networks_textual_inversion.py b/modules/ui_extra_networks_textual_inversion.py
new file mode 100644
index 00000000..dbd23d2d
--- /dev/null
+++ b/modules/ui_extra_networks_textual_inversion.py
@@ -0,0 +1,33 @@
+import json

+import os

+

+from modules import ui_extra_networks, sd_hijack

+

+

+class ExtraNetworksPageTextualInversion(ui_extra_networks.ExtraNetworksPage):

+    def __init__(self):

+        super().__init__('Textual Inversion')

+        self.allow_negative_prompt = True

+

+    def refresh(self):

+        sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings(force_reload=True)

+

+    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 = "./file=" + preview_file.replace('\\', '/') + "?mtime=" + str(os.path.getmtime(preview_file))

+

+            yield {

+                "name": embedding.name,

+                "filename": embedding.filename,

+                "preview": preview,

+                "prompt": json.dumps(embedding.name),

+                "local_preview": path + ".preview.png",

+            }

+

+    def allowed_directories_for_previews(self):

+        return list(sd_hijack.model_hijack.embedding_db.embedding_dirs)

diff --git a/modules/ui_postprocessing.py b/modules/ui_postprocessing.py
new file mode 100644
index 00000000..b418d955
--- /dev/null
+++ b/modules/ui_postprocessing.py
@@ -0,0 +1,57 @@
+import gradio as gr

+from modules import scripts_postprocessing, scripts, shared, gfpgan_model, codeformer_model, ui_common, postprocessing, call_queue

+import modules.generation_parameters_copypaste as parameters_copypaste

+

+

+def create_ui():

+    tab_index = gr.State(value=0)

+

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

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

+

+            submit = gr.Button('Generate', elem_id="extras_generate", variant='primary')

+

+            script_inputs = scripts.scripts_postproc.setup_ui()

+

+        with gr.Column():

+            result_images, html_info_x, html_info, html_log = ui_common.create_output_panel("extras", shared.opts.outdir_extras_samples)

+

+    tab_single.select(fn=lambda: 0, inputs=[], outputs=[tab_index])

+    tab_batch.select(fn=lambda: 1, inputs=[], outputs=[tab_index])

+    tab_batch_dir.select(fn=lambda: 2, inputs=[], outputs=[tab_index])

+

+    submit.click(

+        fn=call_queue.wrap_gradio_gpu_call(postprocessing.run_postprocessing, extra_outputs=[None, '']),

+        inputs=[

+            tab_index,

+            extras_image,

+            image_batch,

+            extras_batch_input_dir,

+            extras_batch_output_dir,

+            show_extras_results,

+            *script_inputs

+        ],

+        outputs=[

+            result_images,

+            html_info_x,

+            html_info,

+        ]

+    )

+

+    parameters_copypaste.add_paste_fields("extras", extras_image, None)

+

+    extras_image.change(

+        fn=scripts.scripts_postproc.image_changed,

+        inputs=[], outputs=[]

+    )

diff --git a/modules/ui_tempdir.py b/modules/ui_tempdir.py
index 07210d14..21945235 100644
--- a/modules/ui_tempdir.py
+++ b/modules/ui_tempdir.py
@@ -1,6 +1,7 @@
 import os

 import tempfile

 from collections import namedtuple

+from pathlib import Path

 

 import gradio as gr

 

@@ -12,10 +13,29 @@ from modules import shared
 Savedfile = namedtuple("Savedfile", ["name"])

 

 

+def register_tmp_file(gradio, filename):

+    if hasattr(gradio, 'temp_file_sets'):  # gradio 3.15

+        gradio.temp_file_sets[0] = gradio.temp_file_sets[0] | {os.path.abspath(filename)}

+

+    if hasattr(gradio, 'temp_dirs'):  # gradio 3.9

+        gradio.temp_dirs = gradio.temp_dirs | {os.path.abspath(os.path.dirname(filename))}

+

+

+def check_tmp_file(gradio, filename):

+    if hasattr(gradio, 'temp_file_sets'):

+        return any([filename in fileset for fileset in gradio.temp_file_sets])

+

+    if hasattr(gradio, 'temp_dirs'):

+        return any(Path(temp_dir).resolve() in Path(filename).resolve().parents for temp_dir in gradio.temp_dirs)

+

+    return False

+

+

 def save_pil_to_file(pil_image, dir=None):

     already_saved_as = getattr(pil_image, 'already_saved_as', None)

     if already_saved_as and os.path.isfile(already_saved_as):

-        shared.demo.temp_dirs = shared.demo.temp_dirs | {os.path.abspath(os.path.dirname(already_saved_as))}

+        register_tmp_file(shared.demo, already_saved_as)

+

         file_obj = Savedfile(already_saved_as)

         return file_obj

 

@@ -44,7 +64,7 @@ def on_tmpdir_changed():
 

     os.makedirs(shared.opts.temp_dir, exist_ok=True)

 

-    shared.demo.temp_dirs = shared.demo.temp_dirs | {os.path.abspath(shared.opts.temp_dir)}

+    register_tmp_file(shared.demo, os.path.join(shared.opts.temp_dir, "x"))

 

 

 def cleanup_tmpdr():

diff --git a/modules/upscaler.py b/modules/upscaler.py
index c4e6e6bd..e2eaa730 100644
--- a/modules/upscaler.py
+++ b/modules/upscaler.py
@@ -11,7 +11,6 @@ from modules import modelloader, shared
 
 LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)
 NEAREST = (Image.Resampling.NEAREST if hasattr(Image, 'Resampling') else Image.NEAREST)
-from modules.paths import models_path
 
 
 class Upscaler:
@@ -39,7 +38,7 @@ class Upscaler:
         self.mod_scale = None
 
         if self.model_path is None and self.name:
-            self.model_path = os.path.join(models_path, self.name)
+            self.model_path = os.path.join(shared.models_path, self.name)
         if self.model_path and create_dirs:
             os.makedirs(self.model_path, exist_ok=True)
 
@@ -53,10 +52,10 @@ class Upscaler:
     def do_upscale(self, img: PIL.Image, selected_model: str):
         return img
 
-    def upscale(self, img: PIL.Image, scale: int, selected_model: str = None):
+    def upscale(self, img: PIL.Image, scale, selected_model: str = None):
         self.scale = scale
-        dest_w = img.width * scale
-        dest_h = img.height * scale
+        dest_w = int(img.width * scale)
+        dest_h = int(img.height * scale)
 
         for i in range(3):
             shape = (img.width, img.height)
@@ -95,6 +94,7 @@ class UpscalerData:
     def __init__(self, name: str, path: str, upscaler: Upscaler = None, scale: int = 4, model=None):
         self.name = name
         self.data_path = path
+        self.local_data_path = path
         self.scaler = upscaler
         self.scale = scale
         self.model = model
@@ -142,4 +142,4 @@ class UpscalerNearest(Upscaler):
     def __init__(self, dirname=None):
         super().__init__(False)
         self.name = "Nearest"
-        self.scalers = [UpscalerData("Nearest", None, self)]
\ No newline at end of file
+        self.scalers = [UpscalerData("Nearest", None, self)]
diff --git a/modules/xlmr.py b/modules/xlmr.py
new file mode 100644
index 00000000..beab3fdf
--- /dev/null
+++ b/modules/xlmr.py
@@ -0,0 +1,137 @@
+from transformers import BertPreTrainedModel,BertModel,BertConfig
+import torch.nn as nn
+import torch
+from transformers.models.xlm_roberta.configuration_xlm_roberta import XLMRobertaConfig
+from transformers import XLMRobertaModel,XLMRobertaTokenizer
+from typing import Optional
+
+class BertSeriesConfig(BertConfig):
+    def __init__(self, vocab_size=30522, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072, hidden_act="gelu", hidden_dropout_prob=0.1, attention_probs_dropout_prob=0.1, max_position_embeddings=512, type_vocab_size=2, initializer_range=0.02, layer_norm_eps=1e-12, pad_token_id=0, position_embedding_type="absolute", use_cache=True, classifier_dropout=None,project_dim=512, pooler_fn="average",learn_encoder=False,model_type='bert',**kwargs):
+
+        super().__init__(vocab_size, hidden_size, num_hidden_layers, num_attention_heads, intermediate_size, hidden_act, hidden_dropout_prob, attention_probs_dropout_prob, max_position_embeddings, type_vocab_size, initializer_range, layer_norm_eps, pad_token_id, position_embedding_type, use_cache, classifier_dropout, **kwargs)
+        self.project_dim = project_dim
+        self.pooler_fn = pooler_fn
+        self.learn_encoder = learn_encoder
+
+class RobertaSeriesConfig(XLMRobertaConfig):
+    def __init__(self, pad_token_id=1, bos_token_id=0, eos_token_id=2,project_dim=512,pooler_fn='cls',learn_encoder=False, **kwargs):
+        super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
+        self.project_dim = project_dim
+        self.pooler_fn = pooler_fn
+        self.learn_encoder = learn_encoder
+
+
+class BertSeriesModelWithTransformation(BertPreTrainedModel):
+
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+    _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"]
+    config_class = BertSeriesConfig
+
+    def __init__(self, config=None, **kargs):
+        # modify initialization for autoloading 
+        if config is None:
+            config = XLMRobertaConfig()
+            config.attention_probs_dropout_prob= 0.1
+            config.bos_token_id=0
+            config.eos_token_id=2
+            config.hidden_act='gelu'
+            config.hidden_dropout_prob=0.1
+            config.hidden_size=1024
+            config.initializer_range=0.02
+            config.intermediate_size=4096
+            config.layer_norm_eps=1e-05
+            config.max_position_embeddings=514
+
+            config.num_attention_heads=16
+            config.num_hidden_layers=24
+            config.output_past=True
+            config.pad_token_id=1
+            config.position_embedding_type= "absolute"
+
+            config.type_vocab_size= 1
+            config.use_cache=True
+            config.vocab_size= 250002
+            config.project_dim = 768
+            config.learn_encoder = False
+        super().__init__(config)
+        self.roberta = XLMRobertaModel(config)
+        self.transformation = nn.Linear(config.hidden_size,config.project_dim)
+        self.pre_LN=nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.tokenizer = XLMRobertaTokenizer.from_pretrained('xlm-roberta-large')
+        self.pooler = lambda x: x[:,0]
+        self.post_init()
+
+    def encode(self,c):
+        device = next(self.parameters()).device
+        text = self.tokenizer(c,
+                        truncation=True,
+                        max_length=77,
+                        return_length=False,
+                        return_overflowing_tokens=False,
+                        padding="max_length",
+                        return_tensors="pt")
+        text["input_ids"] = torch.tensor(text["input_ids"]).to(device)
+        text["attention_mask"] = torch.tensor(
+            text['attention_mask']).to(device)
+        features = self(**text)
+        return features['projection_state'] 
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+    ) :
+        r"""
+        """
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+
+        outputs = self.roberta(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=True,
+            return_dict=return_dict,
+        )
+
+        # last module outputs
+        sequence_output = outputs[0]
+
+
+        # project every module
+        sequence_output_ln = self.pre_LN(sequence_output)
+
+        # pooler
+        pooler_output = self.pooler(sequence_output_ln)
+        pooler_output = self.transformation(pooler_output)
+        projection_state = self.transformation(outputs.last_hidden_state)
+
+        return {
+            'pooler_output':pooler_output,
+            'last_hidden_state':outputs.last_hidden_state,
+            'hidden_states':outputs.hidden_states,
+            'attentions':outputs.attentions,
+            'projection_state':projection_state,
+            'sequence_out': sequence_output
+        }
+
+
+class RobertaSeriesModelWithTransformation(BertSeriesModelWithTransformation):
+    base_model_prefix = 'roberta'
+    config_class= RobertaSeriesConfig
\ No newline at end of file