From 602686a5d27c4c6b71b4dd38d87a4f27d84aace5 Mon Sep 17 00:00:00 2001
From: cmdr2 <secondary.cmdr2@gmail.com>
Date: Wed, 19 Oct 2022 20:27:06 +0530
Subject: [PATCH] Move the current implementation of upscale/redo/double size
 into a custom plugin

---
 ui/media/js/main.js    |  5 -----
 ui/media/js/plugins.js | 51 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/ui/media/js/main.js b/ui/media/js/main.js
index 95b6ade1..c266dd7b 100644
--- a/ui/media/js/main.js
+++ b/ui/media/js/main.js
@@ -369,13 +369,8 @@ function showImages(reqBody, res, outputContainer, livePreview) {
             const buttons = {
                 'imgUseBtn': { text: 'Use as Input', on_click: onUseAsInputClick },
                 'imgSaveBtn': { text: 'Download', on_click: onDownloadImageClick },
-                'imgX2Btn': { text: 'Double Size', on_click: getStartNewTaskHandler('img2img_X2') },
-                'imgRedoBtn': { text: 'Redo', on_click: getStartNewTaskHandler('img2img') },
                 'makeSimilarBtn': { text: 'Make Similar Images', on_click: onMakeSimilarClick },
             }
-            if (!req.use_upscale) {
-                buttons.upscaleBtn = { text: 'Upscale', on_click: getStartNewTaskHandler('upscale') }
-            }
 
             // include the plugins
             Object.assign(buttons, PLUGINS['IMAGE_INFO_BUTTONS'])
diff --git a/ui/media/js/plugins.js b/ui/media/js/plugins.js
index 51c30ecf..92ca296c 100644
--- a/ui/media/js/plugins.js
+++ b/ui/media/js/plugins.js
@@ -26,3 +26,54 @@ const PLUGINS = {
      */
     IMAGE_INFO_BUTTONS: {}
 }
+
+
+PLUGINS['IMAGE_INFO_BUTTONS']['custom_imgX2Btn'] = { text: 'Double Size', on_click: getStartNewTaskHandler('img2img_X2') }
+PLUGINS['IMAGE_INFO_BUTTONS']['custom_imgRedoBtn'] = { text: 'Redo', on_click: getStartNewTaskHandler('img2img') }
+PLUGINS['IMAGE_INFO_BUTTONS']['custom_upscaleBtn'] = { text: 'Upscale', on_click: getStartNewTaskHandler('upscale'), filter: (req, img) => !req.use_upscale }
+
+function getStartNewTaskHandler(mode) {
+    return function(reqBody, img) {
+        const newTaskRequest = getCurrentUserRequest()
+        switch (mode) {
+            case 'img2img':
+            case 'img2img_X2':
+                newTaskRequest.reqBody = Object.assign({}, reqBody, {
+                    num_outputs: 1,
+                    use_cpu: useCPUField.checked,
+                })
+                if (!newTaskRequest.reqBody.init_image || mode === 'img2img_X2') {
+                    newTaskRequest.reqBody.sampler = 'ddim'
+                    newTaskRequest.reqBody.prompt_strength = '0.5'
+                    newTaskRequest.reqBody.init_image = img.src
+                    delete newTaskRequest.reqBody.mask
+                } else {
+                    newTaskRequest.reqBody.seed = 1 + newTaskRequest.reqBody.seed
+                }
+                if (mode === 'img2img_X2') {
+                    newTaskRequest.reqBody.width = reqBody.width * 2
+                    newTaskRequest.reqBody.height = reqBody.height * 2
+                    newTaskRequest.reqBody.num_inference_steps = Math.min(100, reqBody.num_inference_steps * 2)
+                    if (useUpscalingField.checked) {
+                        newTaskRequest.reqBody.use_upscale = upscaleModelField.value
+                    } else {
+                        delete newTaskRequest.reqBody.use_upscale
+                    }
+                }
+                break
+            case 'upscale':
+                newTaskRequest.reqBody = Object.assign({}, reqBody, {
+                    num_outputs: 1,
+                    //use_face_correction: 'GFPGANv1.3',
+                    use_upscale: upscaleModelField.value,
+                })
+                break
+            default:
+                throw new Error("Unknown upscale mode: " + mode)
+        }
+        newTaskRequest.seed = newTaskRequest.reqBody.seed
+        newTaskRequest.numOutputsTotal = 1
+        newTaskRequest.batchCount = 1
+        createTask(newTaskRequest)
+    }
+}