img2tiles mk2.py

import math

import modules.scripts as scripts
import gradio as gr
from PIL import Image, ImageDraw

from modules import processing, images, devices
from modules.processing import Processed
from modules.shared import state

import random
from collections import namedtuple
from PIL import Image, ImageDraw

Tile = namedtuple('Tile', ['x_pos', 'y_pos', 'width', 'height'])
ImageTile = namedtuple('ImageTile', ['tile', 'image'])

#Check if tile is vertical or horizontal or square. returns 0, 1 or 2
def check_tile_type(tile) -> int:
    if tile.width == tile.height:
        return 'square'
    elif tile.width > tile.height:
        return 'vertical'
    else:
        return 'horizontal'

#split tile into two separated tiles. The separation should be across the long side. If tile is a square, choose side randomly Each side of them must be divisible by divider and be bigger or equal than min_tile_size. 
def split_tile(tile, divider, min_tile_size):
    if check_tile_type(tile) == 'square':
        if random.random() > 0.5:
            return split_tile_vertical(tile, divider, min_tile_size)
        else:
            return split_tile_horizontal(tile, divider, min_tile_size)
    elif check_tile_type(tile) == 'vertical':
        return split_tile_vertical(tile, divider, min_tile_size)
    else:
        return split_tile_horizontal(tile, divider, min_tile_size)

#split tile into two separated tiles. The separation should be across the long side. Each side of them must be divisible by divider and be bigger or equal than min_tile_size. 
def split_tile_vertical(tile, divider, min_tile_size):
    if tile.width < min_tile_size * 2:
        return None
    else:
        split_point = random.randint(min_tile_size, tile.width - min_tile_size)
        if split_point % divider != 0:
            split_point = split_point - split_point % divider
        return [Tile(tile.x_pos, tile.y_pos, split_point, tile.height), Tile(tile.x_pos + split_point, tile.y_pos, tile.width - split_point, tile.height)]

#split tile into two separated tiles. The separation should be across the long side. Each side of them must be divisible by divider and be bigger or equal than min_tile_size. 
def split_tile_horizontal(tile, divider, min_tile_size):
    if tile.height < min_tile_size * 2:
        return None
    else:
        split_point = random.randint(min_tile_size, tile.height - min_tile_size)
        if split_point % divider != 0:
            split_point = split_point - split_point % divider
        return [Tile(tile.x_pos, tile.y_pos, tile.width, split_point), Tile(tile.x_pos, tile.y_pos + split_point, tile.width, tile.height - split_point)]

#sorts the tiles in the list by perimeter length
sort_func = lambda x: x.width + x.height

#apply split_tile to the biggest tile in the list recursively
def split_into_tiles(init_tile, divider, min_tile_size):
    assert min_tile_size >= divider
    tiles = [init_tile]
    while True:
        tiles.sort(key=sort_func, reverse=True)
        new_tiles = split_tile(tiles[0], divider, min_tile_size)
        if new_tiles is None:
            break
        tiles.pop(0)
        tiles.extend(new_tiles)
    return tiles

#split PIL image into tiles
def split_image(image, divider, min_tile_size):
    tiles = split_into_tiles(Tile(0, 0, image.width, image.height), divider, min_tile_size)
    image_tiles = []
    for tile in tiles:
        image_tiles.append(ImageTile(tile, image.crop((tile.x_pos, tile.y_pos, tile.x_pos + tile.width, tile.y_pos + tile.height))))
    return image_tiles


#Draw a border of a certain color and width on each ImageTile from the list
def draw_borders(image_tiles, border_width, border_color):
    for image_tile in image_tiles:
        draw = ImageDraw.Draw(image_tile.image)
        draw.rectangle([(0, 0), (image_tile.image.width - 1, image_tile.image.height - 1)], outline=border_color, width=border_width)

#merge tiles from the list into one image
def merge_tiles(image_tiles, image_size):
    image = Image.new('RGB', image_size, color='black')
    for image_tile in image_tiles:
        image.paste(image_tile.image, (image_tile.tile.x_pos, image_tile.tile.y_pos))
    return image

#split image into tiles, draw borders on them and merge them back
def split_draw_borders_and_merge(image, divider, min_tile_size, border_width, border_color):
    image_tiles = split_image(image, divider, min_tile_size)
    draw_borders(image_tiles, border_width, border_color)
    img = merge_tiles(image_tiles, image.size)
    return img

class Script(scripts.Script):
    def title(self):
        return "img2tiles mk2"

    def show(self, is_img2img):
        return is_img2img

    def ui(self, is_img2img):
        tile_size = gr.Slider(minimum=32, maximum=256, step=16,
                              label='Minimal tile size', value=128, visible=True)
        use_random_seeds = gr.Checkbox(value=True, label='Use -1 for seeds', visible=True)
        save_tiles = gr.Checkbox(value=False, label='Save separate tiles', visible=True)

        tile_border_width = gr.Slider(minimum=0, maximum=256, step=1,
                                        label='Tile border width', value=0, visible=True)
        tile_border_color = gr.ColorPicker(label='Tile border color', visible=True)
        return [tile_size, tile_border_width, tile_border_color, use_random_seeds, save_tiles]


    def run(self, p, tile_size, tile_border_width, tile_border_color, use_random_seeds, save_tiles):
        p.do_not_save_samples = not save_tiles

        if use_random_seeds:
            p.seed = -1
        else:
            processing.fix_seed(p)

        if p.seed != -1:
            random.seed(p.seed)

        initial_info = None
        seed = p.seed

        minimal_res_size = p.width
        tile_size_ratio = minimal_res_size // tile_size
        print("img2tiles tile size ratio is ", tile_size_ratio)
        divider = 64 // tile_size_ratio
        #divider = min_tile_size // minimal_res_size // 64

        init_img = p.init_images[0]
        img = init_img

        devices.torch_gc()

        all_tiles = split_image(img, divider, tile_size)

        p.batch_size = 1
        batch_count = math.ceil(len(all_tiles))
        state.job_count = batch_count

        print(
            f"img2tiles will process a total of {len(all_tiles)} tile images in a total of {state.job_count} batches.")

        result_images = []

        for i in range(batch_count):
            p.init_images = [all_tiles[i].image]
            p.width = all_tiles[i].image.width * tile_size_ratio
            p.height = all_tiles[i].image.height * tile_size_ratio
            state.job = f"Batch {i + 1 * batch_count} out of {state.job_count}"
            try:
                processed = processing.process_images(p)
                if initial_info is None:
                    initial_info = processed.info
                all_tiles[i] = ImageTile(
                    Tile(
                        all_tiles[i].tile.x_pos * tile_size_ratio,
                        all_tiles[i].tile.y_pos * tile_size_ratio,
                        all_tiles[i].tile.width * tile_size_ratio,
                        all_tiles[i].tile.height * tile_size_ratio
                    ),
                    processed.images[0])
            except Exception as e:
                print(e)
                print(all_tiles[i].tile)
                all_tiles[i] = ImageTile(
                Tile(
                    all_tiles[i].tile.x_pos * tile_size_ratio,
                    all_tiles[i].tile.y_pos * tile_size_ratio,
                    all_tiles[i].tile.width * tile_size_ratio,
                    all_tiles[i].tile.height * tile_size_ratio
                ),
                Image.new("RGB", (all_tiles[i].tile.width * tile_size_ratio, all_tiles[i].tile.height * tile_size_ratio), color=tile_border_color))


        #shape = (len(grid.tiles[0][2]), len(grid.tiles))
        #image_size = (p.width * shape[1], p.height * shape[0])
        #combined_image = Image.new('RGB', image_size)

        # for row in range(shape[0]):
        #     for col in range(shape[1]):
        #         offset = p.width * col, p.height * row
        #         idx = row * shape[1] + col
        #         w_res = draw_border(work_results[idx], tile_border_color, tile_border_width)
        #         combined_image.paste(w_res, offset)

        #combined_image = draw_border(combined_image, tile_border_color, tile_border_width)

    
        draw_borders(all_tiles, tile_border_width, tile_border_color)
        combined_image = merge_tiles(all_tiles, (init_img.width * tile_size_ratio, init_img.height * tile_size_ratio))

        result_images.append(combined_image)
        images.save_image(combined_image, 'outputs/img2img-grids', basename='grid')
        processed = Processed(p, result_images, seed, initial_info)

        return processed