Skip to content

Add --progress-images to CLI #623

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
kevinschaul wants to merge 1 commit into from
Closed

Add --progress-images to CLI #623

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 21 additions & 2 deletions scripts/dream.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -260,6 +260,20 @@ def main_loop(t2i, outdir, prompt_as_dir, parser, infile):
results = [] # list of filename, prompt pairs
grid_images = dict() # seed -> Image, only used if `do_grid`

step_writer = PngWriter(os.path.join(current_outdir, "intermediates"))
step_index = 1
def image_progress(sample, step):
# since rendering images is moderately expensive, only render every 5th image
# and don't bother with the last one, since it'll render anyway
nonlocal step_index
if opt.progress_images and step % 5 == 0 and step < opt.steps - 1:
image = t2i.sample_to_image(sample)
step_index_padded = str(step_index).rjust(len(str(opt.steps)), '0')
name = f'{prefix}.{opt.seed}.{step_index_padded}.png'
metadata = f'{opt.prompt} -S{opt.seed} [intermediate]'
step_writer.save_image_and_prompt_to_png(image, metadata, name)
step_index += 1

def image_writer(image, seed, upscaled=False):
path = None
if do_grid:
Expand Down Expand Up @@ -296,7 +310,7 @@ def image_writer(image, seed, upscaled=False):
results.append([path, metadata_prompt])
last_results.append([path, seed])

t2i.prompt2image(image_callback=image_writer, **vars(opt))
t2i.prompt2image(step_callback=image_progress, image_callback=image_writer, **vars(opt))

if do_grid and len(grid_images) > 0:
grid_img = make_grid(list(grid_images.values()))
Expand Down Expand Up @@ -544,7 +558,7 @@ def create_cmd_parser():
description='Example: dream> a fantastic alien landscape -W1024 -H960 -s100 -n12'
)
parser.add_argument('prompt')
parser.add_argument('-s', '--steps', type=int, help='Number of steps')
parser.add_argument('-s', '--steps', type=int, default=50, help='Number of steps')
parser.add_argument(
'-S',
'--seed',
Expand Down Expand Up @@ -678,6 +692,11 @@ def create_cmd_parser():
type=str,
help='list of variations to apply, in the format `seed:weight,seed:weight,...'
)
parser.add_argument(
'--progress-images',
action='store_true',
help='Store in-progress images as the image is being rendered.',
)
return parser


Expand Down