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 In Colab
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/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:** 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!')