Skip to content

Commit

Permalink
make postprocess() extension
Browse files Browse the repository at this point in the history
  • Loading branch information
@wkpark
wkpark committed May 12, 2023
1 parent 9238cf8 commit 11fee42
Showing 1 changed file with 161 additions and 103 deletions.
264 changes: 161 additions & 103 deletions scripts/ddetailer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import gradio as gr
import shutil

from copy import copy
from modules import processing, images
from modules import scripts, script_callbacks, shared, devices, modelloader
from modules.processing import Processed, StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img
Expand Down Expand Up @@ -95,100 +96,132 @@ def gr_show(visible=True):
return {"visible": visible, "__type__": "update"}

class DetectionDetailerScript(scripts.Script):
def __init__(self):
super().__init__()

def title(self):
return "Detection Detailer"

def show(self, is_img2img):
return True
return scripts.AlwaysVisible

def ui(self, is_img2img):
import modules.ui

model_list = list_models(dd_models_path)
model_list.insert(0, "None")
if is_img2img:
info = gr.HTML("<p style=\"margin-bottom:0.75em\">Recommended settings: Use from inpaint tab, inpaint at full res ON, denoise <0.5</p>")
else:
info = gr.HTML("")
with gr.Group():
with gr.Accordion("Detection Detailer", open=False):
with gr.Row():
dd_model_a = gr.Dropdown(label="Primary detection model (A)", choices=model_list,value = "None", visible=True, type="value")

with gr.Row():
dd_conf_a = gr.Slider(label='Detection confidence threshold % (A)', minimum=0, maximum=100, step=1, value=30, visible=False)
dd_dilation_factor_a = gr.Slider(label='Dilation factor (A)', minimum=0, maximum=255, step=1, value=4, visible=False)
enabled = gr.Checkbox(label="Enable", value=False, visible=True)

with gr.Row():
dd_offset_x_a = gr.Slider(label='X offset (A)', minimum=-200, maximum=200, step=1, value=0, visible=False)
dd_offset_y_a = gr.Slider(label='Y offset (A)', minimum=-200, maximum=200, step=1, value=0, visible=False)

with gr.Row():
dd_preprocess_b = gr.Checkbox(label='Inpaint model B detections before model A runs', value=False, visible=False)
dd_bitwise_op = gr.Radio(label='Bitwise operation', choices=['None', 'A&B', 'A-B'], value="None", visible=False)

br = gr.HTML("<br>")
model_list = list_models(dd_models_path)
model_list.insert(0, "None")
if is_img2img:
info = gr.HTML("<p style=\"margin-bottom:0.75em\">Recommended settings: Use from inpaint tab, inpaint at full res ON, denoise <0.5</p>")
else:
info = gr.HTML("")
with gr.Group():
with gr.Row():
dd_model_a = gr.Dropdown(label="Primary detection model (A)", choices=model_list,value = "None", visible=True, type="value")

with gr.Row():
dd_conf_a = gr.Slider(label='Detection confidence threshold % (A)', minimum=0, maximum=100, step=1, value=30, visible=False)
dd_dilation_factor_a = gr.Slider(label='Dilation factor (A)', minimum=0, maximum=255, step=1, value=4, visible=False)

with gr.Row():
dd_offset_x_a = gr.Slider(label='X offset (A)', minimum=-200, maximum=200, step=1, value=0, visible=False)
dd_offset_y_a = gr.Slider(label='Y offset (A)', minimum=-200, maximum=200, step=1, value=0, visible=False)

with gr.Row():
dd_preprocess_b = gr.Checkbox(label='Inpaint model B detections before model A runs', value=False, visible=False)
dd_bitwise_op = gr.Radio(label='Bitwise operation', choices=['None', 'A&B', 'A-B'], value="None", visible=False)

br = gr.HTML("<br>")

with gr.Group():
with gr.Row():
dd_model_b = gr.Dropdown(label="Secondary detection model (B) (optional)", choices=model_list,value = "None", visible =False, type="value")

with gr.Row():
dd_conf_b = gr.Slider(label='Detection confidence threshold % (B)', minimum=0, maximum=100, step=1, value=30, visible=False)
dd_dilation_factor_b = gr.Slider(label='Dilation factor (B)', minimum=0, maximum=255, step=1, value=4, visible=False)

with gr.Row():
dd_offset_x_b = gr.Slider(label='X offset (B)', minimum=-200, maximum=200, step=1, value=0, visible=False)
dd_offset_y_b = gr.Slider(label='Y offset (B)', minimum=-200, maximum=200, step=1, value=0, visible=False)

with gr.Group():
with gr.Row():
dd_mask_blur = gr.Slider(label='Mask blur ', minimum=0, maximum=64, step=1, value=4, visible=(not is_img2img))
dd_denoising_strength = gr.Slider(label='Denoising strength (Inpaint)', minimum=0.0, maximum=1.0, step=0.01, value=0.4, visible=(not is_img2img))

with gr.Row():
dd_inpaint_full_res = gr.Checkbox(label='Inpaint at full resolution ', value=True, visible = (not is_img2img))
dd_inpaint_full_res_padding = gr.Slider(label='Inpaint at full resolution padding, pixels ', minimum=0, maximum=256, step=4, value=32, visible=(not is_img2img))

dd_model_a.change(
lambda modelname: {
dd_model_b:gr_show( modelname != "None" ),
dd_conf_a:gr_show( modelname != "None" ),
dd_dilation_factor_a:gr_show( modelname != "None"),
dd_offset_x_a:gr_show( modelname != "None" ),
dd_offset_y_a:gr_show( modelname != "None" )

},
inputs= [dd_model_a],
outputs =[dd_model_b, dd_conf_a, dd_dilation_factor_a, dd_offset_x_a, dd_offset_y_a]
)

dd_model_b.change(
lambda modelname: {
dd_preprocess_b:gr_show( modelname != "None" ),
dd_bitwise_op:gr_show( modelname != "None" ),
dd_conf_b:gr_show( modelname != "None" ),
dd_dilation_factor_b:gr_show( modelname != "None"),
dd_offset_x_b:gr_show( modelname != "None" ),
dd_offset_y_b:gr_show( modelname != "None" )
},
inputs= [dd_model_b],
outputs =[dd_preprocess_b, dd_bitwise_op, dd_conf_b, dd_dilation_factor_b, dd_offset_x_b, dd_offset_y_b]
)

return [info, enabled,
dd_model_a,
dd_conf_a, dd_dilation_factor_a,
dd_offset_x_a, dd_offset_y_a,
dd_preprocess_b, dd_bitwise_op,
br,
dd_model_b,
dd_conf_b, dd_dilation_factor_b,
dd_offset_x_b, dd_offset_y_b,
dd_mask_blur, dd_denoising_strength,
dd_inpaint_full_res, dd_inpaint_full_res_padding
]

def get_seed(self, p) -> tuple[int, int]:
i = p.iteration

if not p.all_seeds:
seed = p.seed
elif i < len(p.all_seeds):
seed = p.all_seeds[i]
else:
j = i % len(p.all_seeds)
seed = p.all_seeds[j]

with gr.Group():
with gr.Row():
dd_model_b = gr.Dropdown(label="Secondary detection model (B) (optional)", choices=model_list,value = "None", visible =False, type="value")
if not p.all_subseeds:
subseed = p.subseed
elif i < len(p.all_subseeds):
subseed = p.all_subseeds[i]
else:
j = i % len(p.all_subseeds)
subseed = p.all_subseeds[j]

with gr.Row():
dd_conf_b = gr.Slider(label='Detection confidence threshold % (B)', minimum=0, maximum=100, step=1, value=30, visible=False)
dd_dilation_factor_b = gr.Slider(label='Dilation factor (B)', minimum=0, maximum=255, step=1, value=4, visible=False)

with gr.Row():
dd_offset_x_b = gr.Slider(label='X offset (B)', minimum=-200, maximum=200, step=1, value=0, visible=False)
dd_offset_y_b = gr.Slider(label='Y offset (B)', minimum=-200, maximum=200, step=1, value=0, visible=False)

with gr.Group():
with gr.Row():
dd_mask_blur = gr.Slider(label='Mask blur ', minimum=0, maximum=64, step=1, value=4, visible=(not is_img2img))
dd_denoising_strength = gr.Slider(label='Denoising strength (Inpaint)', minimum=0.0, maximum=1.0, step=0.01, value=0.4, visible=(not is_img2img))

with gr.Row():
dd_inpaint_full_res = gr.Checkbox(label='Inpaint at full resolution ', value=True, visible = (not is_img2img))
dd_inpaint_full_res_padding = gr.Slider(label='Inpaint at full resolution padding, pixels ', minimum=0, maximum=256, step=4, value=32, visible=(not is_img2img))

dd_model_a.change(
lambda modelname: {
dd_model_b:gr_show( modelname != "None" ),
dd_conf_a:gr_show( modelname != "None" ),
dd_dilation_factor_a:gr_show( modelname != "None"),
dd_offset_x_a:gr_show( modelname != "None" ),
dd_offset_y_a:gr_show( modelname != "None" )

},
inputs= [dd_model_a],
outputs =[dd_model_b, dd_conf_a, dd_dilation_factor_a, dd_offset_x_a, dd_offset_y_a]
)

dd_model_b.change(
lambda modelname: {
dd_preprocess_b:gr_show( modelname != "None" ),
dd_bitwise_op:gr_show( modelname != "None" ),
dd_conf_b:gr_show( modelname != "None" ),
dd_dilation_factor_b:gr_show( modelname != "None"),
dd_offset_x_b:gr_show( modelname != "None" ),
dd_offset_y_b:gr_show( modelname != "None" )
},
inputs= [dd_model_b],
outputs =[dd_preprocess_b, dd_bitwise_op, dd_conf_b, dd_dilation_factor_b, dd_offset_x_b, dd_offset_y_b]
)

return [info,
dd_model_a,
dd_conf_a, dd_dilation_factor_a,
dd_offset_x_a, dd_offset_y_a,
dd_preprocess_b, dd_bitwise_op,
br,
dd_model_b,
dd_conf_b, dd_dilation_factor_b,
dd_offset_x_b, dd_offset_y_b,
dd_mask_blur, dd_denoising_strength,
dd_inpaint_full_res, dd_inpaint_full_res_padding
]
return seed, subseed

def run(self, p, info,
def process(self, p, *args):
if getattr(p, "_disable_ddetailer", False):
return

def postprocess_image(self, p, pp, info, enabled,
dd_model_a,
dd_conf_a, dd_dilation_factor_a,
dd_offset_x_a, dd_offset_y_a,
Expand All @@ -200,21 +233,36 @@ def run(self, p, info,
dd_mask_blur, dd_denoising_strength,
dd_inpaint_full_res, dd_inpaint_full_res_padding):

processing.fix_seed(p)
if getattr(p, "_disable_ddetailer", False):
return

if not enabled:
return

script_mode = False

initial_info = None
seed = p.seed
p.batch_size = 1
ddetail_count = p.n_iter
p.n_iter = 1
p.do_not_save_grid = True
p.do_not_save_samples = True
seed, subseed = self.get_seed(p)
p.seed = seed
p.subseed = subseed

info = ""
if not script_mode:
ddetail_count = 1
else:
ddetail_count = p.n_iter
p.batch_size = 1
p.n_iter = 1
p.do_not_save_grid = True
p.do_not_save_samples = True

is_txt2img = isinstance(p, StableDiffusionProcessingTxt2Img)
p_txt = copy(p)
if (not is_txt2img):
orig_image = p.init_images[0]
else:
p_txt = p
p = StableDiffusionProcessingImg2Img(
init_images = None,
init_images = [pp.image],
resize_mode = 0,
denoising_strength = dd_denoising_strength,
mask = None,
Expand All @@ -235,7 +283,8 @@ def run(self, p, info,
seed_resize_from_h=p_txt.seed_resize_from_h,
seed_resize_from_w=p_txt.seed_resize_from_w,
sampler_name=p_txt.sampler_name,
n_iter=p_txt.n_iter,
batch_size=1,
n_iter=1,
steps=p_txt.steps,
cfg_scale=p_txt.cfg_scale,
width=p_txt.width,
Expand All @@ -244,16 +293,18 @@ def run(self, p, info,
)
p.do_not_save_grid = True
p.do_not_save_samples = True

p._disable_ddetailer = True

output_images = []
state.job_count = ddetail_count
for n in range(ddetail_count):
devices.torch_gc()
start_seed = seed + n
if ( is_txt2img ):
print(f"Processing initial image for output generation {n + 1}.")
p_txt.seed = start_seed
processed = processing.process_images(p_txt)
init_image = processed.images[0]
print(f"Prepare initial image for output generation {p_txt.iteration + 1}.")
init_image = copy(pp.image)
info = processing.create_infotext(p_txt, p_txt.all_prompts, p_txt.all_seeds, p_txt.all_subseeds, None, 0, 0)
else:
init_image = orig_image

Expand All @@ -276,7 +327,7 @@ def run(self, p, info,
images.save_image(segmask_preview_b, opts.outdir_ddetailer_previews, "", start_seed, p.prompt, opts.samples_format, p=p)
gen_count = len(masks_b_pre)
state.job_count += gen_count
print(f"Processing {gen_count} model {label_b_pre} detections for output generation {n + 1}.")
print(f"Processing {gen_count} model {label_b_pre} detections for output generation {p_txt.iteration + 1}.")
p.seed = start_seed
p.init_images = [init_image]

Expand All @@ -285,15 +336,20 @@ def run(self, p, info,
if ( opts.dd_save_masks):
images.save_image(masks_b_pre[i], opts.outdir_ddetailer_masks, "", start_seed, p.prompt, opts.samples_format, p=p)
processed = processing.process_images(p)

if not is_txt2img:
p.prompt = processed.all_prompts[0]
info = processed.info
p.seed = processed.seed + 1
p.subseed = processed.subseed + 1
p.init_images = processed.images

if (gen_count > 0):
output_images[n] = processed.images[0]
init_image = processed.images[0]

else:
print(f"No model B detections for output generation {n} with current settings.")
print(f"No model B detections for output generation {p_txt.iteration + 1} with current settings.")

# Primary run
if (dd_model_a != "None"):
Expand Down Expand Up @@ -335,7 +391,7 @@ def run(self, p, info,
images.save_image(segmask_preview_a, opts.outdir_ddetailer_previews, "", start_seed, p.prompt, opts.samples_format, p=p)
gen_count = len(masks_a)
state.job_count += gen_count
print(f"Processing {gen_count} model {label_a} detections for output generation {n + 1}.")
print(f"Processing {gen_count} model {label_a} detections for output generation {p_txt.iteration + 1}.")
p.seed = start_seed
p.init_images = [init_image]

Expand All @@ -347,21 +403,23 @@ def run(self, p, info,
processed = processing.process_images(p)
if initial_info is None:
initial_info = processed.info
info = processed.info
p.seed = processed.seed + 1
p.subseed = processed.subseed + 1
p.init_images = processed.images

if (gen_count > 0):
output_images[n] = processed.images[0]
if ( opts.samples_save ):
images.save_image(processed.images[0], p.outpath_samples, "", start_seed, p.prompt, opts.samples_format, info=initial_info, p=p)
if ( script_mode and opts.samples_save ):
images.save_image(processed.images[0], p.outpath_samples, "", start_seed, p.prompt, opts.samples_format, info=info, p=p)

else:
print(f"No model {label_a} detections for output generation {n} with current settings.")
state.job = f"Generation {n + 1} out of {state.job_count}"
print(f"No model {label_a} detections for output generation {p_txt.iteration + 1} with current settings.")
state.job = f"Generation {p_txt.iteration + 1} out of {state.job_count}"
if (initial_info is None):
initial_info = "No detections found."

return Processed(p, output_images, seed, initial_info)
pp.image = output_images[0]

def modeldataset(model_shortname):
path = modelpath(model_shortname)
Expand Down

0 comments on commit 11fee42

Please sign in to comment.