From 252cbe6b568af7a88dc5930f0570556bb4fcdd39 Mon Sep 17 00:00:00 2001
From: blessedcoolant <54517381+blessedcoolant@users.noreply.github.com>
Date: Mon, 19 Sep 2022 20:04:26 +1200
Subject: [PATCH 1/2] Update CONTRIBUTORS.md
---
docs/other/CONTRIBUTORS.md | 1 +
1 file changed, 1 insertion(+)
diff --git a/docs/other/CONTRIBUTORS.md b/docs/other/CONTRIBUTORS.md
index 7eba44dbad5..424b06036ae 100644
--- a/docs/other/CONTRIBUTORS.md
+++ b/docs/other/CONTRIBUTORS.md
@@ -54,6 +54,7 @@ We thank them for all of their time and hard work.
- [Matthias Wild](https://github.com/mauwii)
- [Kyle Schouviller](https://github.com/kyle0654)
- [rabidcopy](https://github.com/rabidcopy)
+- [Kevin Schaul](https://github.com/kevinschaul)
## **Original CompVis Authors:**
From 79179d61bc4c463c46be01b376a780adfdb26fe9 Mon Sep 17 00:00:00 2001
From: blessedcoolant <54517381+blessedcoolant@users.noreply.github.com>
Date: Mon, 19 Sep 2022 11:25:33 +1200
Subject: [PATCH 2/2] Show Progress / Save Progress / Make Video CLI
Implementation
---
docs/features/OTHER.md | 54 +++++++++++++++++++++++++-
ldm/dream/args.py | 13 +++++++
ldm/util.py | 19 +++++++++
scripts/dream.py | 87 +++++++++++++++++++++++++++++++++++++++++-
4 files changed, 170 insertions(+), 3 deletions(-)
diff --git a/docs/features/OTHER.md b/docs/features/OTHER.md
index 019ed35b709..c4bec2f50f2 100644
--- a/docs/features/OTHER.md
+++ b/docs/features/OTHER.md
@@ -6,8 +6,7 @@ title: Others
Stable Diffusion AI Notebook:
-
Open and follow instructions to use an isolated environment running
Dream.
@@ -28,6 +27,57 @@ dream> "pond garden with lotus by claude monet" --seamless -s100 -n4
---
+## **Show Progress**
+
+Provides a visual preview of the image generation process.
+
+`-show_progress `
+
+- `step_count`: The number of steps between each progress update. Default: `5`.
+
+- `duration`: The duration (in seconds) for how long you want the final image to be displayed before
+ the preview closes automatically. Default: `2`.
+
+- Set step_count to the same value as your steps to get a preview only when the image is fully
+ generated.
+
+- Enter duration: `0` to keep it open forever until user presses a key. Note that this will block
+ the code from running further until user input.
+
+If you have post processing options, the preview will close after image generation and reopen again
+with the updated changes.
+
+---
+
+## **Save Progress & Make Video**
+
+Allows you to save the intermediate steps during the image generation process and make a video out
+of it.
+
+`-save_progress `
+
+- `step_count`: The number of steps between each intermediate image saved. When no value is given,
+ it defaults to `5`.
+- `video_options`: Allows you to generate a video from the intermediate images. Takes two options:
+ `v` (Video) or `vo` (Video Only)
+
+### **Usage**
+
+`-save_progress`: Saves intermediate frames every 5 steps. No video generation.
+
+`-save_progress 3`: Saves intermediate frames every 3 steps. No video generation.
+
+`-save_progress 3 v`: Saves intermediate frames every 3 steps. Also generates a video from the
+frames at the end.
+
+`-save_progress 3 vo`: Does not save intermediate frames but generates a video of the process every
+3 steps.
+
+`-show_progress 3 -save_progress 3 vo`: Shows a preview of the generation process updating every 3
+seconds while also saving a video of the same.
+
+---
+
## **Shortcuts: Reusing Seeds**
Since it is so common to reuse seeds while refining a prompt, there is now a shortcut as of version
diff --git a/ldm/dream/args.py b/ldm/dream/args.py
index db6d9636455..7c60ac1b042 100644
--- a/ldm/dream/args.py
+++ b/ldm/dream/args.py
@@ -485,6 +485,19 @@ def _create_dream_cmd_parser(self):
type=str,
help='Directory to save generated images and a log of prompts and seeds',
)
+ render_group.add_argument(
+ '-save_progress',
+ '--save_progress',
+ nargs='*',
+ help='Store in-progress images as the image is being rendered. Takes two values : int : v (video) or vo (video only)'
+ )
+ render_group.add_argument(
+ '-show_progress',
+ '--show_progress',
+ nargs='*',
+ type=float,
+ help='Show image generation progress. Takes two values. and '
+ )
img2img_group.add_argument(
'-I',
'--init_img',
diff --git a/ldm/util.py b/ldm/util.py
index d1379cae2be..45a5211af01 100644
--- a/ldm/util.py
+++ b/ldm/util.py
@@ -212,3 +212,22 @@ def parallel_data_prefetch(
return out
else:
return gather_res
+
+def make_video(images, video_location):
+ import cv2
+ width, height = images[0].size
+
+ video = cv2.VideoWriter(video_location, cv2.VideoWriter_fourcc(*'mp4v'), 20.0, (width, height))
+
+ #draw stuff that goes on every frame here
+ for image in images:
+ image = cv2.cvtColor(np.array(image), cv2.COLOR_BGR2RGB)
+ video.write(image)
+
+ video.release()
+
+def show_progress(image):
+ import cv2
+ image = cv2.cvtColor(np.array(image), cv2.COLOR_BGR2RGB)
+ cv2.imshow('Preview', image)
+ cv2.waitKey(1)
\ No newline at end of file
diff --git a/scripts/dream.py b/scripts/dream.py
index 857b5637aa6..de166533e3e 100755
--- a/scripts/dream.py
+++ b/scripts/dream.py
@@ -1,6 +1,7 @@
#!/usr/bin/env python3
# Copyright (c) 2022 Lincoln D. Stein (https://github.com/lstein)
+import cv2
import os
import re
import sys
@@ -12,6 +13,7 @@
from ldm.dream.pngwriter import PngWriter
from ldm.dream.server import DreamServer, ThreadingDreamServer
from ldm.dream.image_util import make_grid
+from ldm.util import make_video, show_progress
from omegaconf import OmegaConf
# Placeholder to be replaced with proper class that tracks the
@@ -220,10 +222,70 @@ def main_loop(gen, opt, infile):
prior_variations = opt.with_variations or []
first_seed = opt.seed
+ if opt.save_progress is not None or opt.show_progress is not None:
+ step_index = 1
+ step_count = 1
+
+ if opt.show_progress is not None:
+ if len(opt.show_progress) < 2:
+ if len(opt.show_progress) == 0:
+ opt.show_progress.extend([5, 2])
+ elif len(opt.show_progress) == 1:
+ opt.show_progress.append(2)
+ step_count = int(opt.show_progress[0])
+
+ if opt.save_progress is not None:
+ if len(opt.save_progress) < 2:
+ if len(opt.save_progress) == 0:
+ opt.save_progress.extend([5, None])
+ elif len(opt.save_progress) == 1:
+ opt.save_progress.append(None)
+
+ step_count, progress_video_type = opt.save_progress
+
+ if progress_video_type is not None:
+ frames_for_video = []
+
+ if progress_video_type != 'vo':
+ step_writer = PngWriter(os.path.join(current_outdir, 'intermediates'))
+
+ def image_progress(sample, step):
+ nonlocal step_index
+ nonlocal step_count
+
+ if step_count == 0:
+ step_count = 5
+
+ if step % int(step_count) == 0 and step < opt.steps - 1:
+ image = gen.sample_to_image(sample)
+
+ if opt.save_progress is not None:
+ nonlocal progress_video_type
+ step_index_padded = str(step_index).rjust(len(str(opt.steps)), '0')
+
+ if progress_video_type != 'vo':
+ interim_seed = '.'
+ if opt.seed is not None:
+ interim_seed = f'.{opt.seed}.'
+ name = f'{prefix}{interim_seed}{step_index_padded}.png'
+ metadata = f'{opt.prompt} -S{interim_seed} [intermediate]'
+ step_writer.save_image_and_prompt_to_png(image, metadata, name)
+
+ if progress_video_type == 'v' or progress_video_type == 'vo':
+ frames_for_video.append(image)
+
+ if opt.show_progress is not None:
+ if step == 0 and int(step_count) == opt.steps:
+ return
+ show_progress(image)
+
+ step_index += 1
+
def image_writer(image, seed, upscaled=False):
path = None
nonlocal first_seed
- nonlocal prior_variations
+ nonlocal prior_variations
+
if opt.grid:
grid_images[seed] = image
else:
@@ -255,8 +317,25 @@ def image_writer(image, seed, upscaled=False):
results.append([path, formatted_dream_prompt])
last_results.append([path, seed])
+ if opt.save_progress is not None:
+ nonlocal progress_video_type
+ if progress_video_type == 'v' or progress_video_type == 'vo':
+ frames_for_video.append(image)
+ make_video(frames_for_video, os.path.join(current_outdir, f'{prefix}.{seed}.mp4'))
+ frames_for_video.clear()
+
+ if opt.show_progress is not None:
+ show_progress(image)
+ cv2.waitKey(1000)
+
catch_ctrl_c = infile is None # if running interactively, we catch keyboard interrupts
+
+ step_callback = None
+ if opt.save_progress is not None or opt.show_progress is not None:
+ step_callback = image_progress
+
gen.prompt2image(
+ step_callback=step_callback,
image_callback=image_writer,
catch_interrupts=catch_ctrl_c,
**vars(opt)
@@ -294,6 +373,12 @@ def image_writer(image, seed, upscaled=False):
print('Outputs:')
log_path = os.path.join(current_outdir, 'dream_log.txt')
write_log_message(results, log_path)
+ if opt.show_progress is not None:
+ if (int(opt.show_progress[1]) == 0):
+ print("Press any key on the preview window to continue ...")
+ cv2.waitKey(int(opt.show_progress[1]) * 1000)
+ cv2.destroyAllWindows()
+ cv2.waitKey(1) # possible fix for window not closing on Macs
print()
print('goodbye!')