Merge branch 'sw-drop' of https://github.com/bakkot/stable-diffusion into bakkot-sw-drop

ldm/dream/pngwriter.py

@@ -73,6 +73,7 @@ class PngWriter:
                     if upscaled:
                         break
                     filename = f'{basecount:06}.{seed}.{series:02}.png'
+                path = os.path.join(self.outdir, filename)
                 finished = not os.path.exists(path)
             return os.path.join(self.outdir, filename)
 

ldm/dream/server.py

@@ -64,7 +64,8 @@ class DreamServer(BaseHTTPRequestHandler):
         upscale_level    = post_data['upscale_level']
         upscale_strength = post_data['upscale_strength']
         upscale = [int(upscale_level),float(upscale_strength)] if upscale_level != '' else None
-        seed = None if int(post_data['seed']) == -1 else int(post_data['seed'])
+        progress_images = 'progress_images' in post_data
+        seed = self.model.seed if int(post_data['seed']) == -1 else int(post_data['seed'])
 
         print(f"Request to generate with prompt: {prompt}")
         # In order to handle upscaled images, the PngWriter needs to maintain state
@@ -116,9 +117,20 @@ class DreamServer(BaseHTTPRequestHandler):
                         {'event':action,'processed_file_cnt':f'{x}/{iterations}'}
                     ) + '\n',"utf-8"))
 
-        def image_progress(image, step):
+        # TODO: refactor PngWriter:
+        # it doesn't need to know if batch_size > 1, just if this is _part of a batch_
+        step_writer = PngWriter('./outputs/intermediates/', prompt, 2)
+        def image_progress(sample, step):
+            url = None
+            # since rendering images is moderately expensive, only render every 5th image
+            # and don't bother with the last one, since it'll render anyway
+            if progress_images and step % 5 == 0 and step < steps - 1:
+                images = self.model._samples_to_images(sample)
+                image = images[0]
+                step_writer.write_image(image, seed) # TODO PngWriter to return path
+                url = step_writer.filepath
             self.wfile.write(bytes(json.dumps(
-                {'event':'step', 'step':step}
+                {'event':'step', 'step':step, 'url': url}
             ) + '\n',"utf-8"))
 
         if initimg is None:

static/dream_web/index.css

@@ -46,7 +46,7 @@ fieldset {
     margin: auto;
     padding-top: 10px;
 }
-img {
+#results img {
     cursor: pointer;
     height: 30vh;
     border-radius: 5px;
@@ -67,3 +67,10 @@ hr {
 label {
     white-space: nowrap;
 }
+#progress-section {
+    display: none;
+}
+#progress-image {
+    width: 30vh;
+    height: 30vh;
+}

static/dream_web/index.html

@@ -67,6 +67,9 @@
           <button type="button" id="reset-seed">&olarr;</button>
           <span>&bull;</span>
           <button type="button" id="reset-all">Reset to Defaults</button>
+          <br>
+          <label for="progress_images">Display in-progress images (slows down generation):</label>
+          <input type="checkbox" name="progress_images" id="progress_images">
           <div id="gfpgan">
             <p><em>The options below require the GFPGAN and ESRGAN packages to be installed</em></p>
             <label title="Strength of the gfpgan (face fixing) algorithm." for="gfpgan_strength">GPFGAN Strength:</label>
@@ -83,12 +86,15 @@
         </fieldset>
       </form>
       <div id="about">For news and support for this web service, visit our <a href="http://github.com/lstein/stable-diffusion">GitHub site</a></div>
+      <div id="progress-section">
+        <progress id="progress-bar" value="0" max="1"></progress>
         <br>
-      <progress id="progress" value="0" max="1"></progress>
+        <img id="progress-image" src='data:image/svg+xml,<svg xmlns="http://www.w3.org/2000/svg"/>'></img>
         <div id="scaling-inprocess-message">
           <i><span>Postprocessing...</span><span id="processing_cnt">1/3</span></i>
         </div>
       </div>
+    </div>
     <div id="results">
       <div id="no-results-message">
         <i><p>No results...</p></i>

static/dream_web/index.js

@@ -53,6 +53,7 @@ function clearFields(form) {
     form.prompt.value = prompt;
 }
 
+const BLANK_IMAGE_URL = 'data:image/svg+xml,<svg xmlns="http://www.w3.org/2000/svg"/>';
 async function generateSubmit(form) {
     const prompt = document.querySelector("#prompt").value;
 
@@ -60,7 +61,14 @@ async function generateSubmit(form) {
     let formData = Object.fromEntries(new FormData(form));
     formData.initimg = formData.initimg.name !== '' ? await toBase64(formData.initimg) : null;
 
-    document.querySelector('progress').setAttribute('max', formData.steps);
+    let progressSectionEle = document.querySelector('#progress-section');
+    progressSectionEle.style.display = 'initial';
+    let progressEle = document.querySelector('#progress-bar');
+    progressEle.setAttribute('max', formData.steps);
+    let progressImageEle = document.querySelector('#progress-image');
+    progressImageEle.src = BLANK_IMAGE_URL;
+
+    progressImageEle.style.display = {}.hasOwnProperty.call(formData, 'progress_images') ? 'initial': 'none';
 
     // Post as JSON, using Fetch streaming to get results
     fetch(form.action, {
@@ -73,7 +81,10 @@ async function generateSubmit(form) {
         while (true) {
             let {value, done} = await reader.read();
             value = new TextDecoder().decode(value);
-            if (done) break;
+            if (done) {
+                progressSectionEle.style.display = 'none';
+                break;
+            }
 
             for (let event of value.split('\n').filter(e => e !== '')) {
                 const data = JSON.parse(event);
@@ -81,14 +92,20 @@ async function generateSubmit(form) {
                 if (data.event == 'result') {
                     noOutputs = false;
                     document.querySelector("#no-results-message")?.remove();
-		    appendOutput(data.files[0],data.files[1],data.config)
+                    appendOutput(data.files[0],data.files[1],data.config);
+                    progressEle.setAttribute('value', 0);
+                    progressEle.setAttribute('max', formData.steps);
+                    progressImageEle.src = BLANK_IMAGE_URL;
                 } else if (data.event == 'upscaling-started') {
                     document.getElementById("processing_cnt").textContent=data.processed_file_cnt;
                     document.getElementById("scaling-inprocess-message").style.display = "block";
                 } else if (data.event == 'upscaling-done') {
                     document.getElementById("scaling-inprocess-message").style.display = "none";
                 } else if (data.event == 'step') {
-                    document.querySelector('progress').setAttribute('value', data.step.toString());
+                    progressEle.setAttribute('value', data.step);
+                    if (data.url) {
+                        progressImageEle.src = data.url;
+                    }
                 }
             }
         }