mirror of
https://github.com/AUTOMATIC1111/stable-diffusion-webui.git
synced 2025-08-04 19:22:32 +00:00
split codebase into multiple files; to anyone this affects negatively: sorry
This commit is contained in:
AUTOMATIC
290
modules/images.py
Normal file
290
modules/images.py
Normal file
@@ -0,0 +1,290 @@
|
||||
import math
|
||||
import os
|
||||
from collections import namedtuple
|
||||
import re
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image, ImageFont, ImageDraw, PngImagePlugin
|
||||
|
||||
from modules.shared import opts
|
||||
|
||||
LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)
|
||||
|
||||
|
||||
def image_grid(imgs, batch_size=1, rows=None):
|
||||
if rows is None:
|
||||
if opts.n_rows > 0:
|
||||
rows = opts.n_rows
|
||||
elif opts.n_rows == 0:
|
||||
rows = batch_size
|
||||
else:
|
||||
rows = math.sqrt(len(imgs))
|
||||
rows = round(rows)
|
||||
|
||||
cols = math.ceil(len(imgs) / rows)
|
||||
|
||||
w, h = imgs[0].size
|
||||
grid = Image.new('RGB', size=(cols * w, rows * h), color='black')
|
||||
|
||||
for i, img in enumerate(imgs):
|
||||
grid.paste(img, box=(i % cols * w, i // cols * h))
|
||||
|
||||
return grid
|
||||
|
||||
|
||||
Grid = namedtuple("Grid", ["tiles", "tile_w", "tile_h", "image_w", "image_h", "overlap"])
|
||||
|
||||
|
||||
def split_grid(image, tile_w=512, tile_h=512, overlap=64):
|
||||
w = image.width
|
||||
h = image.height
|
||||
|
||||
now = tile_w - overlap # non-overlap width
|
||||
noh = tile_h - overlap
|
||||
|
||||
cols = math.ceil((w - overlap) / now)
|
||||
rows = math.ceil((h - overlap) / noh)
|
||||
|
||||
grid = Grid([], tile_w, tile_h, w, h, overlap)
|
||||
for row in range(rows):
|
||||
row_images = []
|
||||
|
||||
y = row * noh
|
||||
|
||||
if y + tile_h >= h:
|
||||
y = h - tile_h
|
||||
|
||||
for col in range(cols):
|
||||
x = col * now
|
||||
|
||||
if x+tile_w >= w:
|
||||
x = w - tile_w
|
||||
|
||||
tile = image.crop((x, y, x + tile_w, y + tile_h))
|
||||
|
||||
row_images.append([x, tile_w, tile])
|
||||
|
||||
grid.tiles.append([y, tile_h, row_images])
|
||||
|
||||
return grid
|
||||
|
||||
|
||||
def combine_grid(grid):
|
||||
def make_mask_image(r):
|
||||
r = r * 255 / grid.overlap
|
||||
r = r.astype(np.uint8)
|
||||
return Image.fromarray(r, 'L')
|
||||
|
||||
mask_w = make_mask_image(np.arange(grid.overlap, dtype=np.float32).reshape((1, grid.overlap)).repeat(grid.tile_h, axis=0))
|
||||
mask_h = make_mask_image(np.arange(grid.overlap, dtype=np.float32).reshape((grid.overlap, 1)).repeat(grid.image_w, axis=1))
|
||||
|
||||
combined_image = Image.new("RGB", (grid.image_w, grid.image_h))
|
||||
for y, h, row in grid.tiles:
|
||||
combined_row = Image.new("RGB", (grid.image_w, h))
|
||||
for x, w, tile in row:
|
||||
if x == 0:
|
||||
combined_row.paste(tile, (0, 0))
|
||||
continue
|
||||
|
||||
combined_row.paste(tile.crop((0, 0, grid.overlap, h)), (x, 0), mask=mask_w)
|
||||
combined_row.paste(tile.crop((grid.overlap, 0, w, h)), (x + grid.overlap, 0))
|
||||
|
||||
if y == 0:
|
||||
combined_image.paste(combined_row, (0, 0))
|
||||
continue
|
||||
|
||||
combined_image.paste(combined_row.crop((0, 0, combined_row.width, grid.overlap)), (0, y), mask=mask_h)
|
||||
combined_image.paste(combined_row.crop((0, grid.overlap, combined_row.width, h)), (0, y + grid.overlap))
|
||||
|
||||
return combined_image
|
||||
|
||||
|
||||
class GridAnnotation:
|
||||
def __init__(self, text='', is_active=True):
|
||||
self.text = text
|
||||
self.is_active = is_active
|
||||
self.size = None
|
||||
|
||||
|
||||
def draw_grid_annotations(im, width, height, hor_texts, ver_texts):
|
||||
def wrap(drawing, text, font, line_length):
|
||||
lines = ['']
|
||||
for word in text.split():
|
||||
line = f'{lines[-1]} {word}'.strip()
|
||||
if drawing.textlength(line, font=font) <= line_length:
|
||||
lines[-1] = line
|
||||
else:
|
||||
lines.append(word)
|
||||
return lines
|
||||
|
||||
def draw_texts(drawing, draw_x, draw_y, lines):
|
||||
for i, line in enumerate(lines):
|
||||
drawing.multiline_text((draw_x, draw_y + line.size[1] / 2), line.text, font=fnt, fill=color_active if line.is_active else color_inactive, anchor="mm", align="center")
|
||||
|
||||
if not line.is_active:
|
||||
drawing.line((draw_x - line.size[0]//2, draw_y + line.size[1]//2, draw_x + line.size[0]//2, draw_y + line.size[1]//2), fill=color_inactive, width=4)
|
||||
|
||||
draw_y += line.size[1] + line_spacing
|
||||
|
||||
fontsize = (width + height) // 25
|
||||
line_spacing = fontsize // 2
|
||||
fnt = ImageFont.truetype(opts.font, fontsize)
|
||||
color_active = (0, 0, 0)
|
||||
color_inactive = (153, 153, 153)
|
||||
|
||||
pad_left = width * 3 // 4 if len(ver_texts) > 0 else 0
|
||||
|
||||
cols = im.width // width
|
||||
rows = im.height // height
|
||||
|
||||
assert cols == len(hor_texts), f'bad number of horizontal texts: {len(hor_texts)}; must be {cols}'
|
||||
assert rows == len(ver_texts), f'bad number of vertical texts: {len(ver_texts)}; must be {rows}'
|
||||
|
||||
calc_img = Image.new("RGB", (1, 1), "white")
|
||||
calc_d = ImageDraw.Draw(calc_img)
|
||||
|
||||
for texts, allowed_width in zip(hor_texts + ver_texts, [width] * len(hor_texts) + [pad_left] * len(ver_texts)):
|
||||
items = [] + texts
|
||||
texts.clear()
|
||||
|
||||
for line in items:
|
||||
wrapped = wrap(calc_d, line.text, fnt, allowed_width)
|
||||
texts += [GridAnnotation(x, line.is_active) for x in wrapped]
|
||||
|
||||
for line in texts:
|
||||
bbox = calc_d.multiline_textbbox((0, 0), line.text, font=fnt)
|
||||
line.size = (bbox[2] - bbox[0], bbox[3] - bbox[1])
|
||||
|
||||
hor_text_heights = [sum([line.size[1] + line_spacing for line in lines]) - line_spacing for lines in hor_texts]
|
||||
ver_text_heights = [sum([line.size[1] + line_spacing for line in lines]) - line_spacing * len(lines) for lines in ver_texts]
|
||||
|
||||
pad_top = max(hor_text_heights) + line_spacing * 2
|
||||
|
||||
result = Image.new("RGB", (im.width + pad_left, im.height + pad_top), "white")
|
||||
result.paste(im, (pad_left, pad_top))
|
||||
|
||||
d = ImageDraw.Draw(result)
|
||||
|
||||
for col in range(cols):
|
||||
x = pad_left + width * col + width / 2
|
||||
y = pad_top / 2 - hor_text_heights[col] / 2
|
||||
|
||||
draw_texts(d, x, y, hor_texts[col])
|
||||
|
||||
for row in range(rows):
|
||||
x = pad_left / 2
|
||||
y = pad_top + height * row + height / 2 - ver_text_heights[row] / 2
|
||||
|
||||
draw_texts(d, x, y, ver_texts[row])
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def draw_prompt_matrix(im, width, height, all_prompts):
|
||||
prompts = all_prompts[1:]
|
||||
boundary = math.ceil(len(prompts) / 2)
|
||||
|
||||
prompts_horiz = prompts[:boundary]
|
||||
prompts_vert = prompts[boundary:]
|
||||
|
||||
hor_texts = [[GridAnnotation(x, is_active=pos & (1 << i) != 0) for i, x in enumerate(prompts_horiz)] for pos in range(1 << len(prompts_horiz))]
|
||||
ver_texts = [[GridAnnotation(x, is_active=pos & (1 << i) != 0) for i, x in enumerate(prompts_vert)] for pos in range(1 << len(prompts_vert))]
|
||||
|
||||
return draw_grid_annotations(im, width, height, hor_texts, ver_texts)
|
||||
|
||||
|
||||
def resize_image(resize_mode, im, width, height):
|
||||
if resize_mode == 0:
|
||||
res = im.resize((width, height), resample=LANCZOS)
|
||||
elif resize_mode == 1:
|
||||
ratio = width / height
|
||||
src_ratio = im.width / im.height
|
||||
|
||||
src_w = width if ratio > src_ratio else im.width * height // im.height
|
||||
src_h = height if ratio <= src_ratio else im.height * width // im.width
|
||||
|
||||
resized = im.resize((src_w, src_h), resample=LANCZOS)
|
||||
res = Image.new("RGB", (width, height))
|
||||
res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2))
|
||||
else:
|
||||
ratio = width / height
|
||||
src_ratio = im.width / im.height
|
||||
|
||||
src_w = width if ratio < src_ratio else im.width * height // im.height
|
||||
src_h = height if ratio >= src_ratio else im.height * width // im.width
|
||||
|
||||
resized = im.resize((src_w, src_h), resample=LANCZOS)
|
||||
res = Image.new("RGB", (width, height))
|
||||
res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2))
|
||||
|
||||
if ratio < src_ratio:
|
||||
fill_height = height // 2 - src_h // 2
|
||||
res.paste(resized.resize((width, fill_height), box=(0, 0, width, 0)), box=(0, 0))
|
||||
res.paste(resized.resize((width, fill_height), box=(0, resized.height, width, resized.height)), box=(0, fill_height + src_h))
|
||||
elif ratio > src_ratio:
|
||||
fill_width = width // 2 - src_w // 2
|
||||
res.paste(resized.resize((fill_width, height), box=(0, 0, 0, height)), box=(0, 0))
|
||||
res.paste(resized.resize((fill_width, height), box=(resized.width, 0, resized.width, height)), box=(fill_width + src_w, 0))
|
||||
|
||||
return res
|
||||
|
||||
|
||||
invalid_filename_chars = '<>:"/\\|?*\n'
|
||||
|
||||
|
||||
def sanitize_filename_part(text):
|
||||
return text.replace(' ', '_').translate({ord(x): '' for x in invalid_filename_chars})[:128]
|
||||
|
||||
|
||||
def save_image(image, path, basename, seed=None, prompt=None, extension='png', info=None, short_filename=False, no_prompt=False):
|
||||
if short_filename or prompt is None or seed is None:
|
||||
file_decoration = ""
|
||||
elif opts.save_to_dirs:
|
||||
file_decoration = f"-{seed}"
|
||||
else:
|
||||
file_decoration = f"-{seed}-{sanitize_filename_part(prompt)[:128]}"
|
||||
|
||||
if extension == 'png' and opts.enable_pnginfo and info is not None:
|
||||
pnginfo = PngImagePlugin.PngInfo()
|
||||
pnginfo.add_text("parameters", info)
|
||||
else:
|
||||
pnginfo = None
|
||||
|
||||
if opts.save_to_dirs and not no_prompt:
|
||||
words = re.findall(r'\w+', prompt or "")
|
||||
if len(words) == 0:
|
||||
words = ["empty"]
|
||||
|
||||
dirname = " ".join(words[0:opts.save_to_dirs_prompt_len])
|
||||
path = os.path.join(path, dirname)
|
||||
|
||||
os.makedirs(path, exist_ok=True)
|
||||
|
||||
filecount = len([x for x in os.listdir(path) if os.path.splitext(x)[1] == '.' + extension])
|
||||
fullfn = "a.png"
|
||||
fullfn_without_extension = "a"
|
||||
for i in range(100):
|
||||
fn = f"{filecount:05}" if basename == '' else f"{basename}-{filecount:04}"
|
||||
fullfn = os.path.join(path, f"{fn}{file_decoration}.{extension}")
|
||||
fullfn_without_extension = os.path.join(path, f"{fn}{file_decoration}")
|
||||
if not os.path.exists(fullfn):
|
||||
break
|
||||
|
||||
image.save(fullfn, quality=opts.jpeg_quality, pnginfo=pnginfo)
|
||||
|
||||
target_side_length = 4000
|
||||
oversize = image.width > target_side_length or image.height > target_side_length
|
||||
if opts.export_for_4chan and (oversize or os.stat(fullfn).st_size > 4 * 1024 * 1024):
|
||||
ratio = image.width / image.height
|
||||
|
||||
if oversize and ratio > 1:
|
||||
image = image.resize((target_side_length, image.height * target_side_length // image.width), LANCZOS)
|
||||
elif oversize:
|
||||
image = image.resize((image.width * target_side_length // image.height, target_side_length), LANCZOS)
|
||||
|
||||
image.save(f"{fullfn_without_extension}.jpg", quality=opts.jpeg_quality, pnginfo=pnginfo)
|
||||
|
||||
if opts.save_txt and info is not None:
|
||||
with open(f"{fullfn_without_extension}.txt", "w", encoding="utf8") as file:
|
||||
file.write(info + "\n")
|
||||
|
||||
Reference in New Issue
Block a user