diff --git a/ldm/dream/pngwriter.py b/ldm/dream/pngwriter.py
index 584d4479ba..8b1fc230ff 100644
--- a/ldm/dream/pngwriter.py
+++ b/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)
diff --git a/ldm/dream/server.py b/ldm/dream/server.py
index e64eefd2c9..3f1126a1ac 100644
--- a/ldm/dream/server.py
+++ b/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:
diff --git a/static/dream_web/index.css b/static/dream_web/index.css
index 8b0f4785f7..beb20e5a74 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 459e78b035..cc478d0ae0 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 @@