Create upscaler.py

upscaler.py

@@ -0,0 +1,128 @@
+from PIL import Image
+from diffusers import StableDiffusionUpscalePipeline
+import torch
+from split_image import split
+import os
+import random
+
+
+def split_image(im, rows, cols, should_square, should_quiet=False):
+    im_width, im_height = im.size
+    row_width = int(im_width / cols)
+    row_height = int(im_height / rows)
+    name = "image"
+    ext = ".png"
+    name = os.path.basename(name)
+    images = []
+    if should_square:
+        min_dimension = min(im_width, im_height)
+        max_dimension = max(im_width, im_height)
+        if not should_quiet:
+            print("Resizing image to a square...")
+            print("Determining background color...")
+        bg_color = split.determine_bg_color(im)
+        if not should_quiet:
+            print("Background color is... " + str(bg_color))
+        im_r = Image.new("RGBA" if ext == "png" else "RGB",
+                         (max_dimension, max_dimension), bg_color)
+        offset = int((max_dimension - min_dimension) / 2)
+        if im_width > im_height:
+            im_r.paste(im, (0, offset))
+        else:
+            im_r.paste(im, (offset, 0))
+        im = im_r
+        row_width = int(max_dimension / cols)
+        row_height = int(max_dimension / rows)
+    n = 0
+    for i in range(0, rows):
+        for j in range(0, cols):
+            box = (j * row_width, i * row_height, j * row_width +
+                   row_width, i * row_height + row_height)
+            outp = im.crop(box)
+            outp_path = name + "_" + str(n) + ext
+            if not should_quiet:
+                print("Exporting image tile: " + outp_path)
+            images.append(outp)
+            n += 1
+    return [img for img in images]
+
+def upscale_image(img, rows, cols,seed,prompt,negative_prompt,xformers,cpu_offload,attention_slicing,enable_custom_sliders=False,guidance=7,iterations=50):
+    model_id = "stabilityai/stable-diffusion-x4-upscaler"
+    try:
+        pipeline = StableDiffusionUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+    except:
+        pipeline = StableDiffusionUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16, local_files_only=True)
+    pipeline = pipeline.to("cuda")
+    if xformers:
+        pipeline.enable_xformers_memory_efficient_attention()
+    else:
+        pipeline.disable_xformers_memory_efficient_attention()
+    if cpu_offload:
+        try:
+            pipeline.enable_sequential_cpu_offload()
+        except:
+            pass
+    if attention_slicing:
+        pipeline.enable_attention_slicing()
+    else:
+        pipeline.disable_attention_slicing()
+    img = Image.fromarray(img)
+    # load model and scheduler
+    if seed==-1:
+        generator = torch.manual_seed(random.randint(0, 9999999))
+    else:
+        generator = torch.manual_seed(seed)
+    
+    original_width, original_height = img.size
+    max_dimension = max(original_width, original_height)
+    tiles = split_image(img, rows, cols, True, False)
+    ups_tiles = []
+    i = 0
+    for x in tiles:
+        i=i+1
+        if enable_custom_sliders:
+            ups_tile = pipeline(prompt=prompt,negative_prompt=negative_prompt,guidance_scale=guidance, num_inference_steps=iterations, image=x.convert("RGB"),generator=generator).images[0]
+        else:
+            ups_tile = pipeline(prompt=prompt,negative_prompt=negative_prompt, image=x.convert("RGB"),generator=generator).images[0]
+        ups_tiles.append(ups_tile)
+        
+    # Determine the size of the merged upscaled image
+    total_width = 0
+    total_height = 0
+    side = 0
+    for ups_tile in ups_tiles:
+        side = ups_tile.width
+        break
+    for x in tiles:
+        tsize = x.width
+        break
+
+    ups_times = abs(side/tsize)
+    new_size = (max_dimension * ups_times, max_dimension * ups_times)
+    total_width = cols*side
+    total_height = rows*side
+
+    # Create a blank image with the calculated size
+    merged_image = Image.new("RGB", (total_width, total_height))
+
+    # Paste each upscaled tile into the blank image
+    current_width = 0
+    current_height = 0
+    maximum_width = cols*side
+    for ups_tile in ups_tiles:
+        merged_image.paste(ups_tile, (current_width, current_height))
+        current_width += ups_tile.width
+        if current_width>=maximum_width:
+            current_width = 0
+            current_height = current_height+side
+
+    # Using the center of the image as pivot, crop the image to the original dimension times four
+    crop_left = (new_size[0] - original_width * ups_times) // 2
+    crop_upper = (new_size[1] - original_height * ups_times) // 2
+    crop_right = crop_left + original_width * ups_times
+    crop_lower = crop_upper + original_height * ups_times
+    final_img = merged_image.crop((crop_left, crop_upper, crop_right, crop_lower))
+
+    # The resulting image should be identical to the original image in proportions / aspect ratio, with no loss of elements.
+    # Save the merged image
+    return final_img