36 files changed, 3033 insertions, 1683 deletions

diff --git a/modules/aesthetic_clip.py b/modules/aesthetic_clip.py
deleted file mode 100644
index 8c828541..00000000
--- a/modules/aesthetic_clip.py
+++ /dev/null
@@ -1,241 +0,0 @@
-import copy
-import itertools
-import os
-from pathlib import Path
-import html
-import gc
-
-import gradio as gr
-import torch
-from PIL import Image
-from torch import optim
-
-from modules import shared
-from transformers import CLIPModel, CLIPProcessor, CLIPTokenizer
-from tqdm.auto import tqdm, trange
-from modules.shared import opts, device
-
-
-def get_all_images_in_folder(folder):
-    return [os.path.join(folder, f) for f in os.listdir(folder) if
-            os.path.isfile(os.path.join(folder, f)) and check_is_valid_image_file(f)]
-
-
-def check_is_valid_image_file(filename):
-    return filename.lower().endswith(('.png', '.jpg', '.jpeg', ".gif", ".tiff", ".webp"))
-
-
-def batched(dataset, total, n=1):
-    for ndx in range(0, total, n):
-        yield [dataset.__getitem__(i) for i in range(ndx, min(ndx + n, total))]
-
-
-def iter_to_batched(iterable, n=1):
-    it = iter(iterable)
-    while True:
-        chunk = tuple(itertools.islice(it, n))
-        if not chunk:
-            return
-        yield chunk
-
-
-def create_ui():
-    import modules.ui
-
-    with gr.Group():
-        with gr.Accordion("Open for Clip Aesthetic!", open=False):
-            with gr.Row():
-                aesthetic_weight = gr.Slider(minimum=0, maximum=1, step=0.01, label="Aesthetic weight",
-                                             value=0.9)
-                aesthetic_steps = gr.Slider(minimum=0, maximum=50, step=1, label="Aesthetic steps", value=5)
-
-            with gr.Row():
-                aesthetic_lr = gr.Textbox(label='Aesthetic learning rate',
-                                          placeholder="Aesthetic learning rate", value="0.0001")
-                aesthetic_slerp = gr.Checkbox(label="Slerp interpolation", value=False)
-                aesthetic_imgs = gr.Dropdown(sorted(shared.aesthetic_embeddings.keys()),
-                                             label="Aesthetic imgs embedding",
-                                             value="None")
-
-                modules.ui.create_refresh_button(aesthetic_imgs, shared.update_aesthetic_embeddings, lambda: {"choices": sorted(shared.aesthetic_embeddings.keys())}, "refresh_aesthetic_embeddings")
-
-            with gr.Row():
-                aesthetic_imgs_text = gr.Textbox(label='Aesthetic text for imgs',
-                                                 placeholder="This text is used to rotate the feature space of the imgs embs",
-                                                 value="")
-                aesthetic_slerp_angle = gr.Slider(label='Slerp angle', minimum=0, maximum=1, step=0.01,
-                                                  value=0.1)
-                aesthetic_text_negative = gr.Checkbox(label="Is negative text", value=False)
-
-    return aesthetic_weight, aesthetic_steps, aesthetic_lr, aesthetic_slerp, aesthetic_imgs, aesthetic_imgs_text, aesthetic_slerp_angle, aesthetic_text_negative
-
-
-aesthetic_clip_model = None
-
-
-def aesthetic_clip():
-    global aesthetic_clip_model
-
-    if aesthetic_clip_model is None or aesthetic_clip_model.name_or_path != shared.sd_model.cond_stage_model.wrapped.transformer.name_or_path:
-        aesthetic_clip_model = CLIPModel.from_pretrained(shared.sd_model.cond_stage_model.wrapped.transformer.name_or_path)
-        aesthetic_clip_model.cpu()
-
-    return aesthetic_clip_model
-
-
-def generate_imgs_embd(name, folder, batch_size):
-    model = aesthetic_clip().to(device)
-    processor = CLIPProcessor.from_pretrained(model.name_or_path)
-
-    with torch.no_grad():
-        embs = []
-        for paths in tqdm(iter_to_batched(get_all_images_in_folder(folder), batch_size),
-                          desc=f"Generating embeddings for {name}"):
-            if shared.state.interrupted:
-                break
-            inputs = processor(images=[Image.open(path) for path in paths], return_tensors="pt").to(device)
-            outputs = model.get_image_features(**inputs).cpu()
-            embs.append(torch.clone(outputs))
-            inputs.to("cpu")
-            del inputs, outputs
-
-        embs = torch.cat(embs, dim=0).mean(dim=0, keepdim=True)
-
-        # The generated embedding will be located here
-        path = str(Path(shared.cmd_opts.aesthetic_embeddings_dir) / f"{name}.pt")
-        torch.save(embs, path)
-
-        model.cpu()
-        del processor
-        del embs
-        gc.collect()
-        torch.cuda.empty_cache()
-        res = f"""
-        Done generating embedding for {name}!
-        Aesthetic embedding saved to {html.escape(path)}
-        """
-        shared.update_aesthetic_embeddings()
-        return gr.Dropdown.update(choices=sorted(shared.aesthetic_embeddings.keys()), label="Imgs embedding",
-                                  value="None"), \
-               gr.Dropdown.update(choices=sorted(shared.aesthetic_embeddings.keys()),
-                                  label="Imgs embedding",
-                                  value="None"), res, ""
-
-
-def slerp(low, high, val):
-    low_norm = low / torch.norm(low, dim=1, keepdim=True)
-    high_norm = high / torch.norm(high, dim=1, keepdim=True)
-    omega = torch.acos((low_norm * high_norm).sum(1))
-    so = torch.sin(omega)
-    res = (torch.sin((1.0 - val) * omega) / so).unsqueeze(1) * low + (torch.sin(val * omega) / so).unsqueeze(1) * high
-    return res
-
-
-class AestheticCLIP:
-    def __init__(self):
-        self.skip = False
-        self.aesthetic_steps = 0
-        self.aesthetic_weight = 0
-        self.aesthetic_lr = 0
-        self.slerp = False
-        self.aesthetic_text_negative = ""
-        self.aesthetic_slerp_angle = 0
-        self.aesthetic_imgs_text = ""
-
-        self.image_embs_name = None
-        self.image_embs = None
-        self.load_image_embs(None)
-
-    def set_aesthetic_params(self, p, aesthetic_lr=0, aesthetic_weight=0, aesthetic_steps=0, image_embs_name=None,
-                             aesthetic_slerp=True, aesthetic_imgs_text="",
-                             aesthetic_slerp_angle=0.15,
-                             aesthetic_text_negative=False):
-        self.aesthetic_imgs_text = aesthetic_imgs_text
-        self.aesthetic_slerp_angle = aesthetic_slerp_angle
-        self.aesthetic_text_negative = aesthetic_text_negative
-        self.slerp = aesthetic_slerp
-        self.aesthetic_lr = aesthetic_lr
-        self.aesthetic_weight = aesthetic_weight
-        self.aesthetic_steps = aesthetic_steps
-        self.load_image_embs(image_embs_name)
-
-        if self.image_embs_name is not None:
-            p.extra_generation_params.update({
-                "Aesthetic LR": aesthetic_lr,
-                "Aesthetic weight": aesthetic_weight,
-                "Aesthetic steps": aesthetic_steps,
-                "Aesthetic embedding": self.image_embs_name,
-                "Aesthetic slerp": aesthetic_slerp,
-                "Aesthetic text": aesthetic_imgs_text,
-                "Aesthetic text negative": aesthetic_text_negative,
-                "Aesthetic slerp angle": aesthetic_slerp_angle,
-            })
-
-    def set_skip(self, skip):
-        self.skip = skip
-
-    def load_image_embs(self, image_embs_name):
-        if image_embs_name is None or len(image_embs_name) == 0 or image_embs_name == "None":
-            image_embs_name = None
-            self.image_embs_name = None
-        if image_embs_name is not None and self.image_embs_name != image_embs_name:
-            self.image_embs_name = image_embs_name
-            self.image_embs = torch.load(shared.aesthetic_embeddings[self.image_embs_name], map_location=device)
-            self.image_embs /= self.image_embs.norm(dim=-1, keepdim=True)
-            self.image_embs.requires_grad_(False)
-
-    def __call__(self, z, remade_batch_tokens):
-        if not self.skip and self.aesthetic_steps != 0 and self.aesthetic_lr != 0 and self.aesthetic_weight != 0 and self.image_embs_name is not None:
-            tokenizer = shared.sd_model.cond_stage_model.tokenizer
-            if not opts.use_old_emphasis_implementation:
-                remade_batch_tokens = [
-                    [tokenizer.bos_token_id] + x[:75] + [tokenizer.eos_token_id] for x in
-                    remade_batch_tokens]
-
-            tokens = torch.asarray(remade_batch_tokens).to(device)
-
-            model = copy.deepcopy(aesthetic_clip()).to(device)
-            model.requires_grad_(True)
-            if self.aesthetic_imgs_text is not None and len(self.aesthetic_imgs_text) > 0:
-                text_embs_2 = model.get_text_features(
-                    **tokenizer([self.aesthetic_imgs_text], padding=True, return_tensors="pt").to(device))
-                if self.aesthetic_text_negative:
-                    text_embs_2 = self.image_embs - text_embs_2
-                    text_embs_2 /= text_embs_2.norm(dim=-1, keepdim=True)
-                img_embs = slerp(self.image_embs, text_embs_2, self.aesthetic_slerp_angle)
-            else:
-                img_embs = self.image_embs
-
-            with torch.enable_grad():
-
-                # We optimize the model to maximize the similarity
-                optimizer = optim.Adam(
-                    model.text_model.parameters(), lr=self.aesthetic_lr
-                )
-
-                for _ in trange(self.aesthetic_steps, desc="Aesthetic optimization"):
-                    text_embs = model.get_text_features(input_ids=tokens)
-                    text_embs = text_embs / text_embs.norm(dim=-1, keepdim=True)
-                    sim = text_embs @ img_embs.T
-                    loss = -sim
-                    optimizer.zero_grad()
-                    loss.mean().backward()
-                    optimizer.step()
-
-                zn = model.text_model(input_ids=tokens, output_hidden_states=-opts.CLIP_stop_at_last_layers)
-                if opts.CLIP_stop_at_last_layers > 1:
-                    zn = zn.hidden_states[-opts.CLIP_stop_at_last_layers]
-                    zn = model.text_model.final_layer_norm(zn)
-                else:
-                    zn = zn.last_hidden_state
-                model.cpu()
-                del model
-                gc.collect()
-                torch.cuda.empty_cache()
-            zn = torch.concat([zn[77 * i:77 * (i + 1)] for i in range(max(z.shape[1] // 77, 1))], 1)
-            if self.slerp:
-                z = slerp(z, zn, self.aesthetic_weight)
-            else:
-                z = z * (1 - self.aesthetic_weight) + zn * self.aesthetic_weight
-
-        return z
diff --git a/modules/api/api.py b/modules/api/api.py
index 5b0c934e..6c06d449 100644
--- a/modules/api/api.py
+++ b/modules/api/api.py
@@ -1,22 +1,32 @@
-from modules.api.processing import StableDiffusionProcessingAPI
-from modules.processing import StableDiffusionProcessingTxt2Img, process_images
-from modules.sd_samplers import all_samplers
-from modules.extras import run_pnginfo
-import modules.shared as shared
+import time
 import uvicorn
-from fastapi import Body, APIRouter, HTTPException
-from fastapi.responses import JSONResponse
-from pydantic import BaseModel, Field, Json
-import json
-import io
-import base64
+from gradio.processing_utils import encode_pil_to_base64, decode_base64_to_file, decode_base64_to_image
+from fastapi import APIRouter, Depends, HTTPException
+import modules.shared as shared
+from modules import devices
+from modules.api.models import *
+from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, process_images
+from modules.sd_samplers import all_samplers
+from modules.extras import run_extras, run_pnginfo
+
+
+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])}")
+
 
 sampler_to_index = lambda name: next(filter(lambda row: name.lower() == row[1].name.lower(), enumerate(all_samplers)), None)
 
-class TextToImageResponse(BaseModel):
-    images: list[str] = Field(default=None, title="Image", description="The generated image in base64 format.")
-    parameters: Json
-    info: Json
+
+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')
+    return reqDict
 
 
 class Api:
@@ -24,16 +34,21 @@ class Api:
         self.router = APIRouter()
         self.app = app
         self.queue_lock = queue_lock
-        self.app.add_api_route("/sdapi/v1/txt2img", self.text2imgapi, methods=["POST"])
+        self.app.add_api_route("/sdapi/v1/txt2img", self.text2imgapi, methods=["POST"], response_model=TextToImageResponse)
+        self.app.add_api_route("/sdapi/v1/img2img", self.img2imgapi, methods=["POST"], response_model=ImageToImageResponse)
+        self.app.add_api_route("/sdapi/v1/extra-single-image", self.extras_single_image_api, methods=["POST"], response_model=ExtrasSingleImageResponse)
+        self.app.add_api_route("/sdapi/v1/extra-batch-images", self.extras_batch_images_api, methods=["POST"], response_model=ExtrasBatchImagesResponse)
+        self.app.add_api_route("/sdapi/v1/png-info", self.pnginfoapi, methods=["POST"], response_model=PNGInfoResponse)
+        self.app.add_api_route("/sdapi/v1/progress", self.progressapi, methods=["GET"], response_model=ProgressResponse)
 
-    def text2imgapi(self, txt2imgreq: StableDiffusionProcessingAPI ):
+    def text2imgapi(self, txt2imgreq: StableDiffusionTxt2ImgProcessingAPI):
         sampler_index = sampler_to_index(txt2imgreq.sampler_index)
-        
+
         if sampler_index is None:
-            raise HTTPException(status_code=404, detail="Sampler not found") 
-        
+            raise HTTPException(status_code=404, detail="Sampler not found")
+
         populate = txt2imgreq.copy(update={ # Override __init__ params
-            "sd_model": shared.sd_model, 
+            "sd_model": shared.sd_model,
             "sampler_index": sampler_index[0],
             "do_not_save_samples": True,
             "do_not_save_grid": True
@@ -41,27 +56,125 @@ class Api:
         )
         p = StableDiffusionProcessingTxt2Img(**vars(populate))
         # Override object param
+
+        shared.state.begin()
+
         with self.queue_lock:
             processed = process_images(p)
-        
-        b64images = []
-        for i in processed.images:
-            buffer = io.BytesIO()
-            i.save(buffer, format="png")
-            b64images.append(base64.b64encode(buffer.getvalue()))
 
-        return TextToImageResponse(images=b64images, parameters=json.dumps(vars(txt2imgreq)), info=json.dumps(processed.info))
-        
-        
+        shared.state.end()
+
+        b64images = list(map(encode_pil_to_base64, processed.images))
+
+        return TextToImageResponse(images=b64images, parameters=vars(txt2imgreq), info=processed.js())
+
+    def img2imgapi(self, img2imgreq: StableDiffusionImg2ImgProcessingAPI):
+        sampler_index = sampler_to_index(img2imgreq.sampler_index)
+
+        if sampler_index is None:
+            raise HTTPException(status_code=404, detail="Sampler not found")
+
+
+        init_images = img2imgreq.init_images
+        if init_images is None:
+            raise HTTPException(status_code=404, detail="Init image not found")
+
+        mask = img2imgreq.mask
+        if mask:
+            mask = decode_base64_to_image(mask)
+
+
+        populate = img2imgreq.copy(update={ # Override __init__ params
+            "sd_model": shared.sd_model,
+            "sampler_index": sampler_index[0],
+            "do_not_save_samples": True,
+            "do_not_save_grid": True,
+            "mask": mask
+            }
+        )
+        p = StableDiffusionProcessingImg2Img(**vars(populate))
+
+        imgs = []
+        for img in init_images:
+            img = decode_base64_to_image(img)
+            imgs = [img] * p.batch_size
+
+        p.init_images = imgs
+
+        shared.state.begin()
+
+        with self.queue_lock:
+            processed = process_images(p)
+
+        shared.state.end()
+
+        b64images = list(map(encode_pil_to_base64, processed.images))
+
+        if (not img2imgreq.include_init_images):
+            img2imgreq.init_images = None
+            img2imgreq.mask = None
+
+        return ImageToImageResponse(images=b64images, parameters=vars(img2imgreq), info=processed.js())
+
+    def extras_single_image_api(self, req: ExtrasSingleImageRequest):
+        reqDict = setUpscalers(req)
+
+        reqDict['image'] = decode_base64_to_image(reqDict['image'])
+
+        with self.queue_lock:
+            result = run_extras(extras_mode=0, image_folder="", input_dir="", output_dir="", **reqDict)
+
+        return ExtrasSingleImageResponse(image=encode_pil_to_base64(result[0][0]), html_info=result[1])
+
+    def extras_batch_images_api(self, req: ExtrasBatchImagesRequest):
+        reqDict = setUpscalers(req)
+
+        def prepareFiles(file):
+            file = decode_base64_to_file(file.data, file_path=file.name)
+            file.orig_name = file.name
+            return file
+
+        reqDict['image_folder'] = list(map(prepareFiles, reqDict['imageList']))
+        reqDict.pop('imageList')
+
+        with self.queue_lock:
+            result = run_extras(extras_mode=1, image="", input_dir="", output_dir="", **reqDict)
+
+        return ExtrasBatchImagesResponse(images=list(map(encode_pil_to_base64, result[0])), html_info=result[1])
+
+    def pnginfoapi(self, req: PNGInfoRequest):
+        if(not req.image.strip()):
+            return PNGInfoResponse(info="")
+
+        result = run_pnginfo(decode_base64_to_image(req.image.strip()))
+
+        return PNGInfoResponse(info=result[1])
+
+    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())
+
+        # avoid dividing zero
+        progress = 0.01
+
+        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_since_start = time.time() - shared.state.time_start
+        eta = (time_since_start/progress)
+        eta_relative = eta-time_since_start
 
-    def img2imgapi(self):
-        raise NotImplementedError
+        progress = min(progress, 1)
 
-    def extrasapi(self):
-        raise NotImplementedError
+        current_image = None
+        if shared.state.current_image and not req.skip_current_image:
+            current_image = encode_pil_to_base64(shared.state.current_image)
 
-    def pnginfoapi(self):
-        raise NotImplementedError
+        return ProgressResponse(progress=progress, eta_relative=eta_relative, state=shared.state.dict(), current_image=current_image)
 
     def launch(self, server_name, port):
         self.app.include_router(self.router)
diff --git a/modules/api/models.py b/modules/api/models.py
new file mode 100644
index 00000000..9ee42a17
--- /dev/null
+++ b/modules/api/models.py
@@ -0,0 +1,167 @@
+import inspect
+from click import prompt
+from pydantic import BaseModel, Field, create_model
+from typing import Any, Optional
+from typing_extensions import Literal
+from inflection import underscore
+from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img
+from modules.shared import sd_upscalers
+
+API_NOT_ALLOWED = [
+    "self",
+    "kwargs",
+    "sd_model",
+    "outpath_samples",
+    "outpath_grids",
+    "sampler_index",
+    "do_not_save_samples",
+    "do_not_save_grid",
+    "extra_generation_params",
+    "overlay_images",
+    "do_not_reload_embeddings",
+    "seed_enable_extras",
+    "prompt_for_display",
+    "sampler_noise_scheduler_override",
+    "ddim_discretize"
+]
+
+class ModelDef(BaseModel):
+    """Assistance Class for Pydantic Dynamic Model Generation"""
+
+    field: str
+    field_alias: str
+    field_type: Any
+    field_value: Any
+    field_exclude: bool = False
+
+
+class PydanticModelGenerator:
+    """
+    Takes in created classes and stubs them out in a way FastAPI/Pydantic is happy about:
+    source_data is a snapshot of the default values produced by the class
+    params are the names of the actual keys required by __init__
+    """
+
+    def __init__(
+        self,
+        model_name: str = None,
+        class_instance = None,
+        additional_fields = None,
+    ):
+        def field_type_generator(k, v):
+            # field_type = str if not overrides.get(k) else overrides[k]["type"]
+            # print(k, v.annotation, v.default)
+            field_type = v.annotation
+
+            return Optional[field_type]
+
+        def merge_class_params(class_):
+            all_classes = list(filter(lambda x: x is not object, inspect.getmro(class_)))
+            parameters = {}
+            for classes in all_classes:
+                parameters = {**parameters, **inspect.signature(classes.__init__).parameters}
+            return parameters
+
+
+        self._model_name = model_name
+        self._class_data = merge_class_params(class_instance)
+        self._model_def = [
+            ModelDef(
+                field=underscore(k),
+                field_alias=k,
+                field_type=field_type_generator(k, v),
+                field_value=v.default
+            )
+            for (k,v) in self._class_data.items() if k not in API_NOT_ALLOWED
+        ]
+
+        for fields in additional_fields:
+            self._model_def.append(ModelDef(
+                field=underscore(fields["key"]),
+                field_alias=fields["key"],
+                field_type=fields["type"],
+                field_value=fields["default"],
+                field_exclude=fields["exclude"] if "exclude" in fields else False))
+
+    def generate_model(self):
+        """
+        Creates a pydantic BaseModel
+        from the json and overrides provided at initialization
+        """
+        fields = {
+            d.field: (d.field_type, Field(default=d.field_value, alias=d.field_alias, exclude=d.field_exclude)) for d in self._model_def
+        }
+        DynamicModel = create_model(self._model_name, **fields)
+        DynamicModel.__config__.allow_population_by_field_name = True
+        DynamicModel.__config__.allow_mutation = True
+        return DynamicModel
+
+StableDiffusionTxt2ImgProcessingAPI = PydanticModelGenerator(
+    "StableDiffusionProcessingTxt2Img",
+    StableDiffusionProcessingTxt2Img,
+    [{"key": "sampler_index", "type": str, "default": "Euler"}]
+).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}]
+).generate_model()
+
+class TextToImageResponse(BaseModel):
+    images: list[str] = Field(default=None, title="Image", description="The generated image in base64 format.")
+    parameters: dict
+    info: str
+
+class ImageToImageResponse(BaseModel):
+    images: list[str] = Field(default=None, title="Image", description="The generated image in base64 format.")
+    parameters: dict
+    info: str
+
+class ExtrasBaseRequest(BaseModel):
+    resize_mode: Literal[0, 1] = Field(default=0, title="Resize Mode", description="Sets the resize mode: 0 to upscale by upscaling_resize amount, 1 to upscale up to upscaling_resize_h x upscaling_resize_w.")
+    show_extras_results: bool = Field(default=True, title="Show results", description="Should the backend return the generated image?")
+    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_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 choosen size?")
+    upscaler_1: str = Field(default="None", title="Main upscaler", description=f"The name of the main upscaler to use, it has to be one of this list: {' , '.join([x.name for x in sd_upscalers])}")
+    upscaler_2: str = Field(default="None", title="Secondary upscaler", description=f"The name of the secondary upscaler to use, it has to be one of this list: {' , '.join([x.name for x in sd_upscalers])}")
+    extras_upscaler_2_visibility: float = Field(default=0, title="Secondary upscaler visibility", ge=0, le=1, allow_inf_nan=False, description="Sets the visibility of secondary upscaler, values should be between 0 and 1.")
+
+class ExtraBaseResponse(BaseModel):
+    html_info: str = Field(title="HTML info", description="A series of HTML tags containing the process info.")
+
+class ExtrasSingleImageRequest(ExtrasBaseRequest):
+    image: str = Field(default="", title="Image", description="Image to work on, must be a Base64 string containing the image's data.")
+
+class ExtrasSingleImageResponse(ExtraBaseResponse):
+    image: str = Field(default=None, title="Image", description="The generated image in base64 format.")
+
+class FileData(BaseModel):
+    data: str = Field(title="File data", description="Base64 representation of the file")
+    name: str = Field(title="File name")
+
+class ExtrasBatchImagesRequest(ExtrasBaseRequest):
+    imageList: list[FileData] = Field(title="Images", description="List of images to work on. Must be Base64 strings")
+
+class ExtrasBatchImagesResponse(ExtraBaseResponse):
+    images: list[str] = Field(title="Images", description="The generated images in base64 format.")
+
+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")
+
+class ProgressRequest(BaseModel):
+    skip_current_image: bool = Field(default=False, title="Skip current image", description="Skip current image serialization")
+
+class ProgressResponse(BaseModel):
+    progress: float = Field(title="Progress", description="The progress with a range of 0 to 1")
+    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.")
diff --git a/modules/api/processing.py b/modules/api/processing.py
deleted file mode 100644
index 4c541241..00000000
--- a/modules/api/processing.py
+++ /dev/null
@@ -1,99 +0,0 @@
-from inflection import underscore
-from typing import Any, Dict, Optional
-from pydantic import BaseModel, Field, create_model
-from modules.processing import StableDiffusionProcessingTxt2Img
-import inspect
-
-
-API_NOT_ALLOWED = [
-    "self",
-    "kwargs",
-    "sd_model",
-    "outpath_samples",
-    "outpath_grids",
-    "sampler_index",
-    "do_not_save_samples",
-    "do_not_save_grid",
-    "extra_generation_params",
-    "overlay_images",
-    "do_not_reload_embeddings",
-    "seed_enable_extras",
-    "prompt_for_display",
-    "sampler_noise_scheduler_override",
-    "ddim_discretize"
-]
-
-class ModelDef(BaseModel):
-    """Assistance Class for Pydantic Dynamic Model Generation"""
-
-    field: str
-    field_alias: str
-    field_type: Any
-    field_value: Any
-
-
-class PydanticModelGenerator:
-    """
-    Takes in created classes and stubs them out in a way FastAPI/Pydantic is happy about:
-    source_data is a snapshot of the default values produced by the class
-    params are the names of the actual keys required by __init__
-    """
-
-    def __init__(
-        self,
-        model_name: str = None,
-        class_instance = None,
-        additional_fields = None,
-    ):
-        def field_type_generator(k, v):
-            # field_type = str if not overrides.get(k) else overrides[k]["type"]
-            # print(k, v.annotation, v.default)
-            field_type = v.annotation
-            
-            return Optional[field_type]
-        
-        def merge_class_params(class_):
-            all_classes = list(filter(lambda x: x is not object, inspect.getmro(class_)))
-            parameters = {}
-            for classes in all_classes:
-                parameters = {**parameters, **inspect.signature(classes.__init__).parameters}
-            return parameters
-            
-                
-        self._model_name = model_name
-        self._class_data = merge_class_params(class_instance)
-        self._model_def = [
-            ModelDef(
-                field=underscore(k),
-                field_alias=k,
-                field_type=field_type_generator(k, v),
-                field_value=v.default
-            )
-            for (k,v) in self._class_data.items() if k not in API_NOT_ALLOWED
-        ]
-        
-        for fields in additional_fields:
-            self._model_def.append(ModelDef(
-                field=underscore(fields["key"]), 
-                field_alias=fields["key"], 
-                field_type=fields["type"],
-                field_value=fields["default"]))
-
-    def generate_model(self):
-        """
-        Creates a pydantic BaseModel
-        from the json and overrides provided at initialization
-        """
-        fields = {
-            d.field: (d.field_type, Field(default=d.field_value, alias=d.field_alias)) for d in self._model_def
-        }
-        DynamicModel = create_model(self._model_name, **fields)
-        DynamicModel.__config__.allow_population_by_field_name = True
-        DynamicModel.__config__.allow_mutation = True
-        return DynamicModel
-    
-StableDiffusionProcessingAPI = PydanticModelGenerator(
-    "StableDiffusionProcessingTxt2Img", 
-    StableDiffusionProcessingTxt2Img,
-    [{"key": "sampler_index", "type": str, "default": "Euler"}]
-).generate_model()
\ No newline at end of file
diff --git a/modules/bsrgan_model.py b/modules/bsrgan_model.py
deleted file mode 100644
index 737e1a76..00000000
--- a/modules/bsrgan_model.py
+++ /dev/null
@@ -1,76 +0,0 @@
-import os.path
-import sys
-import traceback
-
-import PIL.Image
-import numpy as np
-import torch
-from basicsr.utils.download_util import load_file_from_url
-
-import modules.upscaler
-from modules import devices, modelloader
-from modules.bsrgan_model_arch import RRDBNet
-
-
-class UpscalerBSRGAN(modules.upscaler.Upscaler):
-    def __init__(self, dirname):
-        self.name = "BSRGAN"
-        self.model_name = "BSRGAN 4x"
-        self.model_url = "https://github.com/cszn/KAIR/releases/download/v1.0/BSRGAN.pth"
-        self.user_path = dirname
-        super().__init__()
-        model_paths = self.find_models(ext_filter=[".pt", ".pth"])
-        scalers = []
-        if len(model_paths) == 0:
-            scaler_data = modules.upscaler.UpscalerData(self.model_name, self.model_url, self, 4)
-            scalers.append(scaler_data)
-        for file in model_paths:
-            if "http" in file:
-                name = self.model_name
-            else:
-                name = modelloader.friendly_name(file)
-            try:
-                scaler_data = modules.upscaler.UpscalerData(name, file, self, 4)
-                scalers.append(scaler_data)
-            except Exception:
-                print(f"Error loading BSRGAN model: {file}", file=sys.stderr)
-                print(traceback.format_exc(), file=sys.stderr)
-        self.scalers = scalers
-
-    def do_upscale(self, img: PIL.Image, selected_file):
-        torch.cuda.empty_cache()
-        model = self.load_model(selected_file)
-        if model is None:
-            return img
-        model.to(devices.device_bsrgan)
-        torch.cuda.empty_cache()
-        img = np.array(img)
-        img = img[:, :, ::-1]
-        img = np.moveaxis(img, 2, 0) / 255
-        img = torch.from_numpy(img).float()
-        img = img.unsqueeze(0).to(devices.device_bsrgan)
-        with torch.no_grad():
-            output = model(img)
-        output = output.squeeze().float().cpu().clamp_(0, 1).numpy()
-        output = 255. * np.moveaxis(output, 0, 2)
-        output = output.astype(np.uint8)
-        output = output[:, :, ::-1]
-        torch.cuda.empty_cache()
-        return PIL.Image.fromarray(output, 'RGB')
-
-    def load_model(self, path: str):
-        if "http" in path:
-            filename = load_file_from_url(url=self.model_url, model_dir=self.model_path, file_name="%s.pth" % self.name,
-                                          progress=True)
-        else:
-            filename = path
-        if not os.path.exists(filename) or filename is None:
-            print(f"BSRGAN: Unable to load model from {filename}", file=sys.stderr)
-            return None
-        model = RRDBNet(in_nc=3, out_nc=3, nf=64, nb=23, gc=32, sf=4)  # define network
-        model.load_state_dict(torch.load(filename), strict=True)
-        model.eval()
-        for k, v in model.named_parameters():
-            v.requires_grad = False
-        return model
-
diff --git a/modules/bsrgan_model_arch.py b/modules/bsrgan_model_arch.py
deleted file mode 100644
index cb4d1c13..00000000
--- a/modules/bsrgan_model_arch.py
+++ /dev/null
@@ -1,102 +0,0 @@
-import functools
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.nn.init as init
-
-
-def initialize_weights(net_l, scale=1):
-    if not isinstance(net_l, list):
-        net_l = [net_l]
-    for net in net_l:
-        for m in net.modules():
-            if isinstance(m, nn.Conv2d):
-                init.kaiming_normal_(m.weight, a=0, mode='fan_in')
-                m.weight.data *= scale  # for residual block
-                if m.bias is not None:
-                    m.bias.data.zero_()
-            elif isinstance(m, nn.Linear):
-                init.kaiming_normal_(m.weight, a=0, mode='fan_in')
-                m.weight.data *= scale
-                if m.bias is not None:
-                    m.bias.data.zero_()
-            elif isinstance(m, nn.BatchNorm2d):
-                init.constant_(m.weight, 1)
-                init.constant_(m.bias.data, 0.0)
-
-
-def make_layer(block, n_layers):
-    layers = []
-    for _ in range(n_layers):
-        layers.append(block())
-    return nn.Sequential(*layers)
-
-
-class ResidualDenseBlock_5C(nn.Module):
-    def __init__(self, nf=64, gc=32, bias=True):
-        super(ResidualDenseBlock_5C, self).__init__()
-        # gc: growth channel, i.e. intermediate channels
-        self.conv1 = nn.Conv2d(nf, gc, 3, 1, 1, bias=bias)
-        self.conv2 = nn.Conv2d(nf + gc, gc, 3, 1, 1, bias=bias)
-        self.conv3 = nn.Conv2d(nf + 2 * gc, gc, 3, 1, 1, bias=bias)
-        self.conv4 = nn.Conv2d(nf + 3 * gc, gc, 3, 1, 1, bias=bias)
-        self.conv5 = nn.Conv2d(nf + 4 * gc, nf, 3, 1, 1, bias=bias)
-        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
-
-        # initialization
-        initialize_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)
-
-    def forward(self, x):
-        x1 = self.lrelu(self.conv1(x))
-        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
-        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
-        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
-        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
-        return x5 * 0.2 + x
-
-
-class RRDB(nn.Module):
-    '''Residual in Residual Dense Block'''
-
-    def __init__(self, nf, gc=32):
-        super(RRDB, self).__init__()
-        self.RDB1 = ResidualDenseBlock_5C(nf, gc)
-        self.RDB2 = ResidualDenseBlock_5C(nf, gc)
-        self.RDB3 = ResidualDenseBlock_5C(nf, gc)
-
-    def forward(self, x):
-        out = self.RDB1(x)
-        out = self.RDB2(out)
-        out = self.RDB3(out)
-        return out * 0.2 + x
-
-
-class RRDBNet(nn.Module):
-    def __init__(self, in_nc=3, out_nc=3, nf=64, nb=23, gc=32, sf=4):
-        super(RRDBNet, self).__init__()
-        RRDB_block_f = functools.partial(RRDB, nf=nf, gc=gc)
-        self.sf = sf
-
-        self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
-        self.RRDB_trunk = make_layer(RRDB_block_f, nb)
-        self.trunk_conv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
-        #### upsampling
-        self.upconv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
-        if self.sf==4:
-            self.upconv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
-        self.HRconv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
-        self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)
-
-        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
-
-    def forward(self, x):
-        fea = self.conv_first(x)
-        trunk = self.trunk_conv(self.RRDB_trunk(fea))
-        fea = fea + trunk
-
-        fea = self.lrelu(self.upconv1(F.interpolate(fea, scale_factor=2, mode='nearest')))
-        if self.sf==4:
-            fea = self.lrelu(self.upconv2(F.interpolate(fea, scale_factor=2, mode='nearest')))
-        out = self.conv_last(self.lrelu(self.HRconv(fea)))
-
-        return out
\ No newline at end of file
diff --git a/modules/deepbooru.py b/modules/deepbooru.py
index 8914662d..8bbc90a4 100644
--- a/modules/deepbooru.py
+++ b/modules/deepbooru.py
@@ -50,11 +50,12 @@ def create_deepbooru_process(threshold, deepbooru_opts):
     the tags.
     """
     from modules import shared  # prevents circular reference
-    shared.deepbooru_process_manager = multiprocessing.Manager()
+    context = multiprocessing.get_context("spawn")
+    shared.deepbooru_process_manager = context.Manager()
     shared.deepbooru_process_queue = shared.deepbooru_process_manager.Queue()
     shared.deepbooru_process_return = shared.deepbooru_process_manager.dict()
     shared.deepbooru_process_return["value"] = -1
-    shared.deepbooru_process = multiprocessing.Process(target=deepbooru_process, args=(shared.deepbooru_process_queue, shared.deepbooru_process_return, threshold, deepbooru_opts))
+    shared.deepbooru_process = context.Process(target=deepbooru_process, args=(shared.deepbooru_process_queue, shared.deepbooru_process_return, threshold, deepbooru_opts))
     shared.deepbooru_process.start()
 
 
diff --git a/modules/devices.py b/modules/devices.py
index eb422583..7511e1dc 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -1,7 +1,6 @@
+import sys, os, shlex
 import contextlib
-
 import torch
-
 from modules import errors
 
 # has_mps is only available in nightly pytorch (for now), `getattr` for compatibility
@@ -9,10 +8,22 @@ has_mps = getattr(torch, 'has_mps', False)
 
 cpu = torch.device("cpu")
 
+def extract_device_id(args, name):
+    for x in range(len(args)):
+        if name in args[x]: return args[x+1]
+    return None
 
 def get_optimal_device():
     if torch.cuda.is_available():
-        return torch.device("cuda")
+        from modules import shared
+
+        device_id = shared.cmd_opts.device_id
+
+        if device_id is not None:
+            cuda_device = f"cuda:{device_id}"
+            return torch.device(cuda_device)
+        else:
+            return torch.device("cuda")
 
     if has_mps:
         return torch.device("mps")
@@ -34,7 +45,7 @@ def enable_tf32():
 
 errors.run(enable_tf32, "Enabling TF32")
 
-device = device_interrogate = device_gfpgan = device_bsrgan = device_esrgan = device_scunet = device_codeformer = get_optimal_device()
+device = device_interrogate = device_gfpgan = device_swinir = device_esrgan = device_scunet = device_codeformer = None
 dtype = torch.float16
 dtype_vae = torch.float16
 
@@ -70,3 +81,7 @@ def autocast(disable=False):
         return contextlib.nullcontext()
 
     return torch.autocast("cuda")
+
+# MPS workaround for https://github.com/pytorch/pytorch/issues/79383
+def mps_contiguous(input_tensor, device): return input_tensor.contiguous() if device.type == 'mps' else input_tensor
+def mps_contiguous_to(input_tensor, device): return mps_contiguous(input_tensor, device).to(device)
diff --git a/modules/esrgan_model.py b/modules/esrgan_model.py
index 46ad0da3..a13cf6ac 100644
--- a/modules/esrgan_model.py
+++ b/modules/esrgan_model.py
@@ -11,62 +11,109 @@ from modules.upscaler import Upscaler, UpscalerData
 from modules.shared import opts

 

 

-def fix_model_layers(crt_model, pretrained_net):

-    # this code is adapted from https://github.com/xinntao/ESRGAN

-    if 'conv_first.weight' in pretrained_net:

-        return pretrained_net

-

-    if 'model.0.weight' not in pretrained_net:

-        is_realesrgan = "params_ema" in pretrained_net and 'body.0.rdb1.conv1.weight' in pretrained_net["params_ema"]

-        if is_realesrgan:

-            raise Exception("The file is a RealESRGAN model, it can't be used as a ESRGAN model.")

-        else:

-            raise Exception("The file is not a ESRGAN model.")

 

-    crt_net = crt_model.state_dict()

-    load_net_clean = {}

-    for k, v in pretrained_net.items():

-        if k.startswith('module.'):

-            load_net_clean[k[7:]] = v

-        else:

-            load_net_clean[k] = v

-    pretrained_net = load_net_clean

-

-    tbd = []

-    for k, v in crt_net.items():

-        tbd.append(k)

-

-    # directly copy

-    for k, v in crt_net.items():

-        if k in pretrained_net and pretrained_net[k].size() == v.size():

-            crt_net[k] = pretrained_net[k]

-            tbd.remove(k)

-

-    crt_net['conv_first.weight'] = pretrained_net['model.0.weight']

-    crt_net['conv_first.bias'] = pretrained_net['model.0.bias']

-

-    for k in tbd.copy():

-        if 'RDB' in k:

-            ori_k = k.replace('RRDB_trunk.', 'model.1.sub.')

-            if '.weight' in k:

-                ori_k = ori_k.replace('.weight', '.0.weight')

-            elif '.bias' in k:

-                ori_k = ori_k.replace('.bias', '.0.bias')

-            crt_net[k] = pretrained_net[ori_k]

-            tbd.remove(k)

-

-    crt_net['trunk_conv.weight'] = pretrained_net['model.1.sub.23.weight']

-    crt_net['trunk_conv.bias'] = pretrained_net['model.1.sub.23.bias']

-    crt_net['upconv1.weight'] = pretrained_net['model.3.weight']

-    crt_net['upconv1.bias'] = pretrained_net['model.3.bias']

-    crt_net['upconv2.weight'] = pretrained_net['model.6.weight']

-    crt_net['upconv2.bias'] = pretrained_net['model.6.bias']

-    crt_net['HRconv.weight'] = pretrained_net['model.8.weight']

-    crt_net['HRconv.bias'] = pretrained_net['model.8.bias']

-    crt_net['conv_last.weight'] = pretrained_net['model.10.weight']

-    crt_net['conv_last.bias'] = pretrained_net['model.10.bias']

-

-    return crt_net

+def mod2normal(state_dict):

+    # this code is copied from https://github.com/victorca25/iNNfer

+    if 'conv_first.weight' in state_dict:

+        crt_net = {}

+        items = []

+        for k, v in state_dict.items():

+            items.append(k)

+

+        crt_net['model.0.weight'] = state_dict['conv_first.weight']

+        crt_net['model.0.bias'] = state_dict['conv_first.bias']

+

+        for k in items.copy():

+            if 'RDB' in k:

+                ori_k = k.replace('RRDB_trunk.', 'model.1.sub.')

+                if '.weight' in k:

+                    ori_k = ori_k.replace('.weight', '.0.weight')

+                elif '.bias' in k:

+                    ori_k = ori_k.replace('.bias', '.0.bias')

+                crt_net[ori_k] = state_dict[k]

+                items.remove(k)

+

+        crt_net['model.1.sub.23.weight'] = state_dict['trunk_conv.weight']

+        crt_net['model.1.sub.23.bias'] = state_dict['trunk_conv.bias']

+        crt_net['model.3.weight'] = state_dict['upconv1.weight']

+        crt_net['model.3.bias'] = state_dict['upconv1.bias']

+        crt_net['model.6.weight'] = state_dict['upconv2.weight']

+        crt_net['model.6.bias'] = state_dict['upconv2.bias']

+        crt_net['model.8.weight'] = state_dict['HRconv.weight']

+        crt_net['model.8.bias'] = state_dict['HRconv.bias']

+        crt_net['model.10.weight'] = state_dict['conv_last.weight']

+        crt_net['model.10.bias'] = state_dict['conv_last.bias']

+        state_dict = crt_net

+    return state_dict

+

+

+def resrgan2normal(state_dict, nb=23):

+    # this code is copied from https://github.com/victorca25/iNNfer

+    if "conv_first.weight" in state_dict and "body.0.rdb1.conv1.weight" in state_dict:

+        crt_net = {}

+        items = []

+        for k, v in state_dict.items():

+            items.append(k)

+

+        crt_net['model.0.weight'] = state_dict['conv_first.weight']

+        crt_net['model.0.bias'] = state_dict['conv_first.bias']

+

+        for k in items.copy():

+            if "rdb" in k:

+                ori_k = k.replace('body.', 'model.1.sub.')

+                ori_k = ori_k.replace('.rdb', '.RDB')

+                if '.weight' in k:

+                    ori_k = ori_k.replace('.weight', '.0.weight')

+                elif '.bias' in k:

+                    ori_k = ori_k.replace('.bias', '.0.bias')

+                crt_net[ori_k] = state_dict[k]

+                items.remove(k)

+

+        crt_net[f'model.1.sub.{nb}.weight'] = state_dict['conv_body.weight']

+        crt_net[f'model.1.sub.{nb}.bias'] = state_dict['conv_body.bias']

+        crt_net['model.3.weight'] = state_dict['conv_up1.weight']

+        crt_net['model.3.bias'] = state_dict['conv_up1.bias']

+        crt_net['model.6.weight'] = state_dict['conv_up2.weight']

+        crt_net['model.6.bias'] = state_dict['conv_up2.bias']

+        crt_net['model.8.weight'] = state_dict['conv_hr.weight']

+        crt_net['model.8.bias'] = state_dict['conv_hr.bias']

+        crt_net['model.10.weight'] = state_dict['conv_last.weight']

+        crt_net['model.10.bias'] = state_dict['conv_last.bias']

+        state_dict = crt_net

+    return state_dict

+

+

+def infer_params(state_dict):

+    # this code is copied from https://github.com/victorca25/iNNfer

+    scale2x = 0

+    scalemin = 6

+    n_uplayer = 0

+    plus = False

+

+    for block in list(state_dict):

+        parts = block.split(".")

+        n_parts = len(parts)

+        if n_parts == 5 and parts[2] == "sub":

+            nb = int(parts[3])

+        elif n_parts == 3:

+            part_num = int(parts[1])

+            if (part_num > scalemin

+                and parts[0] == "model"

+                and parts[2] == "weight"):

+                scale2x += 1

+            if part_num > n_uplayer:

+                n_uplayer = part_num

+                out_nc = state_dict[block].shape[0]

+        if not plus and "conv1x1" in block:

+            plus = True

+

+    nf = state_dict["model.0.weight"].shape[0]

+    in_nc = state_dict["model.0.weight"].shape[1]

+    out_nc = out_nc

+    scale = 2 ** scale2x

+

+    return in_nc, out_nc, nf, nb, plus, scale

+

 

 class UpscalerESRGAN(Upscaler):

     def __init__(self, dirname):

@@ -109,22 +156,41 @@ class UpscalerESRGAN(Upscaler):
             print("Unable to load %s from %s" % (self.model_path, filename))

             return None

 

-        pretrained_net = torch.load(filename, map_location='cpu' if devices.device_esrgan.type == 'mps' else None)

-        crt_model = arch.RRDBNet(3, 3, 64, 23, gc=32)

+        state_dict = torch.load(filename, map_location='cpu' if devices.device_esrgan.type == 'mps' else None)

+

+        if "params_ema" in state_dict:

+            state_dict = state_dict["params_ema"]

+        elif "params" in state_dict:

+            state_dict = state_dict["params"]

+            num_conv = 16 if "realesr-animevideov3" in filename else 32

+            model = arch.SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=num_conv, upscale=4, act_type='prelu')

+            model.load_state_dict(state_dict)

+            model.eval()

+            return model

+

+        if "body.0.rdb1.conv1.weight" in state_dict and "conv_first.weight" in state_dict:

+            nb = 6 if "RealESRGAN_x4plus_anime_6B" in filename else 23

+            state_dict = resrgan2normal(state_dict, nb)

+        elif "conv_first.weight" in state_dict:

+            state_dict = mod2normal(state_dict)

+        elif "model.0.weight" not in state_dict:

+            raise Exception("The file is not a recognized ESRGAN model.")

+

+        in_nc, out_nc, nf, nb, plus, mscale = infer_params(state_dict)

 

-        pretrained_net = fix_model_layers(crt_model, pretrained_net)

-        crt_model.load_state_dict(pretrained_net)

-        crt_model.eval()

+        model = arch.RRDBNet(in_nc=in_nc, out_nc=out_nc, nf=nf, nb=nb, upscale=mscale, plus=plus)

+        model.load_state_dict(state_dict)

+        model.eval()

 

-        return crt_model

+        return model

 

 

 def upscale_without_tiling(model, img):

     img = np.array(img)

     img = img[:, :, ::-1]

-    img = np.moveaxis(img, 2, 0) / 255

+    img = np.ascontiguousarray(np.transpose(img, (2, 0, 1))) / 255

     img = torch.from_numpy(img).float()

-    img = img.unsqueeze(0).to(devices.device_esrgan)

+    img = devices.mps_contiguous_to(img.unsqueeze(0), devices.device_esrgan)

     with torch.no_grad():

         output = model(img)

     output = output.squeeze().float().cpu().clamp_(0, 1).numpy()

diff --git a/modules/esrgan_model_arch.py b/modules/esrgan_model_arch.py
index e413d36e..bc9ceb2a 100644
--- a/modules/esrgan_model_arch.py
+++ b/modules/esrgan_model_arch.py
@@ -1,80 +1,463 @@
-# this file is taken from https://github.com/xinntao/ESRGAN

+# this file is adapted from https://github.com/victorca25/iNNfer

 

+import math

 import functools

 import torch

 import torch.nn as nn

 import torch.nn.functional as F

 

 

-def make_layer(block, n_layers):

-    layers = []

-    for _ in range(n_layers):

-        layers.append(block())

-    return nn.Sequential(*layers)

+####################

+# RRDBNet Generator

+####################

 

+class RRDBNet(nn.Module):

+    def __init__(self, in_nc, out_nc, nf, nb, nr=3, gc=32, upscale=4, norm_type=None,

+            act_type='leakyrelu', mode='CNA', upsample_mode='upconv', convtype='Conv2D',

+            finalact=None, gaussian_noise=False, plus=False):

+        super(RRDBNet, self).__init__()

+        n_upscale = int(math.log(upscale, 2))

+        if upscale == 3:

+            n_upscale = 1

 

-class ResidualDenseBlock_5C(nn.Module):

-    def __init__(self, nf=64, gc=32, bias=True):

-        super(ResidualDenseBlock_5C, self).__init__()

-        # gc: growth channel, i.e. intermediate channels

-        self.conv1 = nn.Conv2d(nf, gc, 3, 1, 1, bias=bias)

-        self.conv2 = nn.Conv2d(nf + gc, gc, 3, 1, 1, bias=bias)

-        self.conv3 = nn.Conv2d(nf + 2 * gc, gc, 3, 1, 1, bias=bias)

-        self.conv4 = nn.Conv2d(nf + 3 * gc, gc, 3, 1, 1, bias=bias)

-        self.conv5 = nn.Conv2d(nf + 4 * gc, nf, 3, 1, 1, bias=bias)

-        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)

+        self.resrgan_scale = 0

+        if in_nc % 16 == 0:

+            self.resrgan_scale = 1

+        elif in_nc != 4 and in_nc % 4 == 0:

+            self.resrgan_scale = 2

 

-        # initialization

-        # mutil.initialize_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)

+        fea_conv = conv_block(in_nc, nf, kernel_size=3, norm_type=None, act_type=None, convtype=convtype)

+        rb_blocks = [RRDB(nf, nr, kernel_size=3, gc=32, stride=1, bias=1, pad_type='zero',

+            norm_type=norm_type, act_type=act_type, mode='CNA', convtype=convtype,

+            gaussian_noise=gaussian_noise, plus=plus) for _ in range(nb)]

+        LR_conv = conv_block(nf, nf, kernel_size=3, norm_type=norm_type, act_type=None, mode=mode, convtype=convtype)

 

-    def forward(self, x):

-        x1 = self.lrelu(self.conv1(x))

-        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))

-        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))

-        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))

-        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))

-        return x5 * 0.2 + x

+        if upsample_mode == 'upconv':

+            upsample_block = upconv_block

+        elif upsample_mode == 'pixelshuffle':

+            upsample_block = pixelshuffle_block

+        else:

+            raise NotImplementedError('upsample mode [{:s}] is not found'.format(upsample_mode))

+        if upscale == 3:

+            upsampler = upsample_block(nf, nf, 3, act_type=act_type, convtype=convtype)

+        else:

+            upsampler = [upsample_block(nf, nf, act_type=act_type, convtype=convtype) for _ in range(n_upscale)]

+        HR_conv0 = conv_block(nf, nf, kernel_size=3, norm_type=None, act_type=act_type, convtype=convtype)

+        HR_conv1 = conv_block(nf, out_nc, kernel_size=3, norm_type=None, act_type=None, convtype=convtype)

+

+        outact = act(finalact) if finalact else None

+

+        self.model = sequential(fea_conv, ShortcutBlock(sequential(*rb_blocks, LR_conv)),

+            *upsampler, HR_conv0, HR_conv1, outact)

+

+    def forward(self, x, outm=None):

+        if self.resrgan_scale == 1:

+            feat = pixel_unshuffle(x, scale=4)

+        elif self.resrgan_scale == 2:

+            feat = pixel_unshuffle(x, scale=2)

+        else:

+            feat = x

+

+        return self.model(feat)

 

 

 class RRDB(nn.Module):

-    '''Residual in Residual Dense Block'''

+    """

+    Residual in Residual Dense Block

+    (ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks)

+    """

 

-    def __init__(self, nf, gc=32):

+    def __init__(self, nf, nr=3, kernel_size=3, gc=32, stride=1, bias=1, pad_type='zero',

+            norm_type=None, act_type='leakyrelu', mode='CNA', convtype='Conv2D',

+            spectral_norm=False, gaussian_noise=False, plus=False):

         super(RRDB, self).__init__()

-        self.RDB1 = ResidualDenseBlock_5C(nf, gc)

-        self.RDB2 = ResidualDenseBlock_5C(nf, gc)

-        self.RDB3 = ResidualDenseBlock_5C(nf, gc)

+        # This is for backwards compatibility with existing models

+        if nr == 3:

+            self.RDB1 = ResidualDenseBlock_5C(nf, kernel_size, gc, stride, bias, pad_type,

+                    norm_type, act_type, mode, convtype, spectral_norm=spectral_norm,

+                    gaussian_noise=gaussian_noise, plus=plus)

+            self.RDB2 = ResidualDenseBlock_5C(nf, kernel_size, gc, stride, bias, pad_type,

+                    norm_type, act_type, mode, convtype, spectral_norm=spectral_norm,

+                    gaussian_noise=gaussian_noise, plus=plus)

+            self.RDB3 = ResidualDenseBlock_5C(nf, kernel_size, gc, stride, bias, pad_type,

+                    norm_type, act_type, mode, convtype, spectral_norm=spectral_norm,

+                    gaussian_noise=gaussian_noise, plus=plus)

+        else:

+            RDB_list = [ResidualDenseBlock_5C(nf, kernel_size, gc, stride, bias, pad_type,

+                                              norm_type, act_type, mode, convtype, spectral_norm=spectral_norm,

+                                              gaussian_noise=gaussian_noise, plus=plus) for _ in range(nr)]

+            self.RDBs = nn.Sequential(*RDB_list)

 

     def forward(self, x):

-        out = self.RDB1(x)

-        out = self.RDB2(out)

-        out = self.RDB3(out)

+        if hasattr(self, 'RDB1'):

+            out = self.RDB1(x)

+            out = self.RDB2(out)

+            out = self.RDB3(out)

+        else:

+            out = self.RDBs(x)

         return out * 0.2 + x

 

 

-class RRDBNet(nn.Module):

-    def __init__(self, in_nc, out_nc, nf, nb, gc=32):

-        super(RRDBNet, self).__init__()

-        RRDB_block_f = functools.partial(RRDB, nf=nf, gc=gc)

+class ResidualDenseBlock_5C(nn.Module):

+    """

+    Residual Dense Block

+    The core module of paper: (Residual Dense Network for Image Super-Resolution, CVPR 18)

+    Modified options that can be used:

+        - "Partial Convolution based Padding" arXiv:1811.11718

+        - "Spectral normalization" arXiv:1802.05957

+        - "ICASSP 2020 - ESRGAN+ : Further Improving ESRGAN" N. C. 

+            {Rakotonirina} and A. {Rasoanaivo}

+    """

+

+    def __init__(self, nf=64, kernel_size=3, gc=32, stride=1, bias=1, pad_type='zero',

+            norm_type=None, act_type='leakyrelu', mode='CNA', convtype='Conv2D',

+            spectral_norm=False, gaussian_noise=False, plus=False):

+        super(ResidualDenseBlock_5C, self).__init__()

 

-        self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)

-        self.RRDB_trunk = make_layer(RRDB_block_f, nb)

-        self.trunk_conv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)

-        #### upsampling

-        self.upconv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)

-        self.upconv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)

-        self.HRconv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)

-        self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)

+        self.noise = GaussianNoise() if gaussian_noise else None

+        self.conv1x1 = conv1x1(nf, gc) if plus else None

 

-        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)

+        self.conv1 = conv_block(nf, gc, kernel_size, stride, bias=bias, pad_type=pad_type,

+            norm_type=norm_type, act_type=act_type, mode=mode, convtype=convtype,

+            spectral_norm=spectral_norm)

+        self.conv2 = conv_block(nf+gc, gc, kernel_size, stride, bias=bias, pad_type=pad_type,

+            norm_type=norm_type, act_type=act_type, mode=mode, convtype=convtype,

+            spectral_norm=spectral_norm)

+        self.conv3 = conv_block(nf+2*gc, gc, kernel_size, stride, bias=bias, pad_type=pad_type,

+            norm_type=norm_type, act_type=act_type, mode=mode, convtype=convtype,

+            spectral_norm=spectral_norm)

+        self.conv4 = conv_block(nf+3*gc, gc, kernel_size, stride, bias=bias, pad_type=pad_type,

+            norm_type=norm_type, act_type=act_type, mode=mode, convtype=convtype,

+            spectral_norm=spectral_norm)

+        if mode == 'CNA':

+            last_act = None

+        else:

+            last_act = act_type

+        self.conv5 = conv_block(nf+4*gc, nf, 3, stride, bias=bias, pad_type=pad_type,

+            norm_type=norm_type, act_type=last_act, mode=mode, convtype=convtype,

+            spectral_norm=spectral_norm)

 

     def forward(self, x):

-        fea = self.conv_first(x)

-        trunk = self.trunk_conv(self.RRDB_trunk(fea))

-        fea = fea + trunk

+        x1 = self.conv1(x)

+        x2 = self.conv2(torch.cat((x, x1), 1))

+        if self.conv1x1:

+            x2 = x2 + self.conv1x1(x)

+        x3 = self.conv3(torch.cat((x, x1, x2), 1))

+        x4 = self.conv4(torch.cat((x, x1, x2, x3), 1))

+        if self.conv1x1:

+            x4 = x4 + x2

+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))

+        if self.noise:

+            return self.noise(x5.mul(0.2) + x)

+        else:

+            return x5 * 0.2 + x

+

 

-        fea = self.lrelu(self.upconv1(F.interpolate(fea, scale_factor=2, mode='nearest')))

-        fea = self.lrelu(self.upconv2(F.interpolate(fea, scale_factor=2, mode='nearest')))

-        out = self.conv_last(self.lrelu(self.HRconv(fea)))

+####################

+# ESRGANplus

+####################

 

+class GaussianNoise(nn.Module):

+    def __init__(self, sigma=0.1, is_relative_detach=False):

+        super().__init__()

+        self.sigma = sigma

+        self.is_relative_detach = is_relative_detach

+        self.noise = torch.tensor(0, dtype=torch.float)

+

+    def forward(self, x):

+        if self.training and self.sigma != 0:

+            self.noise = self.noise.to(x.device)

+            scale = self.sigma * x.detach() if self.is_relative_detach else self.sigma * x

+            sampled_noise = self.noise.repeat(*x.size()).normal_() * scale

+            x = x + sampled_noise

+        return x 

+

+def conv1x1(in_planes, out_planes, stride=1):

+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)

+

+

+####################

+# SRVGGNetCompact

+####################

+

+class SRVGGNetCompact(nn.Module):

+    """A compact VGG-style network structure for super-resolution.

+    This class is copied from https://github.com/xinntao/Real-ESRGAN

+    """

+

+    def __init__(self, num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu'):

+        super(SRVGGNetCompact, self).__init__()

+        self.num_in_ch = num_in_ch

+        self.num_out_ch = num_out_ch

+        self.num_feat = num_feat

+        self.num_conv = num_conv

+        self.upscale = upscale

+        self.act_type = act_type

+

+        self.body = nn.ModuleList()

+        # the first conv

+        self.body.append(nn.Conv2d(num_in_ch, num_feat, 3, 1, 1))

+        # the first activation

+        if act_type == 'relu':

+            activation = nn.ReLU(inplace=True)

+        elif act_type == 'prelu':

+            activation = nn.PReLU(num_parameters=num_feat)

+        elif act_type == 'leakyrelu':

+            activation = nn.LeakyReLU(negative_slope=0.1, inplace=True)

+        self.body.append(activation)

+

+        # the body structure

+        for _ in range(num_conv):

+            self.body.append(nn.Conv2d(num_feat, num_feat, 3, 1, 1))

+            # activation

+            if act_type == 'relu':

+                activation = nn.ReLU(inplace=True)

+            elif act_type == 'prelu':

+                activation = nn.PReLU(num_parameters=num_feat)

+            elif act_type == 'leakyrelu':

+                activation = nn.LeakyReLU(negative_slope=0.1, inplace=True)

+            self.body.append(activation)

+

+        # the last conv

+        self.body.append(nn.Conv2d(num_feat, num_out_ch * upscale * upscale, 3, 1, 1))

+        # upsample

+        self.upsampler = nn.PixelShuffle(upscale)

+

+    def forward(self, x):

+        out = x

+        for i in range(0, len(self.body)):

+            out = self.body[i](out)

+

+        out = self.upsampler(out)

+        # add the nearest upsampled image, so that the network learns the residual

+        base = F.interpolate(x, scale_factor=self.upscale, mode='nearest')

+        out += base

         return out

+

+

+####################

+# Upsampler

+####################

+

+class Upsample(nn.Module):

+    r"""Upsamples a given multi-channel 1D (temporal), 2D (spatial) or 3D (volumetric) data.

+    The input data is assumed to be of the form

+    `minibatch x channels x [optional depth] x [optional height] x width`.

+    """

+

+    def __init__(self, size=None, scale_factor=None, mode="nearest", align_corners=None):

+        super(Upsample, self).__init__()

+        if isinstance(scale_factor, tuple):

+            self.scale_factor = tuple(float(factor) for factor in scale_factor)

+        else:

+            self.scale_factor = float(scale_factor) if scale_factor else None

+        self.mode = mode

+        self.size = size

+        self.align_corners = align_corners

+

+    def forward(self, x):

+        return nn.functional.interpolate(x, size=self.size, scale_factor=self.scale_factor, mode=self.mode, align_corners=self.align_corners)

+

+    def extra_repr(self):

+        if self.scale_factor is not None:

+            info = 'scale_factor=' + str(self.scale_factor)

+        else:

+            info = 'size=' + str(self.size)

+        info += ', mode=' + self.mode

+        return info

+

+

+def pixel_unshuffle(x, scale):

+    """ Pixel unshuffle.

+    Args:

+        x (Tensor): Input feature with shape (b, c, hh, hw).

+        scale (int): Downsample ratio.

+    Returns:

+        Tensor: the pixel unshuffled feature.

+    """

+    b, c, hh, hw = x.size()

+    out_channel = c * (scale**2)

+    assert hh % scale == 0 and hw % scale == 0

+    h = hh // scale

+    w = hw // scale

+    x_view = x.view(b, c, h, scale, w, scale)

+    return x_view.permute(0, 1, 3, 5, 2, 4).reshape(b, out_channel, h, w)

+

+

+def pixelshuffle_block(in_nc, out_nc, upscale_factor=2, kernel_size=3, stride=1, bias=True,

+                        pad_type='zero', norm_type=None, act_type='relu', convtype='Conv2D'):

+    """

+    Pixel shuffle layer

+    (Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional

+    Neural Network, CVPR17)

+    """

+    conv = conv_block(in_nc, out_nc * (upscale_factor ** 2), kernel_size, stride, bias=bias,

+                        pad_type=pad_type, norm_type=None, act_type=None, convtype=convtype)

+    pixel_shuffle = nn.PixelShuffle(upscale_factor)

+

+    n = norm(norm_type, out_nc) if norm_type else None

+    a = act(act_type) if act_type else None

+    return sequential(conv, pixel_shuffle, n, a)

+

+

+def upconv_block(in_nc, out_nc, upscale_factor=2, kernel_size=3, stride=1, bias=True,

+                pad_type='zero', norm_type=None, act_type='relu', mode='nearest', convtype='Conv2D'):

+    """ Upconv layer """

+    upscale_factor = (1, upscale_factor, upscale_factor) if convtype == 'Conv3D' else upscale_factor

+    upsample = Upsample(scale_factor=upscale_factor, mode=mode)

+    conv = conv_block(in_nc, out_nc, kernel_size, stride, bias=bias,

+                        pad_type=pad_type, norm_type=norm_type, act_type=act_type, convtype=convtype)

+    return sequential(upsample, conv)

+

+

+

+

+

+

+

+

+####################

+# Basic blocks

+####################

+

+

+def make_layer(basic_block, num_basic_block, **kwarg):

+    """Make layers by stacking the same blocks.

+    Args:

+        basic_block (nn.module): nn.module class for basic block. (block)

+        num_basic_block (int): number of blocks. (n_layers)

+    Returns:

+        nn.Sequential: Stacked blocks in nn.Sequential.

+    """

+    layers = []

+    for _ in range(num_basic_block):

+        layers.append(basic_block(**kwarg))

+    return nn.Sequential(*layers)

+

+

+def act(act_type, inplace=True, neg_slope=0.2, n_prelu=1, beta=1.0):

+    """ activation helper """

+    act_type = act_type.lower()

+    if act_type == 'relu':

+        layer = nn.ReLU(inplace)

+    elif act_type in ('leakyrelu', 'lrelu'):

+        layer = nn.LeakyReLU(neg_slope, inplace)

+    elif act_type == 'prelu':

+        layer = nn.PReLU(num_parameters=n_prelu, init=neg_slope)

+    elif act_type == 'tanh':  # [-1, 1] range output

+        layer = nn.Tanh()

+    elif act_type == 'sigmoid':  # [0, 1] range output

+        layer = nn.Sigmoid()

+    else:

+        raise NotImplementedError('activation layer [{:s}] is not found'.format(act_type))

+    return layer

+

+

+class Identity(nn.Module):

+    def __init__(self, *kwargs):

+        super(Identity, self).__init__()

+

+    def forward(self, x, *kwargs):

+        return x

+

+

+def norm(norm_type, nc):

+    """ Return a normalization layer """

+    norm_type = norm_type.lower()

+    if norm_type == 'batch':

+        layer = nn.BatchNorm2d(nc, affine=True)

+    elif norm_type == 'instance':

+        layer = nn.InstanceNorm2d(nc, affine=False)

+    elif norm_type == 'none':

+        def norm_layer(x): return Identity()

+    else:

+        raise NotImplementedError('normalization layer [{:s}] is not found'.format(norm_type))

+    return layer

+

+

+def pad(pad_type, padding):

+    """ padding layer helper """

+    pad_type = pad_type.lower()

+    if padding == 0:

+        return None

+    if pad_type == 'reflect':

+        layer = nn.ReflectionPad2d(padding)

+    elif pad_type == 'replicate':

+        layer = nn.ReplicationPad2d(padding)

+    elif pad_type == 'zero':

+        layer = nn.ZeroPad2d(padding)

+    else:

+        raise NotImplementedError('padding layer [{:s}] is not implemented'.format(pad_type))

+    return layer

+

+

+def get_valid_padding(kernel_size, dilation):

+    kernel_size = kernel_size + (kernel_size - 1) * (dilation - 1)

+    padding = (kernel_size - 1) // 2

+    return padding

+

+

+class ShortcutBlock(nn.Module):

+    """ Elementwise sum the output of a submodule to its input """

+    def __init__(self, submodule):

+        super(ShortcutBlock, self).__init__()

+        self.sub = submodule

+

+    def forward(self, x):

+        output = x + self.sub(x)

+        return output

+

+    def __repr__(self):

+        return 'Identity + \n|' + self.sub.__repr__().replace('\n', '\n|')

+

+

+def sequential(*args):

+    """ Flatten Sequential. It unwraps nn.Sequential. """

+    if len(args) == 1:

+        if isinstance(args[0], OrderedDict):

+            raise NotImplementedError('sequential does not support OrderedDict input.')

+        return args[0]  # No sequential is needed.

+    modules = []

+    for module in args:

+        if isinstance(module, nn.Sequential):

+            for submodule in module.children():

+                modules.append(submodule)

+        elif isinstance(module, nn.Module):

+            modules.append(module)

+    return nn.Sequential(*modules)

+

+

+def conv_block(in_nc, out_nc, kernel_size, stride=1, dilation=1, groups=1, bias=True,

+               pad_type='zero', norm_type=None, act_type='relu', mode='CNA', convtype='Conv2D',

+               spectral_norm=False):

+    """ Conv layer with padding, normalization, activation """

+    assert mode in ['CNA', 'NAC', 'CNAC'], 'Wrong conv mode [{:s}]'.format(mode)

+    padding = get_valid_padding(kernel_size, dilation)

+    p = pad(pad_type, padding) if pad_type and pad_type != 'zero' else None

+    padding = padding if pad_type == 'zero' else 0

+

+    if convtype=='PartialConv2D':

+        c = PartialConv2d(in_nc, out_nc, kernel_size=kernel_size, stride=stride, padding=padding,

+               dilation=dilation, bias=bias, groups=groups)

+    elif convtype=='DeformConv2D':

+        c = DeformConv2d(in_nc, out_nc, kernel_size=kernel_size, stride=stride, padding=padding,

+               dilation=dilation, bias=bias, groups=groups)

+    elif convtype=='Conv3D':

+        c = nn.Conv3d(in_nc, out_nc, kernel_size=kernel_size, stride=stride, padding=padding,

+                dilation=dilation, bias=bias, groups=groups)

+    else:

+        c = nn.Conv2d(in_nc, out_nc, kernel_size=kernel_size, stride=stride, padding=padding,

+                dilation=dilation, bias=bias, groups=groups)

+

+    if spectral_norm:

+        c = nn.utils.spectral_norm(c)

+

+    a = act(act_type) if act_type else None

+    if 'CNA' in mode:

+        n = norm(norm_type, out_nc) if norm_type else None

+        return sequential(p, c, n, a)

+    elif mode == 'NAC':

+        if norm_type is None and act_type is not None:

+            a = act(act_type, inplace=False)

+        n = norm(norm_type, in_nc) if norm_type else None

+        return sequential(n, a, p, c)

diff --git a/modules/extensions.py b/modules/extensions.py
new file mode 100644
index 00000000..897af96e
--- /dev/null
+++ b/modules/extensions.py
@@ -0,0 +1,83 @@
+import os

+import sys

+import traceback

+

+import git

+

+from modules import paths, shared

+

+

+extensions = []

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

+

+

+def active():

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

+

+

+class Extension:

+    def __init__(self, name, path, enabled=True):

+        self.name = name

+        self.path = path

+        self.enabled = enabled

+        self.status = ''

+        self.can_update = False

+

+        repo = None

+        try:

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

+                repo = git.Repo(path)

+        except Exception:

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

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

+

+        if repo is None or repo.bare:

+            self.remote = None

+        else:

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

+            self.status = 'unknown'

+

+    def list_files(self, subdir, extension):

+        from modules import scripts

+

+        dirpath = os.path.join(self.path, subdir)

+        if not os.path.isdir(dirpath):

+            return []

+

+        res = []

+        for filename in sorted(os.listdir(dirpath)):

+            res.append(scripts.ScriptFile(self.path, filename, os.path.join(dirpath, filename)))

+

+        res = [x for x in res if os.path.splitext(x.path)[1].lower() == extension and os.path.isfile(x.path)]

+

+        return res

+

+    def check_updates(self):

+        repo = git.Repo(self.path)

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

+            if fetch.flags != fetch.HEAD_UPTODATE:

+                self.can_update = True

+                self.status = "behind"

+                return

+

+        self.can_update = False

+        self.status = "latest"

+

+    def pull(self):

+        repo = git.Repo(self.path)

+        repo.remotes.origin.pull()

+

+

+def list_extensions():

+    extensions.clear()

+

+    if not os.path.isdir(extensions_dir):

+        return

+

+    for dirname in sorted(os.listdir(extensions_dir)):

+        path = os.path.join(extensions_dir, dirname)

+        if not os.path.isdir(path):

+            continue

+

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

+        extensions.append(extension)

diff --git a/modules/extras.py b/modules/extras.py
index 22c5a1c1..4d51088b 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -1,3 +1,4 @@
+from __future__ import annotations

 import math

 import os

 

@@ -7,6 +8,10 @@ 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

 from modules.shared import opts

 import modules.gfpgan_model

@@ -17,10 +22,38 @@ import piexif.helper
 import gradio as gr

 

 

-cached_images = {}

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

+

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

     devices.torch_gc()

 

     imageArr = []

@@ -39,7 +72,7 @@ def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_
 

         if input_dir == '':

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

-        image_list = [file for file in [os.path.join(input_dir, x) for x in sorted(os.listdir(input_dir))] if os.path.isfile(file)]

+        image_list = shared.listfiles(input_dir)

         for img in image_list:

             try:

                 image = Image.open(img)

@@ -56,72 +89,102 @@ def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_
     else:

         outpath = opts.outdir_samples or opts.outdir_extras_samples

 

-    

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

-        if image is None:

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

-        existing_pnginfo = image.info or {}

+    # Extra operation definitions

 

-        image = image.convert("RGB")

-        info = ""

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

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

 

-            if gfpgan_visibility < 1.0:

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

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

+        return (res, info)

 

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

-            image = res

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

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

 

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

 

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

-            image = res

+    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

+        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=hash(np.array(image.getdata()).tobytes()),

+                                     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 upscaling_resize != 1.0:

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

-                small = image.crop((image.width // 2, image.height // 2, image.width // 2 + 10, image.height // 2 + 10))

-                pixels = tuple(np.array(small).flatten().tolist())

-                key = (resize, scaler_index, image.width, image.height, gfpgan_visibility, codeformer_visibility, codeformer_weight, 

-                       resize_mode, upscaling_resize, upscaling_resize_w, upscaling_resize_h, upscaling_crop) + pixels

-

-                c = cached_images.get(key)

-                if c is None:

-                    upscaler = shared.sd_upscalers[scaler_index]

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

-                    if mode == 1 and crop:

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

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

-                        c = cropped

-                    cached_images[key] = c

-

-                return c

-

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

-            res = upscale(image, extras_upscaler_1, upscaling_resize, resize_mode, upscaling_resize_w, upscaling_resize_h, upscaling_crop)

-

-            if extras_upscaler_2 != 0 and extras_upscaler_2_visibility > 0:

-                res2 = upscale(image, extras_upscaler_2, upscaling_resize, resize_mode, upscaling_resize_w, upscaling_resize_h, upscaling_crop)

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

-                res = Image.blend(res, res2, extras_upscaler_2_visibility)

-

-            image = res

-

-        while len(cached_images) > 2:

-            del cached_images[next(iter(cached_images.keys()))]

-        

         if opts.use_original_name_batch and image_name != None:

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

         else:

@@ -141,6 +204,9 @@ def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_
 

     return outputs, plaintext_to_html(info), ''

 

+def clear_cache():

+    cached_images.clear()

+

 

 def run_pnginfo(image):

     if image is None:

diff --git a/modules/generation_parameters_copypaste.py b/modules/generation_parameters_copypaste.py
index f73647da..985ec95e 100644
--- a/modules/generation_parameters_copypaste.py
+++ b/modules/generation_parameters_copypaste.py
@@ -1,14 +1,25 @@
+import base64

+import io

 import os

 import re

 import gradio as gr

 from modules.shared import script_path

 from modules import shared

+import tempfile

+from PIL import Image

 

 re_param_code = r'\s*([\w ]+):\s*("(?:\\|\"|[^\"])+"|[^,]*)(?:,|$)'

 re_param = re.compile(re_param_code)

 re_params = re.compile(r"^(?:" + re_param_code + "){3,}$")

 re_imagesize = re.compile(r"^(\d+)x(\d+)$")

 type_of_gr_update = type(gr.update())

+paste_fields = {}

+bind_list = []

+

+

+def reset():

+    paste_fields.clear()

+    bind_list.clear()

 

 

 def quote(text):

@@ -20,6 +31,111 @@ def quote(text):
     text = text.replace('"', '\\"')

     return f'"{text}"'

 

+

+def image_from_url_text(filedata):

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

+        filename = filedata["name"]

+        tempdir = os.path.normpath(tempfile.gettempdir())

+        normfn = os.path.normpath(filename)

+        assert normfn.startswith(tempdir), 'trying to open image file not in temporary directory'

+

+        return Image.open(filename)

+

+    if type(filedata) == list:

+        if len(filedata) == 0:

+            return None

+

+        filedata = filedata[0]

+

+    if filedata.startswith("data:image/png;base64,"):

+        filedata = filedata[len("data:image/png;base64,"):]

+

+    filedata = base64.decodebytes(filedata.encode('utf-8'))

+    image = Image.open(io.BytesIO(filedata))

+    return image

+

+

+def add_paste_fields(tabname, init_img, fields):

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

+

+    # backwards compatibility for existing extensions

+    import modules.ui

+    if tabname == 'txt2img':

+        modules.ui.txt2img_paste_fields = fields

+    elif tabname == 'img2img':

+        modules.ui.img2img_paste_fields = fields

+

+

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

+        'sd_model_checkpoint': 'Model hash',

+    }

+    settings_paste_fields = [

+        (component_dict[k], lambda d, k=k, v=v: ui.apply_setting(k, d.get(v, None)))

+        for k, v in settings_map.items()

+    ]

+

+    for tabname, info in paste_fields.items():

+        if info["fields"] is not None:

+            info["fields"] += settings_paste_fields

+

+

+def create_buttons(tabs_list):

+    buttons = {}

+    for tab in tabs_list:

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

+    return buttons

+

+

+#if send_generate_info is a tab name, mean generate_info comes from the params fields of the tab

+def bind_buttons(buttons, send_image, send_generate_info):

+    bind_list.append([buttons, send_image, send_generate_info])

+

+

+def run_bind():

+    for buttons, send_image, 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"]],

+                    )

+                else:

+                    button.click(

+                        fn=lambda x: x,

+                        inputs=[send_image],

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

+                    )

+

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

+                if send_generate_info in paste_fields:

+                    paste_field_names = ['Prompt', 'Negative prompt', 'Steps', 'Face restoration', 'Size-1', 'Size-2'] + (["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],

+                    )

+                else:

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

+

+            button.click(

+                fn=None,

+                _js=f"switch_to_{tab}",

+                inputs=None,

+                outputs=None,

+            )

+

+

 def parse_generation_parameters(x: str):

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

 ```

@@ -68,7 +184,7 @@ Steps: 20, Sampler: Euler a, CFG scale: 7, Seed: 965400086, Size: 512x512, Model
     return res

 

 

-def connect_paste(button, paste_fields, input_comp, js=None):

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

@@ -106,7 +222,9 @@ def connect_paste(button, paste_fields, input_comp, js=None):
 

     button.click(

         fn=paste_func,

-        _js=js,

+        _js=jsfunc,

         inputs=[input_comp],

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

     )

+

+

diff --git a/modules/hypernetworks/hypernetwork.py b/modules/hypernetworks/hypernetwork.py
index 47d91ea5..a11e01d6 100644
--- a/modules/hypernetworks/hypernetwork.py
+++ b/modules/hypernetworks/hypernetwork.py
@@ -1,28 +1,41 @@
+import csv

 import datetime

 import glob

 import html

 import os

 import sys

 import traceback

-import tqdm

-import csv

-

-import torch

+import inspect

 

-from ldm.util import default

-from modules import devices, shared, processing, sd_models

+import modules.textual_inversion.dataset

 import torch

-from torch import einsum

+import tqdm

 from einops import rearrange, repeat

-import modules.textual_inversion.dataset

+from ldm.util import default

+from modules import devices, processing, sd_models, shared

 from modules.textual_inversion import textual_inversion

 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_

+

+from collections import defaultdict, deque

+from statistics import stdev, mean

 

 

 class HypernetworkModule(torch.nn.Module):

     multiplier = 1.0

-

-    def __init__(self, dim, state_dict=None, layer_structure=None, add_layer_norm=False, activation_func=None):

+    activation_dict = {

+        "linear": torch.nn.Identity,

+        "relu": torch.nn.ReLU,

+        "leakyrelu": torch.nn.LeakyReLU,

+        "elu": torch.nn.ELU,

+        "swish": torch.nn.Hardswish,

+        "tanh": torch.nn.Tanh,

+        "sigmoid": torch.nn.Sigmoid,

+    }

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

         super().__init__()

 

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

@@ -31,20 +44,26 @@ class HypernetworkModule(torch.nn.Module):
 

         linears = []

         for i in range(len(layer_structure) - 1):

+

+            # Add a fully-connected layer

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

 

-            if activation_func == "relu":

-                linears.append(torch.nn.ReLU())

-            elif activation_func == "leakyrelu":

-                linears.append(torch.nn.LeakyReLU())

-            elif activation_func == 'linear' or activation_func is None:

+            # Add an activation func

+            if activation_func == "linear" or activation_func is None:

                 pass

+            elif activation_func in self.activation_dict:

+                linears.append(self.activation_dict[activation_func]())

             else:

                 raise RuntimeError(f'hypernetwork uses an unsupported activation function: {activation_func}')

 

+            # Add layer normalization

             if add_layer_norm:

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

 

+            # Add dropout expect last layer

+            if use_dropout and i < len(layer_structure) - 3:

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

+

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

 

         if state_dict is not None:

@@ -53,9 +72,24 @@ class HypernetworkModule(torch.nn.Module):
         else:

             for layer in self.linear:

                 if type(layer) == torch.nn.Linear or type(layer) == torch.nn.LayerNorm:

-                    layer.weight.data.normal_(mean=0.0, std=0.01)

-                    layer.bias.data.zero_()

-

+                    w, b = layer.weight.data, layer.bias.data

+                    if weight_init == "Normal" or type(layer) == torch.nn.LayerNorm:

+                        normal_(w, mean=0.0, std=0.01)

+                        normal_(b, mean=0.0, std=0.005)

+                    elif weight_init == 'XavierUniform':

+                        xavier_uniform_(w)

+                        zeros_(b)

+                    elif weight_init == 'XavierNormal':

+                        xavier_normal_(w)

+                        zeros_(b)

+                    elif weight_init == 'KaimingUniform':

+                        kaiming_uniform_(w, nonlinearity='leaky_relu' if 'leakyrelu' == activation_func else 'relu')

+                        zeros_(b)

+                    elif weight_init == 'KaimingNormal':

+                        kaiming_normal_(w, nonlinearity='leaky_relu' if 'leakyrelu' == activation_func else 'relu')

+                        zeros_(b)

+                    else:

+                        raise KeyError(f"Key {weight_init} is not defined as initialization!")

         self.to(devices.device)

 

     def fix_old_state_dict(self, state_dict):

@@ -93,7 +127,7 @@ class Hypernetwork:
     filename = None

     name = None

 

-    def __init__(self, name=None, enable_sizes=None, layer_structure=None, add_layer_norm=False, activation_func=None):

+    def __init__(self, name=None, enable_sizes=None, layer_structure=None, activation_func=None, weight_init=None, add_layer_norm=False, use_dropout=False):

         self.filename = None

         self.name = name

         self.layers = {}

@@ -101,13 +135,15 @@ class Hypernetwork:
         self.sd_checkpoint = None

         self.sd_checkpoint_name = None

         self.layer_structure = layer_structure

-        self.add_layer_norm = add_layer_norm

         self.activation_func = activation_func

+        self.weight_init = weight_init

+        self.add_layer_norm = add_layer_norm

+        self.use_dropout = use_dropout

 

         for size in enable_sizes or []:

             self.layers[size] = (

-                HypernetworkModule(size, None, self.layer_structure, self.add_layer_norm, self.activation_func),

-                HypernetworkModule(size, None, self.layer_structure, self.add_layer_norm, self.activation_func),

+                HypernetworkModule(size, None, self.layer_structure, self.activation_func, self.weight_init, self.add_layer_norm, self.use_dropout),

+                HypernetworkModule(size, None, self.layer_structure, self.activation_func, self.weight_init, self.add_layer_norm, self.use_dropout),

             )

 

     def weights(self):

@@ -129,8 +165,10 @@ class Hypernetwork:
         state_dict['step'] = self.step

         state_dict['name'] = self.name

         state_dict['layer_structure'] = self.layer_structure

-        state_dict['is_layer_norm'] = self.add_layer_norm

         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

 

@@ -144,14 +182,21 @@ class Hypernetwork:
         state_dict = torch.load(filename, map_location='cpu')

 

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

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

+        print(self.layer_structure)

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

 

         for size, sd in state_dict.items():

             if type(size) == int:

                 self.layers[size] = (

-                    HypernetworkModule(size, sd[0], self.layer_structure, self.add_layer_norm, self.activation_func),

-                    HypernetworkModule(size, sd[1], self.layer_structure, self.add_layer_norm, self.activation_func),

+                    HypernetworkModule(size, sd[0], self.layer_structure, self.activation_func, self.weight_init, self.add_layer_norm, self.use_dropout),

+                    HypernetworkModule(size, sd[1], self.layer_structure, self.activation_func, self.weight_init, self.add_layer_norm, self.use_dropout),

                 )

 

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

@@ -164,13 +209,16 @@ def list_hypernetworks(path):
     res = {}

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

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

-        res[name] = filename

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

+        if name != "None":

+            res[name] = filename

     return res

 

 

 def load_hypernetwork(filename):

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

-    if path is not 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()

@@ -255,11 +303,41 @@ def stack_conds(conds):
     return torch.stack(conds)

 

 

+def statistics(data):

+    if len(data) < 2:

+        std = 0

+    else:

+        std = stdev(data)

+    total_information = f"loss:{mean(data):.3f}" + u"\u00B1" + f"({std/ (len(data) ** 0.5):.3f})"

+    recent_data = data[-32:]

+    if len(recent_data) < 2:

+        std = 0

+    else:

+        std = stdev(recent_data)

+    recent_information = f"recent 32 loss:{mean(recent_data):.3f}" + u"\u00B1" + f"({std / (len(recent_data) ** 0.5):.3f})"

+    return total_information, recent_information

+

+

+def report_statistics(loss_info:dict):

+    keys = sorted(loss_info.keys(), key=lambda x: sum(loss_info[x]) / len(loss_info[x]))

+    for key in keys:

+        try:

+            print("Loss statistics for file " + key)

+            info, recent = statistics(list(loss_info[key]))

+            print(info)

+            print(recent)

+        except Exception as e:

+            print(e)

+

+

+

 def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, data_root, log_directory, training_width, training_height, steps, 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):

     # images allows training previews to have infotext. Importing it at the top causes a circular import problem.

     from modules import images

 

-    assert hypernetwork_name, 'hypernetwork not selected'

+    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, data_root, template_file, steps, save_hypernetwork_every, create_image_every, log_directory, name="hypernetwork")

 

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

     shared.loaded_hypernetwork = Hypernetwork()

@@ -285,35 +363,51 @@ def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, data_root, log
     else:

         images_dir = None

 

+    hypernetwork = shared.loaded_hypernetwork

+    checkpoint = sd_models.select_checkpoint()

+

+    ititial_step = hypernetwork.step or 0

+    if ititial_step >= steps:

+        shared.state.textinfo = f"Model has already been trained beyond specified max steps"

+        return hypernetwork, filename

+

+    scheduler = LearnRateScheduler(learn_rate, steps, ititial_step)

+    

+    # 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)}..."

     with torch.autocast("cuda"):

         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, device=devices.device, template_file=template_file, include_cond=True, batch_size=batch_size)

+

     if unload:

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

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

 

-    hypernetwork = shared.loaded_hypernetwork

+    size = len(ds.indexes)

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

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

+    previous_mean_losses = [0]

+    previous_mean_loss = 0

+    print("Mean loss of {} elements".format(size))

+    

     weights = hypernetwork.weights()

     for weight in weights:

         weight.requires_grad = True

+    # if optimizer == "AdamW": or else Adam / AdamW / SGD, etc...

+    optimizer = torch.optim.AdamW(weights, lr=scheduler.learn_rate)

 

-    losses = torch.zeros((32,))

+    steps_without_grad = 0

 

     last_saved_file = "<none>"

     last_saved_image = "<none>"

     forced_filename = "<none>"

 

-    ititial_step = hypernetwork.step or 0

-    if ititial_step > steps:

-        return hypernetwork, filename

-

-    scheduler = LearnRateScheduler(learn_rate, steps, ititial_step)

-    optimizer = torch.optim.AdamW(weights, lr=scheduler.learn_rate)

-

     pbar = tqdm.tqdm(enumerate(ds), total=steps - ititial_step)

     for i, entries in pbar:

         hypernetwork.step = i + ititial_step

-

+        if len(loss_dict) > 0:

+            previous_mean_losses = [i[-1] for i in loss_dict.values()]

+            previous_mean_loss = mean(previous_mean_losses)

+            

         scheduler.apply(optimizer, hypernetwork.step)

         if scheduler.finished:

             break

@@ -330,28 +424,46 @@ def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, data_root, log
             del c

 

             losses[hypernetwork.step % losses.shape[0]] = loss.item()

-

+            for entry in entries:

+                loss_dict[entry.filename].append(loss.item())

+                

             optimizer.zero_grad()

+            weights[0].grad = None

             loss.backward()

+

+            if weights[0].grad is None:

+                steps_without_grad += 1

+            else:

+                steps_without_grad = 0

+            assert steps_without_grad < 10, 'no gradient found for the trained weight after backward() for 10 steps in a row; this is a bug; training cannot continue'

+

             optimizer.step()

-        mean_loss = losses.mean()

-        if torch.isnan(mean_loss):

+

+        steps_done = hypernetwork.step + 1

+

+        if torch.isnan(losses[hypernetwork.step % losses.shape[0]]): 

             raise RuntimeError("Loss diverged.")

-        pbar.set_description(f"loss: {mean_loss:.7f}")

+        

+        if len(previous_mean_losses) > 1:

+            std = stdev(previous_mean_losses)

+        else:

+            std = 0

+        dataset_loss_info = f"dataset loss:{mean(previous_mean_losses):.3f}" + u"\u00B1" + f"({std / (len(previous_mean_losses) ** 0.5):.3f})"

+        pbar.set_description(dataset_loss_info)

 

-        if hypernetwork.step > 0 and hypernetwork_dir is not None and hypernetwork.step % save_hypernetwork_every == 0:

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

             # Before saving, change name to match current checkpoint.

-            hypernetwork.name = f'{hypernetwork_name}-{hypernetwork.step}'

-            last_saved_file = os.path.join(hypernetwork_dir, f'{hypernetwork.name}.pt')

-            hypernetwork.save(last_saved_file)

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

+            last_saved_file = os.path.join(hypernetwork_dir, f'{hypernetwork_name_every}.pt')

+            save_hypernetwork(hypernetwork, checkpoint, hypernetwork_name, last_saved_file)

 

         textual_inversion.write_loss(log_directory, "hypernetwork_loss.csv", hypernetwork.step, len(ds), {

-            "loss": f"{mean_loss:.7f}",

+            "loss": f"{previous_mean_loss:.7f}",

             "learn_rate": scheduler.learn_rate

         })

 

-        if hypernetwork.step > 0 and images_dir is not None and hypernetwork.step % create_image_every == 0:

-            forced_filename = f'{hypernetwork_name}-{hypernetwork.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)

 

             optimizer.zero_grad()

@@ -388,30 +500,39 @@ def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, data_root, log
 

             if image is not None:

                 shared.state.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)

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

 

         shared.state.job_no = hypernetwork.step

 

         shared.state.textinfo = f"""

 <p>

-Loss: {mean_loss:.7f}<br/>

+Loss: {previous_mean_loss:.7f}<br/>

 Step: {hypernetwork.step}<br/>

 Last prompt: {html.escape(entries[0].cond_text)}<br/>

 Last saved hypernetwork: {html.escape(last_saved_file)}<br/>

 Last saved image: {html.escape(last_saved_image)}<br/>

 </p>

 """

+        

+    report_statistics(loss_dict)

 

-    checkpoint = sd_models.select_checkpoint()

-

-    hypernetwork.sd_checkpoint = checkpoint.hash

-    hypernetwork.sd_checkpoint_name = checkpoint.model_name

-    # Before saving for the last time, change name back to the base name (as opposed to the save_hypernetwork_every step-suffixed naming convention).

-    hypernetwork.name = hypernetwork_name

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

-    hypernetwork.save(filename)

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

+    save_hypernetwork(hypernetwork, checkpoint, hypernetwork_name, filename)

 

     return hypernetwork, filename

 

-

+def save_hypernetwork(hypernetwork, checkpoint, hypernetwork_name, filename):

+    old_hypernetwork_name = hypernetwork.name

+    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_name = checkpoint.model_name

+        hypernetwork.name = hypernetwork_name

+        hypernetwork.save(filename)

+    except:

+        hypernetwork.sd_checkpoint = old_sd_checkpoint

+        hypernetwork.sd_checkpoint_name = old_sd_checkpoint_name

+        hypernetwork.name = old_hypernetwork_name

+        raise

diff --git a/modules/hypernetworks/ui.py b/modules/hypernetworks/ui.py
index e6f50a1f..aad09ffc 100644
--- a/modules/hypernetworks/ui.py
+++ b/modules/hypernetworks/ui.py
@@ -3,14 +3,15 @@ import os
 import re

 

 import gradio as gr

-

-import modules.textual_inversion.textual_inversion

 import modules.textual_inversion.preprocess

-from modules import sd_hijack, shared, devices

+import modules.textual_inversion.textual_inversion

+from modules import devices, sd_hijack, shared

 from modules.hypernetworks import hypernetwork

 

+not_available = ["hardswish", "multiheadattention"]

+keys = ["linear"] + list(x for x in hypernetwork.HypernetworkModule.activation_dict.keys() if x not in not_available)

 

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

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

     # Remove illegal characters from name.

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

 

@@ -25,8 +26,10 @@ def create_hypernetwork(name, enable_sizes, overwrite_old, layer_structure=None,
         name=name,

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

         layer_structure=layer_structure,

-        add_layer_norm=add_layer_norm,

         activation_func=activation_func,

+        weight_init=weight_init,

+        add_layer_norm=add_layer_norm,

+        use_dropout=use_dropout,

     )

     hypernet.save(fn)

 

diff --git a/modules/images.py b/modules/images.py
index b9589563..a0728553 100644
--- a/modules/images.py
+++ b/modules/images.py
@@ -1,4 +1,8 @@
 import datetime

+import sys

+import traceback

+

+import pytz

 import io

 import math

 import os

@@ -12,7 +16,7 @@ from PIL import Image, ImageFont, ImageDraw, PngImagePlugin
 from fonts.ttf import Roboto

 import string

 

-from modules import sd_samplers, shared

+from modules import sd_samplers, shared, script_callbacks

 from modules.shared import opts, cmd_opts

 

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

@@ -273,10 +277,15 @@ invalid_filename_chars = '<>:"/\\|?*\n'
 invalid_filename_prefix = ' '

 invalid_filename_postfix = ' .'

 re_nonletters = re.compile(r'[\s' + string.punctuation + ']+')

+re_pattern = re.compile(r"(.*?)(?:\[([^\[\]]+)\]|$)")

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

 max_filename_part_length = 128

 

 

 def sanitize_filename_part(text, replace_spaces=True):

+    if text is None:

+        return None

+

     if replace_spaces:

         text = text.replace(' ', '_')

 

@@ -286,49 +295,104 @@ def sanitize_filename_part(text, replace_spaces=True):
     return text

 

 

-def apply_filename_pattern(x, p, seed, prompt):

-    max_prompt_words = opts.directories_max_prompt_words

-

-    if seed is not None:

-        x = x.replace("[seed]", str(seed))

+class FilenameGenerator:

+    replacements = {

+        'seed': lambda self: self.seed if self.seed is not None else '',

+        'steps': lambda self:  self.p and self.p.steps,

+        'cfg': lambda self: self.p and self.p.cfg_scale,

+        'width': lambda self: self.image.width,

+        'height': lambda self: self.image.height,

+        'styles': lambda self: self.p and sanitize_filename_part(", ".join([style for style in self.p.styles if not style == "None"]) or "None", replace_spaces=False),

+        'sampler': lambda self: self.p and sanitize_filename_part(sd_samplers.samplers[self.p.sampler_index].name, replace_spaces=False),

+        'model_hash': lambda self: getattr(self.p, "sd_model_hash", shared.sd_model.sd_model_hash),

+        'date': lambda self: datetime.datetime.now().strftime('%Y-%m-%d'),

+        'datetime': lambda self, *args: self.datetime(*args),  # accepts formats: [datetime], [datetime<Format>], [datetime<Format><Time Zone>]

+        'job_timestamp': lambda self: getattr(self.p, "job_timestamp", shared.state.job_timestamp),

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

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

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

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

+    }

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

+

+    def __init__(self, p, seed, prompt, image):

+        self.p = p

+        self.seed = seed

+        self.prompt = prompt

+        self.image = image

+

+    def prompt_no_style(self):

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

+            return None

+

+        prompt_no_style = self.prompt

+        for style in shared.prompt_styles.get_style_prompts(self.p.styles):

+            if len(style) > 0:

+                for part in style.split("{prompt}"):

+                    prompt_no_style = prompt_no_style.replace(part, "").replace(", ,", ",").strip().strip(',')

+

+                prompt_no_style = prompt_no_style.replace(style, "").strip().strip(',').strip()

+

+        return sanitize_filename_part(prompt_no_style, replace_spaces=False)

+

+    def prompt_words(self):

+        words = [x for x in re_nonletters.split(self.prompt or "") if len(x) > 0]

+        if len(words) == 0:

+            words = ["empty"]

+        return sanitize_filename_part(" ".join(words[0:opts.directories_max_prompt_words]), replace_spaces=False)

+

+    def datetime(self, *args):

+        time_datetime = datetime.datetime.now()

+

+        time_format = args[0] if len(args) > 0 and args[0] != "" else self.default_time_format

+        try:

+            time_zone = pytz.timezone(args[1]) if len(args) > 1 else None

+        except pytz.exceptions.UnknownTimeZoneError as _:

+            time_zone = None

+

+        time_zone_time = time_datetime.astimezone(time_zone)

+        try:

+            formatted_time = time_zone_time.strftime(time_format)

+        except (ValueError, TypeError) as _:

+            formatted_time = time_zone_time.strftime(self.default_time_format)

+

+        return sanitize_filename_part(formatted_time, replace_spaces=False)

+

+    def apply(self, x):

+        res = ''

+

+        for m in re_pattern.finditer(x):

+            text, pattern = m.groups()

+            res += text

+

+            if pattern is None:

+                continue

 

-    if p is not None:

-        x = x.replace("[steps]", str(p.steps))

-        x = x.replace("[cfg]", str(p.cfg_scale))

-        x = x.replace("[width]", str(p.width))

-        x = x.replace("[height]", str(p.height))

-        x = x.replace("[styles]", sanitize_filename_part(", ".join([x for x in p.styles if not x == "None"]) or "None", replace_spaces=False))

-        x = x.replace("[sampler]", sanitize_filename_part(sd_samplers.samplers[p.sampler_index].name, replace_spaces=False))

+            pattern_args = []

+            while True:

+                m = re_pattern_arg.match(pattern)

+                if m is None:

+                    break

 

-    x = x.replace("[model_hash]", getattr(p, "sd_model_hash", shared.sd_model.sd_model_hash))

-    x = x.replace("[date]", datetime.date.today().isoformat())

-    x = x.replace("[datetime]", datetime.datetime.now().strftime("%Y%m%d%H%M%S"))

-    x = x.replace("[job_timestamp]", getattr(p, "job_timestamp", shared.state.job_timestamp))

+                pattern, arg = m.groups()

+                pattern_args.insert(0, arg)

 

-    # Apply [prompt] at last. Because it may contain any replacement word.^M

-    if prompt is not None:

-        x = x.replace("[prompt]", sanitize_filename_part(prompt))

-        if "[prompt_no_styles]" in x:

-            prompt_no_style = prompt

-            for style in shared.prompt_styles.get_style_prompts(p.styles):

-                if len(style) > 0:

-                    style_parts = [y for y in style.split("{prompt}")]

-                    for part in style_parts:

-                        prompt_no_style = prompt_no_style.replace(part, "").replace(", ,", ",").strip().strip(',')

-            prompt_no_style = prompt_no_style.replace(style, "").strip().strip(',').strip()

-            x = x.replace("[prompt_no_styles]", sanitize_filename_part(prompt_no_style, replace_spaces=False))

+            fun = self.replacements.get(pattern.lower())

+            if fun is not None:

+                try:

+                    replacement = fun(self, *pattern_args)

+                except Exception:

+                    replacement = None

+                    print(f"Error adding [{pattern}] to filename", file=sys.stderr)

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

 

-        x = x.replace("[prompt_spaces]", sanitize_filename_part(prompt, replace_spaces=False))

-        if "[prompt_words]" in x:

-            words = [x for x in re_nonletters.split(prompt or "") if len(x) > 0]

-            if len(words) == 0:

-                words = ["empty"]

-            x = x.replace("[prompt_words]", sanitize_filename_part(" ".join(words[0:max_prompt_words]), replace_spaces=False))

+                if replacement is not None:

+                    res += str(replacement)

+                    continue

 

-    if cmd_opts.hide_ui_dir_config:

-        x = re.sub(r'^[\\/]+|\.{2,}[\\/]+|[\\/]+\.{2,}', '', x)

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

 

-    return x

+        return res

 

 

 def get_next_sequence_number(path, basename):

@@ -354,7 +418,7 @@ def get_next_sequence_number(path, basename):
 

 

 def save_image(image, path, basename, seed=None, prompt=None, extension='png', info=None, short_filename=False, no_prompt=False, grid=False, pnginfo_section_name='parameters', p=None, existing_info=None, forced_filename=None, suffix="", save_to_dirs=None):

-    '''Save an image.

+    """Save an image.

 

     Args:

         image (`PIL.Image`):

@@ -385,52 +449,57 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
             The full path of the saved imaged.

         txt_fullfn (`str` or None):

             If a text file is saved for this image, this will be its full path. Otherwise None.

-    '''

-    if short_filename or prompt is None or seed is None:

-        file_decoration = ""

-    elif opts.save_to_dirs:

-        file_decoration = opts.samples_filename_pattern or "[seed]"

-    else:

-        file_decoration = opts.samples_filename_pattern or "[seed]-[prompt_spaces]"

-

-    if file_decoration != "":

-        file_decoration = "-" + file_decoration.lower()

-

-    file_decoration = apply_filename_pattern(file_decoration, p, seed, prompt) + suffix

-

-    if extension == 'png' and opts.enable_pnginfo and info is not None:

-        pnginfo = PngImagePlugin.PngInfo()

-

-        if existing_info is not None:

-            for k, v in existing_info.items():

-                pnginfo.add_text(k, str(v))

-

-        pnginfo.add_text(pnginfo_section_name, info)

-    else:

-        pnginfo = None

+    """

+    namegen = FilenameGenerator(p, seed, prompt, image)

 

     if save_to_dirs is None:

         save_to_dirs = (grid and opts.grid_save_to_dirs) or (not grid and opts.save_to_dirs and not no_prompt)

 

     if save_to_dirs:

-        dirname = apply_filename_pattern(opts.directories_filename_pattern or "[prompt_words]", p, seed, prompt).strip('\\ /')

+        dirname = namegen.apply(opts.directories_filename_pattern or "[prompt_words]").lstrip(' ').rstrip('\\ /')

         path = os.path.join(path, dirname)

 

     os.makedirs(path, exist_ok=True)

 

     if forced_filename is None:

-        basecount = get_next_sequence_number(path, basename)

-        fullfn = "a.png"

-        fullfn_without_extension = "a"

-        for i in range(500):

-            fn = f"{basecount + i:05}" if basename == '' else f"{basename}-{basecount + i:04}"

-            fullfn = os.path.join(path, f"{fn}{file_decoration}.{extension}")

-            fullfn_without_extension = os.path.join(path, f"{fn}{file_decoration}")

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

-                break

+        if short_filename or seed is None:

+            file_decoration = ""

+        elif opts.save_to_dirs:

+            file_decoration = opts.samples_filename_pattern or "[seed]"

+        else:

+            file_decoration = opts.samples_filename_pattern or "[seed]-[prompt_spaces]"

+

+        add_number = opts.save_images_add_number or file_decoration == ''

+

+        if file_decoration != "" and add_number:

+            file_decoration = "-" + file_decoration

+

+        file_decoration = namegen.apply(file_decoration) + suffix

+

+        if add_number:

+            basecount = get_next_sequence_number(path, basename)

+            fullfn = None

+            for i in range(500):

+                fn = f"{basecount + i:05}" if basename == '' else f"{basename}-{basecount + i:04}"

+                fullfn = os.path.join(path, f"{fn}{file_decoration}.{extension}")

+                if not os.path.exists(fullfn):

+                    break

+        else:

+            fullfn = os.path.join(path, f"{file_decoration}.{extension}")

     else:

         fullfn = os.path.join(path, f"{forced_filename}.{extension}")

-        fullfn_without_extension = os.path.join(path, forced_filename)

+

+    pnginfo = existing_info or {}

+    if info is not None:

+        pnginfo[pnginfo_section_name] = info

+

+    params = script_callbacks.ImageSaveParams(image, p, fullfn, pnginfo)

+    script_callbacks.before_image_saved_callback(params)

+

+    image = params.image

+    fullfn = params.filename

+    info = params.pnginfo.get(pnginfo_section_name, None)

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

 

     def exif_bytes():

         return piexif.dump({

@@ -439,12 +508,20 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
             },

         })

 

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

+    if extension.lower() == '.png':

+        pnginfo_data = PngImagePlugin.PngInfo()

+        for k, v in params.pnginfo.items():

+            pnginfo_data.add_text(k, str(v))

+

+        image.save(fullfn, quality=opts.jpeg_quality, pnginfo=pnginfo_data)

+

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

         image.save(fullfn, quality=opts.jpeg_quality)

+

         if opts.enable_pnginfo and info is not None:

             piexif.insert(exif_bytes(), fullfn)

     else:

-        image.save(fullfn, quality=opts.jpeg_quality, pnginfo=pnginfo)

+        image.save(fullfn, quality=opts.jpeg_quality)

 

     target_side_length = 4000

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

@@ -467,6 +544,8 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
     else:

         txt_fullfn = None

 

+    script_callbacks.image_saved_callback(params)

+

     return fullfn, txt_fullfn

 

 

diff --git a/modules/images_history.py b/modules/images_history.py
deleted file mode 100644
index 46b23e56..00000000
--- a/modules/images_history.py
+++ /dev/null
@@ -1,183 +0,0 @@
-import os
-import shutil
-import sys
-
-def traverse_all_files(output_dir, image_list, curr_dir=None):
-    curr_path = output_dir if curr_dir is None else os.path.join(output_dir, curr_dir)
-    try:
-        f_list = os.listdir(curr_path)
-    except:
-        if curr_dir[-10:].rfind(".") > 0 and curr_dir[-4:] != ".txt":
-            image_list.append(curr_dir)
-        return image_list
-    for file in f_list:
-        file = file if curr_dir is None else os.path.join(curr_dir, file)
-        file_path = os.path.join(curr_path, file)
-        if file[-4:] == ".txt":
-            pass
-        elif os.path.isfile(file_path) and file[-10:].rfind(".") > 0:
-            image_list.append(file)
-        else:
-            image_list = traverse_all_files(output_dir, image_list, file)
-    return image_list
-
-
-def get_recent_images(dir_name, page_index, step, image_index, tabname):
-    page_index = int(page_index)
-    image_list = []
-    if not os.path.exists(dir_name):
-        pass
-    elif os.path.isdir(dir_name):
-        image_list = traverse_all_files(dir_name, image_list)
-        image_list = sorted(image_list, key=lambda file: -os.path.getctime(os.path.join(dir_name, file)))
-    else:
-        print(f'ERROR: "{dir_name}" is not a directory. Check the path in the settings.', file=sys.stderr)
-    num = 48 if tabname != "extras" else 12
-    max_page_index = len(image_list) // num + 1
-    page_index = max_page_index if page_index == -1 else page_index + step
-    page_index = 1 if page_index < 1 else page_index
-    page_index = max_page_index if page_index > max_page_index else page_index
-    idx_frm = (page_index - 1) * num
-    image_list = image_list[idx_frm:idx_frm + num]
-    image_index = int(image_index)
-    if image_index < 0 or image_index > len(image_list) - 1:
-        current_file = None
-        hidden = None
-    else:
-        current_file = image_list[int(image_index)]
-        hidden = os.path.join(dir_name, current_file)
-    return [os.path.join(dir_name, file) for file in image_list], page_index, image_list, current_file, hidden, ""
-
-
-def first_page_click(dir_name, page_index, image_index, tabname):
-    return get_recent_images(dir_name, 1, 0, image_index, tabname)
-
-
-def end_page_click(dir_name, page_index, image_index, tabname):
-    return get_recent_images(dir_name, -1, 0, image_index, tabname)
-
-
-def prev_page_click(dir_name, page_index, image_index, tabname):
-    return get_recent_images(dir_name, page_index, -1, image_index, tabname)
-
-
-def next_page_click(dir_name, page_index, image_index, tabname):
-    return get_recent_images(dir_name, page_index, 1, image_index, tabname)
-
-
-def page_index_change(dir_name, page_index, image_index, tabname):
-    return get_recent_images(dir_name, page_index, 0, image_index, tabname)
-
-
-def show_image_info(num, image_path, filenames):
-    # print(f"select image {num}")
-    file = filenames[int(num)]
-    return file, num, os.path.join(image_path, file)
-
-
-def delete_image(delete_num, tabname, dir_name, name, page_index, filenames, image_index):
-    if name == "":
-        return filenames, delete_num
-    else:
-        delete_num = int(delete_num)
-        index = list(filenames).index(name)
-        i = 0
-        new_file_list = []
-        for name in filenames:
-            if i >= index and i < index + delete_num:
-                path = os.path.join(dir_name, name)
-                if os.path.exists(path):
-                    print(f"Delete file {path}")
-                    os.remove(path)
-                    txt_file = os.path.splitext(path)[0] + ".txt"
-                    if os.path.exists(txt_file):
-                        os.remove(txt_file)
-                else:
-                    print(f"Not exists file {path}")
-            else:
-                new_file_list.append(name)
-            i += 1
-    return new_file_list, 1
-
-
-def show_images_history(gr, opts, tabname, run_pnginfo, switch_dict):
-    if opts.outdir_samples != "":
-        dir_name = opts.outdir_samples
-    elif tabname == "txt2img":
-        dir_name = opts.outdir_txt2img_samples
-    elif tabname == "img2img":
-        dir_name = opts.outdir_img2img_samples
-    elif tabname == "extras":
-        dir_name = opts.outdir_extras_samples
-    else:
-        return
-    with gr.Row():
-        renew_page = gr.Button('Renew Page', elem_id=tabname + "_images_history_renew_page")
-        first_page = gr.Button('First Page')
-        prev_page = gr.Button('Prev Page')
-        page_index = gr.Number(value=1, label="Page Index")
-        next_page = gr.Button('Next Page')
-        end_page = gr.Button('End Page')
-    with gr.Row(elem_id=tabname + "_images_history"):
-        with gr.Row():
-            with gr.Column(scale=2):
-                history_gallery = gr.Gallery(show_label=False, elem_id=tabname + "_images_history_gallery").style(grid=6)
-                with gr.Row():
-                    delete_num = gr.Number(value=1, interactive=True, label="number of images to delete consecutively next")
-                    delete = gr.Button('Delete', elem_id=tabname + "_images_history_del_button")
-            with gr.Column():
-                with gr.Row():
-                    pnginfo_send_to_txt2img = gr.Button('Send to txt2img')
-                    pnginfo_send_to_img2img = gr.Button('Send to img2img')
-                with gr.Row():
-                    with gr.Column():
-                        img_file_info = gr.Textbox(label="Generate Info", interactive=False)
-                        img_file_name = gr.Textbox(label="File Name", interactive=False)
-                with gr.Row():
-                    # hiden items
-
-                    img_path = gr.Textbox(dir_name.rstrip("/"), visible=False)
-                    tabname_box = gr.Textbox(tabname, visible=False)
-                    image_index = gr.Textbox(value=-1, visible=False)
-                    set_index = gr.Button('set_index', elem_id=tabname + "_images_history_set_index", visible=False)
-                    filenames = gr.State()
-                    hidden = gr.Image(type="pil", visible=False)
-                    info1 = gr.Textbox(visible=False)
-                    info2 = gr.Textbox(visible=False)
-
-    # turn pages
-    gallery_inputs = [img_path, page_index, image_index, tabname_box]
-    gallery_outputs = [history_gallery, page_index, filenames, img_file_name, hidden, img_file_name]
-
-    first_page.click(first_page_click, _js="images_history_turnpage", inputs=gallery_inputs, outputs=gallery_outputs)
-    next_page.click(next_page_click, _js="images_history_turnpage", inputs=gallery_inputs, outputs=gallery_outputs)
-    prev_page.click(prev_page_click, _js="images_history_turnpage", inputs=gallery_inputs, outputs=gallery_outputs)
-    end_page.click(end_page_click, _js="images_history_turnpage", inputs=gallery_inputs, outputs=gallery_outputs)
-    page_index.submit(page_index_change, _js="images_history_turnpage", inputs=gallery_inputs, outputs=gallery_outputs)
-    renew_page.click(page_index_change, _js="images_history_turnpage", inputs=gallery_inputs, outputs=gallery_outputs)
-    # page_index.change(page_index_change, inputs=[tabname_box, img_path,  page_index], outputs=[history_gallery, page_index])
-
-    # other funcitons
-    set_index.click(show_image_info, _js="images_history_get_current_img", inputs=[tabname_box, img_path, filenames], outputs=[img_file_name, image_index, hidden])
-    img_file_name.change(fn=None, _js="images_history_enable_del_buttons", inputs=None, outputs=None)
-    delete.click(delete_image, _js="images_history_delete", inputs=[delete_num, tabname_box, img_path, img_file_name, page_index, filenames, image_index], outputs=[filenames, delete_num])
-    hidden.change(fn=run_pnginfo, inputs=[hidden], outputs=[info1, img_file_info, info2])
-
-    # pnginfo.click(fn=run_pnginfo, inputs=[hidden], outputs=[info1, img_file_info, info2])
-    switch_dict["fn"](pnginfo_send_to_txt2img, switch_dict["t2i"], img_file_info, 'switch_to_txt2img')
-    switch_dict["fn"](pnginfo_send_to_img2img, switch_dict["i2i"], img_file_info, 'switch_to_img2img_img2img')
-
-
-def create_history_tabs(gr, opts, run_pnginfo, switch_dict):
-    with gr.Blocks(analytics_enabled=False) as images_history:
-        with gr.Tabs() as tabs:
-            with gr.Tab("txt2img history"):
-                with gr.Blocks(analytics_enabled=False) as images_history_txt2img:
-                    show_images_history(gr, opts, "txt2img", run_pnginfo, switch_dict)
-            with gr.Tab("img2img history"):
-                with gr.Blocks(analytics_enabled=False) as images_history_img2img:
-                    show_images_history(gr, opts, "img2img", run_pnginfo, switch_dict)
-            with gr.Tab("extras history"):
-                with gr.Blocks(analytics_enabled=False) as images_history_img2img:
-                    show_images_history(gr, opts, "extras", run_pnginfo, switch_dict)
-    return images_history
diff --git a/modules/img2img.py b/modules/img2img.py
index eea5199b..efda26e1 100644
--- a/modules/img2img.py
+++ b/modules/img2img.py
@@ -19,7 +19,7 @@ import modules.scripts
 def process_batch(p, input_dir, output_dir, args):

     processing.fix_seed(p)

 

-    images = [file for file in [os.path.join(input_dir, x) for x in os.listdir(input_dir)] if os.path.isfile(file)]

+    images = shared.listfiles(input_dir)

 

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

 

@@ -39,6 +39,8 @@ def process_batch(p, input_dir, output_dir, args):
             break

 

         img = Image.open(image)

+        # Use the EXIF orientation of photos taken by smartphones.

+        img = ImageOps.exif_transpose(img) 

         p.init_images = [img] * p.batch_size

 

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

@@ -56,24 +58,30 @@ 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_inpaint, init_mask_inpaint, mask_mode, steps: int, sampler_index: int, mask_blur: int, 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, aesthetic_lr=0, aesthetic_weight=0, aesthetic_steps=0, aesthetic_imgs=None, aesthetic_slerp=False, aesthetic_imgs_text="", aesthetic_slerp_angle=0.15, aesthetic_text_negative=False, *args):

+def img2img(mode: int, prompt: str, negative_prompt: str, prompt_style: str, prompt_style2: str, init_img, init_img_with_mask, init_img_inpaint, init_mask_inpaint, mask_mode, steps: int, sampler_index: int, mask_blur: int, 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:

             image = init_img_with_mask['image']

             mask = 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')

+        # Uploaded mask

         else:

             image = init_img_inpaint

             mask = init_mask_inpaint

+    # No mask

     else:

         image = init_img

         mask = None

 

+    # Use the EXIF orientation of photos taken by smartphones.

+    image = ImageOps.exif_transpose(image) 

+

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

 

     p = StableDiffusionProcessingImg2Img(

@@ -109,7 +117,8 @@ def img2img(mode: int, prompt: str, negative_prompt: str, prompt_style: str, pro
         inpainting_mask_invert=inpainting_mask_invert,

     )

 

-    shared.aesthetic_clip.set_aesthetic_params(p, float(aesthetic_lr), float(aesthetic_weight), int(aesthetic_steps), aesthetic_imgs, aesthetic_slerp, aesthetic_imgs_text, aesthetic_slerp_angle, aesthetic_text_negative)

+    p.scripts = modules.scripts.scripts_txt2img

+    p.script_args = args

 

     if shared.cmd_opts.enable_console_prompts:

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

diff --git a/modules/lowvram.py b/modules/lowvram.py
index 7eba1349..f327c3df 100644
--- a/modules/lowvram.py
+++ b/modules/lowvram.py
@@ -1,9 +1,8 @@
 import torch

-from modules.devices import get_optimal_device

+from modules import devices

 

 module_in_gpu = None

 cpu = torch.device("cpu")

-device = gpu = get_optimal_device()

 

 

 def send_everything_to_cpu():

@@ -33,7 +32,7 @@ def setup_for_low_vram(sd_model, use_medvram):
         if module_in_gpu is not None:

             module_in_gpu.to(cpu)

 

-        module.to(gpu)

+        module.to(devices.device)

         module_in_gpu = module

 

     # see below for register_forward_pre_hook;

@@ -51,7 +50,7 @@ def setup_for_low_vram(sd_model, use_medvram):
     # send the model to GPU. Then put modules back. the modules will be in CPU.

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

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

-    sd_model.to(device)

+    sd_model.to(devices.device)

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

 

     # register hooks for those the first two models

@@ -70,7 +69,7 @@ def setup_for_low_vram(sd_model, use_medvram):
         # so that only one of them is in GPU at a time

         stored = diff_model.input_blocks, diff_model.middle_block, diff_model.output_blocks, diff_model.time_embed

         diff_model.input_blocks, diff_model.middle_block, diff_model.output_blocks, diff_model.time_embed = None, None, None, None

-        sd_model.model.to(device)

+        sd_model.model.to(devices.device)

         diff_model.input_blocks, diff_model.middle_block, diff_model.output_blocks, diff_model.time_embed = stored

 

         # install hooks for bits of third model

diff --git a/modules/processing.py b/modules/processing.py
index ff1ec4c9..b1df4918 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -47,6 +47,25 @@ def apply_color_correction(correction, image):
     return image

 

 

+def apply_overlay(image, paste_loc, index, overlays):

+    if overlays is None or index >= len(overlays):

+        return image

+

+    overlay = overlays[index]

+

+    if paste_loc is not None:

+        x, y, w, h = paste_loc

+        base_image = Image.new('RGBA', (overlay.width, overlay.height))

+        image = images.resize_image(1, image, w, h)

+        base_image.paste(image, (x, y))

+        image = base_image

+

+    image = image.convert('RGBA')

+    image.alpha_composite(overlay)

+    image = image.convert('RGB')

+

+    return image

+

 def get_correct_sampler(p):

     if isinstance(p, modules.processing.StableDiffusionProcessingTxt2Img):

         return sd_samplers.samplers

@@ -58,9 +77,8 @@ def get_correct_sampler(p):
 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_index: int=0, 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 = "uniform", s_churn: float = 0.0, s_tmax: float = None, s_tmin: float = 0.0, s_noise: float = 1.0):

+    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_index: int = 0, 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):

         self.sd_model = sd_model

         self.outpath_samples: str = outpath_samples

         self.outpath_grids: str = outpath_grids

@@ -90,13 +108,14 @@ class StableDiffusionProcessing():
         self.do_not_reload_embeddings = do_not_reload_embeddings

         self.paste_to = None

         self.color_corrections = None

-        self.denoising_strength: float = 0

+        self.denoising_strength: float = denoising_strength

         self.sampler_noise_scheduler_override = None

-        self.ddim_discretize = opts.ddim_discretize

+        self.ddim_discretize = ddim_discretize or opts.ddim_discretize

         self.s_churn = s_churn or opts.s_churn

         self.s_tmin = s_tmin or opts.s_tmin

         self.s_tmax = s_tmax or float('inf')  # not representable as a standard ui option

         self.s_noise = s_noise or opts.s_noise

+        self.override_settings = {k: v for k, v in (override_settings or {}).items() if k not in shared.restricted_opts}

 

         if not seed_enable_extras:

             self.subseed = -1

@@ -104,6 +123,78 @@ class StableDiffusionProcessing():
             self.seed_resize_from_h = 0

             self.seed_resize_from_w = 0

 

+        self.scripts = None

+        self.script_args = None

+        self.all_prompts = None

+        self.all_seeds = None

+        self.all_subseeds = None

+

+    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 torch.zeros(

+                x.shape[0], 5, 1, 1, 

+                dtype=x.dtype, 

+                device=x.device

+            )

+

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

+

+        return image_conditioning

+

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

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

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

+            return torch.zeros(

+                latent_image.shape[0], 5, 1, 1,

+                dtype=latent_image.dtype,

+                device=latent_image.device

+            )

+

+        # Handle the different mask inputs

+        if image_mask is not None:

+            if torch.is_tensor(image_mask):

+                conditioning_mask = image_mask

+            else:

+                conditioning_mask = np.array(image_mask.convert("L"))

+                conditioning_mask = conditioning_mask.astype(np.float32) / 255.0

+                conditioning_mask = torch.from_numpy(conditioning_mask[None, None])

+

+                # Inpainting model uses a discretized mask as input, so we round to either 1.0 or 0.0

+                conditioning_mask = torch.round(conditioning_mask)

+        else:

+            conditioning_mask = torch.ones(1, 1, *source_image.shape[-2:])

+

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

+        conditioning_image = torch.lerp(

+            source_image,

+            source_image * (1.0 - conditioning_mask),

+            getattr(self, "inpainting_mask_weight", shared.opts.inpainting_mask_weight)

+        )

+        

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

+

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

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

+        conditioning_mask = conditioning_mask.expand(conditioning_image.shape[0], -1, -1, -1)

+        image_conditioning = torch.cat([conditioning_mask, conditioning_image], dim=1)

+        image_conditioning = image_conditioning.to(shared.device).type(self.sd_model.dtype)

+

+        return image_conditioning

 

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

         pass

@@ -305,6 +396,7 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments, iteration
         "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 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]),

@@ -326,6 +418,22 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments, iteration
 

 

 def process_images(p: StableDiffusionProcessing) -> Processed:

+    stored_opts = {k: opts.data[k] for k in p.override_settings.keys()}

+

+    try:

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

+            opts.data[k] = v  # we don't call onchange for simplicity which makes changing model, hypernet impossible

+

+        res = process_images_inner(p)

+

+    finally:

+        for k, v in stored_opts.items():

+            opts.data[k] = v

+

+    return res

+

+

+def process_images_inner(p: StableDiffusionProcessing) -> Processed:

     """this is the main loop that both txt2img and img2img use; it calls func_init once inside all the scopes and func_sample once per batch"""

 

     if type(p.prompt) == list:

@@ -350,32 +458,35 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
     shared.prompt_styles.apply_styles(p)

 

     if type(p.prompt) == list:

-        all_prompts = p.prompt

+        p.all_prompts = p.prompt

     else:

-        all_prompts = p.batch_size * p.n_iter * [p.prompt]

+        p.all_prompts = p.batch_size * p.n_iter * [p.prompt]

 

     if type(seed) == list:

-        all_seeds = seed

+        p.all_seeds = seed

     else:

-        all_seeds = [int(seed) + (x if p.subseed_strength == 0 else 0) for x in range(len(all_prompts))]

+        p.all_seeds = [int(seed) + (x if p.subseed_strength == 0 else 0) for x in range(len(p.all_prompts))]

 

     if type(subseed) == list:

-        all_subseeds = subseed

+        p.all_subseeds = subseed

     else:

-        all_subseeds = [int(subseed) + x for x in range(len(all_prompts))]

+        p.all_subseeds = [int(subseed) + x for x in range(len(p.all_prompts))]

 

     def infotext(iteration=0, position_in_batch=0):

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

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

 

     if os.path.exists(cmd_opts.embeddings_dir) and not p.do_not_reload_embeddings:

         model_hijack.embedding_db.load_textual_inversion_embeddings()

 

+    if p.scripts is not None:

+        p.scripts.process(p)

+

     infotexts = []

     output_images = []

 

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

         with devices.autocast():

-            p.init(all_prompts, all_seeds, all_subseeds)

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

 

         if state.job_count == -1:

             state.job_count = p.n_iter

@@ -387,15 +498,13 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
             if state.interrupted:

                 break

 

-            prompts = all_prompts[n * p.batch_size:(n + 1) * p.batch_size]

-            seeds = all_seeds[n * p.batch_size:(n + 1) * p.batch_size]

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

+            prompts = p.all_prompts[n * p.batch_size:(n + 1) * p.batch_size]

+            seeds = p.all_seeds[n * p.batch_size:(n + 1) * p.batch_size]

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

 

-            if (len(prompts) == 0):

+            if len(prompts) == 0:

                 break

 

-            #uc = p.sd_model.get_learned_conditioning(len(prompts) * [p.negative_prompt])

-            #c = p.sd_model.get_learned_conditioning(prompts)

             with devices.autocast():

                 uc = prompt_parser.get_learned_conditioning(shared.sd_model, len(prompts) * [p.negative_prompt], p.steps)

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

@@ -442,22 +551,11 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
 

                 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:

-                        images.save_image(image, p.outpath_samples, "", seeds[i], prompts[i], opts.samples_format, info=infotext(n, i), p=p, suffix="-before-color-correction")

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

+                        images.save_image(image_without_cc, p.outpath_samples, "", seeds[i], prompts[i], opts.samples_format, info=infotext(n, i), p=p, suffix="-before-color-correction")

                     image = apply_color_correction(p.color_corrections[i], image)

 

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

-                    overlay = p.overlay_images[i]

-

-                    if p.paste_to is not None:

-                        x, y, w, h = p.paste_to

-                        base_image = Image.new('RGBA', (overlay.width, overlay.height))

-                        image = images.resize_image(1, image, w, h)

-                        base_image.paste(image, (x, y))

-                        image = base_image

-

-                    image = image.convert('RGBA')

-                    image.alpha_composite(overlay)

-                    image = image.convert('RGB')

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

 

                 if opts.samples_save and not p.do_not_save_samples:

                     images.save_image(image, p.outpath_samples, "", seeds[i], prompts[i], opts.samples_format, info=infotext(n, i), p=p)

@@ -490,10 +588,16 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
                 index_of_first_image = 1

 

             if opts.grid_save:

-                images.save_image(grid, p.outpath_grids, "grid", all_seeds[0], all_prompts[0], opts.grid_format, info=infotext(), short_filename=not opts.grid_extended_filename, p=p, grid=True)

+                images.save_image(grid, p.outpath_grids, "grid", p.all_seeds[0], p.all_prompts[0], opts.grid_format, info=infotext(), short_filename=not opts.grid_extended_filename, p=p, grid=True)

 

     devices.torch_gc()

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

+

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

+

+    if p.scripts is not None:

+        p.scripts.postprocess(p, res)

+

+    return res

 

 

 class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):

@@ -515,6 +619,8 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
             else:

                 state.job_count = state.job_count * 2

 

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

+

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

                 desired_pixel_count = 512 * 512

                 actual_pixel_count = self.width * self.height

@@ -536,41 +642,19 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
                     firstphase_width_truncated = self.firstphase_height * self.width / self.height

                     firstphase_height_truncated = self.firstphase_height

 

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

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

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

 

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

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

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

-

-        else:

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

-            image_conditioning = torch.zeros(x.shape[0], 5, 1, 1, dtype=x.dtype, device=x.device)

-

-        return image_conditioning

-

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

         self.sampler = sd_samplers.create_sampler_with_index(sd_samplers.samplers, self.sampler_index, self.sd_model)

 

         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.create_dummy_mask(x))

+            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.create_dummy_mask(x, self.firstphase_width, self.firstphase_height))

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

 

@@ -603,11 +687,13 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 

         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)

 

+        image_conditioning = self.txt2img_image_conditioning(x)

+

         # 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=self.create_dummy_mask(samples))

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

 

         return samples

 

@@ -615,7 +701,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):

     sampler = None

 

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

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

         super().__init__(**kwargs)

 

         self.init_images = init_images

@@ -706,6 +792,10 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
             batch_images = np.expand_dims(imgs[0], axis=0).repeat(self.batch_size, axis=0)

             if self.overlay_images is not None:

                 self.overlay_images = self.overlay_images * self.batch_size

+

+            if self.color_corrections is not None and len(self.color_corrections) == 1:

+                self.color_corrections = self.color_corrections * self.batch_size

+

         elif len(imgs) <= self.batch_size:

             self.batch_size = len(imgs)

             batch_images = np.array(imgs)

@@ -735,33 +825,7 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
             elif self.inpainting_fill == 3:

                 self.init_latent = self.init_latent * self.mask

 

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

-            if self.image_mask is not None:

-                conditioning_mask = np.array(self.image_mask.convert("L"))

-                conditioning_mask = conditioning_mask.astype(np.float32) / 255.0

-                conditioning_mask = torch.from_numpy(conditioning_mask[None, None])

-

-                # Inpainting model uses a discretized mask as input, so we round to either 1.0 or 0.0

-                conditioning_mask = torch.round(conditioning_mask)

-            else:

-                conditioning_mask = torch.ones(1, 1, *image.shape[-2:])

-

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

-            conditioning_mask = conditioning_mask.to(image.device)

-            conditioning_image = image * (1.0 - conditioning_mask)

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

-

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

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

-            conditioning_mask = conditioning_mask.expand(conditioning_image.shape[0], -1, -1, -1)

-            self.image_conditioning = torch.cat([conditioning_mask, conditioning_image], dim=1)

-            self.image_conditioning = self.image_conditioning.to(shared.device).type(self.sd_model.dtype)

-        else:

-            self.image_conditioning = torch.zeros(

-                self.init_latent.shape[0], 5, 1, 1,

-                dtype=self.init_latent.dtype,

-                device=self.init_latent.device

-            )

+        self.image_conditioning = self.img2img_image_conditioning(image, self.init_latent, self.image_mask)

 

 

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

diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py
new file mode 100644
index 00000000..6ea58d61
--- /dev/null
+++ b/modules/script_callbacks.py
@@ -0,0 +1,132 @@
+import sys

+import traceback

+from collections import namedtuple

+import inspect

+

+

+def report_exception(c, job):

+    print(f"Error executing callback {job} for {c.script}", file=sys.stderr)

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

+

+

+class ImageSaveParams:

+    def __init__(self, image, p, filename, pnginfo):

+        self.image = image

+        """the PIL image itself"""

+

+        self.p = p

+        """p object with processing parameters; either StableDiffusionProcessing or an object with same fields"""

+

+        self.filename = filename

+        """name of file that the image would be saved to"""

+

+        self.pnginfo = pnginfo

+        """dictionary with parameters for image's PNG info data; infotext will have the key 'parameters'"""

+

+

+ScriptCallback = namedtuple("ScriptCallback", ["script", "callback"])

+callbacks_model_loaded = []

+callbacks_ui_tabs = []

+callbacks_ui_settings = []

+callbacks_before_image_saved = []

+callbacks_image_saved = []

+

+

+def clear_callbacks():

+    callbacks_model_loaded.clear()

+    callbacks_ui_tabs.clear()

+    callbacks_ui_settings.clear()

+    callbacks_before_image_saved.clear()

+    callbacks_image_saved.clear()

+

+

+def model_loaded_callback(sd_model):

+    for c in callbacks_model_loaded:

+        try:

+            c.callback(sd_model)

+        except Exception:

+            report_exception(c, 'model_loaded_callback')

+

+

+def ui_tabs_callback():

+    res = []

+    

+    for c in callbacks_ui_tabs:

+        try:

+            res += c.callback() or []

+        except Exception:

+            report_exception(c, 'ui_tabs_callback')

+

+    return res

+

+

+def ui_settings_callback():

+    for c in callbacks_ui_settings:

+        try:

+            c.callback()

+        except Exception:

+            report_exception(c, 'ui_settings_callback')

+

+

+def before_image_saved_callback(params: ImageSaveParams):

+    for c in callbacks_image_saved:

+        try:

+            c.callback(params)

+        except Exception:

+            report_exception(c, 'before_image_saved_callback')

+

+

+def image_saved_callback(params: ImageSaveParams):

+    for c in callbacks_image_saved:

+        try:

+            c.callback(params)

+        except Exception:

+            report_exception(c, 'image_saved_callback')

+

+

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

+

+    callbacks.append(ScriptCallback(filename, fun))

+

+

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

+    add_callback(callbacks_model_loaded, callback)

+

+

+def on_ui_tabs(callback):

+    """register a function to be called when the UI is creating new tabs.

+    The function must either return a None, which means no new tabs to be added, or a list, where

+    each element is a tuple:

+        (gradio_component, title, elem_id)

+

+    gradio_component is a gradio component to be used for contents of the tab (usually gr.Blocks)

+    title is tab text displayed to user in the UI

+    elem_id is HTML id for the tab

+    """

+    add_callback(callbacks_ui_tabs, callback)

+

+

+def on_ui_settings(callback):

+    """register a function to be called before UI settings are populated; add your settings

+    by using shared.opts.add_option(shared.OptionInfo(...)) """

+    add_callback(callbacks_ui_settings, callback)

+

+

+def on_before_image_saved(callback):

+    """register a function to be called before an image is saved to a file.

+    The callback is called with one argument:

+        - params: ImageSaveParams - parameters the image is to be saved with. You can change fields in this object.

+    """

+    add_callback(callbacks_before_image_saved, callback)

+

+

+def on_image_saved(callback):

+    """register a function to be called after an image is saved to a file.

+    The callback is called with one argument:

+        - params: ImageSaveParams - parameters the image was saved with. Changing fields in this object does nothing.

+    """

+    add_callback(callbacks_image_saved, callback)

diff --git a/modules/scripts.py b/modules/scripts.py
index 1039fa9c..533db45c 100644
--- a/modules/scripts.py
+++ b/modules/scripts.py
@@ -1,86 +1,171 @@
 import os

 import sys

 import traceback

+from collections import namedtuple

 

 import modules.ui as ui

 import gradio as gr

 

 from modules.processing import StableDiffusionProcessing

-from modules import shared

+from modules import shared, paths, script_callbacks, extensions

+

+AlwaysVisible = object()

+

 

 class Script:

     filename = None

     args_from = None

     args_to = None

+    alwayson = False

+

+    infotext_fields = None

+    """if set in ui(), this is a list of pairs of gradio component + text; the text will be used when

+    parsing infotext to set the value for the component; see ui.py's txt2img_paste_fields for an example

+    """

 

-    # The title of the script. This is what will be displayed in the dropdown menu.

     def title(self):

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

+

         raise NotImplementedError()

 

-    # How the script is displayed in the UI. See https://gradio.app/docs/#components

-    # for the different UI components you can use and how to create them.

-    # Most UI components can return a value, such as a boolean for a checkbox.

-    # The returned values are passed to the run method as parameters.

     def ui(self, is_img2img):

+        """this function should create gradio UI elements. See https://gradio.app/docs/#components

+        The return value should be an array of all components that are used in processing.

+        Values of those returned componenbts will be passed to run() and process() functions.

+        """

+

         pass

 

-    # Determines when the script should be shown in the dropdown menu via the 

-    # returned value. As an example:

-    # is_img2img is True if the current tab is img2img, and False if it is txt2img.

-    # Thus, return is_img2img to only show the script on the img2img tab.

     def show(self, is_img2img):

+        """

+        is_img2img is True if this function is called for the img2img interface, and Fasle otherwise

+

+        This function should return:

+         - False if the script should not be shown in UI at all

+         - True if the script should be shown in UI if it's scelected in the scripts drowpdown

+         - script.AlwaysVisible if the script should be shown in UI at all times

+         """

+

         return True

 

-    # This is where the additional processing is implemented. The parameters include

-    # self, the model object "p" (a StableDiffusionProcessing class, see

-    # processing.py), and the parameters returned by the ui method.

-    # Custom functions can be defined here, and additional libraries can be imported 

-    # to be used in processing. The return value should be a Processed object, which is

-    # what is returned by the process_images method.

-    def run(self, *args):

+    def run(self, p, *args):

+        """

+        This function is called if the script has been selected in the script dropdown.

+        It must do all processing and return the Processed object with results, same as

+        one returned by processing.process_images.

+

+        Usually the processing is done by calling the processing.process_images function.

+

+        args contains all values returned by components from ui()

+        """

+

         raise NotImplementedError()

 

-    # The description method is currently unused.

-    # To add a description that appears when hovering over the title, amend the "titles" 

-    # dict in script.js to include the script title (returned by title) as a key, and 

-    # your description as the value.

+    def process(self, p, *args):

+        """

+        This function is called before processing begins for AlwaysVisible scripts.

+        You can modify the processing object (p) here, inject hooks, etc.

+        args contains all values returned by components from ui()

+        """

+

+        pass

+

+    def postprocess(self, p, processed, *args):

+        """

+        This function is called after processing ends for AlwaysVisible scripts.

+        args contains all values returned by components from ui()

+        """

+

+        pass

+

     def describe(self):

+        """unused"""

         return ""

 

 

+current_basedir = paths.script_path

+

+

+def basedir():

+    """returns the base directory for the current script. For scripts in the main scripts directory,

+    this is the main directory (where webui.py resides), and for scripts in extensions directory

+    (ie extensions/aesthetic/script/aesthetic.py), this is extension's directory (extensions/aesthetic)

+    """

+    return current_basedir

+

+

 scripts_data = []

+ScriptFile = namedtuple("ScriptFile", ["basedir", "filename", "path"])

+ScriptClassData = namedtuple("ScriptClassData", ["script_class", "path", "basedir"])

 

 

-def load_scripts(basedir):

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

-        return

+def list_scripts(scriptdirname, extension):

+    scripts_list = []

 

-    for filename in sorted(os.listdir(basedir)):

-        path = os.path.join(basedir, filename)

+    basedir = os.path.join(paths.script_path, scriptdirname)

+    if os.path.exists(basedir):

+        for filename in sorted(os.listdir(basedir)):

+            scripts_list.append(ScriptFile(paths.script_path, filename, os.path.join(basedir, filename)))

+

+    for ext in extensions.active():

+        scripts_list += ext.list_files(scriptdirname, extension)

+

+    scripts_list = [x for x in scripts_list if os.path.splitext(x.path)[1].lower() == extension and os.path.isfile(x.path)]

+

+    return scripts_list

 

-        if os.path.splitext(path)[1].lower() != '.py':

-            continue

 

-        if not os.path.isfile(path):

+def list_files_with_name(filename):

+    res = []

+

+    dirs = [paths.script_path] + [ext.path for ext in extensions.active()]

+

+    for dirpath in dirs:

+        if not os.path.isdir(dirpath):

             continue

 

+        path = os.path.join(dirpath, filename)

+        if os.path.isfile(filename):

+            res.append(path)

+

+    return res

+

+

+def load_scripts():

+    global current_basedir

+    scripts_data.clear()

+    script_callbacks.clear_callbacks()

+

+    scripts_list = list_scripts("scripts", ".py")

+

+    syspath = sys.path

+

+    for scriptfile in sorted(scripts_list):

         try:

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

+            if scriptfile.basedir != paths.script_path:

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

+            current_basedir = scriptfile.basedir

+

+            with open(scriptfile.path, "r", encoding="utf8") as file:

                 text = file.read()

 

             from types import ModuleType

-            compiled = compile(text, path, 'exec')

-            module = ModuleType(filename)

+            compiled = compile(text, scriptfile.path, 'exec')

+            module = ModuleType(scriptfile.filename)

             exec(compiled, module.__dict__)

 

             for key, script_class in module.__dict__.items():

                 if type(script_class) == type and issubclass(script_class, Script):

-                    scripts_data.append((script_class, path))

+                    scripts_data.append(ScriptClassData(script_class, scriptfile.path, scriptfile.basedir))

 

         except Exception:

-            print(f"Error loading script: {filename}", file=sys.stderr)

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

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

 

+        finally:

+            sys.path = syspath

+            current_basedir = paths.script_path

+

 

 def wrap_call(func, filename, funcname, *args, default=None, **kwargs):

     try:

@@ -96,56 +181,80 @@ def wrap_call(func, filename, funcname, *args, default=None, **kwargs):
 class ScriptRunner:

     def __init__(self):

         self.scripts = []

+        self.selectable_scripts = []

+        self.alwayson_scripts = []

         self.titles = []

+        self.infotext_fields = []

 

     def setup_ui(self, is_img2img):

-        for script_class, path in scripts_data:

+        for script_class, path, basedir in scripts_data:

             script = script_class()

             script.filename = path

 

-            if not script.show(is_img2img):

-                continue

+            visibility = script.show(is_img2img)

 

-            self.scripts.append(script)

+            if visibility == AlwaysVisible:

+                self.scripts.append(script)

+                self.alwayson_scripts.append(script)

+                script.alwayson = True

 

-        self.titles = [wrap_call(script.title, script.filename, "title") or f"{script.filename} [error]" for script in self.scripts]

+            elif visibility:

+                self.scripts.append(script)

+                self.selectable_scripts.append(script)

 

-        dropdown = gr.Dropdown(label="Script", choices=["None"] + self.titles, value="None", type="index")

-        dropdown.save_to_config = True

-        inputs = [dropdown]

+        self.titles = [wrap_call(script.title, script.filename, "title") or f"{script.filename} [error]" for script in self.selectable_scripts]

 

-        for script in self.scripts:

+        inputs = [None]

+        inputs_alwayson = [True]

+

+        def create_script_ui(script, inputs, inputs_alwayson):

             script.args_from = len(inputs)

             script.args_to = len(inputs)

 

             controls = wrap_call(script.ui, script.filename, "ui", is_img2img)

 

             if controls is None:

-                continue

+                return

 

             for control in controls:

                 control.custom_script_source = os.path.basename(script.filename)

-                control.visible = False

+                if not script.alwayson:

+                    control.visible = False

+

+            if script.infotext_fields is not None:

+                self.infotext_fields += script.infotext_fields

 

             inputs += controls

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

             script.args_to = len(inputs)

 

+        for script in self.alwayson_scripts:

+            with gr.Group():

+                create_script_ui(script, inputs, inputs_alwayson)

+

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

+            create_script_ui(script, inputs, inputs_alwayson)

+

         def select_script(script_index):

-            if 0 < script_index <= len(self.scripts):

-                script = self.scripts[script_index-1]

+            if 0 < script_index <= len(self.selectable_scripts):

+                script = self.selectable_scripts[script_index-1]

                 args_from = script.args_from

                 args_to = script.args_to

             else:

                 args_from = 0

                 args_to = 0

 

-            return [ui.gr_show(True if i == 0 else args_from <= i < args_to) for i in range(len(inputs))]

+            return [ui.gr_show(True if i == 0 else args_from <= i < args_to or is_alwayson) for i, is_alwayson in enumerate(inputs_alwayson)]

 

         def init_field(title):

             if title == 'None':

                 return

             script_index = self.titles.index(title)

-            script = self.scripts[script_index]

+            script = self.selectable_scripts[script_index]

             for i in range(script.args_from, script.args_to):

                 inputs[i].visible = True

 

@@ -164,7 +273,7 @@ class ScriptRunner:
         if script_index == 0:

             return None

 

-        script = self.scripts[script_index-1]

+        script = self.selectable_scripts[script_index-1]

 

         if script is None:

             return None

@@ -176,7 +285,25 @@ class ScriptRunner:
 

         return processed

 

-    def reload_sources(self):

+    def process(self, p):

+        for script in self.alwayson_scripts:

+            try:

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

+                script.process(p, *script_args)

+            except Exception:

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

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

+

+    def postprocess(self, p, processed):

+        for script in self.alwayson_scripts:

+            try:

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

+                script.postprocess(p, processed, *script_args)

+            except Exception:

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

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

+

+    def reload_sources(self, cache):

         for si, script in list(enumerate(self.scripts)):

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

                 args_from = script.args_from

@@ -186,9 +313,12 @@ class ScriptRunner:
 

                 from types import ModuleType

 

-                compiled = compile(text, filename, 'exec')

-                module = ModuleType(script.filename)

-                exec(compiled, module.__dict__)

+                module = cache.get(filename, None)

+                if module is None:

+                    compiled = compile(text, filename, 'exec')

+                    module = ModuleType(script.filename)

+                    exec(compiled, module.__dict__)

+                    cache[filename] = module

 

                 for key, script_class in module.__dict__.items():

                     if type(script_class) == type and issubclass(script_class, Script):

@@ -197,19 +327,22 @@ class ScriptRunner:
                         self.scripts[si].args_from = args_from

                         self.scripts[si].args_to = args_to

 

+

 scripts_txt2img = ScriptRunner()

 scripts_img2img = ScriptRunner()

 

+

 def reload_script_body_only():

-    scripts_txt2img.reload_sources()

-    scripts_img2img.reload_sources()

+    cache = {}

+    scripts_txt2img.reload_sources(cache)

+    scripts_img2img.reload_sources(cache)

 

 

-def reload_scripts(basedir):

+def reload_scripts():

     global scripts_txt2img, scripts_img2img

 

-    scripts_data.clear()

-    load_scripts(basedir)

+    load_scripts()

 

     scripts_txt2img = ScriptRunner()

     scripts_img2img = ScriptRunner()

+

diff --git a/modules/scunet_model.py b/modules/scunet_model.py
index 36a996bf..59532274 100644
--- a/modules/scunet_model.py
+++ b/modules/scunet_model.py
@@ -54,9 +54,8 @@ class UpscalerScuNET(modules.upscaler.Upscaler):
         img = img[:, :, ::-1]
         img = np.moveaxis(img, 2, 0) / 255
         img = torch.from_numpy(img).float()
-        img = img.unsqueeze(0).to(device)
+        img = devices.mps_contiguous_to(img.unsqueeze(0), device)
 
-        img = img.to(device)
         with torch.no_grad():
             output = model(img)
         output = output.squeeze().float().cpu().clamp_(0, 1).numpy()
diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index 1f8587d1..0f10828e 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -332,7 +332,6 @@ class FrozenCLIPEmbedderWithCustomWords(torch.nn.Module):
                     multipliers.append([1.0] * 75)

 

             z1 = self.process_tokens(tokens, multipliers)

-            z1 = shared.aesthetic_clip(z1, remade_batch_tokens)

             z = z1 if z is None else torch.cat((z, z1), axis=-2)

 

             remade_batch_tokens = rem_tokens

diff --git a/modules/sd_models.py b/modules/sd_models.py
index d99dbce8..f86dc3ed 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -3,11 +3,12 @@ import os.path
 import sys

 from collections import namedtuple

 import torch

+import re

 from omegaconf import OmegaConf

 

 from ldm.util import instantiate_from_config

 

-from modules import shared, modelloader, devices

+from modules import shared, modelloader, devices, script_callbacks

 from modules.paths import models_path

 from modules.sd_hijack_inpainting import do_inpainting_hijack, should_hijack_inpainting

 

@@ -35,8 +36,10 @@ def setup_model():
     list_models()

 

 

-def checkpoint_tiles():

-    return sorted([x.title for x in checkpoints_list.values()])

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

@@ -170,7 +173,9 @@ def load_model_weights(model, checkpoint_info):
             print(f"Global Step: {pl_sd['global_step']}")

 

         sd = get_state_dict_from_checkpoint(pl_sd)

-        missing, extra = model.load_state_dict(sd, strict=False)

+        del pl_sd

+        model.load_state_dict(sd, strict=False)

+        del sd

 

         if shared.cmd_opts.opt_channelslast:

             model.to(memory_format=torch.channels_last)

@@ -194,9 +199,10 @@ def load_model_weights(model, checkpoint_info):
 

         model.first_stage_model.to(devices.dtype_vae)

 

-        checkpoints_loaded[checkpoint_info] = model.state_dict().copy()

-        while len(checkpoints_loaded) > shared.opts.sd_checkpoint_cache:

-            checkpoints_loaded.popitem(last=False)  # LRU

+        if shared.opts.sd_checkpoint_cache > 0:

+            checkpoints_loaded[checkpoint_info] = model.state_dict().copy()

+            while len(checkpoints_loaded) > shared.opts.sd_checkpoint_cache:

+                checkpoints_loaded.popitem(last=False)  # LRU

     else:

         print(f"Loading weights [{sd_model_hash}] from cache")

         checkpoints_loaded.move_to_end(checkpoint_info)

@@ -238,6 +244,9 @@ def load_model(checkpoint_info=None):
     sd_hijack.model_hijack.hijack(sd_model)

 

     sd_model.eval()

+    shared.sd_model = sd_model

+

+    script_callbacks.model_loaded_callback(sd_model)

 

     print(f"Model loaded.")

     return sd_model

@@ -252,7 +261,7 @@ def reload_model_weights(sd_model, info=None):
 

     if sd_model.sd_checkpoint_info.config != checkpoint_info.config or should_hijack_inpainting(checkpoint_info) != should_hijack_inpainting(sd_model.sd_checkpoint_info):

         checkpoints_loaded.clear()

-        shared.sd_model = load_model(checkpoint_info)

+        load_model(checkpoint_info)

         return shared.sd_model

 

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

@@ -265,6 +274,7 @@ def reload_model_weights(sd_model, info=None):
     load_model_weights(sd_model, checkpoint_info)

 

     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)

diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index f58a29b9..3670b57d 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -7,7 +7,7 @@ import inspect
 import k_diffusion.sampling

 import ldm.models.diffusion.ddim

 import ldm.models.diffusion.plms

-from modules import prompt_parser, devices, processing

+from modules import prompt_parser, devices, processing, images

 

 from modules.shared import opts, cmd_opts, state

 import modules.shared as shared

@@ -71,6 +71,7 @@ sampler_extra_params = {
     'sample_dpm_2': ['s_churn', 's_tmin', 's_tmax', 's_noise'],

 }

 

+

 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

@@ -82,14 +83,22 @@ def setup_img2img_steps(p, steps=None):
     return steps, t_enc

 

 

-def sample_to_image(samples):

-    x_sample = processing.decode_first_stage(shared.sd_model, samples[0:1])[0]

+def single_sample_to_image(sample):

+    x_sample = processing.decode_first_stage(shared.sd_model, sample.unsqueeze(0))[0]

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

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

     x_sample = x_sample.astype(np.uint8)

     return Image.fromarray(x_sample)

 

 

+def sample_to_image(samples):

+    return single_sample_to_image(samples[0])

+

+

+def samples_to_image_grid(samples):

+    return images.image_grid([single_sample_to_image(sample) for sample in samples])

+

+

 def store_latent(decoded):

     state.current_latent = decoded

 

@@ -219,7 +228,7 @@ class VanillaStableDiffusionSampler:
             unconditional_conditioning = {"c_concat": [image_conditioning], "c_crossattn": [unconditional_conditioning]}

             

             

-        samples = self.launch_sampling(steps, lambda: self.sampler.decode(x1, conditioning, t_enc, unconditional_guidance_scale=p.cfg_scale, unconditional_conditioning=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

 

@@ -420,7 +429,7 @@ class KDiffusionSampler:
         self.model_wrap_cfg.init_latent = x

         self.last_latent = x

 

-        samples = self.launch_sampling(steps, lambda: self.func(self.model_wrap_cfg, xi, extra_args={

+        samples = self.launch_sampling(t_enc + 1, lambda: self.func(self.model_wrap_cfg, xi, extra_args={

             'cond': conditioning, 

             'image_cond': image_conditioning, 

             'uncond': unconditional_conditioning, 

diff --git a/modules/shared.py b/modules/shared.py
index ab5a0e9a..01c2942b 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -31,7 +31,6 @@ parser.add_argument("--no-half-vae", action='store_true', help="do not switch th
 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("--aesthetic_embeddings-dir", type=str, default=os.path.join(models_path, 'aesthetic_embeddings'), help="aesthetic_embeddings directory(default: aesthetic_embeddings)")

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

@@ -41,7 +40,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("--share", action='store_true', help="use share=True for gradio and make the UI accessible through their site (doesn't work for me but you might have better luck)")

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

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

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

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

@@ -59,12 +58,13 @@ parser.add_argument("--opt-split-attention", action='store_true', help="force-en
 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("--use-cpu", nargs='+',choices=['all', 'sd', 'interrogate', 'gfpgan', 'bsrgan', 'esrgan', 'scunet', 'codeformer'], help="use CPU as torch device for specified modules", default=[], type=str.lower)

+parser.add_argument("--use-cpu", nargs='+',choices=['all', 'sd', 'interrogate', 'gfpgan', 'swinir', 'esrgan', 'scunet', 'codeformer'], 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("--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("--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)

@@ -80,10 +80,14 @@ parser.add_argument('--vae-path', type=str, help='Path to Variational Autoencode
 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 with the webui")

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

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

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

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

 

 cmd_opts = parser.parse_args()

-restricted_opts = [

+restricted_opts = {

     "samples_filename_pattern",

+    "directories_filename_pattern",

     "outdir_samples",

     "outdir_txt2img_samples",

     "outdir_img2img_samples",

@@ -91,10 +95,12 @@ restricted_opts = [
     "outdir_grids",

     "outdir_txt2img_grids",

     "outdir_save",

-]

+}

 

-devices.device, devices.device_interrogate, devices.device_gfpgan, devices.device_bsrgan, devices.device_esrgan, devices.device_scunet, devices.device_codeformer = \

-(devices.cpu if any(y in cmd_opts.use_cpu for y in [x, 'all']) else devices.get_optimal_device() for x in ['sd', 'interrogate', 'gfpgan', 'bsrgan', 'esrgan', 'scunet', 'codeformer'])

+cmd_opts.disable_extension_access = cmd_opts.share or cmd_opts.listen

+

+devices.device, devices.device_interrogate, devices.device_gfpgan, devices.device_swinir, devices.device_esrgan, devices.device_scunet, devices.device_codeformer = \

+(devices.cpu if any(y in cmd_opts.use_cpu for y in [x, 'all']) else devices.get_optimal_device() for x in ['sd', 'interrogate', 'gfpgan', 'swinir', 'esrgan', 'scunet', 'codeformer'])

 

 device = devices.device

 weight_load_location = None if cmd_opts.lowram else "cpu"

@@ -108,21 +114,6 @@ os.makedirs(cmd_opts.hypernetwork_dir, exist_ok=True)
 hypernetworks = hypernetwork.list_hypernetworks(cmd_opts.hypernetwork_dir)

 loaded_hypernetwork = None

 

-

-os.makedirs(cmd_opts.aesthetic_embeddings_dir, exist_ok=True)

-aesthetic_embeddings = {}

-

-

-def update_aesthetic_embeddings():

-    global aesthetic_embeddings

-    aesthetic_embeddings = {f.replace(".pt", ""): os.path.join(cmd_opts.aesthetic_embeddings_dir, f) for f in

-                            os.listdir(cmd_opts.aesthetic_embeddings_dir) if f.endswith(".pt")}

-    aesthetic_embeddings = OrderedDict(**{"None": None}, **aesthetic_embeddings)

-

-

-update_aesthetic_embeddings()

-

-

 def reload_hypernetworks():

     global hypernetworks

 

@@ -143,6 +134,7 @@ class State:
     current_image = None

     current_image_sampling_step = 0

     textinfo = None

+    need_restart = False

 

     def skip(self):

         self.skipped = True

@@ -155,9 +147,38 @@ class State:
         self.sampling_step = 0

         self.current_image_sampling_step = 0

 

-    def get_job_timestamp(self):

-        return datetime.datetime.now().strftime("%Y%m%d%H%M%S")  # shouldn't this return job_timestamp?

+    def dict(self):

+        obj = {

+            "skipped": self.skipped,

+            "interrupted": self.skipped,

+            "job": self.job,

+            "job_count": self.job_count,

+            "job_no": self.job_no,

+            "sampling_step": self.sampling_step,

+            "sampling_steps": self.sampling_steps,

+        }

+

+        return obj

+

+    def begin(self):

+        self.sampling_step = 0

+        self.job_count = -1

+        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.skipped = False

+        self.interrupted = False

+        self.textinfo = None

+

+        devices.torch_gc()

+

+    def end(self):

+        self.job = ""

+        self.job_count = 0

 

+        devices.torch_gc()

 

 state = State()

 

@@ -179,13 +200,13 @@ def realesrgan_models_names():
 

 

 class OptionInfo:

-    def __init__(self, default=None, label="", component=None, component_args=None, onchange=None, show_on_main_page=False, refresh=None):

+    def __init__(self, default=None, label="", component=None, component_args=None, onchange=None, section=None, refresh=None):

         self.default = default

         self.label = label

         self.component = component

         self.component_args = component_args

         self.onchange = onchange

-        self.section = None

+        self.section = section

         self.refresh = refresh

 

 

@@ -203,7 +224,8 @@ options_templates = {}
 options_templates.update(options_section(('saving-images', "Saving images/grids"), {

     "samples_save": OptionInfo(True, "Always save all generated images"),

     "samples_format": OptionInfo('png', 'File format for images'),

-    "samples_filename_pattern": OptionInfo("", "Images filename pattern"),

+    "samples_filename_pattern": OptionInfo("", "Images filename pattern", component_args=hide_dirs),

+    "save_images_add_number": OptionInfo(True, "Add number to filename when saving", component_args=hide_dirs),

 

     "grid_save": OptionInfo(True, "Always save all generated image grids"),

     "grid_format": OptionInfo('png', 'File format for grids'),

@@ -238,8 +260,8 @@ options_templates.update(options_section(('saving-to-dirs', "Saving to a directo
     "save_to_dirs": OptionInfo(False, "Save images to a subdirectory"),

     "grid_save_to_dirs": OptionInfo(False, "Save grids to a subdirectory"),

     "use_save_to_dirs_for_ui": OptionInfo(False, "When using \"Save\" button, save images to a subdirectory"),

-    "directories_filename_pattern": OptionInfo("", "Directory name pattern"),

-    "directories_max_prompt_words": OptionInfo(8, "Max prompt words for [prompt_words] pattern", gr.Slider, {"minimum": 1, "maximum": 20, "step": 1}),

+    "directories_filename_pattern": OptionInfo("", "Directory name pattern", component_args=hide_dirs),

+    "directories_max_prompt_words": OptionInfo(8, "Max prompt words for [prompt_words] pattern", gr.Slider, {"minimum": 1, "maximum": 20, "step": 1, **hide_dirs}),

 }))

 

 options_templates.update(options_section(('upscaling', "Upscaling"), {

@@ -278,6 +300,7 @@ options_templates.update(options_section(('sd', "Stable Diffusion"), {
     "sd_checkpoint_cache": OptionInfo(0, "Checkpoints to cache in RAM", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}),

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

+    "inpainting_mask_weight": OptionInfo(1.0, "Inpainting conditioning mask strength", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),

     "img2img_color_correction": OptionInfo(False, "Apply color correction to img2img results to match original colors."),

     "save_images_before_color_correction": OptionInfo(False, "Save a copy of image before applying color correction to img2img results"),

     "img2img_fix_steps": OptionInfo(False, "With img2img, do exactly the amount of steps the slider specifies (normally you'd do less with less denoising)."),

@@ -308,11 +331,13 @@ options_templates.update(options_section(('interrogate', "Interrogate Options"),
 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 0 to disable.", gr.Slider, {"minimum": 0, "maximum": 32, "step": 1}),

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

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

     "font": OptionInfo("", "Font for image grids that have text"),

     "js_modal_lightbox": OptionInfo(True, "Enable full page image viewer"),

     "js_modal_lightbox_initially_zoomed": OptionInfo(True, "Show images zoomed in by default in full page image viewer"),

@@ -333,6 +358,12 @@ options_templates.update(options_section(('sampler-params', "Sampler parameters"
     'eta_noise_seed_delta': OptionInfo(0, "Eta noise seed delta", gr.Number, {"precision": 0}),

 }))

 

+options_templates.update(options_section((None, "Hidden options"), {

+    "disabled_extensions": OptionInfo([], "Disable those extensions"),

+}))

+

+options_templates.update()

+

 

 class Options:

     data = None

@@ -344,8 +375,9 @@ class Options:
 

     def __setattr__(self, key, value):

         if self.data is not None:

-            if key in self.data:

+            if key in self.data or key in self.data_labels:

                 self.data[key] = value

+                return

 

         return super(Options, self).__setattr__(key, value)

 

@@ -361,7 +393,7 @@ class Options:
 

     def save(self, filename):

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

-            json.dump(self.data, file)

+            json.dump(self.data, file, indent=4)

 

     def same_type(self, x, y):

         if x is None or y is None:

@@ -396,6 +428,20 @@ class Options:
         d = {k: self.data.get(k, self.data_labels.get(k).default) for k in self.data_labels.keys()}

         return json.dumps(d)

 

+    def add_option(self, key, info):

+        self.data_labels[key] = info

+

+    def reorder(self):

+        """reorder settings so that all items related to section always go together"""

+

+        section_ids = {}

+        settings_items = self.data_labels.items()

+        for k, item in settings_items:

+            if item.section not in section_ids:

+                section_ids[item.section] = len(section_ids)

+

+        self.data_labels = {k: v for k, v in sorted(settings_items, key=lambda x: section_ids[x[1].section])}

+

 

 opts = Options()

 if os.path.exists(config_filename):

@@ -407,9 +453,6 @@ sd_model = None
 

 clip_model = None

 

-from modules.aesthetic_clip import AestheticCLIP

-aesthetic_clip = AestheticCLIP()

-

 progress_print_out = sys.stdout

 

 

@@ -449,3 +492,8 @@ total_tqdm = TotalTQDM()
 

 mem_mon = modules.memmon.MemUsageMonitor("MemMon", device, opts)

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

diff --git a/modules/swinir_model.py b/modules/swinir_model.py
index baa02e3d..4253b66d 100644
--- a/modules/swinir_model.py
+++ b/modules/swinir_model.py
@@ -7,8 +7,8 @@ from PIL import Image
 from basicsr.utils.download_util import load_file_from_url
 from tqdm import tqdm
 
-from modules import modelloader
-from modules.shared import cmd_opts, opts, device
+from modules import modelloader, devices
+from modules.shared import cmd_opts, opts
 from modules.swinir_model_arch import SwinIR as net
 from modules.swinir_model_arch_v2 import Swin2SR as net2
 from modules.upscaler import Upscaler, UpscalerData
@@ -42,7 +42,7 @@ class UpscalerSwinIR(Upscaler):
         model = self.load_model(model_file)
         if model is None:
             return img
-        model = model.to(device)
+        model = model.to(devices.device_swinir)
         img = upscale(img, model)
         try:
             torch.cuda.empty_cache()
@@ -111,7 +111,7 @@ def upscale(
     img = img[:, :, ::-1]
     img = np.moveaxis(img, 2, 0) / 255
     img = torch.from_numpy(img).float()
-    img = img.unsqueeze(0).to(device)
+    img = devices.mps_contiguous_to(img.unsqueeze(0), devices.device_swinir)
     with torch.no_grad(), precision_scope("cuda"):
         _, _, h_old, w_old = img.size()
         h_pad = (h_old // window_size + 1) * window_size - h_old
@@ -139,8 +139,8 @@ def inference(img, model, tile, tile_overlap, window_size, scale):
     stride = tile - tile_overlap
     h_idx_list = list(range(0, h - tile, stride)) + [h - tile]
     w_idx_list = list(range(0, w - tile, stride)) + [w - tile]
-    E = torch.zeros(b, c, h * sf, w * sf, dtype=torch.half, device=device).type_as(img)
-    W = torch.zeros_like(E, dtype=torch.half, device=device)
+    E = torch.zeros(b, c, h * sf, w * sf, dtype=torch.half, device=devices.device_swinir).type_as(img)
+    W = torch.zeros_like(E, dtype=torch.half, device=devices.device_swinir)
 
     with tqdm(total=len(h_idx_list) * len(w_idx_list), desc="SwinIR tiles") as pbar:
         for h_idx in h_idx_list:
diff --git a/modules/textual_inversion/autocrop.py b/modules/textual_inversion/autocrop.py
new file mode 100644
index 00000000..9859974a
--- /dev/null
+++ b/modules/textual_inversion/autocrop.py
@@ -0,0 +1,341 @@
+import cv2

+import requests

+import os

+from collections import defaultdict

+from math import log, sqrt

+import numpy as np

+from PIL import Image, ImageDraw

+

+GREEN = "#0F0"

+BLUE = "#00F"

+RED = "#F00"

+

+

+def crop_image(im, settings):

+  """ Intelligently crop an image to the subject matter """

+

+  scale_by = 1

+  if is_landscape(im.width, im.height):

+    scale_by = settings.crop_height / im.height

+  elif is_portrait(im.width, im.height):

+    scale_by = settings.crop_width / im.width

+  elif is_square(im.width, im.height):

+    if is_square(settings.crop_width, settings.crop_height):

+      scale_by = settings.crop_width / im.width

+    elif is_landscape(settings.crop_width, settings.crop_height):

+      scale_by = settings.crop_width / im.width

+    elif is_portrait(settings.crop_width, settings.crop_height):

+      scale_by = settings.crop_height / im.height

+

+  im = im.resize((int(im.width * scale_by), int(im.height * scale_by)))

+  im_debug = im.copy()

+

+  focus = focal_point(im_debug, settings)

+

+  # take the focal point and turn it into crop coordinates that try to center over the focal

+  # point but then get adjusted back into the frame

+  y_half = int(settings.crop_height / 2)

+  x_half = int(settings.crop_width / 2)

+

+  x1 = focus.x - x_half

+  if x1 < 0:

+      x1 = 0

+  elif x1 + settings.crop_width > im.width:

+      x1 = im.width - settings.crop_width

+

+  y1 = focus.y - y_half

+  if y1 < 0:

+      y1 = 0

+  elif y1 + settings.crop_height > im.height:

+      y1 = im.height - settings.crop_height

+

+  x2 = x1 + settings.crop_width

+  y2 = y1 + settings.crop_height

+

+  crop = [x1, y1, x2, y2]

+

+  results = []

+

+  results.append(im.crop(tuple(crop)))

+

+  if settings.annotate_image:

+    d = ImageDraw.Draw(im_debug)

+    rect = list(crop)

+    rect[2] -= 1

+    rect[3] -= 1

+    d.rectangle(rect, outline=GREEN)

+    results.append(im_debug)

+    if settings.destop_view_image:

+      im_debug.show()

+

+  return results

+

+def focal_point(im, settings):

+    corner_points = image_corner_points(im, settings) if settings.corner_points_weight > 0 else []

+    entropy_points = image_entropy_points(im, settings) if settings.entropy_points_weight > 0 else []

+    face_points = image_face_points(im, settings) if settings.face_points_weight > 0 else []

+

+    pois = []

+

+    weight_pref_total = 0

+    if len(corner_points) > 0:

+      weight_pref_total += settings.corner_points_weight

+    if len(entropy_points) > 0:

+      weight_pref_total += settings.entropy_points_weight

+    if len(face_points) > 0:

+      weight_pref_total += settings.face_points_weight

+

+    corner_centroid = None

+    if len(corner_points) > 0:

+      corner_centroid = centroid(corner_points)

+      corner_centroid.weight = settings.corner_points_weight / weight_pref_total 

+      pois.append(corner_centroid)

+

+    entropy_centroid = None

+    if len(entropy_points) > 0:

+      entropy_centroid = centroid(entropy_points)

+      entropy_centroid.weight = settings.entropy_points_weight / weight_pref_total

+      pois.append(entropy_centroid)

+

+    face_centroid = None

+    if len(face_points) > 0:

+      face_centroid = centroid(face_points)

+      face_centroid.weight = settings.face_points_weight / weight_pref_total 

+      pois.append(face_centroid)

+

+    average_point = poi_average(pois, settings)

+

+    if settings.annotate_image:

+      d = ImageDraw.Draw(im)

+      max_size = min(im.width, im.height) * 0.07

+      if corner_centroid is not None:

+        color = BLUE

+        box = corner_centroid.bounding(max_size * corner_centroid.weight)

+        d.text((box[0], box[1]-15), "Edge: %.02f" % corner_centroid.weight, fill=color)

+        d.ellipse(box, outline=color)

+        if len(corner_points) > 1:

+          for f in corner_points:

+            d.rectangle(f.bounding(4), outline=color)

+      if entropy_centroid is not None:

+        color = "#ff0"

+        box = entropy_centroid.bounding(max_size * entropy_centroid.weight)

+        d.text((box[0], box[1]-15), "Entropy: %.02f" % entropy_centroid.weight, fill=color)

+        d.ellipse(box, outline=color)

+        if len(entropy_points) > 1:

+          for f in entropy_points:

+            d.rectangle(f.bounding(4), outline=color)

+      if face_centroid is not None:

+        color = RED

+        box = face_centroid.bounding(max_size * face_centroid.weight)

+        d.text((box[0], box[1]-15), "Face: %.02f" % face_centroid.weight, fill=color)

+        d.ellipse(box, outline=color)

+        if len(face_points) > 1:

+          for f in face_points:

+            d.rectangle(f.bounding(4), outline=color)

+

+      d.ellipse(average_point.bounding(max_size), outline=GREEN)

+      

+    return average_point

+

+

+def image_face_points(im, settings):

+    if settings.dnn_model_path is not None:

+      detector = cv2.FaceDetectorYN.create(

+          settings.dnn_model_path,

+          "",

+          (im.width, im.height),

+          0.9, # score threshold

+          0.3, # nms threshold

+          5000 # keep top k before nms

+      )

+      faces = detector.detect(np.array(im))

+      results = []

+      if faces[1] is not None:

+        for face in faces[1]:

+          x = face[0]

+          y = face[1]

+          w = face[2]

+          h = face[3]

+          results.append(

+            PointOfInterest(

+              int(x + (w * 0.5)), # face focus left/right is center

+              int(y + (h * 0.33)), # face focus up/down is close to the top of the head

+              size = w,

+              weight = 1/len(faces[1])

+            )

+          )

+      return results

+    else:

+      np_im = np.array(im)

+      gray = cv2.cvtColor(np_im, cv2.COLOR_BGR2GRAY)

+

+      tries = [

+        [ f'{cv2.data.haarcascades}haarcascade_eye.xml', 0.01 ],

+        [ f'{cv2.data.haarcascades}haarcascade_frontalface_default.xml', 0.05 ],

+        [ f'{cv2.data.haarcascades}haarcascade_profileface.xml', 0.05 ],

+        [ f'{cv2.data.haarcascades}haarcascade_frontalface_alt.xml', 0.05 ],

+        [ f'{cv2.data.haarcascades}haarcascade_frontalface_alt2.xml', 0.05 ],

+        [ f'{cv2.data.haarcascades}haarcascade_frontalface_alt_tree.xml', 0.05 ],

+        [ f'{cv2.data.haarcascades}haarcascade_eye_tree_eyeglasses.xml', 0.05 ],

+        [ f'{cv2.data.haarcascades}haarcascade_upperbody.xml', 0.05 ]

+      ]

+      for t in tries:

+        classifier = cv2.CascadeClassifier(t[0])

+        minsize = int(min(im.width, im.height) * t[1]) # at least N percent of the smallest side

+        try:

+          faces = classifier.detectMultiScale(gray, scaleFactor=1.1,

+            minNeighbors=7, minSize=(minsize, minsize), flags=cv2.CASCADE_SCALE_IMAGE)

+        except:

+          continue

+

+        if len(faces) > 0:

+          rects = [[f[0], f[1], f[0] + f[2], f[1] + f[3]] for f in faces]

+          return [PointOfInterest((r[0] +r[2]) // 2, (r[1] + r[3]) // 2, size=abs(r[0]-r[2]), weight=1/len(rects)) for r in rects]

+    return []

+

+

+def image_corner_points(im, settings):

+    grayscale = im.convert("L")

+

+    # naive attempt at preventing focal points from collecting at watermarks near the bottom

+    gd = ImageDraw.Draw(grayscale)

+    gd.rectangle([0, im.height*.9, im.width, im.height], fill="#999")

+

+    np_im = np.array(grayscale)

+

+    points = cv2.goodFeaturesToTrack(

+        np_im,

+        maxCorners=100,

+        qualityLevel=0.04,

+        minDistance=min(grayscale.width, grayscale.height)*0.06,

+        useHarrisDetector=False,

+    )

+

+    if points is None:

+        return []

+

+    focal_points = []

+    for point in points:

+      x, y = point.ravel()

+      focal_points.append(PointOfInterest(x, y, size=4, weight=1/len(points)))

+

+    return focal_points

+

+

+def image_entropy_points(im, settings):

+    landscape = im.height < im.width

+    portrait = im.height > im.width

+    if landscape:

+      move_idx = [0, 2]

+      move_max = im.size[0]

+    elif portrait:

+      move_idx = [1, 3]

+      move_max = im.size[1]

+    else:

+      return []

+

+    e_max = 0

+    crop_current = [0, 0, settings.crop_width, settings.crop_height]

+    crop_best = crop_current

+    while crop_current[move_idx[1]] < move_max:

+        crop = im.crop(tuple(crop_current))

+        e = image_entropy(crop)

+

+        if (e > e_max):

+          e_max = e

+          crop_best = list(crop_current)

+

+        crop_current[move_idx[0]] += 4

+        crop_current[move_idx[1]] += 4

+

+    x_mid = int(crop_best[0] + settings.crop_width/2)

+    y_mid = int(crop_best[1] + settings.crop_height/2)

+

+    return [PointOfInterest(x_mid, y_mid, size=25, weight=1.0)]

+

+

+def image_entropy(im):

+    # greyscale image entropy

+    # band = np.asarray(im.convert("L"))

+    band = np.asarray(im.convert("1"), dtype=np.uint8)

+    hist, _ = np.histogram(band, bins=range(0, 256))

+    hist = hist[hist > 0]

+    return -np.log2(hist / hist.sum()).sum()

+

+def centroid(pois):

+  x = [poi.x for poi in pois]

+  y = [poi.y for poi in pois]

+  return PointOfInterest(sum(x)/len(pois), sum(y)/len(pois))

+

+

+def poi_average(pois, settings):

+    weight = 0.0

+    x = 0.0

+    y = 0.0

+    for poi in pois:

+        weight += poi.weight

+        x += poi.x * poi.weight

+        y += poi.y * poi.weight

+    avg_x = round(x / weight)

+    avg_y = round(y / weight)

+

+    return PointOfInterest(avg_x, avg_y)

+

+

+def is_landscape(w, h):

+  return w > h

+

+

+def is_portrait(w, h):

+  return h > w

+

+

+def is_square(w, h):

+  return w == h

+

+

+def download_and_cache_models(dirname):

+  download_url = 'https://github.com/opencv/opencv_zoo/blob/91fb0290f50896f38a0ab1e558b74b16bc009428/models/face_detection_yunet/face_detection_yunet_2022mar.onnx?raw=true'

+  model_file_name = 'face_detection_yunet.onnx'

+

+  if not os.path.exists(dirname):

+    os.makedirs(dirname)

+

+  cache_file = os.path.join(dirname, model_file_name)

+  if not os.path.exists(cache_file):

+    print(f"downloading face detection model from '{download_url}' to '{cache_file}'")

+    response = requests.get(download_url)

+    with open(cache_file, "wb") as f:

+      f.write(response.content)

+

+  if os.path.exists(cache_file):

+    return cache_file

+  return None

+

+

+class PointOfInterest:

+  def __init__(self, x, y, weight=1.0, size=10):

+    self.x = x

+    self.y = y

+    self.weight = weight

+    self.size = size

+

+  def bounding(self, size):

+    return [

+      self.x - size//2,

+      self.y - size//2,

+      self.x + size//2,

+      self.y + size//2

+    ]

+

+

+class Settings:

+  def __init__(self, crop_width=512, crop_height=512, corner_points_weight=0.5, entropy_points_weight=0.5, face_points_weight=0.5, annotate_image=False, dnn_model_path=None):

+    self.crop_width = crop_width

+    self.crop_height = crop_height

+    self.corner_points_weight = corner_points_weight

+    self.entropy_points_weight = entropy_points_weight

+    self.face_points_weight = face_points_weight

+    self.annotate_image = annotate_image

+    self.destop_view_image = False

+    self.dnn_model_path = dnn_model_path
\ No newline at end of file
diff --git a/modules/textual_inversion/dataset.py b/modules/textual_inversion/dataset.py
index 5b1c5002..ad726577 100644
--- a/modules/textual_inversion/dataset.py
+++ b/modules/textual_inversion/dataset.py
@@ -42,6 +42,8 @@ class PersonalizedBase(Dataset):
         self.lines = lines

 

         assert data_root, 'dataset directory not specified'

+        assert os.path.isdir(data_root), "Dataset directory doesn't exist"

+        assert os.listdir(data_root), "Dataset directory is empty"

 

         cond_model = shared.sd_model.cond_stage_model

 

@@ -86,12 +88,12 @@ class PersonalizedBase(Dataset):
         assert len(self.dataset) > 0, "No images have been found in the dataset."

         self.length = len(self.dataset) * repeats // batch_size

 

-        self.initial_indexes = np.arange(len(self.dataset))

+        self.dataset_length = len(self.dataset)

         self.indexes = None

         self.shuffle()

 

     def shuffle(self):

-        self.indexes = self.initial_indexes[torch.randperm(self.initial_indexes.shape[0]).numpy()]

+        self.indexes = np.random.permutation(self.dataset_length)

 

     def create_text(self, filename_text):

         text = random.choice(self.lines)

diff --git a/modules/textual_inversion/learn_schedule.py b/modules/textual_inversion/learn_schedule.py
index 2062726a..dd0c0ad1 100644
--- a/modules/textual_inversion/learn_schedule.py
+++ b/modules/textual_inversion/learn_schedule.py
@@ -4,30 +4,37 @@ import tqdm
 class LearnScheduleIterator:

     def __init__(self, learn_rate, max_steps, cur_step=0):

         """

-        specify learn_rate as "0.001:100, 0.00001:1000, 1e-5:10000" to have lr of 0.001 until step 100, 0.00001 until 1000, 1e-5:10000 until 10000

+        specify learn_rate as "0.001:100, 0.00001:1000, 1e-5:10000" to have lr of 0.001 until step 100, 0.00001 until 1000, and 1e-5 until 10000

         """

 

         pairs = learn_rate.split(',')

         self.rates = []

         self.it = 0

         self.maxit = 0

-        for i, pair in enumerate(pairs):

-            tmp = pair.split(':')

-            if len(tmp) == 2:

-                step = int(tmp[1])

-                if step > cur_step:

-                    self.rates.append((float(tmp[0]), min(step, max_steps)))

-                    self.maxit += 1

-                    if step > max_steps:

+        try:

+            for i, pair in enumerate(pairs):

+                if not pair.strip():

+                    continue

+                tmp = pair.split(':')

+                if len(tmp) == 2:

+                    step = int(tmp[1])

+                    if step > cur_step:

+                        self.rates.append((float(tmp[0]), min(step, max_steps)))

+                        self.maxit += 1

+                        if step > max_steps:

+                            return

+                    elif step == -1:

+                        self.rates.append((float(tmp[0]), max_steps))

+                        self.maxit += 1

                         return

-                elif step == -1:

+                else:

                     self.rates.append((float(tmp[0]), max_steps))

                     self.maxit += 1

                     return

-            else:

-                self.rates.append((float(tmp[0]), max_steps))

-                self.maxit += 1

-                return

+            assert self.rates

+        except (ValueError, AssertionError):

+            raise Exception('Invalid learning rate schedule. It should be a number or, for example, like "0.001:100, 0.00001:1000, 1e-5:10000" to have lr of 0.001 until step 100, 0.00001 until 1000, and 1e-5 until 10000.')

+

 

     def __iter__(self):

         return self

@@ -52,7 +59,7 @@ class LearnRateScheduler:
         self.finished = False

 

     def apply(self, optimizer, step_number):

-        if step_number <= self.end_step:

+        if step_number < self.end_step:

             return

 

         try:

diff --git a/modules/textual_inversion/preprocess.py b/modules/textual_inversion/preprocess.py
index 33eaddb6..e13b1894 100644
--- a/modules/textual_inversion/preprocess.py
+++ b/modules/textual_inversion/preprocess.py
@@ -7,12 +7,14 @@ import tqdm
 import time

 

 from modules import shared, images

+from modules.paths import models_path

 from modules.shared import opts, cmd_opts

+from modules.textual_inversion import autocrop

 if cmd_opts.deepdanbooru:

     import modules.deepbooru as deepbooru

 

 

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

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

     try:

         if process_caption:

             shared.interrogator.load()

@@ -22,7 +24,7 @@ def preprocess(process_src, process_dst, process_width, process_height, preproce
             db_opts[deepbooru.OPT_INCLUDE_RANKS] = False

             deepbooru.create_deepbooru_process(opts.interrogate_deepbooru_score_threshold, db_opts)

 

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

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

 

     finally:

 

@@ -34,7 +36,7 @@ def preprocess(process_src, process_dst, process_width, process_height, preproce
 

 

 

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

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

     width = process_width

     height = process_height

     src = os.path.abspath(process_src)

@@ -113,6 +115,7 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
                 splitted = image.crop((0, y, to_w, y + to_h))

             yield splitted

 

+

     for index, imagefile in enumerate(tqdm.tqdm(files)):

         subindex = [0]

         filename = os.path.join(src, imagefile)

@@ -137,11 +140,36 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
             ratio = (img.height * width) / (img.width * height)

             inverse_xy = True

 

+        process_default_resize = True

+

         if process_split and ratio < 1.0 and ratio <= split_threshold:

             for splitted in split_pic(img, inverse_xy):

                 save_pic(splitted, index, existing_caption=existing_caption)

-        else:

+            process_default_resize = False

+

+        if process_focal_crop and img.height != img.width:

+

+            dnn_model_path = None

+            try:

+                dnn_model_path = autocrop.download_and_cache_models(os.path.join(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)

+

+            autocrop_settings = autocrop.Settings(

+                crop_width = width,

+                crop_height = height,

+                face_points_weight = process_focal_crop_face_weight,

+                entropy_points_weight = process_focal_crop_entropy_weight,

+                corner_points_weight = process_focal_crop_edges_weight,

+                annotate_image = process_focal_crop_debug,

+                dnn_model_path = dnn_model_path,

+            )

+            for focal in autocrop.crop_image(img, autocrop_settings):

+                save_pic(focal, index, existing_caption=existing_caption)

+            process_default_resize = False

+

+        if process_default_resize:

             img = images.resize_image(1, img, width, height)

             save_pic(img, index, existing_caption=existing_caption)

 

-        shared.state.nextjob()

+        shared.state.nextjob()
\ No newline at end of file
diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py
index 529ed3e2..e0babb46 100644
--- a/modules/textual_inversion/textual_inversion.py
+++ b/modules/textual_inversion/textual_inversion.py
@@ -10,7 +10,7 @@ import csv
 

 from PIL import Image, PngImagePlugin

 

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

+from modules import shared, devices, sd_hijack, processing, sd_models, images

 import modules.textual_inversion.dataset

 from modules.textual_inversion.learn_schedule import LearnRateScheduler

 

@@ -119,7 +119,7 @@ class EmbeddingDatabase:
             vec = emb.detach().to(devices.device, dtype=torch.float32)

             embedding = Embedding(vec, name)

             embedding.step = data.get('step', None)

-            embedding.sd_checkpoint = data.get('hash', None)

+            embedding.sd_checkpoint = data.get('sd_checkpoint', None)

             embedding.sd_checkpoint_name = data.get('sd_checkpoint_name', None)

             self.register_embedding(embedding, shared.sd_model)

 

@@ -157,6 +157,9 @@ def create_embedding(name, num_vectors_per_token, overwrite_old, init_text='*'):
     cond_model = shared.sd_model.cond_stage_model

     embedding_layer = cond_model.wrapped.transformer.text_model.embeddings

 

+    with devices.autocast():

+        cond_model([""])  # will send cond model to GPU if lowvram/medvram is active

+

     ids = cond_model.tokenizer(init_text, max_length=num_vectors_per_token, return_tensors="pt", add_special_tokens=False)["input_ids"]

     embedded = embedding_layer.token_embedding.wrapped(ids.to(devices.device)).squeeze(0)

     vec = torch.zeros((num_vectors_per_token, embedded.shape[1]), device=devices.device)

@@ -164,6 +167,8 @@ def create_embedding(name, num_vectors_per_token, overwrite_old, init_text='*'):
     for i in range(num_vectors_per_token):

         vec[i] = embedded[i * int(embedded.shape[0]) // num_vectors_per_token]

 

+    # Remove illegal characters from name.

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

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

     if not overwrite_old:

         assert not os.path.exists(fn), f"file {fn} already exists"

@@ -179,9 +184,8 @@ def write_loss(log_directory, filename, step, epoch_len, values):
     if shared.opts.training_write_csv_every == 0:

         return

 

-    if step % shared.opts.training_write_csv_every != 0:

+    if (step + 1) % shared.opts.training_write_csv_every != 0:

         return

-

     write_csv_header = False if os.path.exists(os.path.join(log_directory, filename)) else True

 

     with open(os.path.join(log_directory, filename), "a+", newline='') as fout:

@@ -191,18 +195,39 @@ def write_loss(log_directory, filename, step, epoch_len, values):
             csv_writer.writeheader()

 

         epoch = step // epoch_len

-        epoch_step = step - epoch * epoch_len

+        epoch_step = step % epoch_len 

 

         csv_writer.writerow({

             "step": step + 1,

-            "epoch": epoch + 1,

+            "epoch": epoch,

             "epoch_step": epoch_step + 1,

             **values,

         })

 

+def validate_train_inputs(model_name, learn_rate, batch_size, data_root, template_file, 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"

+    assert batch_size > 0, "Batch size must be positive"

+    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 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 isinstance(save_model_every, int), "Save {name} must be integer"

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

+    if save_model_every or create_image_every:

+        assert log_directory, "Log directory is empty"

 

 def train_embedding(embedding_name, learn_rate, batch_size, data_root, log_directory, training_width, training_height, steps, 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):

-    assert embedding_name, 'embedding not selected'

+    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, data_root, template_file, steps, save_embedding_every, create_image_every, log_directory, name="embedding")

 

     shared.state.textinfo = "Initializing textual inversion training..."

     shared.state.job_count = steps

@@ -228,31 +253,36 @@ def train_embedding(embedding_name, learn_rate, batch_size, data_root, log_direc
         os.makedirs(images_embeds_dir, exist_ok=True)

     else:

         images_embeds_dir = None

-        

+

     cond_model = shared.sd_model.cond_stage_model

 

+    hijack = sd_hijack.model_hijack

+

+    embedding = hijack.embedding_db.word_embeddings[embedding_name]

+    checkpoint = sd_models.select_checkpoint()

+

+    ititial_step = embedding.step or 0

+    if ititial_step >= steps:

+        shared.state.textinfo = f"Model has already been trained beyond specified max steps"

+        return embedding, filename

+

+    scheduler = LearnRateScheduler(learn_rate, steps, ititial_step)

+

+    # 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)}..."

     with torch.autocast("cuda"):

         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, device=devices.device, template_file=template_file, batch_size=batch_size)

 

-    hijack = sd_hijack.model_hijack

-

-    embedding = hijack.embedding_db.word_embeddings[embedding_name]

     embedding.vec.requires_grad = True

+    optimizer = torch.optim.AdamW([embedding.vec], lr=scheduler.learn_rate)

 

     losses = torch.zeros((32,))

 

     last_saved_file = "<none>"

     last_saved_image = "<none>"

+    forced_filename = "<none>"

     embedding_yet_to_be_embedded = False

 

-    ititial_step = embedding.step or 0

-    if ititial_step > steps:

-        return embedding, filename

-

-    scheduler = LearnRateScheduler(learn_rate, steps, ititial_step)

-    optimizer = torch.optim.AdamW([embedding.vec], lr=scheduler.learn_rate)

-

     pbar = tqdm.tqdm(enumerate(ds), total=steps-ititial_step)

     for i, entries in pbar:

         embedding.step = i + ititial_step

@@ -276,15 +306,18 @@ def train_embedding(embedding_name, learn_rate, batch_size, data_root, log_direc
             loss.backward()

             optimizer.step()

 

+        steps_done = embedding.step + 1

 

         epoch_num = embedding.step // len(ds)

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

+        epoch_step = embedding.step % len(ds)

 

-        pbar.set_description(f"[Epoch {epoch_num}: {epoch_step}/{len(ds)}]loss: {losses.mean():.7f}")

+        pbar.set_description(f"[Epoch {epoch_num}: {epoch_step+1}/{len(ds)}]loss: {losses.mean():.7f}")

 

-        if embedding.step > 0 and embedding_dir is not None and embedding.step % save_embedding_every == 0:

-            last_saved_file = os.path.join(embedding_dir, f'{embedding_name}-{embedding.step}.pt')

-            embedding.save(last_saved_file)

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

+            save_embedding(embedding, 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, len(ds), {

@@ -292,9 +325,9 @@ def train_embedding(embedding_name, learn_rate, batch_size, data_root, log_direc
             "learn_rate": scheduler.learn_rate

         })

 

-        if embedding.step > 0 and images_dir is not None and embedding.step % create_image_every == 0:

-            last_saved_image = os.path.join(images_dir, f'{embedding_name}-{embedding.step}.png')

-

+        if images_dir is not None and steps_done % create_image_every == 0:

+            forced_filename = f'{embedding_name}-{steps_done}'

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

             p = processing.StableDiffusionProcessingTxt2Img(

                 sd_model=shared.sd_model,

                 do_not_save_grid=True,

@@ -326,7 +359,7 @@ def train_embedding(embedding_name, learn_rate, batch_size, data_root, log_direc
 

             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}-{embedding.step}.png')

+                last_saved_image_chunks = os.path.join(images_embeds_dir, f'{embedding_name}-{steps_done}.png')

 

                 info = PngImagePlugin.PngInfo()

                 data = torch.load(last_saved_file)

@@ -342,7 +375,7 @@ def train_embedding(embedding_name, learn_rate, batch_size, data_root, log_direc
                 checkpoint = sd_models.select_checkpoint()

                 footer_left = checkpoint.model_name

                 footer_mid = '[{}]'.format(checkpoint.hash)

-                footer_right = '{}v {}s'.format(vectorSize, embedding.step)

+                footer_right = '{}v {}s'.format(vectorSize, steps_done)

 

                 captioned_image = caption_image_overlay(image, title, footer_left, footer_mid, footer_right)

                 captioned_image = insert_image_data_embed(captioned_image, data)

@@ -350,8 +383,7 @@ def train_embedding(embedding_name, learn_rate, batch_size, data_root, log_direc
                 captioned_image.save(last_saved_image_chunks, "PNG", pnginfo=info)

                 embedding_yet_to_be_embedded = False

 

-            image.save(last_saved_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}"

 

         shared.state.job_no = embedding.step

@@ -366,11 +398,26 @@ Last saved image: {html.escape(last_saved_image)}<br/>
 </p>

 """

 

-    checkpoint = sd_models.select_checkpoint()

-

-    embedding.sd_checkpoint = checkpoint.hash

-    embedding.sd_checkpoint_name = checkpoint.model_name

-    embedding.cached_checksum = None

-    embedding.save(filename)

+    filename = os.path.join(shared.cmd_opts.embeddings_dir, f'{embedding_name}.pt')

+    save_embedding(embedding, checkpoint, embedding_name, filename, remove_cached_checksum=True)

 

     return embedding, filename

+

+def save_embedding(embedding, 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_name = checkpoint.model_name

+        if remove_cached_checksum:

+            embedding.cached_checksum = None

+        embedding.name = embedding_name

+        embedding.save(filename)

+    except:

+        embedding.sd_checkpoint = old_sd_checkpoint

+        embedding.sd_checkpoint_name = old_sd_checkpoint_name

+        embedding.name = old_embedding_name

+        embedding.cached_checksum = old_cached_checksum

+        raise

diff --git a/modules/txt2img.py b/modules/txt2img.py
index 1761cfa2..c9d5a090 100644
--- a/modules/txt2img.py
+++ b/modules/txt2img.py
@@ -7,7 +7,7 @@ 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, aesthetic_lr=0, aesthetic_weight=0, aesthetic_steps=0, aesthetic_imgs=None, aesthetic_slerp=False, aesthetic_imgs_text="", aesthetic_slerp_angle=0.15, aesthetic_text_negative=False, *args):

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

     p = StableDiffusionProcessingTxt2Img(

         sd_model=shared.sd_model,

         outpath_samples=opts.outdir_samples or opts.outdir_txt2img_samples,

@@ -36,7 +36,8 @@ def txt2img(prompt: str, negative_prompt: str, prompt_style: str, prompt_style2:
         firstphase_height=firstphase_height if enable_hr else None,

     )

 

-    shared.aesthetic_clip.set_aesthetic_params(p, float(aesthetic_lr), float(aesthetic_weight), int(aesthetic_steps), aesthetic_imgs, aesthetic_slerp, aesthetic_imgs_text, aesthetic_slerp_angle, aesthetic_text_negative)

+    p.scripts = modules.scripts.scripts_txt2img

+    p.script_args = args

 

     if cmd_opts.enable_console_prompts:

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

diff --git a/modules/ui.py b/modules/ui.py
index e58f040e..447722cd 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -1,58 +1,53 @@
-import base64

 import html

-import io

 import json

 import math

 import mimetypes

 import os

+import platform

 import random

+import subprocess as sp

 import sys

 import tempfile

 import time

 import traceback

-import platform

-import subprocess as sp

 from functools import partial, reduce

 

+import gradio as gr

+import gradio.routes

+import gradio.utils

 import numpy as np

-import torch

 from PIL import Image, PngImagePlugin

-import piexif

 

-import gradio as gr

-import gradio.utils

-import gradio.routes

 

-from modules import sd_hijack, sd_models, localization

+from modules import sd_hijack, sd_models, localization, script_callbacks, ui_extensions

 from modules.paths import script_path

 

-from modules.shared import opts, cmd_opts, restricted_opts, aesthetic_embeddings

+from modules.shared import opts, cmd_opts, restricted_opts

 

 if cmd_opts.deepdanbooru:

     from modules.deepbooru import get_deepbooru_tags

-import modules.shared as shared

-from modules.sd_samplers import samplers, samplers_for_img2img

-from modules.sd_hijack import model_hijack

+

+import modules.codeformer_model

+import modules.generation_parameters_copypaste as parameters_copypaste

+import modules.gfpgan_model

+import modules.hypernetworks.ui

 import modules.ldsr_model

 import modules.scripts

-import modules.gfpgan_model

-import modules.codeformer_model

+import modules.shared as shared

 import modules.styles

-import modules.generation_parameters_copypaste

+import modules.textual_inversion.ui

 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

 import modules.hypernetworks.ui

-

-import modules.aesthetic_clip as aesthetic_clip

-import modules.images_history as img_his

-

+from modules.generation_parameters_copypaste import image_from_url_text

 

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

 mimetypes.add_type('application/javascript', '.js')

 

-

 if not cmd_opts.share and not cmd_opts.listen:

     # fix gradio phoning home

     gradio.utils.version_check = lambda: None

@@ -95,37 +90,11 @@ def plaintext_to_html(text):
     text = "<p>" + "<br>\n".join([f"{html.escape(x)}" for x in text.split('\n')]) + "</p>"

     return text

 

-

-def image_from_url_text(filedata):

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

-        filename = filedata["name"]

-        tempdir = os.path.normpath(tempfile.gettempdir())

-        normfn = os.path.normpath(filename)

-        assert normfn.startswith(tempdir), 'trying to open image file not in temporary directory'

-

-        return Image.open(filename)

-

-    if type(filedata) == list:

-        if len(filedata) == 0:

-            return None

-

-        filedata = filedata[0]

-

-    if filedata.startswith("data:image/png;base64,"):

-        filedata = filedata[len("data:image/png;base64,"):]

-

-    filedata = base64.decodebytes(filedata.encode('utf-8'))

-    image = Image.open(io.BytesIO(filedata))

-    return image

-

-

 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 = []

@@ -189,7 +158,6 @@ def save_files(js_data, images, do_make_zip, index):
 

     return gr.File.update(value=fullfns, visible=True), '', '', plaintext_to_html(f"Saved: {filenames[0]}")

 

-

 def save_pil_to_file(pil_image, dir=None):

     use_metadata = False

     metadata = PngImagePlugin.PngInfo()

@@ -313,7 +281,10 @@ def check_progress_call(id_part):
         if shared.parallel_processing_allowed:

 

             if shared.state.sampling_step - shared.state.current_image_sampling_step >= opts.show_progress_every_n_steps and shared.state.current_latent is not None:

-                shared.state.current_image = modules.sd_samplers.sample_to_image(shared.state.current_latent)

+                if opts.show_progress_grid:

+                    shared.state.current_image = modules.sd_samplers.samples_to_image_grid(shared.state.current_latent)

+                else:

+                    shared.state.current_image = modules.sd_samplers.sample_to_image(shared.state.current_latent)

                 shared.state.current_image_sampling_step = shared.state.sampling_step

 

         image = shared.state.current_image

@@ -596,6 +567,9 @@ def apply_setting(key, value):
     if value is None:

         return gr.update()

 

+    if shared.cmd_opts.freeze_settings:

+        return gr.update()

+

     # dont allow model to be swapped when model hash exists in prompt

     if key == "sd_model_checkpoint" and opts.disable_weights_auto_swap:

         return gr.update()

@@ -641,10 +615,88 @@ def create_refresh_button(refresh_component, refresh_method, refreshed_args, ele
     return refresh_button

 

 

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

+

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

+

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

+

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

+

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

+

+                        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

+

+

 def create_ui(wrap_gradio_gpu_call):

     import modules.img2img

     import modules.txt2img

 

+    parameters_copypaste.reset()

 

     with gr.Blocks(analytics_enabled=False) as txt2img_interface:

         txt2img_prompt, roll, txt2img_prompt_style, txt2img_negative_prompt, txt2img_prompt_style2, submit, _, _,\

@@ -693,35 +745,11 @@ def create_ui(wrap_gradio_gpu_call):
 

                 seed, reuse_seed, subseed, reuse_subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w, seed_checkbox = create_seed_inputs()

 

-                aesthetic_weight, aesthetic_steps, aesthetic_lr, aesthetic_slerp, aesthetic_imgs, aesthetic_imgs_text, aesthetic_slerp_angle, aesthetic_text_negative = aesthetic_clip.create_ui()

-

                 with gr.Group():

                     custom_inputs = modules.scripts.scripts_txt2img.setup_ui(is_img2img=False)

 

-            with gr.Column(variant='panel'):

-

-                with gr.Group():

-                    txt2img_preview = gr.Image(elem_id='txt2img_preview', visible=False)

-                    txt2img_gallery = gr.Gallery(label='Output', show_label=False, elem_id='txt2img_gallery').style(grid=4)

-

-                with gr.Column():

-                    with gr.Row():

-                        save = gr.Button('Save')

-                        send_to_img2img = gr.Button('Send to img2img')

-                        send_to_inpaint = gr.Button('Send to inpaint')

-                        send_to_extras = gr.Button('Send to extras')

-                        button_id = "hidden_element" if shared.cmd_opts.hide_ui_dir_config else 'open_folder'

-                        open_txt2img_folder = gr.Button(folder_symbol, elem_id=button_id)

-

-                    with gr.Row():

-                        do_make_zip = gr.Checkbox(label="Make Zip when Save?", value=False)

-

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

+            txt2img_gallery, generation_info, html_info = 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)

@@ -750,14 +778,6 @@ def create_ui(wrap_gradio_gpu_call):
                     denoising_strength,

                     firstphase_width,

                     firstphase_height,

-                    aesthetic_lr,

-                    aesthetic_weight,

-                    aesthetic_steps,

-                    aesthetic_imgs,

-                    aesthetic_slerp,

-                    aesthetic_imgs_text,

-                    aesthetic_slerp_angle,

-                    aesthetic_text_negative

                 ] + custom_inputs,

 

                 outputs=[

@@ -788,23 +808,6 @@ def create_ui(wrap_gradio_gpu_call):
                 outputs=[hr_options],

             )

 

-            save.click(

-                fn=wrap_gradio_call(save_files),

-                _js="(x, y, z, w) => [x, y, z, selected_gallery_index()]",

-                inputs=[

-                    generation_info,

-                    txt2img_gallery,

-                    do_make_zip,

-                    html_info,

-                ],

-                outputs=[

-                    download_files,

-                    html_info,

-                    html_info,

-                    html_info,

-                ]

-            )

-

             roll.click(

                 fn=roll_artist,

                 _js="update_txt2img_tokens",

@@ -836,15 +839,9 @@ def create_ui(wrap_gradio_gpu_call):
                 (hr_options, lambda d: gr.Row.update(visible="Denoising strength" in d)),

                 (firstphase_width, "First pass size-1"),

                 (firstphase_height, "First pass size-2"),

-                (aesthetic_lr, "Aesthetic LR"),

-                (aesthetic_weight, "Aesthetic weight"),

-                (aesthetic_steps, "Aesthetic steps"),

-                (aesthetic_imgs, "Aesthetic embedding"),

-                (aesthetic_slerp, "Aesthetic slerp"),

-                (aesthetic_imgs_text, "Aesthetic text"),

-                (aesthetic_text_negative, "Aesthetic text negative"),

-                (aesthetic_slerp_angle, "Aesthetic slerp angle"),

+                *modules.scripts.scripts_txt2img.infotext_fields

             ]

+            parameters_copypaste.add_paste_fields("txt2img", None, txt2img_paste_fields)

 

             txt2img_preview_params = [

                 txt2img_prompt,

@@ -931,35 +928,11 @@ def create_ui(wrap_gradio_gpu_call):
 

                 seed, reuse_seed, subseed, reuse_subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w, seed_checkbox = create_seed_inputs()

 

-                aesthetic_weight_im, aesthetic_steps_im, aesthetic_lr_im, aesthetic_slerp_im, aesthetic_imgs_im, aesthetic_imgs_text_im, aesthetic_slerp_angle_im, aesthetic_text_negative_im = aesthetic_clip.create_ui()

-

                 with gr.Group():

                     custom_inputs = modules.scripts.scripts_img2img.setup_ui(is_img2img=True)

 

-            with gr.Column(variant='panel'):

-

-                with gr.Group():

-                    img2img_preview = gr.Image(elem_id='img2img_preview', visible=False)

-                    img2img_gallery = gr.Gallery(label='Output', show_label=False, elem_id='img2img_gallery').style(grid=4)

-

-                with gr.Column():

-                    with gr.Row():

-                        save = gr.Button('Save')

-                        img2img_send_to_img2img = gr.Button('Send to img2img')

-                        img2img_send_to_inpaint = gr.Button('Send to inpaint')

-                        img2img_send_to_extras = gr.Button('Send to extras')

-                        button_id = "hidden_element" if shared.cmd_opts.hide_ui_dir_config else 'open_folder'

-                        open_img2img_folder = gr.Button(folder_symbol, elem_id=button_id)

-

-                    with gr.Row():

-                        do_make_zip = gr.Checkbox(label="Make Zip when Save?", value=False)

-

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

+            img2img_gallery, generation_info, html_info = 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)

             connect_reuse_seed(subseed, reuse_subseed, generation_info, dummy_component, is_subseed=True)

@@ -1024,14 +997,6 @@ def create_ui(wrap_gradio_gpu_call):
                     inpainting_mask_invert,

                     img2img_batch_input_dir,

                     img2img_batch_output_dir,

-                    aesthetic_lr_im,

-                    aesthetic_weight_im,

-                    aesthetic_steps_im,

-                    aesthetic_imgs_im,

-                    aesthetic_slerp_im,

-                    aesthetic_imgs_text_im,

-                    aesthetic_slerp_angle_im,

-                    aesthetic_text_negative_im,

                 ] + custom_inputs,

                 outputs=[

                     img2img_gallery,

@@ -1055,25 +1020,9 @@ def create_ui(wrap_gradio_gpu_call):
                     fn=interrogate_deepbooru,

                     inputs=[init_img],

                     outputs=[img2img_prompt],

-                )

-

-            save.click(

-                fn=wrap_gradio_call(save_files),

-                _js="(x, y, z, w) => [x, y, z, selected_gallery_index()]",

-                inputs=[

-                    generation_info,

-                    img2img_gallery,

-                    do_make_zip,

-                    html_info,

-                ],

-                outputs=[

-                    download_files,

-                    html_info,

-                    html_info,

-                    html_info,

-                ]

             )

 

+

             roll.click(

                 fn=roll_artist,

                 _js="update_img2img_tokens",

@@ -1107,6 +1056,8 @@ def create_ui(wrap_gradio_gpu_call):
                     outputs=[prompt, negative_prompt, style1, style2],

                 )

 

+            token_button.click(fn=update_token_counter, inputs=[img2img_prompt, steps], outputs=[token_counter])

+

             img2img_paste_fields = [

                 (img2img_prompt, "Prompt"),

                 (img2img_negative_prompt, "Negative prompt"),

@@ -1123,16 +1074,10 @@ def create_ui(wrap_gradio_gpu_call):
                 (seed_resize_from_w, "Seed resize from-1"),

                 (seed_resize_from_h, "Seed resize from-2"),

                 (denoising_strength, "Denoising strength"),

-                (aesthetic_lr_im, "Aesthetic LR"),

-                (aesthetic_weight_im, "Aesthetic weight"),

-                (aesthetic_steps_im, "Aesthetic steps"),

-                (aesthetic_imgs_im, "Aesthetic embedding"),

-                (aesthetic_slerp_im, "Aesthetic slerp"),

-                (aesthetic_imgs_text_im, "Aesthetic text"),

-                (aesthetic_text_negative_im, "Aesthetic text negative"),

-                (aesthetic_slerp_angle_im, "Aesthetic slerp angle"),

+                *modules.scripts.scripts_img2img.infotext_fields

             ]

-            token_button.click(fn=update_token_counter, inputs=[img2img_prompt, steps], outputs=[token_counter])

+            parameters_copypaste.add_paste_fields("img2img", init_img, img2img_paste_fields)

+            parameters_copypaste.add_paste_fields("inpaint", init_img_with_mask, img2img_paste_fields)

 

     with gr.Blocks(analytics_enabled=False) as extras_interface:

         with gr.Row().style(equal_height=False):

@@ -1145,12 +1090,8 @@ def create_ui(wrap_gradio_gpu_call):
                         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."

-                        )

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

 

                 with gr.Tabs(elem_id="extras_resize_mode"):

@@ -1177,17 +1118,12 @@ def create_ui(wrap_gradio_gpu_call):
                     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)

 

-                submit = gr.Button('Generate', elem_id="extras_generate", variant='primary')

+                with gr.Group():

+                    upscale_before_face_fix = gr.Checkbox(label='Upscale Before Restoring Faces', value=False)

 

-            with gr.Column(variant='panel'):

-                result_images = gr.Gallery(label="Result", show_label=False)

-                html_info_x = gr.HTML()

-                html_info = gr.HTML()

-                extras_send_to_img2img = gr.Button('Send to img2img')

-                extras_send_to_inpaint = gr.Button('Send to inpaint')

-                button_id = "hidden_element" if shared.cmd_opts.hide_ui_dir_config else ''

-                open_extras_folder = gr.Button('Open output directory', elem_id=button_id)

+                submit = gr.Button('Generate', elem_id="extras_generate", variant='primary')

 

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

@@ -1210,6 +1146,7 @@ def create_ui(wrap_gradio_gpu_call):
                 extras_upscaler_1,

                 extras_upscaler_2,

                 extras_upscaler_2_visibility,

+                upscale_before_face_fix,

             ],

             outputs=[

                 result_images,

@@ -1217,19 +1154,11 @@ def create_ui(wrap_gradio_gpu_call):
                 html_info,

             ]

         )

+        parameters_copypaste.add_paste_fields("extras", extras_image, None)

 

-        extras_send_to_img2img.click(

-            fn=lambda x: image_from_url_text(x),

-            _js="extract_image_from_gallery_img2img",

-            inputs=[result_images],

-            outputs=[init_img],

-        )

-

-        extras_send_to_inpaint.click(

-            fn=lambda x: image_from_url_text(x),

-            _js="extract_image_from_gallery_inpaint",

-            inputs=[result_images],

-            outputs=[init_img_with_mask],

+        extras_image.change(

+            fn=modules.extras.clear_cache,

+            inputs=[], outputs=[]

         )

 

     with gr.Blocks(analytics_enabled=False) as pnginfo_interface:

@@ -1241,24 +1170,15 @@ def create_ui(wrap_gradio_gpu_call):
                 html = gr.HTML()

                 generation_info = gr.Textbox(visible=False)

                 html2 = gr.HTML()

-

                 with gr.Row():

-                    pnginfo_send_to_txt2img = gr.Button('Send to txt2img')

-                    pnginfo_send_to_img2img = gr.Button('Send to img2img')

+                    buttons = parameters_copypaste.create_buttons(["txt2img", "img2img", "inpaint", "extras"])

+                parameters_copypaste.bind_buttons(buttons, image, generation_info)

 

         image.change(

             fn=wrap_gradio_call(modules.extras.run_pnginfo),

             inputs=[image],

             outputs=[html, generation_info, html2],

         )

-    #images history

-    images_history_switch_dict = {

-        "fn":modules.generation_parameters_copypaste.connect_paste,

-        "t2i":txt2img_paste_fields,

-        "i2i":img2img_paste_fields

-    }

-

-    images_history = img_his.create_history_tabs(gr, opts, wrap_gradio_call(modules.extras.run_pnginfo), images_history_switch_dict)

 

     with gr.Blocks() as modelmerger_interface:

         with gr.Row().style(equal_height=False):

@@ -1300,25 +1220,15 @@ def create_ui(wrap_gradio_gpu_call):
                         with gr.Column():

                             create_embedding = gr.Button(value="Create embedding", variant='primary')

 

-                with gr.Tab(label="Create aesthetic images embedding"):

-

-                    new_embedding_name_ae = gr.Textbox(label="Name")

-                    process_src_ae = gr.Textbox(label='Source directory')

-                    batch_ae = gr.Slider(minimum=1, maximum=1024, step=1, label="Batch size", value=256)

-                    with gr.Row():

-                        with gr.Column(scale=3):

-                            gr.HTML(value="")

-

-                        with gr.Column():

-                            create_embedding_ae = gr.Button(value="Create images embedding", variant='primary')

-

                 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", "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", choices=modules.hypernetworks.ui.keys)

+                    new_hypernetwork_initialization_option = gr.Dropdown(value = "Normal", label="Select Layer weights initialization. relu-like - Kaiming, sigmoid-like - Xavier is recommended", 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_activation_func = gr.Dropdown(value="relu", label="Select activation function of hypernetwork", choices=["linear", "relu", "leakyrelu"])

 

                     with gr.Row():

                         with gr.Column(scale=3):

@@ -1337,6 +1247,7 @@ def create_ui(wrap_gradio_gpu_call):
                     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 if cmd_opts.deepdanbooru else False)

 

@@ -1344,6 +1255,12 @@ def create_ui(wrap_gradio_gpu_call):
                         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)

 

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

+

                     with gr.Row():

                         with gr.Column(scale=3):

                             gr.HTML(value="")

@@ -1357,6 +1274,12 @@ def create_ui(wrap_gradio_gpu_call):
                         outputs=[process_split_extra_row],

                     )

 

+                    process_focal_crop.change(

+                        fn=lambda show: gr_show(show),

+                        inputs=[process_focal_crop],

+                        outputs=[process_focal_crop_row],

+                    )

+

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

@@ -1368,7 +1291,7 @@ def create_ui(wrap_gradio_gpu_call):
                     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)

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

@@ -1411,21 +1334,6 @@ def create_ui(wrap_gradio_gpu_call):
             ]

         )

 

-        create_embedding_ae.click(

-            fn=aesthetic_clip.generate_imgs_embd,

-            inputs=[

-                new_embedding_name_ae,

-                process_src_ae,

-                batch_ae

-            ],

-            outputs=[

-                aesthetic_imgs,

-                aesthetic_imgs_im,

-                ti_output,

-                ti_outcome,

-            ]

-        )

-

         create_hypernetwork.click(

             fn=modules.hypernetworks.ui.create_hypernetwork,

             inputs=[

@@ -1433,8 +1341,10 @@ def create_ui(wrap_gradio_gpu_call):
                 new_hypernetwork_sizes,

                 overwrite_old_hypernetwork,

                 new_hypernetwork_layer_structure,

-                new_hypernetwork_add_layer_norm,

                 new_hypernetwork_activation_func,

+                new_hypernetwork_initialization_option,

+                new_hypernetwork_add_layer_norm,

+                new_hypernetwork_use_dropout

             ],

             outputs=[

                 train_hypernetwork_name,

@@ -1458,6 +1368,11 @@ def create_ui(wrap_gradio_gpu_call):
                 process_caption_deepbooru,

                 process_split_threshold,

                 process_overlap_ratio,

+                process_focal_crop,

+                process_focal_crop_face_weight,

+                process_focal_crop_entropy_weight,

+                process_focal_crop_edges_weight,

+                process_focal_crop_debug,

             ],

             outputs=[

                 ti_output,

@@ -1559,31 +1474,14 @@ def create_ui(wrap_gradio_gpu_call):
     components = []

     component_dict = {}

 

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

+    script_callbacks.ui_settings_callback()

+    opts.reorder()

 

     def run_settings(*args):

         changed = 0

 

+        assert not shared.cmd_opts.freeze_settings, "changing settings is disabled"

+

         for key, value, comp in zip(opts.data_labels.keys(), args, components):

             if comp != dummy_component and not opts.same_type(value, opts.data_labels[key].default):

                 return f"Bad value for setting {key}: {value}; expecting {type(opts.data_labels[key].default).__name__}", opts.dumpjson()

@@ -1613,13 +1511,15 @@ Requested path was: {f}
         return f'{changed} settings changed.', opts.dumpjson()

 

     def run_settings_single(value, key):

+        assert not shared.cmd_opts.freeze_settings, "changing settings is disabled"

+

         if not opts.same_type(value, opts.data_labels[key].default):

             return gr.update(visible=True), opts.dumpjson()

 

+        oldval = opts.data.get(key, None)

         if cmd_opts.hide_ui_dir_config and key in restricted_opts:

             return gr.update(value=oldval), opts.dumpjson()

 

-        oldval = opts.data.get(key, None)

         opts.data[key] = value

 

         if oldval != value:

@@ -1648,8 +1548,9 @@ Requested path was: {f}
         column = None

         with gr.Row(elem_id="settings").style(equal_height=False):

             for i, (k, item) in enumerate(opts.data_labels.items()):

+                section_must_be_skipped = item.section[0] is None

 

-                if previous_section != item.section:

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

@@ -1662,11 +1563,14 @@ Requested path was: {f}
 

                     previous_section = item.section

 

-                    gr.HTML(elem_id="settings_header_text_{}".format(item.section[0]), value='<h1 class="gr-button-lg">{}</h1>'.format(item.section[1]))

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

+                if k in quicksettings_names and not shared.cmd_opts.freeze_settings:

                     quicksettings_list.append((i, k, item))

                     components.append(dummy_component)

+                elif section_must_be_skipped:

+                    components.append(dummy_component)

                 else:

                     component = create_setting_component(k)

                     component_dict[k] = component

@@ -1697,7 +1601,7 @@ Requested path was: {f}
 

         def reload_scripts():

             modules.scripts.reload_script_body_only()

-            reload_javascript() # need to refresh the html page

+            reload_javascript()  # need to refresh the html page

 

         reload_script_bodies.click(

             fn=reload_scripts,

@@ -1708,9 +1612,10 @@ Requested path was: {f}
 

         def request_restart():

             shared.state.interrupt()

-            settings_interface.gradio_ref.do_restart = True

+            shared.state.need_restart = True

 

         restart_gradio.click(

+

             fn=request_restart,

             inputs=[],

             outputs=[],

@@ -1725,30 +1630,40 @@ Requested path was: {f}
         (img2img_interface, "img2img", "img2img"),

         (extras_interface, "Extras", "extras"),

         (pnginfo_interface, "PNG Info", "pnginfo"),

-        (images_history, "History", "images_history"),

         (modelmerger_interface, "Checkpoint Merger", "modelmerger"),

         (train_interface, "Train", "ti"),

-        (settings_interface, "Settings", "settings"),

     ]

 

-    with open(os.path.join(script_path, "style.css"), "r", encoding="utf8") as file:

-        css = file.read()

+    css = ""

+

+    for cssfile in modules.scripts.list_files_with_name("style.css"):

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

+            continue

+

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

+            css += file.read() + "\n"

 

     if os.path.exists(os.path.join(script_path, "user.css")):

         with open(os.path.join(script_path, "user.css"), "r", encoding="utf8") as file:

-            usercss = file.read()

-            css += usercss

+            css += file.read() + "\n"

 

     if not cmd_opts.no_progressbar_hiding:

         css += css_hide_progressbar

 

+    interfaces += script_callbacks.ui_tabs_callback()

+    interfaces += [(settings_interface, "Settings", "settings")]

+

+    extensions_interface = ui_extensions.create_ui()

+    interfaces += [(extensions_interface, "Extensions", "extensions")]

+

     with gr.Blocks(css=css, analytics_enabled=False, title="Stable Diffusion") as demo:

         with gr.Row(elem_id="quicksettings"):

             for i, k, item in quicksettings_list:

                 component = create_setting_component(k, is_quicksettings=True)

                 component_dict[k] = component

 

-        settings_interface.gradio_ref = demo

+        parameters_copypaste.integrate_settings_paste_fields(component_dict)

+        parameters_copypaste.run_bind()

 

         with gr.Tabs(elem_id="tabs") as tabs:

             for interface, label, ifid in interfaces:

@@ -1803,85 +1718,6 @@ Requested path was: {f}
                 component_dict['sd_model_checkpoint'],

             ]

         )

-        paste_field_names = ['Prompt', 'Negative prompt', 'Steps', 'Face restoration', 'Seed', 'Size-1', 'Size-2']

-        txt2img_fields = [field for field,name in txt2img_paste_fields if name in paste_field_names]

-        img2img_fields = [field for field,name in img2img_paste_fields if name in paste_field_names]

-        send_to_img2img.click(

-            fn=lambda img, *args: (image_from_url_text(img),*args),

-            _js="(gallery, ...args) => [extract_image_from_gallery_img2img(gallery), ...args]",

-            inputs=[txt2img_gallery] + txt2img_fields,

-            outputs=[init_img] + img2img_fields,

-        )

-

-        send_to_inpaint.click(

-            fn=lambda x, *args: (image_from_url_text(x), *args),

-            _js="(gallery, ...args) => [extract_image_from_gallery_inpaint(gallery), ...args]",

-            inputs=[txt2img_gallery] + txt2img_fields,

-            outputs=[init_img_with_mask] + img2img_fields,

-        )

-

-        img2img_send_to_img2img.click(

-            fn=lambda x: image_from_url_text(x),

-            _js="extract_image_from_gallery_img2img",

-            inputs=[img2img_gallery],

-            outputs=[init_img],

-        )

-

-        img2img_send_to_inpaint.click(

-            fn=lambda x: image_from_url_text(x),

-            _js="extract_image_from_gallery_inpaint",

-            inputs=[img2img_gallery],

-            outputs=[init_img_with_mask],

-        )

-

-        send_to_extras.click(

-            fn=lambda x: image_from_url_text(x),

-            _js="extract_image_from_gallery_extras",

-            inputs=[txt2img_gallery],

-            outputs=[extras_image],

-        )

-

-        open_txt2img_folder.click(

-            fn=lambda: open_folder(opts.outdir_samples or opts.outdir_txt2img_samples),

-            inputs=[],

-            outputs=[],

-        )

-

-        open_img2img_folder.click(

-            fn=lambda: open_folder(opts.outdir_samples or opts.outdir_img2img_samples),

-            inputs=[],

-            outputs=[],

-        )

-

-        open_extras_folder.click(

-            fn=lambda: open_folder(opts.outdir_samples or opts.outdir_extras_samples),

-            inputs=[],

-            outputs=[],

-        )

-

-        img2img_send_to_extras.click(

-            fn=lambda x: image_from_url_text(x),

-            _js="extract_image_from_gallery_extras",

-            inputs=[img2img_gallery],

-            outputs=[extras_image],

-        )

-

-        settings_map = {

-            'sd_hypernetwork': 'Hypernet',

-            'CLIP_stop_at_last_layers': 'Clip skip',

-            'sd_model_checkpoint': 'Model hash',

-        }

-

-        settings_paste_fields = [

-            (component_dict[k], lambda d, k=k, v=v: apply_setting(k, d.get(v, None)))

-            for k, v in settings_map.items()

-        ]

-

-        modules.generation_parameters_copypaste.connect_paste(txt2img_paste, txt2img_paste_fields + settings_paste_fields, txt2img_prompt)

-        modules.generation_parameters_copypaste.connect_paste(img2img_paste, img2img_paste_fields + settings_paste_fields, img2img_prompt)

-

-        modules.generation_parameters_copypaste.connect_paste(pnginfo_send_to_txt2img, txt2img_paste_fields + settings_paste_fields, generation_info, 'switch_to_txt2img')

-        modules.generation_parameters_copypaste.connect_paste(pnginfo_send_to_img2img, img2img_paste_fields + settings_paste_fields, generation_info, 'switch_to_img2img_img2img')

 

     ui_config_file = cmd_opts.ui_config_file

     ui_settings = {}

@@ -1901,7 +1737,7 @@ Requested path was: {f}
         def apply_field(obj, field, condition=None, init_field=None):

             key = path + "/" + field

 

-            if getattr(obj,'custom_script_source',None) is not None:

+            if getattr(obj, 'custom_script_source', None) is not None:

               key = 'customscript/' + obj.custom_script_source + '/' + key

 

             if getattr(obj, 'do_not_save_to_config', False):

@@ -1960,9 +1796,10 @@ def load_javascript(raw_response):
     with open(os.path.join(script_path, "script.js"), "r", encoding="utf8") as jsfile:

         javascript = f'<script>{jsfile.read()}</script>'

 

-    jsdir = os.path.join(script_path, "javascript")

-    for filename in sorted(os.listdir(jsdir)):

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

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

 

     if cmd_opts.theme is not None:

@@ -1980,6 +1817,5 @@ def load_javascript(raw_response):
     gradio.routes.templates.TemplateResponse = template_response

 

 

-reload_javascript = partial(load_javascript,

-                            gradio.routes.templates.TemplateResponse)

+reload_javascript = partial(load_javascript, gradio.routes.templates.TemplateResponse)

 reload_javascript()

diff --git a/modules/ui_extensions.py b/modules/ui_extensions.py
new file mode 100644
index 00000000..e74b7d68
--- /dev/null
+++ b/modules/ui_extensions.py
@@ -0,0 +1,172 @@
+import json

+import os.path

+import shutil

+import sys

+import time

+import traceback

+

+import git

+

+import gradio as gr

+import html

+

+from modules import extensions, shared, paths

+

+

+def check_access():

+    assert not shared.cmd_opts.disable_extension_access, "extension access disabed because of commandline flags"

+

+

+def apply_and_restart(disable_list, update_list):

+    check_access()

+

+    disabled = json.loads(disable_list)

+    assert type(disabled) == list, f"wrong disable_list data for apply_and_restart: {disable_list}"

+

+    update = json.loads(update_list)

+    assert type(update) == list, f"wrong update_list data for apply_and_restart: {update_list}"

+

+    update = set(update)

+

+    for ext in extensions.extensions:

+        if ext.name not in update:

+            continue

+

+        try:

+            ext.pull()

+        except Exception:

+            print(f"Error pulling updates for {ext.name}:", file=sys.stderr)

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

+

+    shared.opts.disabled_extensions = disabled

+    shared.opts.save(shared.config_filename)

+

+    shared.state.interrupt()

+    shared.state.need_restart = True

+

+

+def check_updates():

+    check_access()

+

+    for ext in extensions.extensions:

+        if ext.remote is None:

+            continue

+

+        try:

+            ext.check_updates()

+        except Exception:

+            print(f"Error checking updates for {ext.name}:", file=sys.stderr)

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

+

+    return extension_table()

+

+

+def extension_table():

+    code = f"""<!-- {time.time()} -->

+    <table id="extensions">

+        <thead>

+            <tr>

+                <th><abbr title="Use checkbox to enable the extension; it will be enabled or disabled when you click apply button">Extension</abbr></th>

+                <th>URL</th>

+                <th><abbr title="Use checkbox to mark the extension for update; it will be updated when you click apply button">Update</abbr></th>

+            </tr>

+        </thead>

+        <tbody>

+    """

+

+    for ext in extensions.extensions:

+        if ext.can_update:

+            ext_status = f"""<label><input class="gr-check-radio gr-checkbox" name="update_{html.escape(ext.name)}" checked="checked" type="checkbox">{html.escape(ext.status)}</label>"""

+        else:

+            ext_status = ext.status

+

+        code += f"""

+            <tr>

+                <td><label><input class="gr-check-radio gr-checkbox" name="enable_{html.escape(ext.name)}" type="checkbox" {'checked="checked"' if ext.enabled else ''}>{html.escape(ext.name)}</label></td>

+                <td><a href="{html.escape(ext.remote or '')}">{html.escape(ext.remote or '')}</a></td>

+                <td{' class="extension_status"' if ext.remote is not None else ''}>{ext_status}</td>

+            </tr>

+    """

+

+    code += """

+        </tbody>

+    </table>

+    """

+

+    return code

+

+

+def install_extension_from_url(dirname, url):

+    check_access()

+

+    assert url, 'No URL specified'

+

+    if dirname is None or dirname == "":

+        *parts, last_part = url.split('/')

+        last_part = last_part.replace(".git", "")

+

+        dirname = last_part

+

+    target_dir = os.path.join(extensions.extensions_dir, dirname)

+    assert not os.path.exists(target_dir), f'Extension directory already exists: {target_dir}'

+

+    assert len([x for x in extensions.extensions if x.remote == url]) == 0, 'Extension with this URL is already installed'

+

+    tmpdir = os.path.join(paths.script_path, "tmp", dirname)

+

+    try:

+        shutil.rmtree(tmpdir, True)

+

+        repo = git.Repo.clone_from(url, tmpdir)

+        repo.remote().fetch()

+

+        os.rename(tmpdir, target_dir)

+

+        extensions.list_extensions()

+        return [extension_table(), html.escape(f"Installed into {target_dir}. Use Installed tab to restart.")]

+    finally:

+        shutil.rmtree(tmpdir, True)

+

+

+def create_ui():

+    import modules.ui

+

+    with gr.Blocks(analytics_enabled=False) as ui:

+        with gr.Tabs(elem_id="tabs_extensions") as tabs:

+            with gr.TabItem("Installed"):

+                extensions_disabled_list = gr.Text(elem_id="extensions_disabled_list", visible=False)

+                extensions_update_list = gr.Text(elem_id="extensions_update_list", visible=False)

+

+                with gr.Row():

+                    apply = gr.Button(value="Apply and restart UI", variant="primary")

+                    check = gr.Button(value="Check for updates")

+

+                extensions_table = gr.HTML(lambda: extension_table())

+

+                apply.click(

+                    fn=apply_and_restart,

+                    _js="extensions_apply",

+                    inputs=[extensions_disabled_list, extensions_update_list],

+                    outputs=[],

+                )

+

+                check.click(

+                    fn=check_updates,

+                    _js="extensions_check",

+                    inputs=[],

+                    outputs=[extensions_table],

+                )

+

+            with gr.TabItem("Install from URL"):

+                install_url = gr.Text(label="URL for extension's git repository")

+                install_dirname = gr.Text(label="Local directory name", placeholder="Leave empty for auto")

+                intall_button = gr.Button(value="Install", variant="primary")

+                intall_result = gr.HTML(elem_id="extension_install_result")

+

+                intall_button.click(

+                    fn=modules.ui.wrap_gradio_call(install_extension_from_url, extra_outputs=[gr.update()]),

+                    inputs=[install_dirname, install_url],

+                    outputs=[extensions_table, intall_result],

+                )

+

+    return ui