Merge branch 'master' of https://github.com/AUTOMATIC1111/stable-diffusion-webui

1 files changed, 37 insertions, 28 deletions

diff --git a/modules/images.py b/modules/images.py
index 848ede75..7870b5b7 100644
--- a/modules/images.py
+++ b/modules/images.py
@@ -16,7 +16,7 @@ from PIL import Image, ImageFont, ImageDraw, PngImagePlugin
 from fonts.ttf import Roboto

 import string

 

-from modules import sd_samplers, shared

+from modules import sd_samplers, shared, script_callbacks

 from modules.shared import opts, cmd_opts

 

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

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

 invalid_filename_postfix = ' .'

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

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

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

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

 max_filename_part_length = 128

 

@@ -343,8 +343,11 @@ class FilenameGenerator:
     def datetime(self, *args):

         time_datetime = datetime.datetime.now()

 

-        time_format = args[0] if len(args) > 0 else self.default_time_format

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

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

+        try:

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

+        except pytz.exceptions.UnknownTimeZoneError as _:

+            time_zone = None

 

         time_zone_time = time_datetime.astimezone(time_zone)

         try:

@@ -359,9 +362,9 @@ class FilenameGenerator:
 

         for m in re_pattern.finditer(x):

             text, pattern = m.groups()

+            res += text

 

             if pattern is None:

-                res += text

                 continue

 

             pattern_args = []

@@ -382,12 +385,9 @@ class FilenameGenerator:
                     print(f"Error adding [{pattern}] to filename", file=sys.stderr)

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

 

-                if replacement is None:

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

-                else:

+                if replacement is not None:

                     res += str(replacement)

-

-                continue

+                    continue

 

             res += f'[{pattern}]'

 

@@ -451,17 +451,6 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
     """

     namegen = FilenameGenerator(p, seed, prompt)

 

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

-        pnginfo = PngImagePlugin.PngInfo()

-

-        if existing_info is not None:

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

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

-

-        pnginfo.add_text(pnginfo_section_name, info)

-    else:

-        pnginfo = None

-

     if save_to_dirs is None:

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

 

@@ -474,8 +463,10 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
     if forced_filename is None:

         if short_filename or seed is None:

             file_decoration = ""

-        else:

+        elif opts.save_to_dirs:

             file_decoration = opts.samples_filename_pattern or "[seed]"

+        else:

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

 

         add_number = opts.save_images_add_number or file_decoration == ''

 

@@ -487,19 +478,27 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
         if add_number:

             basecount = get_next_sequence_number(path, basename)

             fullfn = None

-            fullfn_without_extension = None

             for i in range(500):

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

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

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

                 if not os.path.exists(fullfn):

                     break

         else:

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

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

     else:

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

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

+

+    pnginfo = existing_info or {}

+    if info is not None:

+        pnginfo[pnginfo_section_name] = info

+

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

+    script_callbacks.before_image_saved_callback(params)

+

+    image = params.image

+    fullfn = params.filename

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

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

 

     def exif_bytes():

         return piexif.dump({

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

         })

 

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

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

+        pnginfo_data = PngImagePlugin.PngInfo()

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

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

+

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

+

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

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

+

         if opts.enable_pnginfo and info is not None:

             piexif.insert(exif_bytes(), fullfn)

     else:

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

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

 

     target_side_length = 4000

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

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

         txt_fullfn = None

 

+    script_callbacks.image_saved_callback(params)

+

     return fullfn, txt_fullfn