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 "")