Skip to content

Commit

Permalink
rework saving incomplete images
Browse files Browse the repository at this point in the history
  • Loading branch information
AUTOMATIC1111 committed Aug 6, 2023
1 parent 2844d95 commit d86d12e
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 10 deletions.
18 changes: 9 additions & 9 deletions modules/processing.py
Original file line number Diff line number Diff line change
Expand Up @@ -103,10 +103,6 @@ def txt2img_image_conditioning(sd_model, x, width, height):
return x.new_zeros(x.shape[0], 5, 1, 1, dtype=x.dtype, device=x.device)


def save_images_if_interrupted():
return not (opts.dont_save_interrupted_images and (state.interrupted or state.skipped))


class StableDiffusionProcessing:
"""
The first set of paramaters: sd_models -> do_not_reload_embeddings represent the minimum required to create a StableDiffusionProcessing
Expand Down Expand Up @@ -372,6 +368,10 @@ def setup_conds(self):
def parse_extra_network_prompts(self):
self.prompts, self.extra_network_data = extra_networks.parse_prompts(self.prompts)

def save_samples(self) -> bool:
"""Returns whether generated images need to be written to disk"""
return opts.samples_save and not self.do_not_save_samples and (opts.save_incomplete_images or not state.interrupted and not state.skipped)


class Processed:
def __init__(self, p: StableDiffusionProcessing, images_list, seed=-1, info="", subseed=None, all_prompts=None, all_negative_prompts=None, all_seeds=None, all_subseeds=None, index_of_first_image=0, infotexts=None, comments=""):
Expand Down Expand Up @@ -827,7 +827,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
def infotext(index=0, use_main_prompt=False):
return create_infotext(p, p.prompts, p.seeds, p.subseeds, use_main_prompt=use_main_prompt, index=index, all_negative_prompts=p.negative_prompts)

save_images_if_interrupt = save_images_if_interrupted()
save_samples = p.save_samples()

for i, x_sample in enumerate(x_samples_ddim):
p.batch_index = i
Expand All @@ -836,7 +836,7 @@ def infotext(index=0, use_main_prompt=False):
x_sample = x_sample.astype(np.uint8)

if p.restore_faces:
if opts.save and not p.do_not_save_samples and opts.save_images_before_face_restoration and save_images_if_interrupt:
if save_samples and opts.save_images_before_face_restoration:
images.save_image(Image.fromarray(x_sample), p.outpath_samples, "", p.seeds[i], p.prompts[i], opts.samples_format, info=infotext(i), p=p, suffix="-before-face-restoration")

devices.torch_gc()
Expand All @@ -851,14 +851,14 @@ def infotext(index=0, use_main_prompt=False):
p.scripts.postprocess_image(p, pp)
image = pp.image
if p.color_corrections is not None and i < len(p.color_corrections):
if opts.save and not p.do_not_save_samples and opts.save_images_before_color_correction and save_images_if_interrupt:
if save_samples and opts.save_images_before_color_correction:
image_without_cc = apply_overlay(image, p.paste_to, i, p.overlay_images)
images.save_image(image_without_cc, p.outpath_samples, "", p.seeds[i], p.prompts[i], opts.samples_format, info=infotext(i), p=p, suffix="-before-color-correction")
image = apply_color_correction(p.color_corrections[i], image)

image = apply_overlay(image, p.paste_to, i, p.overlay_images)

if opts.samples_save and not p.do_not_save_samples and save_images_if_interrupt:
if save_samples:
images.save_image(image, p.outpath_samples, "", p.seeds[i], p.prompts[i], opts.samples_format, info=infotext(i), p=p)

text = infotext(i)
Expand Down Expand Up @@ -1096,7 +1096,7 @@ def sample_hr_pass(self, samples, decoded_samples, seeds, subseeds, subseed_stre
def save_intermediate(image, index):
"""saves image before applying hires fix, if enabled in options; takes as an argument either an image or batch with latent space images"""

if not opts.save or self.do_not_save_samples or not opts.save_images_before_highres_fix or not save_images_if_interrupted():
if not self.save_samples() or not opts.save_images_before_highres_fix:
return

if not isinstance(image, Image.Image):
Expand Down
2 changes: 1 addition & 1 deletion modules/shared.py
Original file line number Diff line number Diff line change
Expand Up @@ -385,7 +385,7 @@ def list_samplers():
"temp_dir": OptionInfo("", "Directory for temporary images; leave empty for default"),
"clean_temp_dir_at_start": OptionInfo(False, "Cleanup non-default temporary directory when starting webui"),

"dont_save_interrupted_images": OptionInfo(False, "Don't save incomplete images").info("Don't save images that has been interrupted in mid-generation, they will still show up in webui output."),
"save_incomplete_images": OptionInfo(False, "Save incomplete images").info("save images that has been interrupted in mid-generation; even if not saved, they will still show up in webui output."),
}))

options_templates.update(options_section(('saving-paths', "Paths for saving"), {
Expand Down

0 comments on commit d86d12e

Please sign in to comment.