From e7ca8090fdeb53b922a7ebed4fbeaedf19f60838 Mon Sep 17 00:00:00 2001
From: JeLuF <jf@mormo.org>
Date: Mon, 5 Dec 2022 06:32:33 +0100
Subject: [PATCH] Make JPEG Output quality user controllable (#607)

Add a slider to the image options for the JPEG quality
For PNG images, the slider is hidden.
---
 ui/index.html                  |  3 +++
 ui/media/js/auto-save.js       |  1 +
 ui/media/js/main.js            | 36 ++++++++++++++++++++++++++++++++++
 ui/media/js/utils.js           | 10 ++++++++++
 ui/sd_internal/__init__.py     |  3 +++
 ui/sd_internal/runtime.py      | 26 ++++++++++++++----------
 ui/sd_internal/task_manager.py |  3 +++
 7 files changed, 72 insertions(+), 10 deletions(-)

diff --git a/ui/index.html b/ui/index.html
index 3e7ba8cb..2a99155f 100644
--- a/ui/index.html
+++ b/ui/index.html
@@ -195,6 +195,9 @@
                             <option value="png">png</option>
                         </select>
                     </td></tr>
+                    <tr class="pl-5" id="output_quality_row"><td><label for="output_quality">JPEG Quality:</label></td><td>
+		        <input id="output_quality_slider" name="output_quality" class="editor-slider" value="75" type="range" min="10" max="95"> <input id="output_quality" name="output_quality" size="4" pattern="^[0-9\.]+$" onkeypress="preventNonNumericalInput(event)">
+	            </td></tr>
                     </table></div>
 
                     <div><ul>
diff --git a/ui/media/js/auto-save.js b/ui/media/js/auto-save.js
index f6db831d..2d1543cc 100644
--- a/ui/media/js/auto-save.js
+++ b/ui/media/js/auto-save.js
@@ -21,6 +21,7 @@ const SETTINGS_IDS_LIST = [
     "guidance_scale",
     "prompt_strength",
     "output_format",
+    "output_quality",
     "negative_prompt",
     "stream_image_progress",
     "use_face_correction",
diff --git a/ui/media/js/main.js b/ui/media/js/main.js
index 687e0251..07bd8995 100644
--- a/ui/media/js/main.js
+++ b/ui/media/js/main.js
@@ -16,6 +16,9 @@ let numOutputsParallelField = document.querySelector('#num_outputs_parallel')
 let numInferenceStepsField = document.querySelector('#num_inference_steps')
 let guidanceScaleSlider = document.querySelector('#guidance_scale_slider')
 let guidanceScaleField = document.querySelector('#guidance_scale')
+let outputQualitySlider = document.querySelector('#output_quality_slider')
+let outputQualityField = document.querySelector('#output_quality')
+let outputQualityRow = document.querySelector('#output_quality_row')
 let randomSeedField = document.querySelector("#random_seed")
 let seedField = document.querySelector('#seed')
 let widthField = document.querySelector('#width')
@@ -791,6 +794,7 @@ function getCurrentUserRequest() {
             stream_image_progress: (numOutputsTotal > 50 ? false : streamImageProgressField.checked),
             show_only_filtered_image: showOnlyFilteredImageField.checked,
             output_format: outputFormatField.value,
+            output_quality: outputQualityField.value,
             original_prompt: promptField.value,
             active_tags: (activeTags.map(x => x.name))
         }
@@ -1116,6 +1120,7 @@ document.onkeydown = function(e) {
     }
 }
 
+/********************* Guidance **************************/
 function updateGuidanceScale() {
     guidanceScaleField.value = guidanceScaleSlider.value / 10
     guidanceScaleField.dispatchEvent(new Event("change"))
@@ -1136,6 +1141,7 @@ guidanceScaleSlider.addEventListener('input', updateGuidanceScale)
 guidanceScaleField.addEventListener('input', updateGuidanceScaleSlider)
 updateGuidanceScale()
 
+/********************* Prompt Strength *******************/
 function updatePromptStrength() {
     promptStrengthField.value = promptStrengthSlider.value / 100
     promptStrengthField.dispatchEvent(new Event("change"))
@@ -1156,6 +1162,36 @@ promptStrengthSlider.addEventListener('input', updatePromptStrength)
 promptStrengthField.addEventListener('input', updatePromptStrengthSlider)
 updatePromptStrength()
 
+/********************* JPEG Quality **********************/
+function updateOutputQuality() {
+    outputQualityField.value =  0 | outputQualitySlider.value
+    outputQualityField.dispatchEvent(new Event("change"))
+}
+
+function updateOutputQualitySlider() {
+    if (outputQualityField.value < 10) {
+        outputQualityField.value = 10
+    } else if (outputQualityField.value > 95) {
+        outputQualityField.value = 95
+    }
+
+    outputQualitySlider.value =  0 | outputQualityField.value
+    outputQualitySlider.dispatchEvent(new Event("change"))
+}
+
+outputQualitySlider.addEventListener('input', updateOutputQuality)
+outputQualityField.addEventListener('input', debounce(updateOutputQualitySlider))
+updateOutputQuality()
+
+outputFormatField.addEventListener('change', e => {
+    if (outputFormatField.value == 'jpeg') {
+        outputQualityRow.style.display='table-row'
+    } else {
+        outputQualityRow.style.display='none'
+    }
+})
+
+
 async function getModels() {
     try {
         var sd_model_setting_key = "stable_diffusion_model"
diff --git a/ui/media/js/utils.js b/ui/media/js/utils.js
index a76f030e..109a05d0 100644
--- a/ui/media/js/utils.js
+++ b/ui/media/js/utils.js
@@ -347,6 +347,16 @@ function asyncDelay(timeout) {
     })
 }
 
+/* Simple debounce function, placeholder for the one in engine.js for simple use cases */
+function debounce(func, timeout = 300){
+  let timer;
+  return (...args) => {
+    clearTimeout(timer);
+    timer = setTimeout(() => { func.apply(this, args); }, timeout);
+  };
+}
+
+
 function preventNonNumericalInput(e) {
     e = e || window.event;
     let charCode = (typeof e.which == "undefined") ? e.keyCode : e.which;
diff --git a/ui/sd_internal/__init__.py b/ui/sd_internal/__init__.py
index da5f6b59..18b0a01b 100644
--- a/ui/sd_internal/__init__.py
+++ b/ui/sd_internal/__init__.py
@@ -25,6 +25,7 @@ class Request:
     use_vae_model: str = None
     show_only_filtered_image: bool = False
     output_format: str = "jpeg" # or "png"
+    output_quality: int = 75
 
     stream_progress_updates: bool = False
     stream_image_progress: bool = False
@@ -47,6 +48,7 @@ class Request:
             "use_stable_diffusion_model": self.use_stable_diffusion_model,
             "use_vae_model": self.use_vae_model,
             "output_format": self.output_format,
+            "output_quality": self.output_quality,
         }
 
     def __str__(self):
@@ -70,6 +72,7 @@ class Request:
     use_vae_model: {self.use_vae_model}
     show_only_filtered_image: {self.show_only_filtered_image}
     output_format: {self.output_format}
+    output_quality: {self.output_quality}
 
     stream_progress_updates: {self.stream_progress_updates}
     stream_image_progress: {self.stream_image_progress}'''
diff --git a/ui/sd_internal/runtime.py b/ui/sd_internal/runtime.py
index ed1d69a7..0f491865 100644
--- a/ui/sd_internal/runtime.py
+++ b/ui/sd_internal/runtime.py
@@ -675,12 +675,12 @@ def do_mk_img(req: Request, data_queue: queue.Queue, task_temp_images: list, ste
                         if req.save_to_disk_path is not None:
                             if return_orig_img:
                                 img_out_path = get_base_path(req.save_to_disk_path, req.session_id, prompts[0], img_id, req.output_format)
-                                save_image(img, img_out_path)
+                                save_image(img, img_out_path, req.output_format, req.output_quality)
                             meta_out_path = get_base_path(req.save_to_disk_path, req.session_id, prompts[0], img_id, 'txt')
                             save_metadata(meta_out_path, req, prompts[0], opt_seed)
 
                         if return_orig_img:
-                            img_buffer = img_to_buffer(img, req.output_format)
+                            img_buffer = img_to_buffer(img, req.output_format, req.output_quality)
                             img_str = buffer_to_base64_str(img_buffer, req.output_format)
                             res_image_orig = ResponseImage(data=img_str, seed=opt_seed)
                             res.images.append(res_image_orig)
@@ -700,14 +700,14 @@ def do_mk_img(req: Request, data_queue: queue.Queue, task_temp_images: list, ste
                                 filters_applied.append(req.use_upscale)
                             if (len(filters_applied) > 0):
                                 filtered_image = Image.fromarray(img_data[i])
-                                filtered_buffer = img_to_buffer(filtered_image, req.output_format)
+                                filtered_buffer = img_to_buffer(filtered_image, req.output_format, req.output_quality)
                                 filtered_img_data = buffer_to_base64_str(filtered_buffer, req.output_format)
                                 response_image = ResponseImage(data=filtered_img_data, seed=opt_seed)
                                 res.images.append(response_image)
                                 task_temp_images[i] = filtered_buffer
                                 if req.save_to_disk_path is not None:
                                     filtered_img_out_path = get_base_path(req.save_to_disk_path, req.session_id, prompts[0], img_id, req.output_format, "_".join(filters_applied))
-                                    save_image(filtered_image, filtered_img_out_path)
+                                    save_image(filtered_image, filtered_img_out_path, req.output_format, req.output_quality)
                                     response_image.path_abs = filtered_img_out_path
                                 del filtered_image
                         # Filter Applied, move to next seed
@@ -728,9 +728,12 @@ def do_mk_img(req: Request, data_queue: queue.Queue, task_temp_images: list, ste
 
     return res
 
-def save_image(img, img_out_path):
+def save_image(img, img_out_path, output_format="", output_quality=75):
     try:
-        img.save(img_out_path)
+        if output_format.upper() == "JPEG":
+            img.save(img_out_path, quality=output_quality)
+        else:
+            img.save(img_out_path)
     except:
         print('could not save the file', traceback.format_exc())
 
@@ -938,13 +941,16 @@ def load_mask(mask_str, h0, w0, newH, newW, invert=False):
     return image
 
 # https://stackoverflow.com/a/61114178
-def img_to_base64_str(img, output_format="PNG"):
-    buffered = img_to_buffer(img, output_format)
+def img_to_base64_str(img, output_format="PNG", output_quality=75):
+    buffered = img_to_buffer(img, output_format, quality=output_quality)
     return buffer_to_base64_str(buffered, output_format)
 
-def img_to_buffer(img, output_format="PNG"):
+def img_to_buffer(img, output_format="PNG", output_quality=75):
     buffered = BytesIO()
-    img.save(buffered, format=output_format)
+    if ( output_format.upper() == "JPEG" ):
+        img.save(buffered, format=output_format, quality=output_quality)
+    else:
+        img.save(buffered, format=output_format)
     buffered.seek(0)
     return buffered
 
diff --git a/ui/sd_internal/task_manager.py b/ui/sd_internal/task_manager.py
index ff6cbb4c..d21f2877 100644
--- a/ui/sd_internal/task_manager.py
+++ b/ui/sd_internal/task_manager.py
@@ -79,6 +79,7 @@ class ImageRequest(BaseModel):
     use_vae_model: str = None
     show_only_filtered_image: bool = False
     output_format: str = "jpeg" # or "png"
+    output_quality: int = 75
 
     stream_progress_updates: bool = False
     stream_image_progress: bool = False
@@ -95,6 +96,7 @@ class FilterRequest(BaseModel):
     render_device: str = None
     use_full_precision: bool = False
     output_format: str = "jpeg" # or "png"
+    output_quality: int = 75
 
 # Temporary cache to allow to query tasks results for a short time after they are completed.
 class TaskCache():
@@ -504,6 +506,7 @@ def render(req : ImageRequest):
     r.use_vae_model = req.use_vae_model
     r.show_only_filtered_image = req.show_only_filtered_image
     r.output_format = req.output_format
+    r.output_quality = req.output_quality
 
     r.stream_progress_updates = True # the underlying implementation only supports streaming
     r.stream_image_progress = req.stream_image_progress