From aadfe586b282925a0e4b2ef9deea63bbbc7f3dc7 Mon Sep 17 00:00:00 2001 From: tesseractcat Date: Fri, 26 Aug 2022 21:10:13 -0400 Subject: [PATCH 1/5] Add step_callback to prompt2image --- ldm/simplet2i.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/ldm/simplet2i.py b/ldm/simplet2i.py index 50816ec403e..cae57af345c 100644 --- a/ldm/simplet2i.py +++ b/ldm/simplet2i.py @@ -202,6 +202,7 @@ def prompt2image( ddim_eta=None, skip_normalize=False, image_callback=None, + step_callback=None, # these are specific to txt2img width=None, height=None, @@ -228,9 +229,14 @@ def prompt2image( gfpgan_strength // strength for GFPGAN. 0.0 preserves image exactly, 1.0 replaces it completely ddim_eta // image randomness (eta=0.0 means the same seed always produces the same image) variants // if >0, the 1st generated image will be passed back to img2img to generate the requested number of variants + step_callback // a function or method that will be called each step image_callback // a function or method that will be called each time an image is generated - To use the callback, define a function of method that receives two arguments, an Image object + To use the step callback, define a function that receives two arguments: + - Image GPU data + - The step number + + To use the image callback, define a function of method that receives two arguments, an Image object and the seed. You can then do whatever you like with the image, including converting it to different formats and manipulating it. For example: @@ -285,6 +291,7 @@ def process_image(image,seed): skip_normalize=skip_normalize, init_img=init_img, strength=strength, + callback=step_callback, ) else: images_iterator = self._txt2img( @@ -297,6 +304,7 @@ def process_image(image,seed): skip_normalize=skip_normalize, width=width, height=height, + callback=step_callback, ) with scope(self.device.type), self.model.ema_scope(): @@ -348,6 +356,7 @@ def _txt2img( skip_normalize, width, height, + callback, ): """ An infinite iterator of images from the prompt. @@ -371,6 +380,7 @@ def _txt2img( unconditional_guidance_scale=cfg_scale, unconditional_conditioning=uc, eta=ddim_eta, + img_callback=callback ) yield self._samples_to_images(samples) @@ -386,6 +396,7 @@ def _img2img( skip_normalize, init_img, strength, + callback, # Currently not implemented for img2img ): """ An infinite iterator of images from the prompt and the initial image From bd417b846e0230735cfe66b7449611409ca4ba74 Mon Sep 17 00:00:00 2001 From: tesseractcat Date: Fri, 26 Aug 2022 21:10:25 -0400 Subject: [PATCH 2/5] Switch to http fetch streaming technique --- scripts/dream_web.py | 73 +++++++++++++++++++++++++-------------- static/dream_web/index.js | 44 ++++++++++++++--------- 2 files changed, 74 insertions(+), 43 deletions(-) diff --git a/scripts/dream_web.py b/scripts/dream_web.py index 227e8acdc78..830890754a3 100644 --- a/scripts/dream_web.py +++ b/scripts/dream_web.py @@ -7,6 +7,7 @@ print("Loading model...") from ldm.simplet2i import T2I +from ldm.dream.pngwriter import PngWriter model = T2I(sampler_name='k_lms') # to get rid of annoying warning messages from pytorch @@ -56,16 +57,44 @@ def do_POST(self): print(f"Request to generate with prompt: {prompt}") - outputs = [] + def image_done(image, seed): + config = post_data.copy() # Shallow copy + config['initimg'] = '' + + # Write PNGs + pngwriter = PngWriter( + "./outputs/img-samples/", config['prompt'], 1 + ) + # metadata_str = f'prompt2png({json.dumps(config)} seed={seed}' # gets written into the PNG + pngwriter.write_image(image, seed) + + # Append post_data to log + with open("./outputs/img-samples/dream_web_log.txt", "a") as log: + for file_path, _ in pngwriter.files_written: + log.write(f"{file_path}: {json.dumps(config)}\n") + + self.wfile.write(bytes(json.dumps( + {'event':'result', 'files':pngwriter.files_written, 'config':config} + ) + '\n',"utf-8")) + + def image_progress(image, step): + self.wfile.write(bytes(json.dumps( + {'event':'step', 'step':step} + ) + '\n',"utf-8")) + + # outputs = [] if initimg is None: # Run txt2img - outputs = model.txt2img(prompt, - iterations=iterations, - cfg_scale = cfgscale, - width = width, - height = height, - seed = seed, - steps = steps) + model.prompt2image(prompt, + iterations=iterations, + cfg_scale = cfgscale, + width = width, + height = height, + seed = seed, + steps = steps, + + step_callback=image_progress, + image_callback=image_done) else: # Decode initimg as base64 to temp file with open("./img2img-tmp.png", "wb") as f: @@ -73,27 +102,19 @@ def do_POST(self): f.write(base64.b64decode(initimg)) # Run img2img - outputs = model.img2img(prompt, - init_img = "./img2img-tmp.png", - iterations = iterations, - cfg_scale = cfgscale, - seed = seed, - steps = steps) + model.prompt2image(prompt, + init_img = "./img2img-tmp.png", + iterations = iterations, + cfg_scale = cfgscale, + seed = seed, + steps = steps, + + step_callback=image_progress, + image_callback=image_done) # Remove the temp file os.remove("./img2img-tmp.png") - print(f"Prompt generated with output: {outputs}") - - post_data['initimg'] = '' # Don't send init image back - - # Append post_data to log - with open("./outputs/img-samples/dream_web_log.txt", "a") as log: - for output in outputs: - log.write(f"{output[0]}: {json.dumps(post_data)}\n") - - outputs = [x + [post_data] for x in outputs] # Append config to each output - result = {'outputs': outputs} - self.wfile.write(bytes(json.dumps(result), "utf-8")) + print(f"Prompt generated!") if __name__ == "__main__": # Change working directory to the stable-diffusion directory diff --git a/static/dream_web/index.js b/static/dream_web/index.js index 3b99deecf49..aa0a157a999 100644 --- a/static/dream_web/index.js +++ b/static/dream_web/index.js @@ -7,12 +7,11 @@ function toBase64(file) { }); } -function appendOutput(output) { +function appendOutput(src, seed, config) { let outputNode = document.createElement("img"); - outputNode.src = output[0]; + outputNode.src = src; - let outputConfig = output[2]; - let altText = output[1].toString() + " | " + outputConfig.prompt; + let altText = seed.toString() + " | " + config.prompt; outputNode.alt = altText; outputNode.title = altText; @@ -20,9 +19,9 @@ function appendOutput(output) { outputNode.addEventListener('click', () => { let form = document.querySelector("#generate-form"); for (const [k, v] of new FormData(form)) { - form.querySelector(`*[name=${k}]`).value = outputConfig[k]; + form.querySelector(`*[name=${k}]`).value = config[k]; } - document.querySelector("#seed").value = output[1]; + document.querySelector("#seed").value = seed; saveFields(document.querySelector("#generate-form")); }); @@ -30,12 +29,6 @@ function appendOutput(output) { document.querySelector("#results").prepend(outputNode); } -function appendOutputs(outputs) { - for (const output of outputs) { - appendOutput(output); - } -} - function saveFields(form) { for (const [k, v] of new FormData(form)) { if (typeof v !== 'object') { // Don't save 'file' type @@ -59,20 +52,37 @@ async function generateSubmit(form) { let formData = Object.fromEntries(new FormData(form)); formData.initimg = formData.initimg.name !== '' ? await toBase64(formData.initimg) : null; - // Post as JSON + // Post as JSON, using Fetch streaming to get results fetch(form.action, { method: form.method, body: JSON.stringify(formData), - }).then(async (result) => { - let data = await result.json(); + }).then(async (response) => { + const reader = response.body.pipeThrough(new TextDecoderStream()).getReader(); + + let noOutputs = true; + while (true) { + const {value, done} = await reader.read(); + if (done) break; + + for (let event of value.split('\n').filter(e => e !== '')) { + const data = JSON.parse(event); + + if (data.event == 'result') { + noOutputs = false; + + for (let [file, seed] of data.files) { + appendOutput(file, seed, data.config) + } + } + } + } // Re-enable form, remove no-results-message form.querySelector('fieldset').removeAttribute('disabled'); document.querySelector("#prompt").value = prompt; - if (data.outputs.length != 0) { + if (!noOutputs) { document.querySelector("#no-results-message")?.remove(); - appendOutputs(data.outputs); } else { alert("Error occurred while generating."); } From 87ee4b4363330ed072595cdd457c2ae49db2400e Mon Sep 17 00:00:00 2001 From: tesseractcat Date: Fri, 26 Aug 2022 21:28:21 -0400 Subject: [PATCH 3/5] Fix minor bug, add progress bar --- static/dream_web/index.html | 3 ++- static/dream_web/index.js | 12 ++++++++---- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/static/dream_web/index.html b/static/dream_web/index.html index 4dccd05fb3b..ed76a682809 100644 --- a/static/dream_web/index.html +++ b/static/dream_web/index.html @@ -55,8 +55,9 @@
For news and support for this web service, visit our GitHub site
+
+ -

No results...

 diff --git a/static/dream_web/index.js b/static/dream_web/index.js index aa0a157a999..a0f79a14819 100644 --- a/static/dream_web/index.js +++ b/static/dream_web/index.js @@ -52,6 +52,8 @@ 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); + // Post as JSON, using Fetch streaming to get results fetch(form.action, { method: form.method, @@ -69,10 +71,13 @@ async function generateSubmit(form) { if (data.event == 'result') { noOutputs = false; + document.querySelector("#no-results-message")?.remove(); for (let [file, seed] of data.files) { - appendOutput(file, seed, data.config) + appendOutput(file, seed, data.config); } + } else if (data.event == 'step') { + document.querySelector('progress').setAttribute('value', data.step.toString()); } } } @@ -80,10 +85,9 @@ async function generateSubmit(form) { // Re-enable form, remove no-results-message form.querySelector('fieldset').removeAttribute('disabled'); document.querySelector("#prompt").value = prompt; + document.querySelector('progress').setAttribute('value', '0'); - if (!noOutputs) { - document.querySelector("#no-results-message")?.remove(); - } else { + if (noOutputs) { alert("Error occurred while generating."); } }); From 1add47bb0d1db8948a223d649989b0aaabaa0fbe Mon Sep 17 00:00:00 2001 From: tesseractcat Date: Fri, 26 Aug 2022 22:02:21 -0400 Subject: [PATCH 4/5] Fixed pipeThrough compatibility issue with Firefox android --- static/dream_web/index.js | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/static/dream_web/index.js b/static/dream_web/index.js index a0f79a14819..3952201b736 100644 --- a/static/dream_web/index.js +++ b/static/dream_web/index.js @@ -59,11 +59,12 @@ async function generateSubmit(form) { method: form.method, body: JSON.stringify(formData), }).then(async (response) => { - const reader = response.body.pipeThrough(new TextDecoderStream()).getReader(); + const reader = response.body.getReader(); let noOutputs = true; while (true) { - const {value, done} = await reader.read(); + let {value, done} = await reader.read(); + value = new TextDecoder().decode(value); if (done) break; for (let event of value.split('\n').filter(e => e !== '')) { From 8bf037c9da7352ee928887f6f769e395c363bf3d Mon Sep 17 00:00:00 2001 From: tesseractcat Date: Sat, 27 Aug 2022 01:11:59 -0400 Subject: [PATCH 5/5] Add img_callback support to KSampler --- ldm/models/diffusion/ksampler.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/ldm/models/diffusion/ksampler.py b/ldm/models/diffusion/ksampler.py index 1da81eee5a7..6d463fa6ba2 100644 --- a/ldm/models/diffusion/ksampler.py +++ b/ldm/models/diffusion/ksampler.py @@ -61,6 +61,9 @@ def sample( # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ... **kwargs, ): + def route_callback(k_callback_values): + if img_callback is not None: + img_callback(k_callback_values['x'], k_callback_values['i']) sigmas = self.model.get_sigmas(S) if x_T: @@ -78,7 +81,8 @@ def sample( } return ( K.sampling.__dict__[f'sample_{self.schedule}']( - model_wrap_cfg, x, sigmas, extra_args=extra_args + model_wrap_cfg, x, sigmas, extra_args=extra_args, + callback=route_callback ), None, )