474 lines
19 KiB
Python
474 lines
19 KiB
Python
import math
|
|
from PIL import Image, ImageDraw, ImageOps
|
|
from .processing import StableDiffusionProcessing
|
|
from .processing import Processed
|
|
from enum import Enum
|
|
from . import shared
|
|
from . import images
|
|
from . import devices
|
|
from . import processing
|
|
|
|
class USDUMode(Enum):
|
|
LINEAR = 0
|
|
CHESS = 1
|
|
NONE = 2
|
|
|
|
class USDUSFMode(Enum):
|
|
NONE = 0
|
|
BAND_PASS = 1
|
|
HALF_TILE = 2
|
|
HALF_TILE_PLUS_INTERSECTIONS = 3
|
|
|
|
class USDUpscaler():
|
|
|
|
def __init__(self, p, image, upscaler_index:int, save_redraw, save_seams_fix, tile_width, tile_height) -> None:
|
|
self.p:StableDiffusionProcessing = p
|
|
self.image:Image = image
|
|
self.scale_factor = math.ceil(max(p.width, p.height) / max(image.width, image.height))
|
|
self.upscaler = shared.sd_upscalers[upscaler_index]
|
|
self.redraw = USDURedraw()
|
|
self.redraw.save = save_redraw
|
|
self.redraw.tile_width = tile_width if tile_width > 0 else tile_height
|
|
self.redraw.tile_height = tile_height if tile_height > 0 else tile_width
|
|
self.seams_fix = USDUSeamsFix()
|
|
self.seams_fix.save = save_seams_fix
|
|
self.seams_fix.tile_width = tile_width if tile_width > 0 else tile_height
|
|
self.seams_fix.tile_height = tile_height if tile_height > 0 else tile_width
|
|
self.initial_info = None
|
|
self.rows = math.ceil(self.p.height / 4 / self.redraw.tile_height)
|
|
self.cols = math.ceil(self.p.width / self.redraw.tile_width)
|
|
|
|
def get_factor(self, num):
|
|
# Its just return, don't need elif
|
|
if num == 1:
|
|
return 2
|
|
if num % 4 == 0:
|
|
return 4
|
|
if num % 3 == 0:
|
|
return 3
|
|
if num % 2 == 0:
|
|
return 2
|
|
return 0
|
|
|
|
def get_factors(self):
|
|
scales = []
|
|
current_scale = 1
|
|
current_scale_factor = self.get_factor(self.scale_factor)
|
|
while current_scale_factor == 0:
|
|
self.scale_factor += 1
|
|
current_scale_factor = self.get_factor(self.scale_factor)
|
|
while current_scale < self.scale_factor:
|
|
current_scale_factor = self.get_factor(self.scale_factor // current_scale)
|
|
scales.append(current_scale_factor)
|
|
current_scale = current_scale * current_scale_factor
|
|
if current_scale_factor == 0:
|
|
break
|
|
self.scales = enumerate(scales)
|
|
|
|
def upscale(self):
|
|
# Log info
|
|
print(f"Canva size: {self.p.width}x{self.p.height}")
|
|
print(f"Image size: {self.image.width}x{self.image.height}")
|
|
print(f"Scale factor: {self.scale_factor}")
|
|
# Check upscaler is not empty
|
|
if self.upscaler.name == "None":
|
|
self.image = self.image.resize((self.p.width, self.p.height), resample=Image.LANCZOS)
|
|
return
|
|
# Get list with scale factors
|
|
self.get_factors()
|
|
# Upscaling image over all factors
|
|
for index, value in self.scales:
|
|
print(f"Upscaling iteration {index+1} with scale factor {value}")
|
|
self.image = self.upscaler.scaler.upscale(self.image, value, self.upscaler.data_path)
|
|
# Resize image to set values
|
|
self.image = self.image.resize((self.p.width, self.p.height), resample=Image.LANCZOS)
|
|
shared.batch = [self.image] + [img.resize((self.p.width, self.p.height), resample=Image.LANCZOS) for img in shared.batch[1:]]
|
|
|
|
|
|
def setup_redraw(self, redraw_mode, padding, mask_blur):
|
|
self.redraw.mode = USDUMode(redraw_mode)
|
|
self.redraw.enabled = self.redraw.mode != USDUMode.NONE
|
|
self.redraw.padding = padding
|
|
self.p.mask_blur = mask_blur
|
|
|
|
|
|
def setup_seams_fix(self, padding, denoise, mask_blur, width, mode):
|
|
self.seams_fix.padding = padding
|
|
self.seams_fix.denoise = denoise
|
|
self.seams_fix.mask_blur = mask_blur
|
|
self.seams_fix.width = width
|
|
self.seams_fix.mode = USDUSFMode(mode)
|
|
self.seams_fix.enabled = self.seams_fix.mode != USDUSFMode.NONE
|
|
|
|
|
|
def save_image(self):
|
|
if type(self.p.prompt) != list:
|
|
images.save_image(self.image, self.p.outpath_samples, "", self.p.seed, self.p.prompt, shared.opts.samples_format, info=self.initial_info, p=self.p)
|
|
else:
|
|
images.save_image(self.image, self.p.outpath_samples, "", self.p.seed, self.p.prompt[0], shared.opts.samples_format, info=self.initial_info, p=self.p)
|
|
|
|
def calc_jobs_count(self):
|
|
redraw_job_count = (self.rows * self.cols) if self.redraw.enabled else 0
|
|
seams_job_count = 0
|
|
if self.seams_fix.mode == USDUSFMode.BAND_PASS:
|
|
seams_job_count = self.rows + self.cols - 2
|
|
elif self.seams_fix.mode == USDUSFMode.HALF_TILE:
|
|
seams_job_count = self.rows * (self.cols - 1) + (self.rows - 1) * self.cols
|
|
elif self.seams_fix.mode == USDUSFMode.HALF_TILE_PLUS_INTERSECTIONS:
|
|
seams_job_count = self.rows * (self.cols - 1) + (self.rows - 1) * self.cols + (self.rows - 1) * (self.cols - 1)
|
|
|
|
shared.state.job_count = redraw_job_count + seams_job_count
|
|
|
|
def print_info(self):
|
|
print(f"Tile size: {self.redraw.tile_width}x{self.redraw.tile_height}")
|
|
print(f"Tiles amount: {self.rows * self.cols}")
|
|
print(f"Grid: {self.rows}x{self.cols}")
|
|
print(f"Redraw enabled: {self.redraw.enabled}")
|
|
print(f"Seams fix mode: {self.seams_fix.mode.name}")
|
|
|
|
def add_extra_info(self):
|
|
self.p.extra_generation_params["Ultimate SD upscale upscaler"] = self.upscaler.name
|
|
self.p.extra_generation_params["Ultimate SD upscale tile_width"] = self.redraw.tile_width
|
|
self.p.extra_generation_params["Ultimate SD upscale tile_height"] = self.redraw.tile_height
|
|
self.p.extra_generation_params["Ultimate SD upscale mask_blur"] = self.p.mask_blur
|
|
self.p.extra_generation_params["Ultimate SD upscale padding"] = self.redraw.padding
|
|
|
|
def process(self):
|
|
shared.state.begin()
|
|
self.calc_jobs_count()
|
|
self.result_images = []
|
|
if self.redraw.enabled:
|
|
self.image = self.redraw.start(self.p, self.image, self.rows, self.cols)
|
|
self.initial_info = self.redraw.initial_info
|
|
self.result_images.append(self.image)
|
|
if self.redraw.save:
|
|
self.save_image()
|
|
|
|
if self.seams_fix.enabled:
|
|
self.image = self.seams_fix.start(self.p, self.image, self.rows, self.cols)
|
|
self.initial_info = self.seams_fix.initial_info
|
|
self.result_images.append(self.image)
|
|
if self.seams_fix.save:
|
|
self.save_image()
|
|
shared.state.end()
|
|
|
|
class USDURedraw():
|
|
|
|
def init_draw(self, p, width, height):
|
|
p.inpaint_full_res = True
|
|
p.inpaint_full_res_padding = self.padding
|
|
p.width = math.ceil((self.tile_width+self.padding) / 8) * 8
|
|
p.height = math.ceil((self.tile_height+self.padding) / 8) * 8
|
|
mask = Image.new("L", (width, height), "black")
|
|
draw = ImageDraw.Draw(mask)
|
|
return mask, draw
|
|
|
|
def calc_rectangle(self, xi, yi):
|
|
x1 = xi * self.tile_width
|
|
y1 = yi * self.tile_height
|
|
x2 = xi * self.tile_width + self.tile_width
|
|
y2 = yi * self.tile_height + self.tile_height
|
|
|
|
return x1, y1, x2, y2
|
|
|
|
def linear_process(self, p, image, rows, cols):
|
|
mask, draw = self.init_draw(p, image.width, image.height)
|
|
for yi in range(rows):
|
|
for xi in range(cols):
|
|
if shared.state.interrupted:
|
|
break
|
|
draw.rectangle(self.calc_rectangle(xi, yi), fill="white")
|
|
p.init_images = [image]
|
|
p.image_mask = mask
|
|
processed = processing.process_images(p)
|
|
draw.rectangle(self.calc_rectangle(xi, yi), fill="black")
|
|
if (len(processed.images) > 0):
|
|
image = processed.images[0]
|
|
|
|
p.width = image.width
|
|
p.height = image.height
|
|
self.initial_info = processed.infotext(p, 0)
|
|
|
|
return image
|
|
|
|
def chess_process(self, p, image, rows, cols):
|
|
mask, draw = self.init_draw(p, image.width, image.height)
|
|
tiles = []
|
|
# calc tiles colors
|
|
for yi in range(rows):
|
|
for xi in range(cols):
|
|
if shared.state.interrupted:
|
|
break
|
|
if xi == 0:
|
|
tiles.append([])
|
|
color = xi % 2 == 0
|
|
if yi > 0 and yi % 2 != 0:
|
|
color = not color
|
|
tiles[yi].append(color)
|
|
|
|
for yi in range(len(tiles)):
|
|
for xi in range(len(tiles[yi])):
|
|
if shared.state.interrupted:
|
|
break
|
|
if not tiles[yi][xi]:
|
|
tiles[yi][xi] = not tiles[yi][xi]
|
|
continue
|
|
tiles[yi][xi] = not tiles[yi][xi]
|
|
draw.rectangle(self.calc_rectangle(xi, yi), fill="white")
|
|
p.init_images = [image]
|
|
p.image_mask = mask
|
|
processed = processing.process_images(p)
|
|
draw.rectangle(self.calc_rectangle(xi, yi), fill="black")
|
|
if (len(processed.images) > 0):
|
|
image = processed.images[0]
|
|
|
|
for yi in range(len(tiles)):
|
|
for xi in range(len(tiles[yi])):
|
|
if shared.state.interrupted:
|
|
break
|
|
if not tiles[yi][xi]:
|
|
continue
|
|
draw.rectangle(self.calc_rectangle(xi, yi), fill="white")
|
|
p.init_images = [image]
|
|
p.image_mask = mask
|
|
processed = processing.process_images(p)
|
|
draw.rectangle(self.calc_rectangle(xi, yi), fill="black")
|
|
if (len(processed.images) > 0):
|
|
image = processed.images[0]
|
|
|
|
p.width = image.width
|
|
p.height = image.height
|
|
self.initial_info = processed.infotext(p, 0)
|
|
|
|
return image
|
|
|
|
def start(self, p, image, rows, cols):
|
|
self.initial_info = None
|
|
if self.mode == USDUMode.LINEAR:
|
|
return self.linear_process(p, image, rows, cols)
|
|
if self.mode == USDUMode.CHESS:
|
|
return self.chess_process(p, image, rows, cols)
|
|
|
|
class USDUSeamsFix():
|
|
|
|
def init_draw(self, p):
|
|
self.initial_info = None
|
|
p.width = math.ceil((self.tile_width+self.padding) / 8) * 8
|
|
p.height = math.ceil((self.tile_height+self.padding) / 8) * 8
|
|
|
|
def half_tile_process(self, p, image, rows, cols):
|
|
|
|
self.init_draw(p)
|
|
processed = None
|
|
|
|
gradient = Image.linear_gradient("L")
|
|
row_gradient = Image.new("L", (self.tile_width, self.tile_height), "black")
|
|
row_gradient.paste(gradient.resize(
|
|
(self.tile_width, self.tile_height//2), resample=Image.BICUBIC), (0, 0))
|
|
row_gradient.paste(gradient.rotate(180).resize(
|
|
(self.tile_width, self.tile_height//2), resample=Image.BICUBIC),
|
|
(0, self.tile_height//2))
|
|
col_gradient = Image.new("L", (self.tile_width, self.tile_height), "black")
|
|
col_gradient.paste(gradient.rotate(90).resize(
|
|
(self.tile_width//2, self.tile_height), resample=Image.BICUBIC), (0, 0))
|
|
col_gradient.paste(gradient.rotate(270).resize(
|
|
(self.tile_width//2, self.tile_height), resample=Image.BICUBIC), (self.tile_width//2, 0))
|
|
|
|
p.denoising_strength = self.denoise
|
|
p.mask_blur = self.mask_blur
|
|
|
|
for yi in range(rows-1):
|
|
for xi in range(cols):
|
|
if shared.state.interrupted:
|
|
break
|
|
p.width = self.tile_width
|
|
p.height = self.tile_height
|
|
p.inpaint_full_res = True
|
|
p.inpaint_full_res_padding = self.padding
|
|
mask = Image.new("L", (image.width, image.height), "black")
|
|
mask.paste(row_gradient, (xi*self.tile_width, yi*self.tile_height + self.tile_height//2))
|
|
|
|
p.init_images = [image]
|
|
p.image_mask = mask
|
|
processed = processing.process_images(p)
|
|
if (len(processed.images) > 0):
|
|
image = processed.images[0]
|
|
|
|
for yi in range(rows):
|
|
for xi in range(cols-1):
|
|
if shared.state.interrupted:
|
|
break
|
|
p.width = self.tile_width
|
|
p.height = self.tile_height
|
|
p.inpaint_full_res = True
|
|
p.inpaint_full_res_padding = self.padding
|
|
mask = Image.new("L", (image.width, image.height), "black")
|
|
mask.paste(col_gradient, (xi*self.tile_width+self.tile_width//2, yi*self.tile_height))
|
|
|
|
p.init_images = [image]
|
|
p.image_mask = mask
|
|
processed = processing.process_images(p)
|
|
if (len(processed.images) > 0):
|
|
image = processed.images[0]
|
|
|
|
p.width = image.width
|
|
p.height = image.height
|
|
if processed is not None:
|
|
self.initial_info = processed.infotext(p, 0)
|
|
|
|
return image
|
|
|
|
def half_tile_process_corners(self, p, image, rows, cols):
|
|
fixed_image = self.half_tile_process(p, image, rows, cols)
|
|
processed = None
|
|
self.init_draw(p)
|
|
gradient = Image.radial_gradient("L").resize(
|
|
(self.tile_width, self.tile_height), resample=Image.BICUBIC)
|
|
gradient = ImageOps.invert(gradient)
|
|
p.denoising_strength = self.denoise
|
|
#p.mask_blur = 0
|
|
p.mask_blur = self.mask_blur
|
|
|
|
for yi in range(rows-1):
|
|
for xi in range(cols-1):
|
|
if shared.state.interrupted:
|
|
break
|
|
p.width = self.tile_width
|
|
p.height = self.tile_height
|
|
p.inpaint_full_res = True
|
|
p.inpaint_full_res_padding = 0
|
|
mask = Image.new("L", (fixed_image.width, fixed_image.height), "black")
|
|
mask.paste(gradient, (xi*self.tile_width + self.tile_width//2,
|
|
yi*self.tile_height + self.tile_height//2))
|
|
|
|
p.init_images = [fixed_image]
|
|
p.image_mask = mask
|
|
processed = processing.process_images(p)
|
|
if (len(processed.images) > 0):
|
|
fixed_image = processed.images[0]
|
|
|
|
p.width = fixed_image.width
|
|
p.height = fixed_image.height
|
|
if processed is not None:
|
|
self.initial_info = processed.infotext(p, 0)
|
|
|
|
return fixed_image
|
|
|
|
def band_pass_process(self, p, image, cols, rows):
|
|
|
|
self.init_draw(p)
|
|
processed = None
|
|
|
|
p.denoising_strength = self.denoise
|
|
p.mask_blur = 0
|
|
|
|
gradient = Image.linear_gradient("L")
|
|
mirror_gradient = Image.new("L", (256, 256), "black")
|
|
mirror_gradient.paste(gradient.resize((256, 128), resample=Image.BICUBIC), (0, 0))
|
|
mirror_gradient.paste(gradient.rotate(180).resize((256, 128), resample=Image.BICUBIC), (0, 128))
|
|
|
|
row_gradient = mirror_gradient.resize((image.width, self.width), resample=Image.BICUBIC)
|
|
col_gradient = mirror_gradient.rotate(90).resize((self.width, image.height), resample=Image.BICUBIC)
|
|
|
|
for xi in range(1, rows):
|
|
if shared.state.interrupted:
|
|
break
|
|
p.width = self.width + self.padding * 2
|
|
p.height = image.height
|
|
p.inpaint_full_res = True
|
|
p.inpaint_full_res_padding = self.padding
|
|
mask = Image.new("L", (image.width, image.height), "black")
|
|
mask.paste(col_gradient, (xi * self.tile_width - self.width // 2, 0))
|
|
|
|
p.init_images = [image]
|
|
p.image_mask = mask
|
|
processed = processing.process_images(p)
|
|
if (len(processed.images) > 0):
|
|
image = processed.images[0]
|
|
for yi in range(1, cols):
|
|
if shared.state.interrupted:
|
|
break
|
|
p.width = image.width
|
|
p.height = self.width + self.padding * 2
|
|
p.inpaint_full_res = True
|
|
p.inpaint_full_res_padding = self.padding
|
|
mask = Image.new("L", (image.width, image.height), "black")
|
|
mask.paste(row_gradient, (0, yi * self.tile_height - self.width // 2))
|
|
|
|
p.init_images = [image]
|
|
p.image_mask = mask
|
|
processed = processing.process_images(p)
|
|
if (len(processed.images) > 0):
|
|
image = processed.images[0]
|
|
|
|
p.width = image.width
|
|
p.height = image.height
|
|
if processed is not None:
|
|
self.initial_info = processed.infotext(p, 0)
|
|
|
|
return image
|
|
|
|
def start(self, p, image, rows, cols):
|
|
if USDUSFMode(self.mode) == USDUSFMode.BAND_PASS:
|
|
return self.band_pass_process(p, image, rows, cols)
|
|
elif USDUSFMode(self.mode) == USDUSFMode.HALF_TILE:
|
|
return self.half_tile_process(p, image, rows, cols)
|
|
elif USDUSFMode(self.mode) == USDUSFMode.HALF_TILE_PLUS_INTERSECTIONS:
|
|
return self.half_tile_process_corners(p, image, rows, cols)
|
|
else:
|
|
return image
|
|
|
|
class Script():
|
|
def title(self):
|
|
return "Ultimate SD upscale"
|
|
|
|
def show(self, is_img2img):
|
|
return is_img2img
|
|
|
|
def run(self, p, _, tile_width, tile_height, mask_blur, padding, seams_fix_width, seams_fix_denoise, seams_fix_padding,
|
|
upscaler_index, save_upscaled_image, redraw_mode, save_seams_fix_image, seams_fix_mask_blur,
|
|
seams_fix_type, target_size_type, custom_width, custom_height, custom_scale):
|
|
|
|
# Init
|
|
processing.fix_seed(p)
|
|
devices.torch_gc()
|
|
|
|
p.do_not_save_grid = True
|
|
p.do_not_save_samples = True
|
|
p.inpaint_full_res = False
|
|
|
|
p.inpainting_fill = 1
|
|
p.n_iter = 1
|
|
p.batch_size = 1
|
|
|
|
seed = p.seed
|
|
|
|
# Init image
|
|
init_img = p.init_images[0]
|
|
if init_img == None:
|
|
return Processed(p, [], seed, "Empty image")
|
|
init_img = images.flatten(init_img, shared.opts.img2img_background_color)
|
|
|
|
#override size
|
|
if target_size_type == 1:
|
|
p.width = custom_width
|
|
p.height = custom_height
|
|
if target_size_type == 2:
|
|
p.width = math.ceil((init_img.width * custom_scale) / 8) * 8
|
|
p.height = math.ceil((init_img.height * custom_scale) / 8) * 8
|
|
|
|
# Upscaling
|
|
upscaler = USDUpscaler(p, init_img, upscaler_index, save_upscaled_image, save_seams_fix_image, tile_width, tile_height)
|
|
upscaler.upscale()
|
|
|
|
# Drawing
|
|
upscaler.setup_redraw(redraw_mode, padding, mask_blur)
|
|
upscaler.setup_seams_fix(seams_fix_padding, seams_fix_denoise, seams_fix_mask_blur, seams_fix_width, seams_fix_type)
|
|
upscaler.print_info()
|
|
upscaler.add_extra_info()
|
|
upscaler.process()
|
|
result_images = upscaler.result_images
|
|
|
|
return Processed(p, result_images, seed, upscaler.initial_info if upscaler.initial_info is not None else "")
|
|
|