25 files changed, 442 insertions, 245 deletions

diff --git a/modules/api/api.py b/modules/api/api.py
index 54ee7cb0..33845045 100644
--- a/modules/api/api.py
+++ b/modules/api/api.py
@@ -67,10 +67,10 @@ def encode_pil_to_base64(image):
 class Api:
     def __init__(self, app: FastAPI, queue_lock: Lock):
         if shared.cmd_opts.api_auth:
-            self.credenticals = dict()
+            self.credentials = dict()
             for auth in shared.cmd_opts.api_auth.split(","):
                 user, password = auth.split(":")
-                self.credenticals[user] = password
+                self.credentials[user] = password
 
         self.router = APIRouter()
         self.app = app
@@ -93,7 +93,7 @@ class Api:
         self.add_api_route("/sdapi/v1/hypernetworks", self.get_hypernetworks, methods=["GET"], response_model=List[HypernetworkItem])
         self.add_api_route("/sdapi/v1/face-restorers", self.get_face_restorers, methods=["GET"], response_model=List[FaceRestorerItem])
         self.add_api_route("/sdapi/v1/realesrgan-models", self.get_realesrgan_models, methods=["GET"], response_model=List[RealesrganItem])
-        self.add_api_route("/sdapi/v1/prompt-styles", self.get_promp_styles, methods=["GET"], response_model=List[PromptStyleItem])
+        self.add_api_route("/sdapi/v1/prompt-styles", self.get_prompt_styles, methods=["GET"], response_model=List[PromptStyleItem])
         self.add_api_route("/sdapi/v1/artist-categories", self.get_artists_categories, methods=["GET"], response_model=List[str])
         self.add_api_route("/sdapi/v1/artists", self.get_artists, methods=["GET"], response_model=List[ArtistItem])
 
@@ -102,9 +102,9 @@ class Api:
             return self.app.add_api_route(path, endpoint, dependencies=[Depends(self.auth)], **kwargs)
         return self.app.add_api_route(path, endpoint, **kwargs)
 
-    def auth(self, credenticals: HTTPBasicCredentials = Depends(HTTPBasic())):
-        if credenticals.username in self.credenticals:
-            if compare_digest(credenticals.password, self.credenticals[credenticals.username]):
+    def auth(self, credentials: HTTPBasicCredentials = Depends(HTTPBasic())):
+        if credentials.username in self.credentials:
+            if compare_digest(credentials.password, self.credentials[credentials.username]):
                 return True
 
         raise HTTPException(status_code=401, detail="Incorrect username or password", headers={"WWW-Authenticate": "Basic"})
@@ -157,12 +157,7 @@ class Api:
         args.pop('include_init_images', None)  # this is meant to be done by "exclude": True in model, but it's for a reason that I cannot determine.
         p = StableDiffusionProcessingImg2Img(**args)
 
-        imgs = []
-        for img in init_images:
-            img = decode_base64_to_image(img)
-            imgs = [img] * p.batch_size
-
-        p.init_images = imgs
+        p.init_images = [decode_base64_to_image(x) for x in init_images]
 
         shared.state.begin()
 
@@ -244,7 +239,7 @@ class Api:
     def interrogateapi(self, interrogatereq: InterrogateRequest):
         image_b64 = interrogatereq.image
         if image_b64 is None:
-            raise HTTPException(status_code=404, detail="Image not found") 
+            raise HTTPException(status_code=404, detail="Image not found")
 
         img = decode_base64_to_image(image_b64)
         img = img.convert('RGB')
@@ -257,7 +252,7 @@ class Api:
                 processed = deepbooru.model.tag(img)
             else:
                 raise HTTPException(status_code=404, detail="Model not found")
-        
+
         return InterrogateResponse(caption=processed)
 
     def interruptapi(self):
@@ -313,7 +308,7 @@ class Api:
     def get_realesrgan_models(self):
         return [{"name":x.name,"path":x.data_path, "scale":x.scale} for x in get_realesrgan_models(None)]
 
-    def get_promp_styles(self):
+    def get_prompt_styles(self):
         styleList = []
         for k in shared.prompt_styles.styles:
             style = shared.prompt_styles.styles[k]
diff --git a/modules/api/models.py b/modules/api/models.py
index f77951fc..a22bc6b3 100644
--- a/modules/api/models.py
+++ b/modules/api/models.py
@@ -128,7 +128,7 @@ class ExtrasBaseRequest(BaseModel):
     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?")
+    upscaling_crop: bool = Field(default=True, title="Crop to fit", description="Should the upscaler crop the image to fit in the chosen 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.")
diff --git a/modules/devices.py b/modules/devices.py
index f8cffae1..800510b7 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -125,7 +125,16 @@ def layer_norm_fix(*args, **kwargs):
     return orig_layer_norm(*args, **kwargs)
 
 
+# MPS workaround for https://github.com/pytorch/pytorch/issues/90532
+orig_tensor_numpy = torch.Tensor.numpy
+def numpy_fix(self, *args, **kwargs):
+    if self.requires_grad:
+        self = self.detach()
+    return orig_tensor_numpy(self, *args, **kwargs)
+
+
 # PyTorch 1.13 doesn't need these fixes but unfortunately is slower and has regressions that prevent training from working
 if has_mps() and version.parse(torch.__version__) < version.parse("1.13"):
     torch.Tensor.to = tensor_to_fix
     torch.nn.functional.layer_norm = layer_norm_fix
+    torch.Tensor.numpy = numpy_fix
diff --git a/modules/extras.py b/modules/extras.py
index bc349d5e..6fa7d856 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -62,7 +62,7 @@ def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_
     # Also keep track of original file names

     imageNameArr = []

     outputs = []

-    

+

     if extras_mode == 1:

         #convert file to pillow image

         for img in image_folder:

@@ -188,13 +188,19 @@ def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_
         for op in extras_ops:

             image, info = op(image, info)

 

-        if opts.use_original_name_batch and image_name != None:

+        if opts.use_original_name_batch and image_name is not None:

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

         else:

             basename = ''

 

+        # Add upscaler name as a suffix.

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

+        # Add second upscaler if applicable.

+        if suffix and extras_upscaler_2 and extras_upscaler_2_visibility:

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

+

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

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

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

 

         if opts.enable_pnginfo:

             image.info = existing_pnginfo

@@ -234,7 +240,7 @@ def run_pnginfo(image):
     return '', geninfo, info

 

 

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

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

     def weighted_sum(theta0, theta1, alpha):

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

 

@@ -246,30 +252,25 @@ def run_modelmerger(primary_model_name, secondary_model_name, teritary_model_nam
 

     primary_model_info = sd_models.checkpoints_list[primary_model_name]

     secondary_model_info = sd_models.checkpoints_list[secondary_model_name]

-    teritary_model_info = sd_models.checkpoints_list.get(teritary_model_name, None)

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

     result_is_inpainting_model = False

 

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

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

-

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

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

-

-    if teritary_model_info is not None:

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

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

-    else:

-        theta_2 = None

-

     theta_funcs = {

         "Weighted sum": (None, weighted_sum),

         "Add difference": (get_difference, add_difference),

     }

     theta_func1, theta_func2 = theta_funcs[interp_method]

 

-    print(f"Merging...")

+    if theta_func1 and not tertiary_model_info:

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

+

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

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

 

     if theta_func1:

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

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

+

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

             if 'model' in key:

                 if key in theta_2:

@@ -277,7 +278,12 @@ def run_modelmerger(primary_model_name, secondary_model_name, teritary_model_nam
                     theta_1[key] = theta_func1(theta_1[key], t2)

                 else:

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

-    del theta_2

+        del theta_2

+

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

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

+

+    print("Merging...")

 

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

         if 'model' in key and key in theta_1:

@@ -307,6 +313,7 @@ def run_modelmerger(primary_model_name, secondary_model_name, teritary_model_nam
             theta_0[key] = theta_1[key]

             if save_as_half:

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

+    del theta_1

 

     ckpt_dir = shared.cmd_opts.ckpt_dir or sd_models.model_path

 

@@ -332,5 +339,5 @@ def run_modelmerger(primary_model_name, secondary_model_name, teritary_model_nam
 

     sd_models.list_models()

 

-    print(f"Checkpoint saved.")

+    print("Checkpoint saved.")

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

diff --git a/modules/generation_parameters_copypaste.py b/modules/generation_parameters_copypaste.py
index 44fe1a6c..565e342d 100644
--- a/modules/generation_parameters_copypaste.py
+++ b/modules/generation_parameters_copypaste.py
@@ -77,6 +77,7 @@ def integrate_settings_paste_fields(component_dict):
         'inpainting_mask_weight': 'Conditional mask weight',

         'sd_model_checkpoint': 'Model hash',

         'eta_noise_seed_delta': 'ENSD',

+        'initial_noise_multiplier': 'Noise multiplier',

     }

     settings_paste_fields = [

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

@@ -121,8 +122,7 @@ def run_bind():
 

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

-

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

                     button.click(

                         fn=lambda *x: x,

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

diff --git a/modules/images.py b/modules/images.py
index 08a72e67..809ad9f7 100644
--- a/modules/images.py
+++ b/modules/images.py
@@ -136,8 +136,19 @@ def draw_grid_annotations(im, width, height, hor_texts, ver_texts):
                 lines.append(word)

         return lines

 

-    def draw_texts(drawing, draw_x, draw_y, lines):

+    def get_font(fontsize):

+        try:

+            return ImageFont.truetype(opts.font or Roboto, fontsize)

+        except Exception:

+            return ImageFont.truetype(Roboto, fontsize)

+

+    def draw_texts(drawing, draw_x, draw_y, lines, initial_fnt, initial_fontsize):

         for i, line in enumerate(lines):

+            fnt = initial_fnt

+            fontsize = initial_fontsize

+            while drawing.multiline_textsize(line.text, font=fnt)[0] > line.allowed_width and fontsize > 0:

+                fontsize -= 1

+                fnt = get_font(fontsize)

             drawing.multiline_text((draw_x, draw_y + line.size[1] / 2), line.text, font=fnt, fill=color_active if line.is_active else color_inactive, anchor="mm", align="center")

 

             if not line.is_active:

@@ -148,10 +159,7 @@ def draw_grid_annotations(im, width, height, hor_texts, ver_texts):
     fontsize = (width + height) // 25

     line_spacing = fontsize // 2

 

-    try:

-        fnt = ImageFont.truetype(opts.font or Roboto, fontsize)

-    except Exception:

-        fnt = ImageFont.truetype(Roboto, fontsize)

+    fnt = get_font(fontsize)

 

     color_active = (0, 0, 0)

     color_inactive = (153, 153, 153)

@@ -178,6 +186,7 @@ def draw_grid_annotations(im, width, height, hor_texts, ver_texts):
         for line in texts:

             bbox = calc_d.multiline_textbbox((0, 0), line.text, font=fnt)

             line.size = (bbox[2] - bbox[0], bbox[3] - bbox[1])

+            line.allowed_width = allowed_width

 

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

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

@@ -194,13 +203,13 @@ def draw_grid_annotations(im, width, height, hor_texts, ver_texts):
         x = pad_left + width * col + width / 2

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

 

-        draw_texts(d, x, y, hor_texts[col])

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

 

     for row in range(rows):

         x = pad_left / 2

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

 

-        draw_texts(d, x, y, ver_texts[row])

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

 

     return result

 

@@ -429,7 +438,7 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
             The directory to save the image. Note, the option `save_to_dirs` will make the image to be saved into a sub directory.

         basename (`str`):

             The base filename which will be applied to `filename pattern`.

-        seed, prompt, short_filename, 

+        seed, prompt, short_filename,

         extension (`str`):

             Image file extension, default is `png`.

         pngsectionname (`str`):

@@ -501,30 +510,39 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
     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({

-            "Exif": {

-                piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(info or "", encoding="unicode")

-            },

-        })

+    def _atomically_save_image(image_to_save, filename_without_extension, extension):

+        # save image with .tmp extension to avoid race condition when another process detects new image in the directory

+        temp_file_path = filename_without_extension + ".tmp"

+        image_format = Image.registered_extensions()[extension]

 

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

-        pnginfo_data = PngImagePlugin.PngInfo()

-        if opts.enable_pnginfo:

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

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

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

+            pnginfo_data = PngImagePlugin.PngInfo()

+            if opts.enable_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)

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

 

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

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

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

+            image_to_save.save(temp_file_path, format=image_format, 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)

+            if opts.enable_pnginfo and info is not None:

+                exif_bytes = piexif.dump({

+                    "Exif": {

+                        piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(info or "", encoding="unicode")

+                    },

+                })

+

+                piexif.insert(exif_bytes, temp_file_path)

+        else:

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

+

+        # atomically rename the file with correct extension

+        os.replace(temp_file_path, filename_without_extension + extension)

+

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

+    _atomically_save_image(image, fullfn_without_extension, extension)

 

     image.already_saved_as = fullfn

 

@@ -538,9 +556,7 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
         elif oversize:

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

 

-        image.save(fullfn_without_extension + ".jpg", quality=opts.jpeg_quality)

-        if opts.enable_pnginfo and info is not None:

-            piexif.insert(exif_bytes(), fullfn_without_extension + ".jpg")

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

 

     if opts.save_txt and info is not None:

         txt_fullfn = f"{fullfn_without_extension}.txt"

@@ -583,7 +599,7 @@ def read_info_from_image(image):
 Negative prompt: {json_info["uc"]}

 Steps: {json_info["steps"]}, Sampler: {sampler}, CFG scale: {json_info["scale"]}, Seed: {json_info["seed"]}, Size: {image.width}x{image.height}, Clip skip: 2, ENSD: 31337"""

         except Exception:

-            print(f"Error parsing NovelAI iamge generation parameters:", file=sys.stderr)

+            print(f"Error parsing NovelAI image generation parameters:", file=sys.stderr)

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

 

     return geninfo, items

@@ -606,3 +622,14 @@ def image_data(data):
         pass

 

     return '', None

+

+

+def flatten(img, bgcolor):

+    """replaces transparency with bgcolor (example: "#ffffff"), returning an RGB mode image with no transparency"""

+

+    if img.mode == "RGBA":

+        background = Image.new('RGBA', img.size, bgcolor)

+        background.paste(img, mask=img)

+        img = background

+

+    return img.convert('RGB')

diff --git a/modules/import_hook.py b/modules/import_hook.py
new file mode 100644
index 00000000..28c67dfa
--- /dev/null
+++ b/modules/import_hook.py
@@ -0,0 +1,5 @@
+import sys
+
+# this will break any attempt to import xformers which will prevent stability diffusion repo from trying to use it
+if "--xformers" not in "".join(sys.argv):
+    sys.modules["xformers"] = None
diff --git a/modules/lowvram.py b/modules/lowvram.py
index aa464a95..042a0254 100644
--- a/modules/lowvram.py
+++ b/modules/lowvram.py
@@ -55,18 +55,20 @@ def setup_for_low_vram(sd_model, use_medvram):
     if hasattr(sd_model.cond_stage_model, 'model'):

         sd_model.cond_stage_model.transformer = sd_model.cond_stage_model.model

 

-    # remove three big modules, cond, first_stage, and unet from the model and then

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

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

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

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

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

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

     sd_model.to(devices.device)

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

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

 

-    # register hooks for those the first two models

+    # register hooks for those the first three models

     sd_model.cond_stage_model.transformer.register_forward_pre_hook(send_me_to_gpu)

     sd_model.first_stage_model.register_forward_pre_hook(send_me_to_gpu)

     sd_model.first_stage_model.encode = first_stage_model_encode_wrap

     sd_model.first_stage_model.decode = first_stage_model_decode_wrap

+    if sd_model.depth_model:

+        sd_model.depth_model.register_forward_pre_hook(send_me_to_gpu)

     parents[sd_model.cond_stage_model.transformer] = sd_model.cond_stage_model

 

     if hasattr(sd_model.cond_stage_model, 'model'):

diff --git a/modules/ngrok.py b/modules/ngrok.py
index 64c9a3c2..3df2c06b 100644
--- a/modules/ngrok.py
+++ b/modules/ngrok.py
@@ -2,7 +2,7 @@ from pyngrok import ngrok, conf, exception
 
 def connect(token, port, region):
     account = None
-    if token == None:
+    if token is None:
         token = 'None'
     else:
         if ':' in token:
@@ -14,7 +14,7 @@ def connect(token, port, region):
         auth_token=token, region=region
     )
     try:
-        if account == None:
+        if account is None:
             public_url = ngrok.connect(port, pyngrok_config=config, bind_tls=True).public_url
         else:
             public_url = ngrok.connect(port, pyngrok_config=config, bind_tls=True, auth=account).public_url
diff --git a/modules/processing.py b/modules/processing.py
index ab5a34d0..75b0067c 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -13,15 +13,21 @@ from skimage import exposure
 from typing import Any, Dict, List, Optional

 

 import modules.sd_hijack

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

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

 from modules.sd_hijack import model_hijack

 from modules.shared import opts, cmd_opts, state

 import modules.shared as shared

 import modules.face_restoration

 import modules.images as images

 import modules.styles

+import modules.sd_models as sd_models

+import modules.sd_vae as sd_vae

 import logging

+from ldm.data.util import AddMiDaS

+from ldm.models.diffusion.ddpm import LatentDepth2ImageDiffusion

 

+from einops import repeat, rearrange

+from blendmodes.blend import blendLayers, BlendType

 

 # some of those options should not be changed at all because they would break the model, so I removed them from options.

 opt_C = 4

@@ -34,17 +40,19 @@ def setup_color_correction(image):
     return correction_target

 

 

-def apply_color_correction(correction, image):

+def apply_color_correction(correction, original_image):

     logging.info("Applying color correction.")

     image = Image.fromarray(cv2.cvtColor(exposure.match_histograms(

         cv2.cvtColor(

-            np.asarray(image),

+            np.asarray(original_image),

             cv2.COLOR_RGB2LAB

         ),

         correction,

         channel_axis=2

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

-

+    

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

+    

     return image

 

 

@@ -72,7 +80,7 @@ class StableDiffusionProcessing():
     """

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

     """

-    def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prompt: str = "", styles: List[str] = None, seed: int = -1, subseed: int = -1, subseed_strength: float = 0, seed_resize_from_h: int = -1, seed_resize_from_w: int = -1, seed_enable_extras: bool = True, sampler_name: str = None, batch_size: int = 1, n_iter: int = 1, steps: int = 50, cfg_scale: float = 7.0, width: int = 512, height: int = 512, restore_faces: bool = False, tiling: bool = False, do_not_save_samples: bool = False, do_not_save_grid: bool = False, extra_generation_params: Dict[Any, Any] = None, overlay_images: Any = None, negative_prompt: str = None, eta: float = None, do_not_reload_embeddings: bool = False, denoising_strength: float = 0, ddim_discretize: str = None, s_churn: float = 0.0, s_tmax: float = None, s_tmin: float = 0.0, s_noise: float = 1.0, override_settings: Dict[str, Any] = None, sampler_index: int = None):

+    def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prompt: str = "", styles: List[str] = None, seed: int = -1, subseed: int = -1, subseed_strength: float = 0, seed_resize_from_h: int = -1, seed_resize_from_w: int = -1, seed_enable_extras: bool = True, sampler_name: str = None, batch_size: int = 1, n_iter: int = 1, steps: int = 50, cfg_scale: float = 7.0, width: int = 512, height: int = 512, restore_faces: bool = False, tiling: bool = False, do_not_save_samples: bool = False, do_not_save_grid: bool = False, extra_generation_params: Dict[Any, Any] = None, overlay_images: Any = None, negative_prompt: str = None, eta: float = None, do_not_reload_embeddings: bool = False, denoising_strength: float = 0, ddim_discretize: str = None, s_churn: float = 0.0, s_tmax: float = None, s_tmin: float = 0.0, s_noise: float = 1.0, override_settings: Dict[str, Any] = None, override_settings_restore_afterwards: bool = True, sampler_index: int = None):

         if sampler_index is not None:

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

 

@@ -113,6 +121,7 @@ class StableDiffusionProcessing():
         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}

+        self.override_settings_restore_afterwards = override_settings_restore_afterwards

         self.is_using_inpainting_conditioning = False

 

         if not seed_enable_extras:

@@ -142,19 +151,34 @@ class StableDiffusionProcessing():
 

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

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

+        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 latent_image.new_zeros(latent_image.shape[0], 5, 1, 1)

+    def depth2img_image_conditioning(self, source_image):

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

+        transformer = AddMiDaS(model_type="dpt_hybrid")

+        transformed = transformer({"jpg": rearrange(source_image[0], "c h w -> h w c")})

+        midas_in = torch.from_numpy(transformed["midas_in"][None, ...]).to(device=shared.device)

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

+

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

+        conditioning = torch.nn.functional.interpolate(

+            self.sd_model.depth_model(midas_in),

+            size=conditioning_image.shape[2:],

+            mode="bicubic",

+            align_corners=False,

+        )

 

+        (depth_min, depth_max) = torch.aminmax(conditioning)

+        conditioning = 2. * (conditioning - depth_min) / (depth_max - depth_min) - 1.

+        return conditioning

+

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

         self.is_using_inpainting_conditioning = True

 

         # Handle the different mask inputs

@@ -179,7 +203,7 @@ class StableDiffusionProcessing():
             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))

 

@@ -191,6 +215,18 @@ class StableDiffusionProcessing():
 

         return image_conditioning

 

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

+        # HACK: Using introspection as the Depth2Image model doesn't appear to uniquely

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

+        if isinstance(self.sd_model, LatentDepth2ImageDiffusion):

+            return self.depth2img_image_conditioning(source_image)

+

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

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

+

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

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

+

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

         pass

 

@@ -424,15 +460,21 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
 

     try:

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

-            setattr(opts, k, v)  # we don't call onchange for simplicity which makes changing model impossible

-            if k == 'sd_hypernetwork': shared.reload_hypernetworks()  # make onchange call for changing hypernet since it is relatively fast to load on-change, while SD models are not

+            setattr(opts, k, v)

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

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

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

 

         res = process_images_inner(p)

 

-    finally:  # restore opts to original state

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

-            setattr(opts, k, v)

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

+    finally:

+        # restore opts to original state

+        if p.override_settings_restore_afterwards:

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

+                setattr(opts, k, v)

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

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

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

 

     return res

 

@@ -501,7 +543,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
         for n in range(p.n_iter):

             if state.skipped:

                 state.skipped = False

-            

+

             if state.interrupted:

                 break

 

@@ -541,9 +583,8 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
 

             devices.torch_gc()

 

-            if opts.filter_nsfw:

-                import modules.safety as safety

-                x_samples_ddim = modules.safety.censor_batch(x_samples_ddim)

+            if p.scripts is not None:

+                p.scripts.postprocess_batch(p, x_samples_ddim, batch_number=n)

 

             for i, x_sample in enumerate(x_samples_ddim):

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

@@ -577,7 +618,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                     image.info["parameters"] = text

                 output_images.append(image)

 

-            del x_samples_ddim 

+            del x_samples_ddim

 

             devices.torch_gc()

 

@@ -669,7 +710,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 

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

 

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

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

         def save_intermediate(image, index):

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

                 return

@@ -685,7 +726,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 

             samples = torch.nn.functional.interpolate(samples, size=(self.height // opt_f, self.width // opt_f), mode="bilinear")

 

-            # Avoid making the inpainting conditioning unless necessary as 

+            # Avoid making the inpainting conditioning unless necessary as

             # this does need some extra compute to decode / encode the image again.

             if getattr(self, "inpainting_mask_weight", shared.opts.inpainting_mask_weight) < 1.0:

                 image_conditioning = self.img2img_image_conditioning(decode_first_stage(self.sd_model, samples), samples)

@@ -734,7 +775,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):

     sampler = None

 

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

         super().__init__(**kwargs)

 

         self.init_images = init_images

@@ -749,6 +790,7 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
         self.inpaint_full_res = inpaint_full_res

         self.inpaint_full_res_padding = inpaint_full_res_padding

         self.inpainting_mask_invert = inpainting_mask_invert

+        self.initial_noise_multiplier = opts.initial_noise_multiplier if initial_noise_multiplier is None else initial_noise_multiplier

         self.mask = None

         self.nmask = None

         self.image_conditioning = None

@@ -793,7 +835,7 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
             self.color_corrections = []

         imgs = []

         for img in self.init_images:

-            image = img.convert("RGB")

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

 

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

                 image = images.resize_image(self.resize_mode, image, self.width, self.height)

@@ -866,6 +908,10 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
     def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts):

         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)

 

+        if self.initial_noise_multiplier != 1.0:

+            self.extra_generation_params["Noise multiplier"] = self.initial_noise_multiplier

+            x *= self.initial_noise_multiplier

+

         samples = self.sampler.sample_img2img(self, self.init_latent, x, conditioning, unconditional_conditioning, image_conditioning=self.image_conditioning)

 

         if self.mask is not None:

diff --git a/modules/safe.py b/modules/safe.py
index 10460ad0..479c8b86 100644
--- a/modules/safe.py
+++ b/modules/safe.py
@@ -37,16 +37,16 @@ class RestrictedUnpickler(pickle.Unpickler):
 

         if module == 'collections' and name == 'OrderedDict':

             return getattr(collections, name)

-        if module == 'torch._utils' and name in ['_rebuild_tensor_v2', '_rebuild_parameter']:

+        if module == 'torch._utils' and name in ['_rebuild_tensor_v2', '_rebuild_parameter', '_rebuild_device_tensor_from_numpy']:

             return getattr(torch._utils, name)

-        if module == 'torch' and name in ['FloatStorage', 'HalfStorage', 'IntStorage', 'LongStorage', 'DoubleStorage', 'ByteStorage']:

+        if module == 'torch' and name in ['FloatStorage', 'HalfStorage', 'IntStorage', 'LongStorage', 'DoubleStorage', 'ByteStorage', 'float32']:

             return getattr(torch, name)

         if module == 'torch.nn.modules.container' and name in ['ParameterDict']:

             return getattr(torch.nn.modules.container, name)

-        if module == 'numpy.core.multiarray' and name == 'scalar':

-            return numpy.core.multiarray.scalar

-        if module == 'numpy' and name == 'dtype':

-            return numpy.dtype

+        if module == 'numpy.core.multiarray' and name in ['scalar', '_reconstruct']:

+            return getattr(numpy.core.multiarray, name)

+        if module == 'numpy' and name in ['dtype', 'ndarray']:

+            return getattr(numpy, name)

         if module == '_codecs' and name == 'encode':

             return encode

         if module == "pytorch_lightning.callbacks" and name == 'model_checkpoint':

@@ -80,7 +80,7 @@ def check_pt(filename, extra_handler):
         # new pytorch format is a zip file

         with zipfile.ZipFile(filename) as z:

             check_zip_filenames(filename, z.namelist())

-            

+

             # find filename of data.pkl in zip file: '<directory name>/data.pkl'

             data_pkl_filenames = [f for f in z.namelist() if data_pkl_re.match(f)]

             if len(data_pkl_filenames) == 0:

@@ -108,7 +108,7 @@ def load(filename, *args, **kwargs):
 

 def load_with_extra(filename, extra_handler=None, *args, **kwargs):

     """

-    this functon is intended to be used by extensions that want to load models with

+    this function is intended to be used by extensions that want to load models with

     some extra classes in them that the usual unpickler would find suspicious.

 

     Use the extra_handler argument to specify a function that takes module and field name as text,

diff --git a/modules/safety.py b/modules/safety.py
deleted file mode 100644
index cff4b278..00000000
--- a/modules/safety.py
+++ /dev/null
@@ -1,42 +0,0 @@
-import torch

-from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker

-from transformers import AutoFeatureExtractor

-from PIL import Image

-

-import modules.shared as shared

-

-safety_model_id = "CompVis/stable-diffusion-safety-checker"

-safety_feature_extractor = None

-safety_checker = None

-

-def numpy_to_pil(images):

-    """

-    Convert a numpy image or a batch of images to a PIL image.

-    """

-    if images.ndim == 3:

-        images = images[None, ...]

-    images = (images * 255).round().astype("uint8")

-    pil_images = [Image.fromarray(image) for image in images]

-

-    return pil_images

-

-# check and replace nsfw content

-def check_safety(x_image):

-    global safety_feature_extractor, safety_checker

-

-    if safety_feature_extractor is None:

-        safety_feature_extractor = AutoFeatureExtractor.from_pretrained(safety_model_id)

-        safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_model_id)

-

-    safety_checker_input = safety_feature_extractor(numpy_to_pil(x_image), return_tensors="pt")

-    x_checked_image, has_nsfw_concept = safety_checker(images=x_image, clip_input=safety_checker_input.pixel_values)

-

-    return x_checked_image, has_nsfw_concept

-

-

-def censor_batch(x):

-    x_samples_ddim_numpy = x.cpu().permute(0, 2, 3, 1).numpy()

-    x_checked_image, has_nsfw_concept = check_safety(x_samples_ddim_numpy)

-    x = torch.from_numpy(x_checked_image).permute(0, 3, 1, 2)

-

-    return x

diff --git a/modules/scripts.py b/modules/scripts.py
index b934d881..722f8685 100644
--- a/modules/scripts.py
+++ b/modules/scripts.py
@@ -36,7 +36,7 @@ class Script:
     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.

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

         """

 

         pass

@@ -47,7 +47,7 @@ class Script:
 

         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

+         - True if the script should be shown in UI if it's selected in the scripts dropdown

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

          """

 

@@ -88,6 +88,17 @@ class Script:
 

         pass

 

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

+        """

+        Same as process_batch(), but called for every batch after it has been generated.

+

+        **kwargs will have same items as process_batch, and also:

+          - batch_number - index of current batch, from 0 to number of batches-1

+          - images - torch tensor with all generated images, with values ranging from 0 to 1;

+        """

+

+        pass

+

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

         """

         This function is called after processing ends for AlwaysVisible scripts.

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

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

 

+    def postprocess_batch(self, p, images, **kwargs):

+        for script in self.alwayson_scripts:

+            try:

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

+                script.postprocess_batch(p, *script_args, images=images, **kwargs)

+            except Exception:

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

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

+

     def before_component(self, component, **kwargs):

         for script in self.scripts:

             try:

diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index 95a17093..690a9ec2 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -1,17 +1,11 @@
-import math

-import os

-import sys

-import traceback

 import torch

-import numpy as np

-from torch import einsum

 from torch.nn.functional import silu

 

 import modules.textual_inversion.textual_inversion

-from modules import prompt_parser, devices, sd_hijack_optimizations, shared, sd_hijack_checkpoint

+from modules import devices, sd_hijack_optimizations, shared, sd_hijack_checkpoint

 from modules.hypernetworks import hypernetwork

-from modules.shared import opts, device, cmd_opts

-from modules import sd_hijack_clip, sd_hijack_open_clip

+from modules.shared import cmd_opts

+from modules import sd_hijack_clip, sd_hijack_open_clip, sd_hijack_unet

 

 from modules.sd_hijack_optimizations import invokeAI_mps_available

 

@@ -35,10 +29,12 @@ ldm.modules.attention.BasicTransformerBlock.ATTENTION_MODES["softmax-xformers"]
 ldm.modules.attention.print = lambda *args: None

 ldm.modules.diffusionmodules.model.print = lambda *args: None

 

+

 def apply_optimizations():

     undo_optimizations()

 

     ldm.modules.diffusionmodules.model.nonlinearity = silu

+    ldm.modules.diffusionmodules.openaimodel.th = sd_hijack_unet.th

 

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

         print("Applying xformers cross attention optimization.")

diff --git a/modules/sd_hijack_inpainting.py b/modules/sd_hijack_inpainting.py
index 938f9a58..85e7281f 100644
--- a/modules/sd_hijack_inpainting.py
+++ b/modules/sd_hijack_inpainting.py
@@ -209,7 +209,7 @@ def p_sample_plms(self, x, c, t, index, repeat_noise=False, use_original_steps=F
         else:
             x_in = torch.cat([x] * 2)
             t_in = torch.cat([t] * 2)
-            
+
             if isinstance(c, dict):
                 assert isinstance(unconditional_conditioning, dict)
                 c_in = dict()
@@ -278,7 +278,7 @@ def p_sample_plms(self, x, c, t, index, repeat_noise=False, use_original_steps=F
     x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t_prime, index)
 
     return x_prev, pred_x0, e_t
-    
+
 # =================================================================================================
 # Monkey patch LatentInpaintDiffusion to load the checkpoint with a proper config.
 # Adapted from:
@@ -324,12 +324,11 @@ def should_hijack_inpainting(checkpoint_info):
 
 def do_inpainting_hijack():
     # most of this stuff seems to no longer be needed because it is already included into SD2.0
-    # LatentInpaintDiffusion remains because SD2.0's LatentInpaintDiffusion can't be loaded without specifying a checkpoint
     # p_sample_plms is needed because PLMS can't work with dicts as conditionings
-    # this file should be cleaned up later if weverything tuens out to work fine
+    # this file should be cleaned up later if everything turns out to work fine
 
     # ldm.models.diffusion.ddpm.get_unconditional_conditioning = get_unconditional_conditioning
-    ldm.models.diffusion.ddpm.LatentInpaintDiffusion = LatentInpaintDiffusion
+    # ldm.models.diffusion.ddpm.LatentInpaintDiffusion = LatentInpaintDiffusion
 
     # ldm.models.diffusion.ddim.DDIMSampler.p_sample_ddim = p_sample_ddim
     # ldm.models.diffusion.ddim.DDIMSampler.sample = sample_ddim
diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 98123fbf..02c87f40 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -127,7 +127,7 @@ def check_for_psutil():
 

 invokeAI_mps_available = check_for_psutil()

 

-# -- Taken from https://github.com/invoke-ai/InvokeAI --

+# -- Taken from https://github.com/invoke-ai/InvokeAI and modified --

 if invokeAI_mps_available:

     import psutil

     mem_total_gb = psutil.virtual_memory().total // (1 << 30)

@@ -152,14 +152,16 @@ def einsum_op_slice_1(q, k, v, slice_size):
     return r

 

 def einsum_op_mps_v1(q, k, v):

-    if q.shape[1] <= 4096: # (512x512) max q.shape[1]: 4096

+    if q.shape[0] * q.shape[1] <= 2**16: # (512x512) max q.shape[1]: 4096

         return einsum_op_compvis(q, k, v)

     else:

         slice_size = math.floor(2**30 / (q.shape[0] * q.shape[1]))

+        if slice_size % 4096 == 0:

+            slice_size -= 1

         return einsum_op_slice_1(q, k, v, slice_size)

 

 def einsum_op_mps_v2(q, k, v):

-    if mem_total_gb > 8 and q.shape[1] <= 4096:

+    if mem_total_gb > 8 and q.shape[0] * q.shape[1] <= 2**16:

         return einsum_op_compvis(q, k, v)

     else:

         return einsum_op_slice_0(q, k, v, 1)

@@ -188,7 +190,7 @@ def einsum_op(q, k, v):
         return einsum_op_cuda(q, k, v)

 

     if q.device.type == 'mps':

-        if mem_total_gb >= 32:

+        if mem_total_gb >= 32 and q.shape[0] % 32 != 0 and q.shape[0] * q.shape[1] < 2**18:

             return einsum_op_mps_v1(q, k, v)

         return einsum_op_mps_v2(q, k, v)

 

diff --git a/modules/sd_hijack_unet.py b/modules/sd_hijack_unet.py
new file mode 100644
index 00000000..18daf8c1
--- /dev/null
+++ b/modules/sd_hijack_unet.py
@@ -0,0 +1,30 @@
+import torch

+

+

+class TorchHijackForUnet:

+    """

+    This is torch, but with cat that resizes tensors to appropriate dimensions if they do not match;

+    this makes it possible to create pictures with dimensions that are multiples of 8 rather than 64

+    """

+

+    def __getattr__(self, item):

+        if item == 'cat':

+            return self.cat

+

+        if hasattr(torch, item):

+            return getattr(torch, item)

+

+        raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, item))

+

+    def cat(self, tensors, *args, **kwargs):

+        if len(tensors) == 2:

+            a, b = tensors

+            if a.shape[-2:] != b.shape[-2:]:

+                a = torch.nn.functional.interpolate(a, b.shape[-2:], mode="nearest")

+

+            tensors = (a, b)

+

+        return torch.cat(tensors, *args, **kwargs)

+

+

+th = TorchHijackForUnet()

diff --git a/modules/sd_models.py b/modules/sd_models.py
index 283cf1cd..f36b299f 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -7,6 +7,9 @@ import torch
 import re

 import safetensors.torch

 from omegaconf import OmegaConf

+from os import mkdir

+from urllib import request

+import ldm.modules.midas as midas

 

 from ldm.util import instantiate_from_config

 

@@ -36,6 +39,7 @@ def setup_model():
         os.makedirs(model_path)

 

     list_models()

+    enable_midas_autodownload()

 

 

 def checkpoint_tiles(): 

@@ -223,10 +227,54 @@ def load_model_weights(model, checkpoint_info, vae_file="auto"):
     model.sd_model_checkpoint = checkpoint_file

     model.sd_checkpoint_info = checkpoint_info

 

+    sd_vae.delete_base_vae()

+    sd_vae.clear_loaded_vae()

     vae_file = sd_vae.resolve_vae(checkpoint_file, vae_file=vae_file)

     sd_vae.load_vae(model, vae_file)

 

 

+def enable_midas_autodownload():

+    """

+    Gives the ldm.modules.midas.api.load_model function automatic downloading.

+

+    When the 512-depth-ema model, and other future models like it, is loaded,

+    it calls midas.api.load_model to load the associated midas depth model.

+    This function applies a wrapper to download the model to the correct

+    location automatically.

+    """

+

+    midas_path = os.path.join(models_path, 'midas')

+

+    # stable-diffusion-stability-ai hard-codes the midas model path to

+    # a location that differs from where other scripts using this model look.

+    # HACK: Overriding the path here.

+    for k, v in midas.api.ISL_PATHS.items():

+        file_name = os.path.basename(v)

+        midas.api.ISL_PATHS[k] = os.path.join(midas_path, file_name)

+

+    midas_urls = {

+        "dpt_large": "https://github.com/intel-isl/DPT/releases/download/1_0/dpt_large-midas-2f21e586.pt",

+        "dpt_hybrid": "https://github.com/intel-isl/DPT/releases/download/1_0/dpt_hybrid-midas-501f0c75.pt",

+        "midas_v21": "https://github.com/AlexeyAB/MiDaS/releases/download/midas_dpt/midas_v21-f6b98070.pt",

+        "midas_v21_small": "https://github.com/AlexeyAB/MiDaS/releases/download/midas_dpt/midas_v21_small-70d6b9c8.pt",

+    }

+

+    midas.api.load_model_inner = midas.api.load_model

+

+    def load_model_wrapper(model_type):

+        path = midas.api.ISL_PATHS[model_type]

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

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

+                mkdir(midas_path)

+    

+            print(f"Downloading midas model weights for {model_type} to {path}")

+            request.urlretrieve(midas_urls[model_type], path)

+            print(f"{model_type} downloaded")

+

+        return midas.api.load_model_inner(model_type)

+

+    midas.api.load_model = load_model_wrapper

+

 def load_model(checkpoint_info=None):

     from modules import lowvram, sd_hijack

     checkpoint_info = checkpoint_info or select_checkpoint()

@@ -245,13 +293,16 @@ def load_model(checkpoint_info=None):
     if should_hijack_inpainting(checkpoint_info):

         # Hardcoded config for now...

         sd_config.model.target = "ldm.models.diffusion.ddpm.LatentInpaintDiffusion"

-        sd_config.model.params.use_ema = False

         sd_config.model.params.conditioning_key = "hybrid"

         sd_config.model.params.unet_config.params.in_channels = 9

+        sd_config.model.params.finetune_keys = None

 

         # Create a "fake" config with a different name so that we know to unload it when switching models.

         checkpoint_info = checkpoint_info._replace(config=checkpoint_info.config.replace(".yaml", "-inpainting.yaml"))

 

+    if not hasattr(sd_config.model.params, "use_ema"):

+        sd_config.model.params.use_ema = False

+

     do_inpainting_hijack()

 

     if shared.cmd_opts.no_half:

diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index 4c123d3b..d26e48dc 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -23,16 +23,16 @@ samplers_k_diffusion = [
     ('Euler', 'sample_euler', ['k_euler'], {}),

     ('LMS', 'sample_lms', ['k_lms'], {}),

     ('Heun', 'sample_heun', ['k_heun'], {}),

-    ('DPM2', 'sample_dpm_2', ['k_dpm_2'], {}),

-    ('DPM2 a', 'sample_dpm_2_ancestral', ['k_dpm_2_a'], {}),

+    ('DPM2', 'sample_dpm_2', ['k_dpm_2'], {'discard_next_to_last_sigma': True}),

+    ('DPM2 a', 'sample_dpm_2_ancestral', ['k_dpm_2_a'], {'discard_next_to_last_sigma': True}),

     ('DPM++ 2S a', 'sample_dpmpp_2s_ancestral', ['k_dpmpp_2s_a'], {}),

     ('DPM++ 2M', 'sample_dpmpp_2m', ['k_dpmpp_2m'], {}),

     ('DPM++ SDE', 'sample_dpmpp_sde', ['k_dpmpp_sde'], {}),

     ('DPM fast', 'sample_dpm_fast', ['k_dpm_fast'], {}),

     ('DPM adaptive', 'sample_dpm_adaptive', ['k_dpm_ad'], {}),

     ('LMS Karras', 'sample_lms', ['k_lms_ka'], {'scheduler': 'karras'}),

-    ('DPM2 Karras', 'sample_dpm_2', ['k_dpm_2_ka'], {'scheduler': 'karras'}),

-    ('DPM2 a Karras', 'sample_dpm_2_ancestral', ['k_dpm_2_a_ka'], {'scheduler': 'karras'}),

+    ('DPM2 Karras', 'sample_dpm_2', ['k_dpm_2_ka'], {'scheduler': 'karras', 'discard_next_to_last_sigma': True}),

+    ('DPM2 a Karras', 'sample_dpm_2_ancestral', ['k_dpm_2_a_ka'], {'scheduler': 'karras', 'discard_next_to_last_sigma': True}),

     ('DPM++ 2S a Karras', 'sample_dpmpp_2s_ancestral', ['k_dpmpp_2s_a_ka'], {'scheduler': 'karras'}),

     ('DPM++ 2M Karras', 'sample_dpmpp_2m', ['k_dpmpp_2m_ka'], {'scheduler': 'karras'}),

     ('DPM++ SDE Karras', 'sample_dpmpp_sde', ['k_dpmpp_sde_ka'], {'scheduler': 'karras'}),

@@ -444,9 +444,7 @@ class KDiffusionSampler:
 

         return extra_params_kwargs

 

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

-        steps, t_enc = setup_img2img_steps(p, steps)

-

+    def get_sigmas(self, p, steps):

         if p.sampler_noise_scheduler_override:

             sigmas = p.sampler_noise_scheduler_override(steps)

         elif self.config is not None and self.config.options.get('scheduler', None) == 'karras':

@@ -454,6 +452,16 @@ class KDiffusionSampler:
         else:

             sigmas = self.model_wrap.get_sigmas(steps)

 

+        if self.config is not None and self.config.options.get('discard_next_to_last_sigma', False):

+            sigmas = torch.cat([sigmas[:-2], sigmas[-1:]])

+

+        return sigmas

+

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

+        steps, t_enc = setup_img2img_steps(p, steps)

+

+        sigmas = self.get_sigmas(p, steps)

+

         sigma_sched = sigmas[steps - t_enc - 1:]

         xi = x + noise * sigma_sched[0]

         

@@ -485,12 +493,7 @@ class KDiffusionSampler:
     def sample(self, p, x, conditioning, unconditional_conditioning, steps=None, image_conditioning = None):

         steps = steps or p.steps

 

-        if p.sampler_noise_scheduler_override:

-            sigmas = p.sampler_noise_scheduler_override(steps)

-        elif self.config is not None and self.config.options.get('scheduler', None) == 'karras':

-            sigmas = k_diffusion.sampling.get_sigmas_karras(n=steps, sigma_min=0.1, sigma_max=10, device=shared.device)

-        else:

-            sigmas = self.model_wrap.get_sigmas(steps)

+        sigmas = self.get_sigmas(p, steps)

 

         x = x * sigmas[0]

 

diff --git a/modules/sd_vae.py b/modules/sd_vae.py
index 9c120975..25638a83 100644
--- a/modules/sd_vae.py
+++ b/modules/sd_vae.py
@@ -4,6 +4,7 @@ from collections import namedtuple
 from modules import shared, devices, script_callbacks
 from modules.paths import models_path
 import glob
+from copy import deepcopy
 
 
 model_dir = "Stable-diffusion"
@@ -15,7 +16,7 @@ vae_path = os.path.abspath(os.path.join(models_path, vae_dir))
 vae_ignore_keys = {"model_ema.decay", "model_ema.num_updates"}
 
 
-default_vae_dict = {"auto": "auto", "None": "None"}
+default_vae_dict = {"auto": "auto", "None": None, None: None}
 default_vae_list = ["auto", "None"]
 
 
@@ -39,7 +40,8 @@ def get_base_vae(model):
 def store_base_vae(model):
     global base_vae, checkpoint_info
     if checkpoint_info != model.sd_checkpoint_info:
-        base_vae = model.first_stage_model.state_dict().copy()
+        assert not loaded_vae_file, "Trying to store non-base VAE!"
+        base_vae = deepcopy(model.first_stage_model.state_dict())
         checkpoint_info = model.sd_checkpoint_info
 
 
@@ -50,9 +52,11 @@ def delete_base_vae():
 
 
 def restore_base_vae(model):
-    global base_vae, checkpoint_info
+    global loaded_vae_file
     if base_vae is not None and checkpoint_info == model.sd_checkpoint_info:
-        load_vae_dict(model, base_vae)
+        print("Restoring base VAE")
+        _load_vae_dict(model, base_vae)
+        loaded_vae_file = None
     delete_base_vae()
 
 
@@ -148,9 +152,10 @@ def load_vae(model, vae_file=None):
     if vae_file:
         assert os.path.isfile(vae_file), f"VAE file doesn't exist: {vae_file}"
         print(f"Loading VAE weights from: {vae_file}")
+        store_base_vae(model)
         vae_ckpt = torch.load(vae_file, map_location=shared.weight_load_location)
         vae_dict_1 = {k: v for k, v in vae_ckpt["state_dict"].items() if k[0:4] != "loss" and k not in vae_ignore_keys}
-        load_vae_dict(model, vae_dict_1)
+        _load_vae_dict(model, vae_dict_1)
 
         # If vae used is not in dict, update it
         # It will be removed on refresh though
@@ -158,30 +163,22 @@ def load_vae(model, vae_file=None):
         if vae_opt not in vae_dict:
             vae_dict[vae_opt] = vae_file
             vae_list.append(vae_opt)
+    elif loaded_vae_file:
+        restore_base_vae(model)
 
     loaded_vae_file = vae_file
 
-    """
-    # Save current VAE to VAE settings, maybe? will it work?
-    if save_settings:
-        if vae_file is None:
-            vae_opt = "None"
-
-        # shared.opts.sd_vae = vae_opt
-    """
-
     first_load = False
 
 
 # don't call this from outside
-def load_vae_dict(model, vae_dict_1=None):
-    if vae_dict_1:
-        store_base_vae(model)
-        model.first_stage_model.load_state_dict(vae_dict_1)
-    else:
-        restore_base_vae()
+def _load_vae_dict(model, vae_dict_1):
+    model.first_stage_model.load_state_dict(vae_dict_1)
     model.first_stage_model.to(devices.dtype_vae)
 
+def clear_loaded_vae():
+    global loaded_vae_file
+    loaded_vae_file = None
 
 def reload_vae_weights(sd_model=None, vae_file="auto"):
     from modules import lowvram, devices, sd_hijack
diff --git a/modules/shared.py b/modules/shared.py
index dc45fcaa..8ea3b441 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -5,6 +5,7 @@ import os
 import sys

 import time

 

+from PIL import Image

 import gradio as gr

 import tqdm

 

@@ -293,6 +294,7 @@ options_templates.update(options_section(('saving-images', "Saving images/grids"
     "export_for_4chan": OptionInfo(True, "If PNG image is larger than 4MB or any dimension is larger than 4000, downscale and save copy as JPG"),

 

     "use_original_name_batch": OptionInfo(False, "Use original name for output filename during batch process in extras tab"),

+    "use_upscaler_name_as_suffix": OptionInfo(False, "Use upscaler name as filename suffix in the extras tab"),

     "save_selected_only": OptionInfo(True, "When using 'Save' button, only save a single selected image"),

     "do_not_add_watermark": OptionInfo(False, "Do not add watermark to images"),

 

@@ -359,14 +361,15 @@ options_templates.update(options_section(('sd', "Stable Diffusion"), {
     "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}),

+    "initial_noise_multiplier": OptionInfo(1.0, "Noise multiplier for img2img", gr.Slider, {"minimum": 0.5, "maximum": 1.5, "step": 0.01 }),

     "img2img_color_correction": OptionInfo(False, "Apply color correction to img2img results to match original colors."),

     "img2img_fix_steps": OptionInfo(False, "With img2img, do exactly the amount of steps the slider specifies (normally you'd do less with less denoising)."),

+    "img2img_background_color": OptionInfo("#ffffff", "With img2img, fill image's transparent parts with this color.", gr.ColorPicker, {}),

     "enable_quantization": OptionInfo(False, "Enable quantization in K samplers for sharper and cleaner results. This may change existing seeds. Requires restart to apply."),

     "enable_emphasis": OptionInfo(True, "Emphasis: use (text) to make model pay more attention to text and [text] to make it pay less attention"),

     "use_old_emphasis_implementation": OptionInfo(False, "Use old emphasis implementation. Can be useful to reproduce old seeds."),

     "enable_batch_seeds": OptionInfo(True, "Make K-diffusion samplers produce same images in a batch as when making a single image"),

     "comma_padding_backtrack": OptionInfo(20, "Increase coherency by padding from the last comma within n tokens when using more than 75 tokens", gr.Slider, {"minimum": 0, "maximum": 74, "step": 1 }),

-    "filter_nsfw": OptionInfo(False, "Filter NSFW content"),

     'CLIP_stop_at_last_layers': OptionInfo(1, "Clip skip", gr.Slider, {"minimum": 1, "maximum": 12, "step": 1}),

     "random_artist_categories": OptionInfo([], "Allowed categories for random artists selection when using the Roll button", gr.CheckboxGroup, {"choices": artist_db.categories()}),

 }))

@@ -395,6 +398,7 @@ options_templates.update(options_section(('ui', "User interface"), {
     "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"),

+    "send_size": OptionInfo(True, "Send size when sending prompt or image to another interface"),

     "font": OptionInfo("", "Font for image grids that have text"),

     "js_modal_lightbox": OptionInfo(True, "Enable full page image viewer"),

     "js_modal_lightbox_initially_zoomed": OptionInfo(True, "Show images zoomed in by default in full page image viewer"),

diff --git a/modules/textual_inversion/dataset.py b/modules/textual_inversion/dataset.py
index 2dc64c3c..88d68c76 100644
--- a/modules/textual_inversion/dataset.py
+++ b/modules/textual_inversion/dataset.py
@@ -28,9 +28,9 @@ class DatasetEntry:
 

 

 class PersonalizedBase(Dataset):

-    def __init__(self, data_root, width, height, repeats, flip_p=0.5, placeholder_token="*", model=None, cond_model=None, device=None, template_file=None, include_cond=False, batch_size=1, gradient_step=1, shuffle_tags=False, tag_drop_out=0, latent_sampling_method='once'):        

+    def __init__(self, data_root, width, height, repeats, flip_p=0.5, placeholder_token="*", model=None, cond_model=None, device=None, template_file=None, include_cond=False, batch_size=1, gradient_step=1, shuffle_tags=False, tag_drop_out=0, latent_sampling_method='once'):

         re_word = re.compile(shared.opts.dataset_filename_word_regex) if len(shared.opts.dataset_filename_word_regex) > 0 else None

-        

+

         self.placeholder_token = placeholder_token

 

         self.width = width

@@ -50,14 +50,14 @@ class PersonalizedBase(Dataset):
 

         self.image_paths = [os.path.join(data_root, file_path) for file_path in os.listdir(data_root)]

 

-        

+

         self.shuffle_tags = shuffle_tags

         self.tag_drop_out = tag_drop_out

 

         print("Preparing dataset...")

         for path in tqdm.tqdm(self.image_paths):

             if shared.state.interrupted:

-                raise Exception("inturrupted")

+                raise Exception("interrupted")

             try:

                 image = Image.open(path).convert('RGB').resize((self.width, self.height), PIL.Image.BICUBIC)

             except Exception:

@@ -144,7 +144,7 @@ class PersonalizedDataLoader(DataLoader):
             self.collate_fn = collate_wrapper_random

         else:

             self.collate_fn = collate_wrapper

-        

+

 

 class BatchLoader:

     def __init__(self, data):

diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py
index e28c357a..daf3997b 100644
--- a/modules/textual_inversion/textual_inversion.py
+++ b/modules/textual_inversion/textual_inversion.py
@@ -133,7 +133,7 @@ class EmbeddingDatabase:
 

                 process_file(fullfn, fn)

             except Exception:

-                print(f"Error loading emedding {fn}:", file=sys.stderr)

+                print(f"Error loading embedding {fn}:", file=sys.stderr)

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

                 continue

 

@@ -194,7 +194,7 @@ def write_loss(log_directory, filename, step, epoch_len, values):
             csv_writer.writeheader()

 

         epoch = (step - 1) // epoch_len

-        epoch_step = (step - 1) % epoch_len 

+        epoch_step = (step - 1) % epoch_len

 

         csv_writer.writerow({

             "step": step,

@@ -270,9 +270,9 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
    # dataset loading may take a while, so input validations and early returns should be done before this

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

     old_parallel_processing_allowed = shared.parallel_processing_allowed

-    

+

     pin_memory = shared.opts.pin_memory

-    

+

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

 

     latent_sampling_method = ds.latent_sampling_method

@@ -295,12 +295,12 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
     loss_step = 0

     _loss_step = 0 #internal

 

-    

+

     last_saved_file = "<none>"

     last_saved_image = "<none>"

     forced_filename = "<none>"

     embedding_yet_to_be_embedded = False

-    

+

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

     try:

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

@@ -327,10 +327,10 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
                     c = shared.sd_model.cond_stage_model(batch.cond_text)

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

                     del x

-                    

+

                     _loss_step += loss.item()

                 scaler.scale(loss).backward()

-                

+

                 # go back until we reach gradient accumulation steps

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

                     continue

diff --git a/modules/ui.py b/modules/ui.py
index fe4abe05..faba69a4 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -49,10 +49,14 @@ if not cmd_opts.share and not cmd_opts.listen:
     gradio.utils.version_check = lambda: None

     gradio.utils.get_local_ip_address = lambda: '127.0.0.1'

 

-if cmd_opts.ngrok != None:

+if cmd_opts.ngrok is not None:

     import modules.ngrok as ngrok

     print('ngrok authtoken detected, trying to connect...')

-    ngrok.connect(cmd_opts.ngrok, cmd_opts.port if cmd_opts.port != None else 7860, cmd_opts.ngrok_region)

+    ngrok.connect(

+        cmd_opts.ngrok,

+        cmd_opts.port if cmd_opts.port is not None else 7860,

+        cmd_opts.ngrok_region

+        )

 

 

 def gr_show(visible=True):

@@ -82,6 +86,7 @@ folder_symbol = '\U0001f4c2'  # 📂
 refresh_symbol = '\U0001f504'  # 🔄

 save_style_symbol = '\U0001f4be'  # 💾

 apply_style_symbol = '\U0001f4cb'  # 📋

+clear_prompt_symbol = '\U0001F5D1'  # 🗑️

 

 

 def plaintext_to_html(text):

@@ -302,8 +307,8 @@ def create_seed_inputs():
 

     with gr.Row(visible=False) as seed_extra_row_2:

         seed_extras.append(seed_extra_row_2)

-        seed_resize_from_w = gr.Slider(minimum=0, maximum=2048, step=64, label="Resize seed from width", value=0)

-        seed_resize_from_h = gr.Slider(minimum=0, maximum=2048, step=64, label="Resize seed from height", value=0)

+        seed_resize_from_w = gr.Slider(minimum=0, maximum=2048, step=8, label="Resize seed from width", value=0)

+        seed_resize_from_h = gr.Slider(minimum=0, maximum=2048, step=8, label="Resize seed from height", value=0)

 

     random_seed.click(fn=lambda: -1, show_progress=False, inputs=[], outputs=[seed])

     random_subseed.click(fn=lambda: -1, show_progress=False, inputs=[], outputs=[subseed])

@@ -316,6 +321,17 @@ def create_seed_inputs():
     return seed, reuse_seed, subseed, reuse_subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w, seed_checkbox

 

 

+

+def connect_clear_prompt(button):

+    """Given clear button, prompt, and token_counter objects, setup clear prompt button click event"""

+    button.click(

+        _js="clear_prompt",

+        fn=None,

+        inputs=[],

+        outputs=[],

+    )

+

+

 def connect_reuse_seed(seed: gr.Number, reuse_seed: gr.Button, generation_info: gr.Textbox, dummy_component, is_subseed):

     """ Connects a 'reuse (sub)seed' button's click event so that it copies last used

         (sub)seed value from generation info the to the seed field. If copying subseed and subseed strength

@@ -391,10 +407,17 @@ def create_toprow(is_img2img):
             paste = gr.Button(value=paste_symbol, elem_id="paste")

             save_style = gr.Button(value=save_style_symbol, elem_id="style_create")

             prompt_style_apply = gr.Button(value=apply_style_symbol, elem_id="style_apply")

-

+            clear_prompt_button = gr.Button(value=clear_prompt_symbol, elem_id=f"{id_part}_clear_prompt")

             token_counter = gr.HTML(value="<span></span>", elem_id=f"{id_part}_token_counter")

             token_button = gr.Button(visible=False, elem_id=f"{id_part}_token_button")

 

+            clear_prompt_button.click(

+                fn=lambda *x: x,

+                _js="confirm_clear_prompt",

+                inputs=[prompt, negative_prompt],

+                outputs=[prompt, negative_prompt],

+            )

+

         button_interrogate = None

         button_deepbooru = None

         if is_img2img:

@@ -616,10 +639,14 @@ def create_ui():
     modules.scripts.scripts_txt2img.initialize_scripts(is_img2img=False)

 

     with gr.Blocks(analytics_enabled=False) as txt2img_interface:

-        txt2img_prompt, roll, txt2img_prompt_style, txt2img_negative_prompt, txt2img_prompt_style2, submit, _, _, txt2img_prompt_style_apply, txt2img_save_style, txt2img_paste, token_counter, token_button = create_toprow(is_img2img=False)

+        txt2img_prompt, roll, txt2img_prompt_style, txt2img_negative_prompt, txt2img_prompt_style2, submit, _, _,txt2img_prompt_style_apply, txt2img_save_style, txt2img_paste, token_counter, token_button = create_toprow(is_img2img=False)

+

         dummy_component = gr.Label(visible=False)

         txt_prompt_img = gr.File(label="", elem_id="txt2img_prompt_image", file_count="single", type="bytes", visible=False)

 

+

+

+

         with gr.Row(elem_id='txt2img_progress_row'):

             with gr.Column(scale=1):

                 pass

@@ -635,8 +662,8 @@ def create_ui():
                 sampler_index = gr.Radio(label='Sampling method', elem_id="txt2img_sampling", choices=[x.name for x in samplers], value=samplers[0].name, type="index")

 

                 with gr.Group():

-                    width = gr.Slider(minimum=64, maximum=2048, step=64, label="Width", value=512)

-                    height = gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512)

+                    width = gr.Slider(minimum=64, maximum=2048, step=8, label="Width", value=512)

+                    height = gr.Slider(minimum=64, maximum=2048, step=8, label="Height", value=512)

 

                 with gr.Row():

                     restore_faces = gr.Checkbox(label='Restore faces', value=False, visible=len(shared.face_restorers) > 1)

@@ -644,8 +671,8 @@ def create_ui():
                     enable_hr = gr.Checkbox(label='Highres. fix', value=False)

 

                 with gr.Row(visible=False) as hr_options:

-                    firstphase_width = gr.Slider(minimum=0, maximum=1024, step=64, label="Firstpass width", value=0)

-                    firstphase_height = gr.Slider(minimum=0, maximum=1024, step=64, label="Firstpass height", value=0)

+                    firstphase_width = gr.Slider(minimum=0, maximum=1024, step=8, label="Firstpass width", value=0)

+                    firstphase_height = gr.Slider(minimum=0, maximum=1024, step=8, label="Firstpass height", value=0)

                     denoising_strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Denoising strength', value=0.7)

 

                 with gr.Row(equal_height=True):

@@ -770,7 +797,8 @@ def create_ui():
     modules.scripts.scripts_img2img.initialize_scripts(is_img2img=True)

 

     with gr.Blocks(analytics_enabled=False) as img2img_interface:

-        img2img_prompt, roll, img2img_prompt_style, img2img_negative_prompt, img2img_prompt_style2, submit, img2img_interrogate, img2img_deepbooru, img2img_prompt_style_apply, img2img_save_style, img2img_paste, token_counter, token_button = create_toprow(is_img2img=True)

+        img2img_prompt, roll, img2img_prompt_style, img2img_negative_prompt, img2img_prompt_style2, submit, img2img_interrogate, img2img_deepbooru, img2img_prompt_style_apply, img2img_save_style, img2img_paste,token_counter, token_button = create_toprow(is_img2img=True)

+

 

         with gr.Row(elem_id='img2img_progress_row'):

             img2img_prompt_img = gr.File(label="", elem_id="img2img_prompt_image", file_count="single", type="bytes", visible=False)

@@ -788,7 +816,7 @@ def create_ui():
 

                 with gr.Tabs(elem_id="mode_img2img") as tabs_img2img_mode:

                     with gr.TabItem('img2img', id='img2img'):

-                        init_img = gr.Image(label="Image for img2img", elem_id="img2img_image", show_label=False, source="upload", interactive=True, type="pil", tool=cmd_opts.gradio_img2img_tool).style(height=480)

+                        init_img = gr.Image(label="Image for img2img", elem_id="img2img_image", show_label=False, source="upload", interactive=True, type="pil", tool=cmd_opts.gradio_img2img_tool, image_mode="RGBA").style(height=480)

 

                     with gr.TabItem('Inpaint', id='inpaint'):

                         init_img_with_mask = gr.Image(label="Image for inpainting with mask", show_label=False, elem_id="img2maskimg", source="upload", interactive=True, type="pil", tool=cmd_opts.gradio_inpaint_tool, image_mode="RGBA").style(height=480)

@@ -835,8 +863,8 @@ def create_ui():
                 sampler_index = gr.Radio(label='Sampling method', choices=[x.name for x in samplers_for_img2img], value=samplers_for_img2img[0].name, type="index")

 

                 with gr.Group():

-                    width = gr.Slider(minimum=64, maximum=2048, step=64, label="Width", value=512, elem_id="img2img_width")

-                    height = gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512, elem_id="img2img_height")

+                    width = gr.Slider(minimum=64, maximum=2048, step=8, label="Width", value=512, elem_id="img2img_width")

+                    height = gr.Slider(minimum=64, maximum=2048, step=8, label="Height", value=512, elem_id="img2img_height")

 

                 with gr.Row():

                     restore_faces = gr.Checkbox(label='Restore faces', value=False, visible=len(shared.face_restorers) > 1)

@@ -1171,8 +1199,8 @@ def create_ui():
                 with gr.Tab(label="Preprocess images"):

                     process_src = gr.Textbox(label='Source directory')

                     process_dst = gr.Textbox(label='Destination directory')

-                    process_width = gr.Slider(minimum=64, maximum=2048, step=64, label="Width", value=512)

-                    process_height = gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512)

+                    process_width = gr.Slider(minimum=64, maximum=2048, step=8, label="Width", value=512)

+                    process_height = gr.Slider(minimum=64, maximum=2048, step=8, label="Height", value=512)

                     preprocess_txt_action = gr.Dropdown(label='Existing Caption txt Action', value="ignore", choices=["ignore", "copy", "prepend", "append"])

 

                     with gr.Row():

@@ -1230,8 +1258,8 @@ def create_ui():
                     dataset_directory = gr.Textbox(label='Dataset directory', placeholder="Path to directory with input images")

                     log_directory = gr.Textbox(label='Log directory', placeholder="Path to directory where to write outputs", value="textual_inversion")

                     template_file = gr.Textbox(label='Prompt template file', value=os.path.join(script_path, "textual_inversion_templates", "style_filewords.txt"))

-                    training_width = gr.Slider(minimum=64, maximum=2048, step=64, label="Width", value=512)

-                    training_height = gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512)

+                    training_width = gr.Slider(minimum=64, maximum=2048, step=8, label="Width", value=512)

+                    training_height = gr.Slider(minimum=64, maximum=2048, step=8, label="Height", value=512)

                     steps = gr.Number(label='Max steps', value=100000, precision=0)

                     create_image_every = gr.Number(label='Save an image to log directory every N steps, 0 to disable', value=500, precision=0)

                     save_embedding_every = gr.Number(label='Save a copy of embedding to log directory every N steps, 0 to disable', value=500, precision=0)

@@ -1450,7 +1478,7 @@ def create_ui():
             opts.save(shared.config_filename)

         except RuntimeError:

             return opts.dumpjson(), f'{len(changed)} settings changed without save: {", ".join(changed)}.'

-        return opts.dumpjson(), f'{len(changed)} settings changed: {", ".join(changed)}.'

+        return opts.dumpjson(), f'{len(changed)} settings changed{": " if len(changed) > 0 else ""}{", ".join(changed)}.'

 

     def run_settings_single(value, key):

         if not opts.same_type(value, opts.data_labels[key].default):

diff --git a/modules/ui_extensions.py b/modules/ui_extensions.py
index b487ac25..eec9586f 100644
--- a/modules/ui_extensions.py
+++ b/modules/ui_extensions.py
@@ -9,6 +9,8 @@ import git
 

 import gradio as gr

 import html

+import shutil

+import errno

 

 from modules import extensions, shared, paths

 

@@ -138,7 +140,18 @@ def install_extension_from_url(dirname, url):
         repo = git.Repo.clone_from(url, tmpdir)

         repo.remote().fetch()

 

-        os.rename(tmpdir, target_dir)

+        try:

+            os.rename(tmpdir, target_dir)

+        except OSError as err:

+            # TODO what does this do on windows? I think it'll be a different error code but I don't have a system to check it

+            # Shouldn't cause any new issues at least but we probably want to handle it there too.

+            if err.errno == errno.EXDEV:

+                # Cross device link, typical in docker or when tmp/ and extensions/ are on different file systems

+                # Since we can't use a rename, do the slower but more versitile shutil.move()

+                shutil.move(tmpdir, target_dir)

+            else:

+                # Something else, not enough free space, permissions, etc.  rethrow it so that it gets handled.

+                raise(err)

 

         import launch

         launch.run_extension_installer(target_dir)

@@ -206,12 +219,13 @@ def refresh_available_extensions_from_data(hide_tags):
         if url is None:

             continue

 

+        existing = installed_extension_urls.get(normalize_git_url(url), None)

+        extension_tags = extension_tags + ["installed"] if existing else extension_tags

+

         if len([x for x in extension_tags if x in tags_to_hide]) > 0:

             hidden += 1

             continue

 

-        existing = installed_extension_urls.get(normalize_git_url(url), None)

-

         install_code = f"""<input onclick="install_extension_from_index(this, '{html.escape(url)}')" type="button" value="{"Install" if not existing else "Installed"}" {"disabled=disabled" if existing else ""} class="gr-button gr-button-lg gr-button-secondary">"""

 

         tags_text = ", ".join([f"<span class='extension-tag' title='{tags.get(x, '')}'>{x}</span>" for x in extension_tags])

@@ -222,7 +236,11 @@ def refresh_available_extensions_from_data(hide_tags):
                 <td>{html.escape(description)}</td>

                 <td>{install_code}</td>

             </tr>

-    """

+        

+        """

+

+        for tag in [x for x in extension_tags if x not in tags]:

+            tags[tag] = tag

 

     code += """

         </tbody>

@@ -272,7 +290,7 @@ def create_ui():
                     install_extension_button = gr.Button(elem_id="install_extension_button", visible=False)

 

                 with gr.Row():

-                    hide_tags = gr.CheckboxGroup(value=["ads", "localization"], label="Hide extensions with tags", choices=["script", "ads", "localization"])

+                    hide_tags = gr.CheckboxGroup(value=["ads", "localization", "installed"], label="Hide extensions with tags", choices=["script", "ads", "localization", "installed"])

 

                 install_result = gr.HTML()

                 available_extensions_table = gr.HTML()