From 96a6b11ab4d9af11dd90c90903efd091896efb11 Mon Sep 17 00:00:00 2001
From: cmdr2 <secondary.cmdr2@gmail.com>
Date: Mon, 5 Sep 2022 13:10:19 +0530
Subject: [PATCH] Show a warning if the initial image is larger than 768x768,
 can cause out of memory errors

---
 ui/index.html | 53 +++++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 47 insertions(+), 6 deletions(-)

diff --git a/ui/index.html b/ui/index.html
index 4578ec04..3cff9ebd 100644
--- a/ui/index.html
+++ b/ui/index.html
@@ -384,6 +384,8 @@ const USE_FULL_PRECISION_KEY = "useFullPrecision"
 const USE_TURBO_MODE_KEY = "useTurboMode"
 const HEALTH_PING_INTERVAL = 5 // seconds
 
+const IMAGE_REGEX = new RegExp('data:image/[A-Za-z]+;base64')
+
 let promptField = document.querySelector('#prompt')
 let numOutputsTotalField = document.querySelector('#num_outputs_total')
 let numOutputsParallelField = document.querySelector('#num_outputs_parallel')
@@ -495,8 +497,21 @@ function setStatus(statusType, msg, msgType) {
     }
 }
 
+function logMsg(msg, level) {
+    if (level === 'error') {
+        outputMsg.innerHTML = '<span style="color: red">Error: ' + msg + '</span>'
+    } else if (level === 'warn') {
+        outputMsg.innerHTML = '<span style="color: orange">Warning: ' + msg + '</span>'
+    } else {
+        outputMsg.innerHTML = msg
+    }
+
+    console.log(level, msg)
+}
+
 function logError(msg, res) {
-    outputMsg.innerHTML = '<span style="color: red">Error: ' + msg + '</span>'
+    logMsg(msg, 'error')
+
     console.log('request error', res)
     setStatus('request', 'error', 'error')
 }
@@ -643,20 +658,46 @@ async function doMakeImage(reqBody) {
     return true
 }
 
+function validateInput() {
+    const MAX_INIT_IMAGE_DIMENSION = 768
+
+    let width = parseInt(widthField.value)
+    let height = parseInt(heightField.value)
+
+    if (IMAGE_REGEX.test(initImagePreview.src)) {
+        if (initImagePreview.naturalWidth > MAX_INIT_IMAGE_DIMENSION || initImagePreview.naturalHeight > MAX_INIT_IMAGE_DIMENSION) {
+            return {'isValid': false, 'warning': `The dimensions of your initial image are very large, and can cause 'Out of Memory' errors! Please ensure that its dimensions are equal (or smaller) than the desired output image.
+                    <br/><br/>
+                    Your initial image size is ${initImagePreview.naturalWidth}x${initImagePreview.naturalHeight} pixels. Please try to keep it smaller than ${MAX_INIT_IMAGE_DIMENSION}x${MAX_INIT_IMAGE_DIMENSION}.`}
+        }
+    }
+
+    return {'isValid': true}
+}
+
 async function makeImage() {
     if (serverStatus !== 'online') {
         logError('The server is still starting up..')
         return
     }
 
+    let validation = validateInput()
+    if (validation['isValid']) {
+        outputMsg.innerHTML = 'Fetching..'
+    } else {
+        if (validation['error']) {
+            logError(validation['error'])
+            return
+        } else if (validation['warning']) {
+            logMsg(validation['warning'], 'warn')
+        }
+    }
+
     setStatus('request', 'fetching..')
 
     makeImageBtn.innerHTML = 'Processing..'
     makeImageBtn.disabled = true
 
-    outputMsg.innerHTML = 'Fetching..'
-
-    const imageRegex = new RegExp('data:image/[A-Za-z]+;base64')
     let seed = (randomSeedField.checked ? Math.floor(Math.random() * 10000) : parseInt(seedField.value))
     let numOutputsTotal = parseInt(numOutputsTotalField.value)
     let numOutputsParallel = parseInt(numOutputsParallelField.value)
@@ -685,11 +726,11 @@ async function makeImage() {
         use_full_precision: useFullPrecisionField.checked
     }
 
-    if (imageRegex.test(initImagePreview.src)) {
+    if (IMAGE_REGEX.test(initImagePreview.src)) {
         reqBody['init_image'] = initImagePreview.src
         reqBody['prompt_strength'] = parseInt(promptStrengthField.value) / 10
 
-        // if (imageRegex.test(maskImagePreview.src)) {
+        // if (IMAGE_REGEX.test(maskImagePreview.src)) {
         //     reqBody['mask'] = maskImagePreview.src
         // }
     }