CLI: Show Progress / Save Progress / Generate Progress Video

docs/features/OTHER.md

@@ -6,8 +6,7 @@ title: Others
 
 Stable Diffusion AI Notebook: <a
 href="https://colab.research.google.com/github/lstein/stable-diffusion/blob/main/notebooks/Stable_Diffusion_AI_Notebook.ipynb"
-target="_parent">
-<img
+target="_parent"> <img
 src="https://colab.research.google.com/assets/colab-badge.svg"
 alt="Open In Colab"/></a> <br> Open and follow instructions to use an isolated environment running
 Dream.<br>
@@ -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: int> <duration: float>`
+
+- `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: int default=5> <video_options: v | vo default: None>`
+
+- `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

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:**
 

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 <step_count> : int <progress_video_type>: 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. <step_count: int> and <final_display_time: float>'
+        )
         img2img_group.add_argument(
             '-I',
             '--init_img',

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)

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!')