mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 09:14:35 +02:00
refactor internal post loop
Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
+102
-99
@@ -250,6 +250,100 @@ def process_init(p: StableDiffusionProcessing):
|
||||
p.prompts, _ = extra_networks.parse_prompts(p.prompts)
|
||||
|
||||
|
||||
def process_samples(p: StableDiffusionProcessing, samples):
|
||||
out_images = []
|
||||
out_infotexts = []
|
||||
for i, sample in enumerate(samples):
|
||||
debug(f'Processing result: index={i+1}/{len(samples)}')
|
||||
p.batch_index = i
|
||||
info = create_infotext(p, p.prompts, p.seeds, p.subseeds, index=i)
|
||||
if isinstance(sample, Image.Image) or (isinstance(sample, list) and isinstance(sample[0], Image.Image)):
|
||||
image = sample
|
||||
sample = np.array(sample)
|
||||
else:
|
||||
sample = validate_sample(sample)
|
||||
image = Image.fromarray(sample)
|
||||
|
||||
if p.restore_faces:
|
||||
p.ops.append('restore')
|
||||
if not p.do_not_save_samples and shared.opts.save_images_before_detailer:
|
||||
images.save_image(Image.fromarray(sample), path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-restore")
|
||||
sample = face_restoration.restore_faces(sample, p)
|
||||
if sample is not None:
|
||||
image = Image.fromarray(sample)
|
||||
|
||||
if p.detailer_enabled:
|
||||
p.ops.append('detailer')
|
||||
if not p.do_not_save_samples and shared.opts.save_images_before_detailer:
|
||||
images.save_image(Image.fromarray(sample), path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-detailer")
|
||||
sample = detailer.detail(sample, p)
|
||||
if sample is not None:
|
||||
image = Image.fromarray(sample)
|
||||
|
||||
if p.color_corrections is not None and i < len(p.color_corrections):
|
||||
p.ops.append('color')
|
||||
if not p.do_not_save_samples and shared.opts.save_images_before_color_correction:
|
||||
orig = p.color_corrections
|
||||
p.color_corrections = None
|
||||
p.color_corrections = orig
|
||||
image_without_cc = apply_overlay(image, p.paste_to, i, p.overlay_images)
|
||||
images.save_image(image_without_cc, path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-color-correct")
|
||||
image = apply_color_correction(p.color_corrections[i], image)
|
||||
|
||||
if p.scripts is not None and isinstance(p.scripts, scripts_manager.ScriptRunner):
|
||||
pp = scripts_manager.PostprocessImageArgs(image)
|
||||
p.scripts.postprocess_image(p, pp)
|
||||
if pp.image is not None:
|
||||
image = pp.image
|
||||
|
||||
if shared.opts.mask_apply_overlay:
|
||||
image = apply_overlay(image, p.paste_to, i, p.overlay_images)
|
||||
|
||||
if hasattr(p, 'mask_for_overlay') and p.mask_for_overlay and any([shared.opts.save_mask, shared.opts.save_mask_composite, shared.opts.return_mask, shared.opts.return_mask_composite]):
|
||||
image_mask = p.mask_for_overlay.convert('RGB')
|
||||
image1 = image.convert('RGBA').convert('RGBa')
|
||||
image2 = Image.new('RGBa', image.size)
|
||||
mask = images.resize_image(3, p.mask_for_overlay, image.width, image.height).convert('L')
|
||||
image_mask_composite = Image.composite(image1, image2, mask).convert('RGBA')
|
||||
if shared.opts.save_mask:
|
||||
images.save_image(image_mask, p.outpath_samples, "", p.all_seeds[i], p.all_prompts[i], shared.opts.samples_format, info=info, p=p, suffix="-mask")
|
||||
if shared.opts.save_mask_composite:
|
||||
images.save_image(image_mask_composite, p.outpath_samples, "", p.all_seeds[i], p.all_prompts[i], shared.opts.samples_format, info=info, p=p, suffix="-mask-composite")
|
||||
if shared.opts.return_mask:
|
||||
out_infotexts.append(info)
|
||||
out_images.append(image_mask)
|
||||
if shared.opts.return_mask_composite:
|
||||
out_infotexts.append(info)
|
||||
out_images.append(image_mask_composite)
|
||||
|
||||
if shared.opts.include_mask:
|
||||
if shared.opts.mask_apply_overlay and p.overlay_images is not None and len(p.overlay_images) > 0:
|
||||
p.image_mask = create_binary_mask(p.overlay_images[0])
|
||||
p.image_mask = ImageOps.invert(p.image_mask)
|
||||
out_infotexts.append(info)
|
||||
out_images.append(p.image_mask)
|
||||
elif getattr(p, 'image_mask', None) is not None and isinstance(p.image_mask, Image.Image):
|
||||
if getattr(p, 'mask_for_detailer', None) is not None:
|
||||
out_infotexts.append(info)
|
||||
out_images.append(p.mask_for_detailer)
|
||||
else:
|
||||
out_infotexts.append(info)
|
||||
out_images.append(p.image_mask)
|
||||
|
||||
if p.selected_scale_tab_after == 1:
|
||||
p.width_after, p.height_after = int(image.width * p.scale_by_after), int(image.height * p.scale_by_after)
|
||||
if p.resize_mode_after != 0 and p.resize_name_after != 'None':
|
||||
image = images.resize_image(p.resize_mode_after, image, p.width_after, p.height_after, p.resize_name_after, context=p.resize_context_after)
|
||||
|
||||
if shared.opts.samples_save and not p.do_not_save_samples and p.outpath_samples is not None:
|
||||
images.save_image(image, p.outpath_samples, "", p.all_seeds[i], p.all_prompts[i], shared.opts.samples_format, info=info, p=p) # main save image
|
||||
|
||||
image.info["parameters"] = info
|
||||
out_infotexts.append(info)
|
||||
out_images.append(image)
|
||||
return out_images, out_infotexts
|
||||
|
||||
|
||||
def process_images_inner(p: StableDiffusionProcessing) -> Processed:
|
||||
"""this is the main loop that both txt2img and img2img use; it calls func_init once inside all the scopes and func_sample once per batch"""
|
||||
if type(p.prompt) == list:
|
||||
@@ -302,7 +396,9 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
|
||||
processed = p.scripts.process_images(p)
|
||||
if processed is not None:
|
||||
samples = processed.images
|
||||
infotexts += processed.infotexts
|
||||
for script_image, script_infotext in zip(processed.images, processed.infotexts):
|
||||
output_images.append(script_image)
|
||||
infotexts.append(script_infotext)
|
||||
if samples is None:
|
||||
from modules.processing_diffusers import process_diffusers
|
||||
samples = process_diffusers(p)
|
||||
@@ -320,92 +416,10 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
|
||||
p.scripts.postprocess_batch_list(p, batch_params, batch_number=n)
|
||||
samples = batch_params.images
|
||||
|
||||
for i, sample in enumerate(samples):
|
||||
debug(f'Processing result: index={i+1}/{len(samples)} iteration={n+1}/{p.n_iter}')
|
||||
p.batch_index = i
|
||||
if isinstance(sample, Image.Image) or (isinstance(sample, list) and isinstance(sample[0], Image.Image)):
|
||||
image = sample
|
||||
sample = np.array(sample)
|
||||
else:
|
||||
sample = validate_sample(sample)
|
||||
image = Image.fromarray(sample)
|
||||
if p.restore_faces:
|
||||
p.ops.append('restore')
|
||||
if not p.do_not_save_samples and shared.opts.save_images_before_detailer:
|
||||
info = create_infotext(p, p.prompts, p.seeds, p.subseeds, index=i)
|
||||
images.save_image(Image.fromarray(sample), path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-restore")
|
||||
sample = face_restoration.restore_faces(sample, p)
|
||||
if sample is not None:
|
||||
image = Image.fromarray(sample)
|
||||
if p.detailer_enabled:
|
||||
p.ops.append('detailer')
|
||||
if not p.do_not_save_samples and shared.opts.save_images_before_detailer:
|
||||
info = create_infotext(p, p.prompts, p.seeds, p.subseeds, index=i)
|
||||
images.save_image(Image.fromarray(sample), path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-detailer")
|
||||
sample = detailer.detail(sample, p)
|
||||
if sample is not None:
|
||||
image = Image.fromarray(sample)
|
||||
if p.color_corrections is not None and i < len(p.color_corrections):
|
||||
p.ops.append('color')
|
||||
if not p.do_not_save_samples and shared.opts.save_images_before_color_correction:
|
||||
orig = p.color_corrections
|
||||
p.color_corrections = None
|
||||
p.color_corrections = orig
|
||||
image_without_cc = apply_overlay(image, p.paste_to, i, p.overlay_images)
|
||||
info = create_infotext(p, p.prompts, p.seeds, p.subseeds, index=i)
|
||||
images.save_image(image_without_cc, path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-color-correct")
|
||||
image = apply_color_correction(p.color_corrections[i], image)
|
||||
if p.scripts is not None and isinstance(p.scripts, scripts_manager.ScriptRunner):
|
||||
pp = scripts_manager.PostprocessImageArgs(image)
|
||||
p.scripts.postprocess_image(p, pp)
|
||||
if pp.image is not None:
|
||||
image = pp.image
|
||||
if shared.opts.mask_apply_overlay:
|
||||
image = apply_overlay(image, p.paste_to, i, p.overlay_images)
|
||||
|
||||
info = create_infotext(p, p.prompts, p.seeds, p.subseeds, index=i, all_negative_prompts=p.negative_prompts)
|
||||
infotexts.append(info)
|
||||
if isinstance(image, list):
|
||||
for img in image:
|
||||
img.info["parameters"] = info
|
||||
output_images = image
|
||||
else:
|
||||
image.info["parameters"] = info
|
||||
output_images.append(image)
|
||||
devices.torch_gc()
|
||||
del samples
|
||||
|
||||
for i, image in enumerate(output_images):
|
||||
is_grid = len(output_images) == p.batch_size * p.n_iter + 1 and i == 0
|
||||
# resize after
|
||||
if p.selected_scale_tab_after == 1:
|
||||
p.width_after, p.height_after = int(image.width * p.scale_by_after), int(image.height * p.scale_by_after)
|
||||
if p.resize_mode_after != 0 and p.resize_name_after != 'None' and not is_grid:
|
||||
image = images.resize_image(p.resize_mode_after, image, p.width_after, p.height_after, p.resize_name_after, context=p.resize_context_after)
|
||||
|
||||
# save images
|
||||
if shared.opts.samples_save and not p.do_not_save_samples and p.outpath_samples is not None:
|
||||
info = create_infotext(p, p.prompts, p.seeds, p.subseeds, index=i)
|
||||
if isinstance(image, list):
|
||||
for img in image:
|
||||
images.save_image(img, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=info, p=p) # main save image
|
||||
else:
|
||||
images.save_image(image, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=info, p=p) # main save image
|
||||
|
||||
if hasattr(p, 'mask_for_overlay') and p.mask_for_overlay and any([shared.opts.save_mask, shared.opts.save_mask_composite, shared.opts.return_mask, shared.opts.return_mask_composite]):
|
||||
image_mask = p.mask_for_overlay.convert('RGB')
|
||||
image1 = image.convert('RGBA').convert('RGBa')
|
||||
image2 = Image.new('RGBa', image.size)
|
||||
mask = images.resize_image(3, p.mask_for_overlay, image.width, image.height).convert('L')
|
||||
image_mask_composite = Image.composite(image1, image2, mask).convert('RGBA')
|
||||
if shared.opts.save_mask:
|
||||
images.save_image(image_mask, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=info, p=p, suffix="-mask")
|
||||
if shared.opts.save_mask_composite:
|
||||
images.save_image(image_mask_composite, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=info, p=p, suffix="-mask-composite")
|
||||
if shared.opts.return_mask:
|
||||
output_images.append(image_mask)
|
||||
if shared.opts.return_mask_composite:
|
||||
output_images.append(image_mask_composite)
|
||||
batch_images, batch_infotexts = process_samples(p, samples)
|
||||
for batch_image, batch_infotext in zip(batch_images, batch_infotexts):
|
||||
output_images.append(batch_image)
|
||||
infotexts.append(batch_infotext)
|
||||
|
||||
if shared.cmd_opts.lowvram:
|
||||
devices.torch_gc(force=True, reason='lowvram')
|
||||
@@ -431,18 +445,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
|
||||
output_images.insert(0, grid)
|
||||
index_of_first_image = 1
|
||||
if shared.opts.grid_save:
|
||||
images.save_image(grid, p.outpath_grids, "", p.all_seeds[0], p.all_prompts[0], shared.opts.grid_format, info=grid_info, p=p, grid=True, suffix="-grid") # main save grid
|
||||
|
||||
if shared.opts.include_mask:
|
||||
if shared.opts.mask_apply_overlay and p.overlay_images is not None and len(p.overlay_images) > 0:
|
||||
p.image_mask = create_binary_mask(p.overlay_images[0])
|
||||
p.image_mask = ImageOps.invert(p.image_mask)
|
||||
output_images.append(p.image_mask)
|
||||
elif getattr(p, 'image_mask', None) is not None and isinstance(p.image_mask, Image.Image):
|
||||
if getattr(p, 'mask_for_detailer', None) is not None:
|
||||
output_images.append(p.mask_for_detailer)
|
||||
else:
|
||||
output_images.append(p.image_mask)
|
||||
images.save_image(grid, p.outpath_grids, "", p.all_seeds[0], p.all_prompts[0], shared.opts.grid_format, info=grid_info, p=p, grid=True) # main save grid
|
||||
|
||||
processed = Processed(
|
||||
p,
|
||||
|
||||
Reference in New Issue
Block a user