Skip to content

Commit 0173d3a

Browse files
bakkotbakkot
committed
stream images
1 parent e7658b9 commit 0173d3a

4 files changed

Lines changed: 59 additions & 17 deletions

File tree

‎ldm/dream/server.py‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,8 @@ def do_POST(self):
6464
upscale_level = post_data['upscale_level']
6565
upscale_strength = post_data['upscale_strength']
6666
upscale = [int(upscale_level),float(upscale_strength)] if upscale_level != '' else None
67-
seed = None if int(post_data['seed']) == -1 else int(post_data['seed'])
67+
progress_images = 'progress_images' in post_data
68+
seed = self.model.seed if int(post_data['seed']) == -1 else int(post_data['seed'])
6869

6970
print(f"Request to generate with prompt: {prompt}")
7071
# In order to handle upscaled images, the PngWriter needs to maintain state
@@ -116,9 +117,20 @@ def image_done(image, seed, upscaled=False):
116117
{'event':action,'processed_file_cnt':f'{x}/{iterations}'}
117118
) + '\n',"utf-8"))
118119

119-
def image_progress(image, step):
120+
# TODO: refactor PngWriter:
121+
# it doesn't need to know if batch_size > 1, just if this is _part of a batch_
122+
step_writer = PngWriter('./outputs/intermediates/', prompt, 2)
123+
def image_progress(sample, step):
124+
url = None
125+
# since rendering images is moderately expensive, only render every 5th image
126+
# and don't bother with the last one, since it'll render anyway
127+
if progress_images and step % 5 == 0 and step < steps - 1:
128+
images = self.model._samples_to_images(sample)
129+
image = images[0]
130+
step_writer.write_image(image, seed) # TODO PngWriter to return path
131+
url = step_writer.filepath
120132
self.wfile.write(bytes(json.dumps(
121-
{'event':'step', 'step':step}
133+
{'event':'step', 'step':step, 'url': url}
122134
) + '\n',"utf-8"))
123135

124136
if initimg is None:

‎static/dream_web/index.css‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ fieldset {
4646
margin: auto;
4747
padding-top: 10px;
4848
}
49-
img {
49+
#results img {
5050
cursor: pointer;
5151
height: 30vh;
5252
border-radius: 5px;
@@ -67,3 +67,10 @@ hr {
6767
label {
6868
white-space: nowrap;
6969
}
70+
#progress-section {
71+
display: none;
72+
}
73+
#progress-image {
74+
width: 30vh;
75+
height: 30vh;
76+
}

‎static/dream_web/index.html‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,9 @@ <h2 id="header">Stable Diffusion Dream Server</h2>
6767
<button type="button" id="reset-seed">&olarr;</button>
6868
<span>&bull;</span>
6969
<button type="button" id="reset-all">Reset to Defaults</button>
70+
<br>
71+
<label for="progress_images">Display in-progress images (slows down generation):</label>
72+
<input type="checkbox" name="progress_images" id="progress_images">
7073
<div id="gfpgan">
7174
<p><em>The options below require the GFPGAN and ESRGAN packages to be installed</em></p>
7275
<label title="Strength of the gfpgan (face fixing) algorithm." for="gfpgan_strength">GPFGAN Strength:</label>
@@ -83,10 +86,13 @@ <h2 id="header">Stable Diffusion Dream Server</h2>
8386
</fieldset>
8487
</form>
8588
<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>
86-
<br>
87-
<progress id="progress" value="0" max="1"></progress>
88-
<div id="scaling-inprocess-message">
89-
<i><span>Postprocessing...</span><span id="processing_cnt">1/3</span></i>
89+
<div id="progress-section">
90+
<progress id="progress-bar" value="0" max="1"></progress>
91+
<br>
92+
<img id="progress-image" src='data:image/svg+xml,<svg xmlns="http://www.w3.org/2000/svg"/>'></img>
93+
<div id="scaling-inprocess-message">
94+
<i><span>Postprocessing...</span><span id="processing_cnt">1/3</span></i>
95+
</div>
9096
</div>
9197
</div>
9298
<div id="results">

‎static/dream_web/index.js‎

Lines changed: 26 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -53,14 +53,22 @@ function clearFields(form) {
5353
form.prompt.value = prompt;
5454
}
5555

56+
const BLANK_IMAGE_URL = 'data:image/svg+xml,<svg xmlns="http://www.w3.org/2000/svg"/>';
5657
async function generateSubmit(form) {
5758
const prompt = document.querySelector("#prompt").value;
5859

5960
// Convert file data to base64
6061
let formData = Object.fromEntries(new FormData(form));
6162
formData.initimg = formData.initimg.name !== '' ? await toBase64(formData.initimg) : null;
6263

63-
document.querySelector('progress').setAttribute('max', formData.steps);
64+
let progressSectionEle = document.querySelector('#progress-section');
65+
progressSectionEle.style.display = 'initial';
66+
let progressEle = document.querySelector('#progress-bar');
67+
progressEle.setAttribute('max', formData.steps);
68+
let progressImageEle = document.querySelector('#progress-image');
69+
progressImageEle.src = BLANK_IMAGE_URL;
70+
71+
progressImageEle.style.display = {}.hasOwnProperty.call(formData, 'progress_images') ? 'initial': 'none';
6472

6573
// Post as JSON, using Fetch streaming to get results
6674
fetch(form.action, {
@@ -73,22 +81,31 @@ async function generateSubmit(form) {
7381
while (true) {
7482
let {value, done} = await reader.read();
7583
value = new TextDecoder().decode(value);
76-
if (done) break;
84+
if (done) {
85+
progressSectionEle.style.display = 'none';
86+
break;
87+
}
7788

7889
for (let event of value.split('\n').filter(e => e !== '')) {
7990
const data = JSON.parse(event);
8091

8192
if (data.event == 'result') {
8293
noOutputs = false;
8394
document.querySelector("#no-results-message")?.remove();
84-
appendOutput(data.files[0],data.files[1],data.config)
85-
} else if (data.event == 'upscaling-started') {
86-
document.getElementById("processing_cnt").textContent=data.processed_file_cnt;
87-
document.getElementById("scaling-inprocess-message").style.display = "block";
88-
} else if (data.event == 'upscaling-done') {
89-
document.getElementById("scaling-inprocess-message").style.display = "none";
95+
appendOutput(data.files[0],data.files[1],data.config);
96+
progressEle.setAttribute('value', 0);
97+
progressEle.setAttribute('max', formData.steps);
98+
progressImageEle.src = BLANK_IMAGE_URL;
99+
} else if (data.event == 'upscaling-started') {
100+
document.getElementById("processing_cnt").textContent=data.processed_file_cnt;
101+
document.getElementById("scaling-inprocess-message").style.display = "block";
102+
} else if (data.event == 'upscaling-done') {
103+
document.getElementById("scaling-inprocess-message").style.display = "none";
90104
} else if (data.event == 'step') {
91-
document.querySelector('progress').setAttribute('value', data.step.toString());
105+
progressEle.setAttribute('value', data.step);
106+
if (data.url) {
107+
progressImageEle.src = data.url;
108+
}
92109
}
93110
}
94111
}

0 commit comments

Comments
 (0)