enhance image saving

CHANGELOG.md

@@ -2,8 +2,9 @@
 
 ## Update for 05/31/2023
 
-- add pause option next to stop/skip
 - redesign action box to be uniform accross all themes
+- add pause option next to stop/skip
+- redesign progress bar
 
 ## Update for 05/30/2023
 

javascript/edit-attention.js

@@ -79,16 +79,12 @@ function keyupEditAttention(event) {
   weight = parseFloat(weight.toPrecision(12));
   if (String(weight).length === 1) weight += '.0';
 
-  console.log('HERE', closeCharacter, weight);
   if (closeCharacter == ')' && weight == 1) {
-    console.log('HERE2');
     text = text.slice(0, selectionStart - 1) + text.slice(selectionStart, selectionEnd) + text.slice(selectionEnd + 5);
     selectionStart--;
     selectionEnd--;
-    console.log('HERE2', text);
   } else {
     text = text.slice(0, selectionEnd + 1) + weight + text.slice(selectionEnd + 1 + end - 1);
-    console.log('HERE3', text);
   }
 
   target.focus();

javascript/style.css

@@ -101,6 +101,8 @@ footer { display: none !important; }
 }
 
 #txt2img_gallery img, #img2img_gallery img, #extras_gallery img { object-fit: scale-down; width: -webkit-fill-available !important; }
+
+#txt2img_generate_box, #img2img_generate_box { gap: 0.5em; flex-wrap: wrap-reverse; }
 #txt2img_actions_column, #img2img_actions_column { gap: 0.5em; }
 #txt2img_generate_box > button, #img2img_generate_box > button { height: 2.2em; line-height: 0; }
 #txt2img_generate_line2, #img2img_generate_line2 { display: flex; }

modules/api/api.py

@@ -60,28 +60,39 @@ def decode_base64_to_image(encoding):
         shared.log.warning(f'API cannot decode image: {e}')
         raise HTTPException(status_code=500, detail="Invalid encoded image") from e
 
-def encode_pil_to_base64(image):
-    with io.BytesIO() as output_bytes:
-        if shared.opts.samples_format.lower() == 'png':
-            use_metadata = False
-            encoded_metadata = PngImagePlugin.PngInfo()
-            for k, v in image.info.items():
-                if isinstance(k, str) and isinstance(v, str):
-                    encoded_metadata.add_text(k, v)
-                    use_metadata = True
-            image.save(output_bytes, format="PNG", pnginfo=(encoded_metadata if use_metadata else None), quality=shared.opts.jpeg_quality)
 
-        elif shared.opts.samples_format.lower() in ("jpg", "jpeg", "webp"):
-            parameters = image.info.get('parameters', None)
-            exif_bytes = piexif.dump({
-                "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(parameters or "", encoding="unicode") }
-            })
-            if shared.opts.samples_format.lower() in ("jpg", "jpeg"):
-                image.save(output_bytes, format="JPEG", exif = exif_bytes, quality=shared.opts.jpeg_quality)
-            else:
-                image.save(output_bytes, format="WEBP", exif = exif_bytes, quality=shared.opts.jpeg_quality)
-        else:
-            raise HTTPException(status_code=500, detail="Invalid image format")
+def save_image(image, fn, ext):
+    # actual save
+    parameters = image.info.get('parameters', None)
+    image_format = Image.registered_extensions()[f'.{ext}']
+    if image_format == 'PNG':
+        pnginfo_data = PngImagePlugin.PngInfo()
+        for k, v in image.info.items():
+            pnginfo_data.add_text(k, str(v))
+        image.save(fn, format=image_format, quality=shared.opts.jpeg_quality, pnginfo=pnginfo_data)
+    elif image_format == 'JPEG':
+        if image.mode == 'RGBA':
+            shared.log.warning('Saving RGBA image as JPEG: Alpha channel will be lost')
+            image = image.convert("RGB")
+        elif image.mode == 'I;16':
+            image = image.point(lambda p: p * 0.0038910505836576).convert("L")
+        exif_bytes = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(parameters or "", encoding="unicode") } })
+        image.save(fn, format=image_format, quality=shared.opts.jpeg_quality, exif=exif_bytes)
+    elif image_format == 'WEBP':
+        if image.mode == 'I;16':
+            image = image.point(lambda p: p * 0.0038910505836576).convert("RGB")
+        exif_bytes = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(parameters or "", encoding="unicode") } })
+        image.save(fn, format=image_format, quality=shared.opts.jpeg_quality, lossless=shared.opts.webp_lossless, exif=exif_bytes)
+    else:
+        # shared.log.warning(f'Unrecognized image format: {extension} attempting save as {image_format}')
+        image.save(fn, format=image_format, quality=shared.opts.jpeg_quality)
+
+
+def encode_pil_to_base64(image):
+    # TODO jpeg
+    print('HERE1', vars(image))
+    with io.BytesIO() as output_bytes:
+        save_image(image, output_bytes, shared.opts.samples_format)
         bytes_data = output_bytes.getvalue()
     return base64.b64encode(bytes_data)
 

modules/images.py

@@ -431,6 +431,7 @@ def atomically_save_image():
     while True:
         image, filename, extension, params, exifinfo_data, txt_fullfn = save_queue.get()
         fn = filename + extension
+        filename = filename.strip()
         image_format = Image.registered_extensions()[extension]
         shared.log.debug(f'Saving image: {image_format} {fn} {image.size}')
         # actual save
@@ -453,7 +454,7 @@ def atomically_save_image():
             exif_bytes = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo_data or "", encoding="unicode") } })
             image.save(fn, format=image_format, quality=shared.opts.jpeg_quality, lossless=shared.opts.webp_lossless, exif=exif_bytes)
         else:
-            shared.log.warning(f'Unrecognized image format: {extension} attempting save as {image_format}')
+            # shared.log.warning(f'Unrecognized image format: {extension} attempting save as {image_format}')
             image.save(fn, format=image_format, quality=shared.opts.jpeg_quality)
         # additional metadata saved in files
         if shared.opts.save_txt and len(exifinfo_data) > 0:

modules/img2img.py

@@ -123,8 +123,6 @@ def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_s
             assert image, "Can't scale by because no image is selected"
             width = int(image.width * scale_by)
             height = int(image.height * scale_by)
-    else:
-        return
 
     assert 0. <= denoising_strength <= 1., 'can only work with strength in [0.0, 1.0]'
 

modules/processing.py

@@ -713,7 +713,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                         p.restore_faces = False
                         info=infotext(n, i)
                         p.restore_faces = orig
-                        images.save_image(Image.fromarray(x_sample), p.outpath_samples, "", seeds[i], prompts[i], opts.samples_format, info=info, p=p, suffix="-before-face-restoration")
+                        images.save_image(Image.fromarray(x_sample), path=p.outpath_samples, basename="", seed=seeds[i], prompt=prompts[i], extension=opts.samples_format, info=info, p=p, suffix="-before-face-restoration")
                     x_sample = modules.face_restoration.restore_faces(x_sample)
                 image = Image.fromarray(x_sample)
                 if p.scripts is not None:
@@ -727,7 +727,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                         info=infotext(n, i)
                         p.color_corrections = orig
                         image_without_cc = apply_overlay(image, p.paste_to, i, p.overlay_images)
-                        images.save_image(image_without_cc, p.outpath_samples, "", seeds[i], prompts[i], opts.samples_format, info=info, p=p, suffix="-before-color-correction")
+                        images.save_image(image_without_cc, path=p.outpath_samples, basename="", seed=seeds[i], prompt=prompts[i], extension=opts.samples_format, info=info, p=p, suffix="-before-color-correction")
                     image = apply_color_correction(p.color_corrections[i], image)
                 image = apply_overlay(image, p.paste_to, i, p.overlay_images)
                 if opts.samples_save and not p.do_not_save_samples:

modules/progress.py

@@ -77,17 +77,11 @@ def progressapi(req: ProgressRequest):
     predicted_duration = elapsed_since_start / progress if progress > 0 else None
     eta = predicted_duration - elapsed_since_start if predicted_duration is not None else None
     id_live_preview = req.id_live_preview
+    live_preview = None
     shared.state.set_current_image()
-    if shared.opts.live_previews_enable and shared.state.id_live_preview != req.id_live_preview:
-        image = shared.state.current_image
-        if image is not None:
-            buffered = io.BytesIO()
-            fmt = 'jpeg' if shared.opts.samples_format == 'jpg' else shared.opts.samples_format
-            image.save(buffered, format=fmt)
-            live_preview = f'data:image/{fmt};base64,{base64.b64encode(buffered.getvalue()).decode("ascii")}'
-            id_live_preview = shared.state.id_live_preview
-        else:
-            live_preview = None
-    else:
-        live_preview = None
+    if shared.opts.live_previews_enable and (shared.state.id_live_preview != req.id_live_preview) and (shared.state.current_image is not None):
+        buffered = io.BytesIO()
+        shared.state.current_image.save(buffered, format='jpeg')
+        live_preview = f'data:image/jpeg;base64,{base64.b64encode(buffered.getvalue()).decode("ascii")}'
+        id_live_preview = shared.state.id_live_preview
     return InternalProgressResponse(active=active, queued=queued, paused=paused, completed=completed, progress=progress, eta=eta, live_preview=live_preview, id_live_preview=id_live_preview, textinfo=shared.state.textinfo)

modules/shared.py

@@ -319,11 +319,11 @@ options_templates.update(options_section(('system-paths', "System Paths"), {
 
 options_templates.update(options_section(('saving-images', "Image Options"), {
     "samples_save": OptionInfo(True, "Always save all generated images"),
-    "samples_format": OptionInfo('jpg', 'File format for images'),
+    "samples_format": OptionInfo('jpg', 'File format for generated images', gr.Dropdown, lambda: {"choices": ["jpg", "png", "webp", "tiff", "jp2", "psd"]}),
     "samples_filename_pattern": OptionInfo("[seed]-[prompt_spaces]", "Images filename pattern", component_args=hide_dirs),
     "save_images_add_number": OptionInfo(True, "Add number to filename when saving", component_args=hide_dirs),
     "grid_save": OptionInfo(True, "Always save all generated image grids"),
-    "grid_format": OptionInfo('jpg', 'File format for grids'),
+    "grid_format": OptionInfo('jpg', 'File format for grids', gr.Dropdown, lambda: {"choices": ["jpg", "png", "webp", "tiff", "jp2", "psd"]}),
     "grid_extended_filename": OptionInfo(True, "Add extended info (seed, prompt) to filename when saving grid"),
     "grid_only_if_multiple": OptionInfo(True, "Do not save grids consisting of one picture"),
     "grid_prevent_empty_spots": OptionInfo(True, "Prevent empty spots in grid (when set to autodetect)"),

modules/ui_extra_networks.py

@@ -3,11 +3,8 @@ import html
 import os.path
 import urllib.parse
 from pathlib import Path
-from PIL import PngImagePlugin
 import gradio as gr
-
 from modules import shared
-from modules.images import read_info_from_image
 from modules.generation_parameters_copypaste import image_from_url_text
 from modules.ui_components import ToolButton
 
@@ -156,9 +153,7 @@ class ExtraNetworksPage:
         """
         Find a preview PNG for a given path (without extension) and call link_preview on it.
         """
-        preview_extensions = ["png", "jpg", "webp"]
-        if shared.opts.samples_format not in preview_extensions:
-            preview_extensions.append(shared.opts.samples_format)
+        preview_extensions = ["jpg", "png", "webp", "tiff", "jp2", "psd"]
         potential_files = sum([[path + "." + ext, path + ".preview." + ext] for ext in preview_extensions], [])
         for file in potential_files:
             if os.path.isfile(file):
@@ -265,19 +260,13 @@ def setup_ui(ui, gallery):
         index = len(images) - 1 if index >= len(images) else index
         img_info = images[index if index >= 0 else 0]
         image = image_from_url_text(img_info)
-        geninfo, _items = read_info_from_image(image)
         is_allowed = False
         for extra_page in ui.stored_extra_pages:
             if any([path_is_parent(x, filename) for x in extra_page.allowed_directories_for_previews()]):
                 is_allowed = True
                 break
         assert is_allowed, f'writing to {filename} is not allowed'
-        if geninfo:
-            pnginfo_data = PngImagePlugin.PngInfo()
-            pnginfo_data.add_text('parameters', geninfo)
-            image.save(filename, pnginfo=pnginfo_data)
-        else:
-            image.save(filename)
+        image.save(filename)
         return [page.create_html(ui.tabname) for page in ui.stored_extra_pages]
 
     ui.button_save_preview.click(