From 5fa3a7ca441faa90179fe40763a2161872d43402 Mon Sep 17 00:00:00 2001
From: cmdr2 <secondary.cmdr2@gmail.com>
Date: Wed, 19 Oct 2022 19:50:05 +0530
Subject: [PATCH] UI-side plugin system; Use PLUGINS['IMAGE_INFO_BUTTONS'] to
 add additional buttons on the generated images

---
 ui/index.html          |  1 +
 ui/media/js/main.js    | 37 ++++++++++++++++++-------------------
 ui/media/js/plugins.js | 23 +++++++++++++++++++++++
 3 files changed, 42 insertions(+), 19 deletions(-)
 create mode 100644 ui/media/js/plugins.js

diff --git a/ui/index.html b/ui/index.html
index 07064615..07a1876c 100644
--- a/ui/index.html
+++ b/ui/index.html
@@ -252,6 +252,7 @@
 </div>
 </body>
 
+<script src="media/js/plugins.js?v=1"></script>
 <script src="media/js/utils.js?v=2"></script>
 <script src="media/js/inpainting-editor.js?v=1"></script>
 <script src="media/js/image-modifiers.js"></script>
diff --git a/ui/media/js/main.js b/ui/media/js/main.js
index 78969875..e0f77636 100644
--- a/ui/media/js/main.js
+++ b/ui/media/js/main.js
@@ -376,6 +376,10 @@ function showImages(reqBody, res, outputContainer, livePreview) {
             if (!req.use_upscale) {
                 buttons.upscaleBtn = { text: 'Upscale', on_click: getStartNewTaskHandler('upscale') }
             }
+
+            // include the plugins
+            Object.assign(buttons, PLUGINS['IMAGE_INFO_BUTTONS'])
+
             const imgItemInfo = imageItemElem.querySelector('.imgItemInfo')
             const createButton = function(name, btnInfo) {
                 const newButton = document.createElement('button')
@@ -383,7 +387,8 @@ function showImages(reqBody, res, outputContainer, livePreview) {
                 newButton.classList.add('tasksBtns')
                 newButton.innerText = btnInfo.text
                 newButton.addEventListener('click', function() {
-                    btnInfo.on_click(req, imageItemElem)
+                    let img = imageItemElem.querySelector('img')
+                    btnInfo.on_click(req, img)
                 })
                 imgItemInfo.appendChild(newButton)
             }
@@ -392,9 +397,8 @@ function showImages(reqBody, res, outputContainer, livePreview) {
     })
 }
 
-function onUseAsInputClick(req, imageItemElem) {
-    const imageElem = imageItemElem.querySelector('img')
-    const imgData = imageElem.src
+function onUseAsInputClick(req, img) {
+    const imgData = img.src
 
     initImageSelector.value = null
     initImagePreview.src = imgData
@@ -406,13 +410,12 @@ function onUseAsInputClick(req, imageItemElem) {
     samplerSelectionContainer.style.display = 'none'
 }
 
-function onDownloadImageClick(req, imageItemElem) {
-    const imageElem = imageItemElem.querySelector('img')
-    const imgData = imageElem.src
-    const imageSeed = imageElem.getAttribute('data-seed')
-    const imagePrompt = imageElem.getAttribute('data-prompt')
-    const imageInferenceSteps = imageElem.getAttribute('data-steps')
-    const imageGuidanceScale = imageElem.getAttribute('data-guidance')
+function onDownloadImageClick(req, img) {
+    const imgData = img.src
+    const imageSeed = img.getAttribute('data-seed')
+    const imagePrompt = img.getAttribute('data-prompt')
+    const imageInferenceSteps = img.getAttribute('data-steps')
+    const imageGuidanceScale = img.getAttribute('data-guidance')
 
     const imgDownload = document.createElement('a')
     imgDownload.download = createFileName(imagePrompt, imageSeed, imageInferenceSteps, imageGuidanceScale, req['output_format'])
@@ -421,12 +424,11 @@ function onDownloadImageClick(req, imageItemElem) {
 }
 
 function getStartNewTaskHandler(mode) {
-    return function(reqBody, imageItemElem) {
+    return function(reqBody, img) {
         if (!isServerAvailable()) {
             alert('The server is not available.')
             return
         }
-        const imageElem = imageItemElem.querySelector('img')
         const newTaskRequest = getCurrentUserRequest()
         switch (mode) {
             case 'img2img':
@@ -438,7 +440,7 @@ function getStartNewTaskHandler(mode) {
                 if (!newTaskRequest.reqBody.init_image || mode === 'img2img_X2') {
                     newTaskRequest.reqBody.sampler = 'ddim'
                     newTaskRequest.reqBody.prompt_strength = '0.5'
-                    newTaskRequest.reqBody.init_image = imageElem.src
+                    newTaskRequest.reqBody.init_image = img.src
                     delete newTaskRequest.reqBody.mask
                 } else {
                     newTaskRequest.reqBody.seed = 1 + newTaskRequest.reqBody.seed
@@ -471,10 +473,7 @@ function getStartNewTaskHandler(mode) {
     }
 }
 
-function onMakeSimilarClick(req, imageItemElem) {
-    const similarImagesCount = 5
-    const imageElem = imageItemElem.querySelector('img')
-
+function onMakeSimilarClick(req, img) {
     let newTaskRequest = getCurrentUserRequest()
 
     newTaskRequest.reqBody = Object.assign({}, req, {
@@ -483,7 +482,7 @@ function onMakeSimilarClick(req, imageItemElem) {
         num_inference_steps: 50,
         guidance_scale: 7.5,
         prompt_strength: 0.7,
-        init_image: imageElem.src,
+        init_image: img.src,
         seed: Math.floor(Math.random() * 10000000)
     })
 
diff --git a/ui/media/js/plugins.js b/ui/media/js/plugins.js
new file mode 100644
index 00000000..91d16c3c
--- /dev/null
+++ b/ui/media/js/plugins.js
@@ -0,0 +1,23 @@
+const PLUGIN_API_VERSION = "1.0"
+
+const PLUGINS = {
+    /**
+     * Register new buttons to show on each output image.
+     * 
+     * Example:
+     * PLUGINS['IMAGE_INFO_BUTTONS']['myCustomVariationButton'] = {
+     *   text: 'Make a Similar Image',
+     *   on_click: function(origRequest, image) {
+     *     let newTaskRequest = getCurrentUserRequest()
+     *     newTaskRequest.reqBody = Object.assign({}, origRequest, {
+     *       init_image: image.src,
+     *       prompt_strength: 0.7,
+     *       seed: Math.floor(Math.random() * 10000000)
+     *     })
+     *     newTaskRequest.seed = newTaskRequest.reqBody.seed
+     *     createTask(newTaskRequest)
+     *   }
+     * }
+     */
+    IMAGE_INFO_BUTTONS: {}
+}