Merge branch 'dev' into better-status-reporting-1

7 files changed, 84 insertions, 51 deletions

diff --git a/modules/api/api.py b/modules/api/api.py
index 8b79495d..9d33b9a9 100644
--- a/modules/api/api.py
+++ b/modules/api/api.py
@@ -30,6 +30,7 @@ from modules import devices
 from typing import Dict, List, Any
 import piexif
 import piexif.helper
+from contextlib import closing
 
 
 def script_name_to_index(name, scripts):
@@ -77,6 +78,8 @@ def encode_pil_to_base64(image):
             image.save(output_bytes, format="PNG", pnginfo=(metadata if use_metadata else None), quality=opts.jpeg_quality)
 
         elif opts.samples_format.lower() in ("jpg", "jpeg", "webp"):
+            if image.mode == "RGBA":
+                image = image.convert("RGB")
             parameters = image.info.get('parameters', None)
             exif_bytes = piexif.dump({
                 "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(parameters or "", encoding="unicode") }
@@ -322,19 +325,19 @@ class Api:
         args.pop('save_images', None)
 
         with self.queue_lock:
-            p = StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **args)
-            p.scripts = script_runner
-            p.outpath_grids = opts.outdir_txt2img_grids
-            p.outpath_samples = opts.outdir_txt2img_samples
-
-            shared.state.begin(job="scripts_txt2img")
-            if selectable_scripts is not None:
-                p.script_args = script_args
-                processed = scripts.scripts_txt2img.run(p, *p.script_args) # Need to pass args as list here
-            else:
-                p.script_args = tuple(script_args) # Need to pass args as tuple here
-                processed = process_images(p)
-            shared.state.end()
+            with closing(StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **args)) as p:
+                p.scripts = script_runner
+                p.outpath_grids = opts.outdir_txt2img_grids
+                p.outpath_samples = opts.outdir_txt2img_samples
+
+                shared.state.begin(job="scripts_txt2img")
+                if selectable_scripts is not None:
+                    p.script_args = script_args
+                    processed = scripts.scripts_txt2img.run(p, *p.script_args) # Need to pass args as list here
+                else:
+                    p.script_args = tuple(script_args) # Need to pass args as tuple here
+                    processed = process_images(p)
+                shared.state.end()
 
         b64images = list(map(encode_pil_to_base64, processed.images)) if send_images else []
 
@@ -378,20 +381,20 @@ class Api:
         args.pop('save_images', None)
 
         with self.queue_lock:
-            p = StableDiffusionProcessingImg2Img(sd_model=shared.sd_model, **args)
-            p.init_images = [decode_base64_to_image(x) for x in init_images]
-            p.scripts = script_runner
-            p.outpath_grids = opts.outdir_img2img_grids
-            p.outpath_samples = opts.outdir_img2img_samples
-
-            shared.state.begin(job="scripts_img2img")
-            if selectable_scripts is not None:
-                p.script_args = script_args
-                processed = scripts.scripts_img2img.run(p, *p.script_args) # Need to pass args as list here
-            else:
-                p.script_args = tuple(script_args) # Need to pass args as tuple here
-                processed = process_images(p)
-            shared.state.end()
+            with closing(StableDiffusionProcessingImg2Img(sd_model=shared.sd_model, **args)) as p:
+                p.init_images = [decode_base64_to_image(x) for x in init_images]
+                p.scripts = script_runner
+                p.outpath_grids = opts.outdir_img2img_grids
+                p.outpath_samples = opts.outdir_img2img_samples
+
+                shared.state.begin(job="scripts_img2img")
+                if selectable_scripts is not None:
+                    p.script_args = script_args
+                    processed = scripts.scripts_img2img.run(p, *p.script_args) # Need to pass args as list here
+                else:
+                    p.script_args = tuple(script_args) # Need to pass args as tuple here
+                    processed = process_images(p)
+                shared.state.end()
 
         b64images = list(map(encode_pil_to_base64, processed.images)) if send_images else []
 
diff --git a/modules/images.py b/modules/images.py
index 1906e2ab..fdccec09 100644
--- a/modules/images.py
+++ b/modules/images.py
@@ -1,3 +1,5 @@
+from __future__ import annotations

+

 import datetime

 

 import pytz

@@ -497,13 +499,23 @@ def get_next_sequence_number(path, basename):
     return result + 1

 

 

-def save_image_with_geninfo(image, geninfo, filename, extension=None, existing_pnginfo=None):

+def save_image_with_geninfo(image, geninfo, filename, extension=None, existing_pnginfo=None, pnginfo_section_name='parameters'):

+    """

+    Saves image to filename, including geninfo as text information for generation info.

+    For PNG images, geninfo is added to existing pnginfo dictionary using the pnginfo_section_name argument as key.

+    For JPG images, there's no dictionary and geninfo just replaces the EXIF description.

+    """

+

     if extension is None:

         extension = os.path.splitext(filename)[1]

 

     image_format = Image.registered_extensions()[extension]

 

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

+        existing_pnginfo = existing_pnginfo or {}

+        if opts.enable_pnginfo:

+            existing_pnginfo[pnginfo_section_name] = geninfo

+

         if opts.enable_pnginfo:

             pnginfo_data = PngImagePlugin.PngInfo()

             for k, v in (existing_pnginfo or {}).items():

@@ -622,7 +634,7 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
         """

         temp_file_path = f"{filename_without_extension}.tmp"

 

-        save_image_with_geninfo(image_to_save, info, temp_file_path, extension, params.pnginfo)

+        save_image_with_geninfo(image_to_save, info, temp_file_path, extension, existing_pnginfo=params.pnginfo, pnginfo_section_name=pnginfo_section_name)

 

         os.replace(temp_file_path, filename_without_extension + extension)

 

@@ -639,12 +651,18 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
     oversize = image.width > opts.target_side_length or image.height > opts.target_side_length

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

         ratio = image.width / image.height

-

+        resize_to = None

         if oversize and ratio > 1:

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

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

         elif oversize:

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

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

 

+        if resize_to is not None:

+            try:

+                # Resizing image with LANCZOS could throw an exception if e.g. image mode is I;16

+                image = image.resize(resize_to, LANCZOS)

+            except:

+                image = image.resize(resize_to)

         try:

             _atomically_save_image(image, fullfn_without_extension, ".jpg")

         except Exception as e:

@@ -662,8 +680,15 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
     return fullfn, txt_fullfn

 

 

-def read_info_from_image(image):

-    items = image.info or {}

+IGNORED_INFO_KEYS = {

+    'jfif', 'jfif_version', 'jfif_unit', 'jfif_density', 'dpi', 'exif',

+    'loop', 'background', 'timestamp', 'duration', 'progressive', 'progression',

+    'icc_profile', 'chromaticity', 'photoshop',

+}

+

+

+def read_info_from_image(image: Image.Image) -> tuple[str | None, dict]:

+    items = (image.info or {}).copy()

 

     geninfo = items.pop('parameters', None)

 

@@ -679,9 +704,7 @@ def read_info_from_image(image):
             items['exif comment'] = exif_comment

             geninfo = exif_comment

 

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

-                    'loop', 'background', 'timestamp', 'duration', 'progressive', 'progression',

-                    'icc_profile', 'chromaticity']:

+    for field in IGNORED_INFO_KEYS:

         items.pop(field, None)

 

     if items.get("Software", None) == "NovelAI":

diff --git a/modules/mac_specific.py b/modules/mac_specific.py
index d74c6b95..735847f5 100644
--- a/modules/mac_specific.py
+++ b/modules/mac_specific.py
@@ -4,16 +4,21 @@ from modules.sd_hijack_utils import CondFunc
 from packaging import version
 
 
-# has_mps is only available in nightly pytorch (for now) and macOS 12.3+.
-# check `getattr` and try it for compatibility
+# before torch version 1.13, has_mps is only available in nightly pytorch and macOS 12.3+,
+# use check `getattr` and try it for compatibility.
+# in torch version 1.13, backends.mps.is_available() and backends.mps.is_built() are introduced in to check mps availabilty,
+# since torch 2.0.1+ nightly build, getattr(torch, 'has_mps', False) was deprecated, see https://github.com/pytorch/pytorch/pull/103279
 def check_for_mps() -> bool:
-    if not getattr(torch, 'has_mps', False):
-        return False
-    try:
-        torch.zeros(1).to(torch.device("mps"))
-        return True
-    except Exception:
-        return False
+    if version.parse(torch.__version__) <= version.parse("2.0.1"):
+        if not getattr(torch, 'has_mps', False):
+            return False
+        try:
+            torch.zeros(1).to(torch.device("mps"))
+            return True
+        except Exception:
+            return False
+    else:
+        return torch.backends.mps.is_available() and torch.backends.mps.is_built()
 has_mps = check_for_mps()
 
 
diff --git a/modules/postprocessing.py b/modules/postprocessing.py
index 544b2f72..136e9c88 100644
--- a/modules/postprocessing.py
+++ b/modules/postprocessing.py
@@ -53,7 +53,9 @@ def run_postprocessing(extras_mode, image, image_folder, input_dir, output_dir,
     for image, name in zip(image_data, image_names):

         shared.state.textinfo = name

 

-        existing_pnginfo = image.info or {}

+        parameters, existing_pnginfo = images.read_info_from_image(image)

+        if parameters:

+            existing_pnginfo["parameters"] = parameters

 

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

 

diff --git a/modules/processing.py b/modules/processing.py
index 85e104aa..89ad1bd0 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -575,7 +575,7 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iter
         "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(':', '')),

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

         "Variation seed strength": (None if p.subseed_strength == 0 else p.subseed_strength),

-        "Seed resize from": (None if p.seed_resize_from_w == 0 or p.seed_resize_from_h == 0 else f"{p.seed_resize_from_w}x{p.seed_resize_from_h}"),

+        "Seed resize from": (None if p.seed_resize_from_w <= 0 or p.seed_resize_from_h <= 0 else f"{p.seed_resize_from_w}x{p.seed_resize_from_h}"),

         "Denoising strength": getattr(p, 'denoising_strength', None),

         "Conditional mask weight": getattr(p, "inpainting_mask_weight", shared.opts.inpainting_mask_weight) if p.is_using_inpainting_conditioning else None,

         "Clip skip": None if clip_skip <= 1 else clip_skip,

diff --git a/modules/ui.py b/modules/ui.py
index e3c8e43a..d4dc7910 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -155,7 +155,7 @@ def process_interrogate(interrogation_function, mode, ii_input_dir, ii_output_di
             img = Image.open(image)

             filename = os.path.basename(image)

             left, _ = os.path.splitext(filename)

-            print(interrogation_function(img), file=open(os.path.join(ii_output_dir, f"{left}.txt"), 'a'))

+            print(interrogation_function(img), file=open(os.path.join(ii_output_dir, f"{left}.txt"), 'a', encoding='utf-8'))

 

         return [gr.update(), None]

 

diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py
index a7d3bc79..1efd00b0 100644
--- a/modules/ui_extra_networks.py
+++ b/modules/ui_extra_networks.py
@@ -30,8 +30,8 @@ def fetch_file(filename: str = ""):
         raise ValueError(f"File cannot be fetched: {filename}. Must be in one of directories registered by extra pages.")

 

     ext = os.path.splitext(filename)[1].lower()

-    if ext not in (".png", ".jpg", ".jpeg", ".webp"):

-        raise ValueError(f"File cannot be fetched: {filename}. Only png and jpg and webp.")

+    if ext not in (".png", ".jpg", ".jpeg", ".webp", ".gif"):

+        raise ValueError(f"File cannot be fetched: {filename}. Only png, jpg, webp, and gif.")

 

     # would profit from returning 304

     return FileResponse(filename, headers={"Accept-Ranges": "bytes"})