Some fixes

This commit is contained in:
Антон Антонов
2023-01-02 16:59:28 +03:00
parent e2fbb03976
commit 9302f48783

View File

@@ -44,6 +44,72 @@ def upscale(p, init_img, upscaler_index, tileSize, padding):
return upscaled_img.resize((p.width, p.height), resample=Image.LANCZOS)
def redraw_image(p, upscaled_img, rows, cols, tileSize, padding):
for yi in range(rows):
for xi in range(cols):
p.width = tileSize
p.height = tileSize
p.inpaint_full_res = True
p.inpaint_full_res_padding = padding
mask = Image.new("L", (upscaled_img.width, upscaled_img.height), "black")
draw = ImageDraw.Draw(mask)
draw.rectangle((
xi * tileSize,
yi * tileSize,
xi * tileSize + tileSize,
yi * tileSize + tileSize
), fill="white")
p.init_images = [upscaled_img]
p.image_mask = mask
processed = processing.process_images(p)
initial_info = processed.info
if (len(processed.images)>0):
upscaled_img = processed.images[0]
return upscaled_img
def seam_draw(p, upscaled_img, seam_pass_width, seam_pass_padding, seam_pass_denoise, padding, tileSize):
for xi in range(1, cols):
p.width = seam_pass_width + seam_pass_padding*2
p.height = upscaled_img.height
p.inpaint_full_res = True
p.inpaint_full_res_padding = padding
mask = Image.new("L", (upscaled_img.width, upscaled_img.height), "black")
draw = ImageDraw.Draw(mask)
draw.rectangle((
xi * tileSize - seam_pass_width//2,
0,
xi * tileSize + seam_pass_width//2,
mask.height
), fill="white")
p.init_images = [upscaled_img]
p.image_mask = mask
processed = processing.process_images(p)
if (len(processed.images) > 0):
upscaled_img = processed.images[0]
for yi in range(1, rows):
p.width = upscaled_img.width
p.height = seam_pass_width + seam_pass_padding*2
p.inpaint_full_res = True
p.inpaint_full_res_padding = padding
mask = Image.new("L", (upscaled_img.width, upscaled_img.height), "black")
draw = ImageDraw.Draw(mask)
draw.rectangle((
0,
yi * tileSize - seam_pass_width//2,
mask.width,
yi * tileSize + seam_pass_width//2
), fill="white")
p.init_images = [upscaled_img]
p.image_mask = mask
processed = processing.process_images(p)
if (len(processed.images) > 0):
upscaled_img = processed.images[0]
return upscaled_img
class Script(scripts.Script):
def title(self):
return "Ultimate SD upscale"
@@ -74,8 +140,10 @@ class Script(scripts.Script):
init_img = p.init_images[0]
init_img = images.flatten(init_img, opts.img2img_background_color)
#Upscaling
upscaled_img = upscale(p, init_img, upscaler_index, tileSize, padding)
#Drawing
devices.torch_gc()
p.inpaint_full_res = False
@@ -83,79 +151,23 @@ class Script(scripts.Script):
rows = math.ceil(p.height / tileSize)
cols = math.ceil(p.width / tileSize)
print(f"Tiles amount: {rows*cols}")
print(f"Grid: {rows}x{cols}")
print(f"Seam path: {seam_pass_enabled}")
seams = 0
if seam_pass_enabled:
seams = rows-1 + cols - 1
result_images = []
state.job_count = rows*cols + seams
for yi in range(rows):
for xi in range(cols):
p.width = tileSize
p.height = tileSize
p.inpaint_full_res = True
p.inpaint_full_res_padding = padding
mask = Image.new("L", (upscaled_img.width, upscaled_img.height), "black")
draw = ImageDraw.Draw(mask)
draw.rectangle((
xi * tileSize,
yi * tileSize,
xi * tileSize + tileSize,
yi * tileSize + tileSize
), fill="white")
p.init_images = [upscaled_img]
p.image_mask = mask
processed = processing.process_images(p)
initial_info = processed.info
if (len(processed.images)>0):
upscaled_img = processed.images[0]
result_images.append(upscaled_img)
result_images.append(redraw_image(p, upscaled_img, rows, cols, tileSize, padding))
if seam_pass_enabled:
for xi in range(1, cols):
p.width = seam_pass_width + seam_pass_padding*2
p.height = upscaled_img.height
p.inpaint_full_res = True
p.inpaint_full_res_padding = padding
mask = Image.new("L", (upscaled_img.width, upscaled_img.height), "black")
draw = ImageDraw.Draw(mask)
draw.rectangle((
xi * tileSize - seam_pass_width//2,
0,
xi * tileSize + seam_pass_width//2,
mask.height
), fill="white")
print(f"Starting seam path drawing")
result_images.append(seam_draw(p, upscaled_img, seam_pass_width, seam_pass_padding, seam_pass_denoise, padding, tileSize))
p.init_images = [upscaled_img]
p.image_mask = mask
processed = processing.process_images(p)
if (len(processed.images) > 0):
upscaled_img = processed.images[0]
for yi in range(1, rows):
p.width = upscaled_img.width
p.height = seam_pass_width + seam_pass_padding*2
p.inpaint_full_res = True
p.inpaint_full_res_padding = padding
mask = Image.new("L", (upscaled_img.width, upscaled_img.height), "black")
draw = ImageDraw.Draw(mask)
draw.rectangle((
0,
yi * tileSize - seam_pass_width//2,
mask.width,
yi * tileSize + seam_pass_width//2
), fill="white")
p.init_images = [upscaled_img]
p.image_mask = mask
processed = processing.process_images(p)
if (len(processed.images) > 0):
upscaled_img = processed.images[0]
result_images.append(upscaled_img)
processed = Processed(p, result_images, seed, initial_info)
return processed