Allow configuring advanced settings like seed, width, height, num_outputs, num_inference_steps and guidance_scale

README.md

@@ -7,7 +7,7 @@ Two things:
 
 All the processing will happen on your local computer, it does not transmit your prompts or process on any remote server.
 
-![Screenshot of tool](shot.jpg?raw=true)
+![Screenshot of tool](shot-v1.jpg?raw=true)
 
 # System Requirements
 1. Requires [Docker](https://docs.docker.com/engine/install/) and [Python](https://www.python.org/downloads/) (3.6 or higher).
@@ -27,6 +27,13 @@ All the processing will happen on your local computer, it does not transmit your
 3. Press `Make Image`. This will take a while, depending on your system's processing power.
 4. See the image generated using your prompt. If there's an error, the status message at the top will show 'error' in red.
 
+# Advanced Settings
+You can also set the configuration like `seed`, `width`, `height`, `num_outputs`, `num_inference_steps` and `guidance_scale` using the 'show' button next to 'Advanced settings'.
+
+Use the same `seed` number to get the same image for a certain prompt. This is useful for refining a prompt without losing the basic image design. Use a seed of `-1` to get random images.
+
+![Screenshot of advanced settings](config-v1.jpg?raw=true)
+
 # Behind the scenes
 This project is a quick way to get started with Stable Diffusion. You do not need to have Stable Diffusion already installed, and do not need any API keys. This project will automatically download Stable Diffusion's docker image, the first time it is run.
 

index.html

@@ -15,9 +15,18 @@
             width: 95%;
         }
     }
-    #image {
+    #configHeader {
-        width: 512px;
+        margin-top: 5px;
-        height: 512px;
+        margin-bottom: 5px;
+        font-size: 10pt;
+    }
+    #config {
+        font-size: 9pt;
+        margin-bottom: 5px;
+        padding-left: 10px;
+    }
+    #outputMsg {
+        font-size: small;
     }
     #footer {
         margin-top: 5px;
@@ -36,9 +45,21 @@
     <b>Prompt:</b><br/>
     <textarea id="prompt">a photograph of an astronaut riding a horse</textarea><br/>
 
+    <div id="configHeader"><b>Advanced settings:</b> [<a id="configToggleBtn" href="#">show</a>]</div>
+    <div id="config">
+        <label for="seed">Seed:</label> <input id="seed" name="seed" value="30000"> <small>(use -1 for random seed)</small><br/>
+        <label for="num_outputs">Number of outputs:</label> <select id="num_outputs" name="num_outputs" value="1"><option value="1" selected>1</option><option value="4">4</option></select><br/>
+        <label for="width">Width:</label> <select id="width" name="width" value="512"><option value="128">128</option><option value="256">256</option><option value="512" selected>512</option><option value="768">768</option><option value="1024">1024</option></select><br/>
+        <label for="height">Height:</label> <select id="height" name="height" value="512"><option value="128">128</option><option value="256">256</option><option value="512" selected>512</option><option value="768">768</option></select><br/>
+        <label for="num_inference_steps">Number of inference steps:</label> <input id="num_inference_steps" name="num_inference_steps" value="50"><br/>
+        <label for="guidance_scale">Guidance Scale:</label> <input id="guidance_scale" name="guidance_scale" value="7.5"><br/>
+    </div>
+
     <button id="makeImage">Make Image</button> <br/><br/>
 
-    <img id="image" />
+    <div id="outputMsg"></div>
+
+    <div id="images"></div>
 
     <div id="footer">
         <p>Please feel free to <a href="https://github.com/cmdr2/stable-diffusion-ui/issues/new">file an issue</a> if you have any problems or suggestions in using this interface.</p>
@@ -96,8 +117,19 @@ async function makeImage() {
     btn.innerHTML = 'Processing..'
     btn.disabled = true;
 
-    let prompt = document.querySelector('#prompt').value
+    let outputMsg = document.querySelector('#outputMsg')
+
+    let reqBody = {
+        prompt: document.querySelector('#prompt').value,
+        num_outputs: document.querySelector('#num_outputs').value,
+        num_inference_steps: document.querySelector('#num_inference_steps').value,
+        guidance_scale: document.querySelector('#guidance_scale').value,
+        width: document.querySelector('#width').value,
+        height: document.querySelector('#height').value,
+        seed: document.querySelector('#seed').value,
+    }
     let res = ''
+    let time = new Date().getTime()
 
     try {
         res = await fetch('/image', {
@@ -105,13 +137,21 @@ async function makeImage() {
             headers: {
                 'Content-Type': 'application/json'
             },
-            body: JSON.stringify({
+            body: JSON.stringify(reqBody)
-                prompt: prompt,
-                width: 512,
-                height: 512
-            })
         })
         res = await res.json()
+
+        if (res.status !== 'succeeded') {
+            if (res.detail !== undefined) {
+                outputMsg.innerHTML = 'Error: ' + res.detail[0].msg + " in " + JSON.stringify(res.detail[0].loc)
+            } else {
+                outputMsg.innerHTML = 'Error: ' + res
+            }
+            outputMsg.innerHTML = '<span style="color: red">' + outputMsg.innerHTML + '</span>'
+            console.log('request error', res)
+            setStatus('request', 'error', 'error')
+            res = undefined
+        }
     } catch (e) {
         console.log('request error', e)
         setStatus('request', 'error', 'error')
@@ -126,24 +166,47 @@ async function makeImage() {
         return
     }
 
+    time = new Date().getTime() - time
+    time /= 1000
+
+    outputMsg.innerHTML = 'Processed in ' + time + ' seconds'
+
+    let images = document.querySelector('#images')
+    images.innerHTML = ''
+
+    for (let idx in res.output) {
         let imgBody = ''
 
         try {
-        imgBody = res.output[0]
+            imgBody = res.output[idx]
         } catch (e) {
             console.log(imgBody)
             setStatus('request', 'invalid image', 'error')
             return
         }
 
-    let img = document.querySelector('#image')
+        let img = document.createElement('img')
+        img.width = parseInt(reqBody.width)
+        img.height = parseInt(reqBody.height)
         img.src = imgBody
 
+        images.appendChild(img)
+    }
+
     setStatus('request', 'done', 'success')
 }
 
 document.querySelector('#makeImage').addEventListener('click', makeImage)
 
+let config = document.querySelector('#config')
+config.style.display = 'none'
+
+document.querySelector('#configToggleBtn').addEventListener('click', function() {
+    config.style.display = (config.style.display === 'none' ? 'block' : 'none')
+    document.querySelector('#configToggleBtn').innerHTML = (config.style.display === 'none' ? 'show' : 'hide')
+    return false
+})
+
 setInterval(healthCheck, HEALTH_PING_INTERVAL * 1000)
 </script>
 

main.py

@@ -9,10 +9,15 @@ PREDICT_URL = LOCAL_SERVER_URL + '/predictions'
 
 app = FastAPI()
 
+# defaults from https://huggingface.co/blog/stable_diffusion
 class ImageRequest(BaseModel):
     prompt: str
-    width: int = 512
+    num_outputs: str = "1"
-    height: int = 512
+    num_inference_steps: str = "50"
+    guidance_scale: str = "7.5"
+    width: str = "512"
+    height: str = "512"
+    seed: str = "30000"
 
 @app.get('/')
 def read_root():
@@ -28,13 +33,22 @@ async def ping():
 
 @app.post('/image')
 async def image(req : ImageRequest):
-    res = requests.post(PREDICT_URL, json={
+    data = {
         "input": {
             "prompt": req.prompt,
-            "width": str(req.width),
+            "num_outputs": req.num_outputs,
-            "height": str(req.height),
+            "num_inference_steps": req.num_inference_steps,
+            "width": req.width,
+            "height": req.height,
+            "seed": req.seed,
         }
-    })
+    }
+
+    if req.seed == "-1":
+        del data['input']['seed']
+
+    res = requests.post(PREDICT_URL, json=data)
+    print(res)
     return res.json()
 
 @app.get('/ding.mp3')