diff --git a/ldm/dream/pngwriter.py b/ldm/dream/pngwriter.py index 584d4479ba8..8b1fc230ff8 100644 --- a/ldm/dream/pngwriter.py +++ b/ldm/dream/pngwriter.py @@ -73,6 +73,7 @@ def unique_filename(self, seed, upscaled=False, previouspath=None): 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) diff --git a/ldm/dream/server.py b/ldm/dream/server.py index e64eefd2c95..3f1126a1ac0 100644 --- a/ldm/dream/server.py +++ b/ldm/dream/server.py @@ -64,7 +64,8 @@ def do_POST(self): 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 @@ def image_done(image, seed, upscaled=False): {'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: diff --git a/static/dream_web/index.css b/static/dream_web/index.css index 8b0f4785f71..beb20e5a742 100644 --- a/static/dream_web/index.css +++ b/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; +} diff --git a/static/dream_web/index.html b/static/dream_web/index.html index 459e78b035b..cc478d0ae04 100644 --- a/static/dream_web/index.html +++ b/static/dream_web/index.html @@ -67,6 +67,9 @@ • +
+ +

The options below require the GFPGAN and ESRGAN packages to be installed

@@ -83,10 +86,13 @@
For news and support for this web service, visit our GitHub site
-
- -
- Postprocessing...1/3 +
+ +
+ +
+ Postprocessing...1/3 +
diff --git a/static/dream_web/index.js b/static/dream_web/index.js index 0ac3374b93d..37899cf41b0 100644 --- a/static/dream_web/index.js +++ b/static/dream_web/index.js @@ -53,6 +53,7 @@ function clearFields(form) { form.prompt.value = prompt; } +const BLANK_IMAGE_URL = 'data:image/svg+xml,'; 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) - } 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"; + 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; + } } } }