Skip to content

Commit

Permalink
rework code/UI for #15293
Browse files Browse the repository at this point in the history
  • Loading branch information
AUTOMATIC1111 committed Mar 31, 2024
1 parent 7f3ce06 commit dcd4f88
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 41 deletions.
1 change: 0 additions & 1 deletion modules/shared_options.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,6 @@
"DAT_tile_overlap": OptionInfo(8, "Tile overlap for DAT upscalers.", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}).info("Low values = visible seam"),
"upscaler_for_img2img": OptionInfo(None, "Upscaler for img2img", gr.Dropdown, lambda: {"choices": [x.name for x in shared.sd_upscalers]}),
"set_scale_by_when_changing_upscaler": OptionInfo(False, "Automatically set the Scale by factor based on the name of the selected Upscaler."),
"show_limit_target_resolution_in_extras_upscale": OptionInfo(False, 'Show "Limit target resolution" slider in "Upscale" extras script. Useful for batches where can be big images.').needs_reload_ui(),
}))

options_templates.update(options_section(('face-restoration', "Face restoration", "postprocessing"), {
Expand Down
75 changes: 35 additions & 40 deletions scripts/postprocessing_upscale.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,17 +12,15 @@
upscale_cache = {}


def limitSizeByOneDemention(size: tuple, limit: int):
w, h = size
if h > w:
if h > limit:
w = limit / h * w
h = limit
else:
if w > limit:
h = limit / w * h
w = limit
return (int(w), int(h))
def limit_size_by_one_dimention(w, h, limit):
if h > w and h > limit:
w = limit * w // h
h = limit
elif w > h and w > limit:
h = limit * h // w
w = limit

return int(w), int(h)


class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing):
Expand All @@ -43,12 +41,11 @@ def ui(self):
with FormRow():
with gr.Tabs(elem_id="extras_resize_mode"):
with gr.TabItem('Scale by', elem_id="extras_scale_by_tab") as tab_scale_by:
upscaling_resize = gr.Slider(minimum=1.0, maximum=8.0, step=0.05, label="Resize", value=4, elem_id="extras_upscaling_resize")
if shared.opts.show_limit_target_resolution_in_extras_upscale:
limit_target_resolution = gr.Slider(minimum=0, maximum=10000, step=8, label="Limit target resolution", value=8000, elem_id="extras_upscale_limit_target_resolution",
tooltip="0 = no limit. Limit maximal target resolution by the biggest demension. Useful for batches where can be big images.")
else:
limit_target_resolution = gr.Number(0, visible=False)
with gr.Row():
with gr.Column(scale=3):
upscaling_resize = gr.Slider(minimum=1.0, maximum=8.0, step=0.05, label="Resize", value=4, elem_id="extras_upscaling_resize")
with gr.Column(scale=1):
max_side_length = gr.Number(label="Max side length", value=0, elem_id="extras_upscale_max_side_length", tooltip="If any of two sides of the image ends up larger than specified, will downscale it to fit. 0 = no limit.")

with gr.TabItem('Scale to', elem_id="extras_scale_to_tab") as tab_scale_to:
with FormRow():
Expand Down Expand Up @@ -79,7 +76,7 @@ def on_selected_upscale_method(upscale_method):
"upscale_enabled": upscale_enabled,
"upscale_mode": selected_tab,
"upscale_by": upscaling_resize,
"limit_target_resolution": limit_target_resolution,
"max_side_length": max_side_length,
"upscale_to_width": upscaling_resize_w,
"upscale_to_height": upscaling_resize_h,
"upscale_crop": upscaling_crop,
Expand All @@ -88,18 +85,18 @@ def on_selected_upscale_method(upscale_method):
"upscaler_2_visibility": extras_upscaler_2_visibility,
}

def upscale(self, image, info, upscaler, upscale_mode, upscale_by, limit_target_resolution, upscale_to_width, upscale_to_height, upscale_crop):
def upscale(self, image, info, upscaler, upscale_mode, upscale_by, max_side_length, upscale_to_width, upscale_to_height, upscale_crop):
if upscale_mode == 1:
upscale_by = max(upscale_to_width/image.width, upscale_to_height/image.height)
info["Postprocess upscale to"] = f"{upscale_to_width}x{upscale_to_height}"
else:
info["Postprocess upscale by"] = upscale_by
if limit_target_resolution != 0 and max(*image.size)*upscale_by > limit_target_resolution:
if max_side_length != 0 and max(*image.size)*upscale_by > max_side_length:
upscale_mode = 1
upscale_crop = False
upscale_to_width, upscale_to_height = limitSizeByOneDemention((image.width*upscale_by, image.height*upscale_by), limit_target_resolution)
upscale_to_width, upscale_to_height = limit_size_by_one_dimention(image.width*upscale_by, image.height*upscale_by, max_side_length)
upscale_by = max(upscale_to_width/image.width, upscale_to_height/image.height)
info["Limit target resolution"] = limit_target_resolution
info["Max side length"] = max_side_length

cache_key = (hash(np.array(image.getdata()).tobytes()), upscaler.name, upscale_mode, upscale_by, upscale_to_width, upscale_to_height, upscale_crop)
cached_image = upscale_cache.pop(cache_key, None)
Expand All @@ -121,21 +118,21 @@ def upscale(self, image, info, upscaler, upscale_mode, upscale_by, limit_target_

return image

def process_firstpass(self, pp: scripts_postprocessing.PostprocessedImage, **args):
if args['upscale_mode'] == 1:
pp.shared.target_width = args['upscale_to_width']
pp.shared.target_height = args['upscale_to_height']
def process_firstpass(self, pp: scripts_postprocessing.PostprocessedImage, upscale_enabled=True, upscale_mode=1, upscale_by=2.0, max_side_length=0, upscale_to_width=None, upscale_to_height=None, upscale_crop=False, upscaler_1_name=None, upscaler_2_name=None, upscaler_2_visibility=0.0):
if upscale_mode == 1:
pp.shared.target_width = upscale_to_width
pp.shared.target_height = upscale_to_height
else:
pp.shared.target_width = int(pp.image.width * args['upscale_by'])
pp.shared.target_height = int(pp.image.height * args['upscale_by'])
if args['limit_target_resolution'] != 0:
pp.shared.target_width, pp.shared.target_height = limitSizeByOneDemention((pp.shared.target_width, pp.shared.target_height), args['limit_target_resolution'])
pp.shared.target_width = int(pp.image.width * upscale_by)
pp.shared.target_height = int(pp.image.height * upscale_by)

pp.shared.target_width, pp.shared.target_height = limit_size_by_one_dimention(pp.shared.target_width, pp.shared.target_height, max_side_length)

def process(self, pp: scripts_postprocessing.PostprocessedImage, **args):
if not args['upscale_enabled']:
def process(self, pp: scripts_postprocessing.PostprocessedImage, upscale_enabled=True, upscale_mode=1, upscale_by=2.0, max_side_length=0, upscale_to_width=None, upscale_to_height=None, upscale_crop=False, upscaler_1_name=None, upscaler_2_name=None, upscaler_2_visibility=0.0):
if not upscale_enabled:
return

upscaler_1_name = args['upscaler_1_name']
upscaler_1_name = upscaler_1_name
if upscaler_1_name == "None":
upscaler_1_name = None

Expand All @@ -145,21 +142,19 @@ def process(self, pp: scripts_postprocessing.PostprocessedImage, **args):
if not upscaler1:
return

upscaler_2_name = args['upscaler_2_name']
upscaler_2_name = upscaler_2_name
if upscaler_2_name == "None":
upscaler_2_name = None

upscaler2 = next(iter([x for x in shared.sd_upscalers if x.name == upscaler_2_name and x.name != "None"]), None)
assert upscaler2 or (upscaler_2_name is None), f'could not find upscaler named {upscaler_2_name}'

upscaled_image = self.upscale(pp.image, pp.info, upscaler1, args['upscale_mode'], args['upscale_by'], args['limit_target_resolution'], args['upscale_to_width'],
args['upscale_to_height'], args['upscale_crop'])
upscaled_image = self.upscale(pp.image, pp.info, upscaler1, upscale_mode, upscale_by, max_side_length, upscale_to_width, upscale_to_height, upscale_crop)
pp.info["Postprocess upscaler"] = upscaler1.name

if upscaler2 and args['upscaler_2_visibility'] > 0:
second_upscale = self.upscale(pp.image, pp.info, upscaler2, args['upscale_mode'], args['upscale_by'], args['upscale_to_width'],
args['upscale_to_height'], args['upscale_crop'])
upscaled_image = Image.blend(upscaled_image, second_upscale, args['upscaler_2_visibility'])
if upscaler2 and upscaler_2_visibility > 0:
second_upscale = self.upscale(pp.image, pp.info, upscaler2, upscale_mode, upscale_by, upscale_to_width, upscale_to_height, upscale_crop)
upscaled_image = Image.blend(upscaled_image, second_upscale, upscaler_2_visibility)

pp.info["Postprocess upscaler 2"] = upscaler2.name

Expand Down

1 comment on commit dcd4f88

@light-and-ray
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍🏻

Please sign in to comment.