From cec203a6dbe5f4e40adb7b35a3ad02f0f2f6752a Mon Sep 17 00:00:00 2001
From: cmdr2 <secondary.cmdr2@gmail.com>
Date: Thu, 6 Oct 2022 15:05:34 +0530
Subject: [PATCH] Allow saving in jpeg or png format; Default to saving in jpeg

---
 ui/index.html              |  8 +++++++-
 ui/media/main.js           | 14 ++++++++------
 ui/sd_internal/__init__.py |  3 +++
 ui/sd_internal/runtime.py  | 10 +++++-----
 ui/server.py               |  2 ++
 5 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/ui/index.html b/ui/index.html
index f824efed..9d681e97 100644
--- a/ui/index.html
+++ b/ui/index.html
@@ -155,6 +155,12 @@
                     <li class="pl-5"><label for="num_inference_steps">Number of inference steps:</label> <input id="num_inference_steps" name="num_inference_steps" size="4" value="50"></li>
                     <li class="pl-5"><label for="guidance_scale_slider">Guidance Scale:</label> <input id="guidance_scale_slider" name="guidance_scale_slider" class="editor-slider" value="75" type="range" min="10" max="500"> <input id="guidance_scale" name="guidance_scale" size="4"></li>
                     <li class="pl-5"><span id="prompt_strength_container"><label for="prompt_strength_slider">Prompt Strength:</label> <input id="prompt_strength_slider" name="prompt_strength_slider" class="editor-slider" value="80" type="range" min="0" max="99"> <input id="prompt_strength" name="prompt_strength" size="4"><br/></span></li>
+                    <li class="pl-5"><label for="output_format">Output format:</label>
+                        <select id="output_format" name="output_format">
+                            <option value="jpeg" selected>jpeg</option>
+                            <option value="png">png</option>
+                        </select>
+                    </li>
 
                     <br/>
 
@@ -218,7 +224,7 @@
 </div>
 </body>
 
-<script src="media/main.js?v=30"></script>
+<script src="media/main.js?v=31"></script>
 <script>
 async function init() {
     await loadModifiers()
diff --git a/ui/media/main.js b/ui/media/main.js
index 913374ab..bb6b0d9a 100644
--- a/ui/media/main.js
+++ b/ui/media/main.js
@@ -47,6 +47,7 @@ let useFaceCorrectionField = document.querySelector("#use_face_correction")
 let useUpscalingField = document.querySelector("#use_upscale")
 let upscaleModelField = document.querySelector("#upscale_model")
 let stableDiffusionModelField = document.querySelector('#stable_diffusion_model')
+let outputFormatField = document.querySelector('#output_format')
 let showOnlyFilteredImageField = document.querySelector("#show_only_filtered_image")
 let updateBranchLabel = document.querySelector("#updateBranchLabel")
 let streamImageProgressField = document.querySelector("#stream_image_progress")
@@ -300,7 +301,7 @@ function showImages(req, res, outputContainer, livePreview) {
                 saveImageBtn = imageItemElem.querySelector('.imgSaveBtn');
 
             useAsInputBtn.addEventListener('click', getUseAsInputHandler(imageItemElem))
-            saveImageBtn.addEventListener('click', getSaveImageHandler(imageItemElem))
+            saveImageBtn.addEventListener('click', getSaveImageHandler(imageItemElem, req['output_format']))
 
             outputContainer.appendChild(imageItemElem)
         }
@@ -343,14 +344,14 @@ function getUseAsInputHandler(imageItemElem) {
     }
 }
 
-function getSaveImageHandler(imageItemElem) {
+function getSaveImageHandler(imageItemElem, outputFormat) {
     return function() {
         const imageElem = imageItemElem.querySelector('img')
         const imgData = imageElem.src
         const imageSeed = imageElem.getAttribute('data-seed')
 
         const imgDownload = document.createElement('a')
-        imgDownload.download = createFileName(imageSeed)
+        imgDownload.download = createFileName(imageSeed, outputFormat)
         imgDownload.href = imgData
         imgDownload.click()
     }
@@ -628,7 +629,8 @@ async function makeImage() {
         use_stable_diffusion_model: stableDiffusionModelField.value,
         stream_progress_updates: true,
         stream_image_progress: streamImageProgress,
-        show_only_filtered_image: showOnlyFilteredImageField.checked
+        show_only_filtered_image: showOnlyFilteredImageField.checked,
+        output_format: outputFormatField.value
     }
 
     if (IMAGE_REGEX.test(initImagePreview.src)) {
@@ -733,7 +735,7 @@ async function makeImage() {
 
 // create a file name with embedded prompt and metadata
 // for easier cateloging and comparison
-function createFileName(seed) {
+function createFileName(seed, outputFormat) {
 
     // Most important information is the prompt
     let underscoreName = lastPromptUsed.replace(/[^a-zA-Z0-9]/g, '_')
@@ -761,7 +763,7 @@ function createFileName(seed) {
     // fileName += `${tagString}`
 
     // add the file extension
-    fileName += `.png`
+    fileName += '.' + (outputFormat === 'png' ? 'png' : 'jpeg')
 
     return fileName
 }
diff --git a/ui/sd_internal/__init__.py b/ui/sd_internal/__init__.py
index c2548cfd..f2a9901f 100644
--- a/ui/sd_internal/__init__.py
+++ b/ui/sd_internal/__init__.py
@@ -24,6 +24,7 @@ class Request:
     use_upscale: str = None # or "RealESRGAN_x4plus" or "RealESRGAN_x4plus_anime_6B"
     use_stable_diffusion_model: str = "sd-v1-4"
     show_only_filtered_image: bool = False
+    output_format: str = "jpeg" # or "png"
 
     stream_progress_updates: bool = False
     stream_image_progress: bool = False
@@ -44,6 +45,7 @@ class Request:
             "use_face_correction": self.use_face_correction,
             "use_upscale": self.use_upscale,
             "use_stable_diffusion_model": self.use_stable_diffusion_model,
+            "output_format": self.output_format,
         }
 
     def to_string(self):
@@ -66,6 +68,7 @@ class Request:
     use_upscale: {self.use_upscale}
     use_stable_diffusion_model: {self.use_stable_diffusion_model}
     show_only_filtered_image: {self.show_only_filtered_image}
+    output_format: {self.output_format}
 
     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 ec567c05..7c20eff9 100644
--- a/ui/sd_internal/runtime.py
+++ b/ui/sd_internal/runtime.py
@@ -284,7 +284,7 @@ def do_mk_img(req: Request):
     opt_use_face_correction = req.use_face_correction
     opt_use_upscale = req.use_upscale
     opt_show_only_filtered = req.show_only_filtered_image
-    opt_format = 'png'
+    opt_format = req.output_format
     opt_sampler_name = req.sampler
 
     print(req.to_string(), '\n    device', device)
@@ -457,7 +457,7 @@ def do_mk_img(req: Request):
                             save_metadata(meta_out_path, prompts, opt_seed, opt_W, opt_H, opt_ddim_steps, opt_scale, opt_strength, opt_use_face_correction, opt_use_upscale, opt_sampler_name, req.negative_prompt, ckpt_file)
 
                         if return_orig_img:
-                            img_data = img_to_base64_str(img)
+                            img_data = img_to_base64_str(img, opt_format)
                             res_image_orig = ResponseImage(data=img_data, seed=opt_seed)
                             res.images.append(res_image_orig)
 
@@ -484,7 +484,7 @@ def do_mk_img(req: Request):
 
                             filtered_image = Image.fromarray(x_sample)
 
-                            filtered_img_data = img_to_base64_str(filtered_image)
+                            filtered_img_data = img_to_base64_str(filtered_image, opt_format)
                             res_image_filtered = ResponseImage(data=filtered_img_data, seed=opt_seed)
                             res.images.append(res_image_filtered)
 
@@ -652,9 +652,9 @@ def load_mask(mask_str, h0, w0, newH, newW, invert=False):
     return image
 
 # https://stackoverflow.com/a/61114178
-def img_to_base64_str(img):
+def img_to_base64_str(img, output_format="PNG"):
     buffered = BytesIO()
-    img.save(buffered, format="PNG")
+    img.save(buffered, format=output_format)
     buffered.seek(0)
     img_byte = buffered.getvalue()
     img_str = "data:image/png;base64," + base64.b64encode(img_byte).decode()
diff --git a/ui/server.py b/ui/server.py
index 452d85c4..5d7f1bfd 100644
--- a/ui/server.py
+++ b/ui/server.py
@@ -60,6 +60,7 @@ class ImageRequest(BaseModel):
     use_upscale: str = None # or "RealESRGAN_x4plus" or "RealESRGAN_x4plus_anime_6B"
     use_stable_diffusion_model: str = "sd-v1-4"
     show_only_filtered_image: bool = False
+    output_format: str = "jpeg" # or "png"
 
     stream_progress_updates: bool = False
     stream_image_progress: bool = False
@@ -163,6 +164,7 @@ def image(req : ImageRequest):
     r.use_upscale: str = req.use_upscale
     r.use_face_correction = req.use_face_correction
     r.show_only_filtered_image = req.show_only_filtered_image
+    r.output_format = req.output_format
 
     r.stream_progress_updates = True # the underlying implementation only supports streaming
     r.stream_image_progress = req.stream_image_progress