Spaces:
Runtime error
Runtime error
File size: 4,924 Bytes
d692ff9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
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
|