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

This commit is contained in:
Lincoln Stein 2022-08-30 09:17:07 -04:00
commit 186d0f9d10
5 changed files with 60 additions and 17 deletions

View File

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

View File

@ -64,7 +64,8 @@ class DreamServer(BaseHTTPRequestHandler):
upscale_level = post_data['upscale_level'] upscale_level = post_data['upscale_level']
upscale_strength = post_data['upscale_strength'] upscale_strength = post_data['upscale_strength']
upscale = [int(upscale_level),float(upscale_strength)] if upscale_level != '' else None 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}") print(f"Request to generate with prompt: {prompt}")
# In order to handle upscaled images, the PngWriter needs to maintain state # 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}'} {'event':action,'processed_file_cnt':f'{x}/{iterations}'}
) + '\n',"utf-8")) ) + '\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( self.wfile.write(bytes(json.dumps(
{'event':'step', 'step':step} {'event':'step', 'step':step, 'url': url}
) + '\n',"utf-8")) ) + '\n',"utf-8"))
if initimg is None: if initimg is None:

View File

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

View File

@ -67,6 +67,9 @@
<button type="button" id="reset-seed">&olarr;</button> <button type="button" id="reset-seed">&olarr;</button>
<span>&bull;</span> <span>&bull;</span>
<button type="button" id="reset-all">Reset to Defaults</button> <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"> <div id="gfpgan">
<p><em>The options below require the GFPGAN and ESRGAN packages to be installed</em></p> <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> <label title="Strength of the gfpgan (face fixing) algorithm." for="gfpgan_strength">GPFGAN Strength:</label>
@ -83,10 +86,13 @@
</fieldset> </fieldset>
</form> </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="about">For news and support for this web service, visit our <a href="http://github.com/lstein/stable-diffusion">GitHub site</a></div>
<br> <div id="progress-section">
<progress id="progress" value="0" max="1"></progress> <progress id="progress-bar" value="0" max="1"></progress>
<div id="scaling-inprocess-message"> <br>
<i><span>Postprocessing...</span><span id="processing_cnt">1/3</span></i> <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> </div>
<div id="results"> <div id="results">

View File

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