Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 17:24:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

add callback

Browse Source
This commit is contained in:
ljleb
2023-08-02 15:26:04 -04:00
parent 246989129f
commit 88cc6f63e1
2 changed files with 61 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/processing.py
+27 -14
View File
@@ -444,8 +444,13 @@ def fix_seed(p):
p.subseed = get_fixed_seed(p.subseed)
def create_infotext(p: StableDiffusionProcessing, all_prompts, all_seeds, all_subseeds, comments=None, iteration=0, position_in_batch=0): # pylint: disable=unused-argument
index = position_in_batch + iteration * p.batch_size
def create_infotext(p: StableDiffusionProcessing, all_prompts, all_seeds, all_subseeds, comments=None, iteration=0, position_in_batch=0, index=None, all_negative_prompts=None): # pylint: disable=unused-argument
if index is None:
index = position_in_batch + iteration * p.batch_size
if all_negative_prompts is None:
all_negative_prompts = p.all_negative_prompts
generation_params = {
"Steps": p.steps,
"Seed": all_seeds[index],
@@ -487,7 +492,7 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts, all_seeds, all_su
generation_params['Token merging ratio hr'] = token_merging_ratio_hr if token_merging_ratio_hr != 0 else None
generation_params.update(p.extra_generation_params)
generation_params_text = ", ".join([k if k == v else f'{k}: {generation_parameters_copypaste.quote(v)}' for k, v in generation_params.items() if v is not None])
negative_prompt_text = f"\nNegative prompt: {p.all_negative_prompts[index]}" if p.all_negative_prompts[index] else ""
negative_prompt_text = f"\nNegative prompt: {all_negative_prompts[index]}" if all_negative_prompts[index] else ""
return f"{all_prompts[index]}{negative_prompt_text}\n{generation_params_text}".strip()
@@ -599,9 +604,6 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
else:
p.all_subseeds = [int(subseed) + x for x in range(len(p.all_prompts))]
def infotext(iteration=0, position_in_batch=0):
return create_infotext(p, p.all_prompts, p.all_seeds, p.all_subseeds, comments, iteration, position_in_batch)
if os.path.exists(shared.opts.embeddings_dir) and not p.do_not_reload_embeddings:
model_hijack.embedding_db.load_textual_inversion_embeddings()
if p.scripts is not None:
@@ -709,6 +711,17 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
devices.torch_gc()
if p.scripts is not None:
p.scripts.postprocess_batch(p, x_samples_ddim, batch_number=n)
p.seeds = seeds
p.subseeds = subseeds
if p.scripts is not None:
p.prompts = p.all_prompts[n * p.batch_size:(n + 1) * p.batch_size]
p.negative_prompts = p.all_negative_prompts[n * p.batch_size:(n + 1) * p.batch_size]
batch_params = scripts.PostprocessBatchListArgs(list(x_samples_ddim))
p.scripts.postprocess_batch_list(p, batch_params, batch_number=n)
x_samples_ddim = batch_params.images
def infotext(index=0):
return create_infotext(p, p.prompts, p.seeds, p.subseeds, index=index, all_negative_prompts=p.negative_prompts)
for i, x_sample in enumerate(x_samples_ddim):
p.batch_index = i
@@ -721,9 +734,9 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
if shared.opts.save and not p.do_not_save_samples and shared.opts.save_images_before_face_restoration:
orig = p.restore_faces
p.restore_faces = False
info=infotext(n, i)
info = infotext(i)
p.restore_faces = orig
images.save_image(Image.fromarray(x_sample), path=p.outpath_samples, basename="", seed=seeds[i], prompt=prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-face-restoration")
images.save_image(Image.fromarray(x_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-face-restoration")
p.ops.append('face')
x_sample = modules.face_restoration.restore_faces(x_sample)
image = Image.fromarray(x_sample)
@@ -735,16 +748,16 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
if shared.opts.save and not p.do_not_save_samples and shared.opts.save_images_before_color_correction:
orig = p.color_corrections
p.color_corrections = None
info=infotext(n, i)
info = infotext(i)
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=seeds[i], prompt=prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-color-correction")
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-correction")
p.ops.append('color')
image = apply_color_correction(p.color_corrections[i], image)
image = apply_overlay(image, p.paste_to, i, p.overlay_images)
if shared.opts.samples_save and not p.do_not_save_samples:
images.save_image(image, p.outpath_samples, "", seeds[i], prompts[i], shared.opts.samples_format, info=infotext(n, i), p=p)
text = infotext(n, i)
images.save_image(image, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=infotext(i), p=p)
text = infotext(i)
infotexts.append(text)
image.info["parameters"] = text
output_images.append(image)
@@ -752,9 +765,9 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
image_mask = p.mask_for_overlay.convert('RGB')
image_mask_composite = Image.composite(image.convert('RGBA').convert('RGBa'), Image.new('RGBa', image.size), images.resize_image(3, p.mask_for_overlay, image.width, image.height).convert('L')).convert('RGBA')
if shared.opts.save_mask:
images.save_image(image_mask, p.outpath_samples, "", seeds[i], prompts[i], shared.opts.samples_format, info=infotext(n, i), p=p, suffix="-mask")
images.save_image(image_mask, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=infotext(i), p=p, suffix="-mask")
if shared.opts.save_mask_composite:
images.save_image(image_mask_composite, p.outpath_samples, "", seeds[i], prompts[i], shared.opts.samples_format, info=infotext(n, i), p=p, suffix="-mask-composite")
images.save_image(image_mask_composite, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=infotext(i), p=p, suffix="-mask-composite")
if shared.opts.return_mask:
output_images.append(image_mask)
if shared.opts.return_mask_composite:
modules/scripts.py
+34
View File
@@ -18,6 +18,11 @@ class PostprocessImageArgs:
self.image = image
class PostprocessBatchListArgs:
def __init__(self, images):
self.images = images
class Script:
name = None
filename = None
@@ -108,6 +113,23 @@ class Script:
"""
pass # pylint: disable=unnecessary-pass
def postprocess_batch_list(self, p, pp: PostprocessBatchListArgs, *args, **kwargs):
"""
Same as postprocess_batch(), but receives batch images as a list of 3D tensors instead of a 4D tensor.
This is useful when you want to update the entire batch instead of individual images.
You can modify the postprocessing object (pp) to update the images in the batch, remove images, add images, etc.
If the number of images is different from the batch size when returning,
then the script has the responsibility to also update the following attributes in the processing object (p):
- p.prompts
- p.negative_prompts
- p.seeds
- p.subseeds
**kwargs will have same items as process_batch, and also:
- batch_number - index of current batch, from 0 to number of batches-1
"""
pass # pylint: disable=unnecessary-pass
def postprocess(self, p, processed, *args):
"""
This function is called after processing ends for AlwaysVisible scripts.
@@ -457,6 +479,18 @@ class ScriptRunner:
errors.display(e, f'Running script before postprocess batch: {script.filename}')
log.debug(f'Script postprocess-batch: {s}')
def postprocess_batch_list(self, p, pp: PostprocessBatchListArgs, **kwargs):
s = []
for script in self.alwayson_scripts:
try:
t0 = time.time()
args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
script.postprocess_batch_list(p, pp, *args, **kwargs)
s.append(f'{script.title()}:{round(time.time()-t0, 2)}s')
except Exception as e:
errors.display(e, f'Running script before postprocess batch list: {script.filename}')
log.debug(f'Script postprocess-batch-list: {s}')
def postprocess_image(self, p, pp: PostprocessImageArgs):
s = []
for script in self.alwayson_scripts: