2022-09-13 05:44:08 +08:00
import datetime
2022-09-03 17:08:45 +08:00
import math
import os
from collections import namedtuple
import re
import numpy as np
2022-09-12 19:40:02 +08:00
import piexif
import piexif . helper
2022-09-03 17:08:45 +08:00
from PIL import Image , ImageFont , ImageDraw , PngImagePlugin
2022-09-14 00:53:42 +08:00
from fonts . ttf import Roboto
2022-09-10 13:45:16 +08:00
import string
2022-09-03 17:08:45 +08:00
2022-09-13 01:47:46 +08:00
from modules import sd_samplers , shared
2022-09-15 09:05:00 +08:00
from modules . shared import opts , cmd_opts
2022-09-03 17:08:45 +08:00
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
2022-09-04 06:29:43 +08:00
non_overlap_width = tile_w - overlap
non_overlap_height = tile_h - overlap
2022-09-03 17:08:45 +08:00
2022-09-04 06:29:43 +08:00
cols = math . ceil ( ( w - overlap ) / non_overlap_width )
rows = math . ceil ( ( h - overlap ) / non_overlap_height )
2022-09-30 06:46:23 +08:00
dx = ( w - tile_w ) / ( cols - 1 ) if cols > 1 else 0
dy = ( h - tile_h ) / ( rows - 1 ) if rows > 1 else 0
2022-09-03 17:08:45 +08:00
grid = Grid ( [ ] , tile_w , tile_h , w , h , overlap )
for row in range ( rows ) :
row_images = [ ]
2022-09-04 23:54:12 +08:00
y = int ( row * dy )
2022-09-03 17:08:45 +08:00
if y + tile_h > = h :
y = h - tile_h
for col in range ( cols ) :
2022-09-04 23:54:12 +08:00
x = int ( col * dx )
2022-09-03 17:08:45 +08:00
2022-09-30 06:46:23 +08:00
if x + tile_w > = w :
2022-09-03 17:08:45 +08:00
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 ' )
2022-09-30 16:42:40 +08:00
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 ) )
2022-09-03 17:08:45 +08:00
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 ) :
2022-09-30 16:42:40 +08:00
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 " )
2022-09-03 17:08:45 +08:00
if not line . is_active :
2022-09-30 16:42:40 +08:00
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 )
2022-09-03 17:08:45 +08:00
draw_y + = line . size [ 1 ] + line_spacing
fontsize = ( width + height ) / / 25
line_spacing = fontsize / / 2
2022-09-13 00:17:02 +08:00
try :
fnt = ImageFont . truetype ( opts . font or Roboto , fontsize )
except Exception :
fnt = ImageFont . truetype ( Roboto , fontsize )
2022-09-03 17:08:45 +08:00
color_active = ( 0 , 0 , 0 )
color_inactive = ( 153 , 153 , 153 )
2022-09-09 22:54:04 +08:00
pad_left = 0 if sum ( [ sum ( [ len ( line . text ) for line in lines ] ) for lines in ver_texts ] ) == 0 else width * 3 / / 4
2022-09-03 17:08:45 +08:00
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 ]
2022-09-30 06:46:23 +08:00
ver_text_heights = [ sum ( [ line . size [ 1 ] + line_spacing for line in lines ] ) - line_spacing * len ( lines ) for lines in
ver_texts ]
2022-09-03 17:08:45 +08:00
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 : ]
2022-09-30 16:42:40 +08:00
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 ) ) ]
2022-09-03 17:08:45 +08:00
return draw_grid_annotations ( im , width , height , hor_texts , ver_texts )
def resize_image ( resize_mode , im , width , height ) :
2022-09-23 22:37:47 +08:00
def resize ( im , w , h ) :
2022-09-24 04:29:53 +08:00
if opts . upscaler_for_img2img is None or opts . upscaler_for_img2img == " None " or im . mode == ' L ' :
2022-09-23 22:37:47 +08:00
return im . resize ( ( w , h ) , resample = LANCZOS )
2022-09-30 15:38:48 +08:00
scale = max ( w / im . width , h / im . height )
2022-09-30 19:23:41 +08:00
if scale > 1.0 :
upscalers = [ x for x in shared . sd_upscalers if x . name == opts . upscaler_for_img2img ]
assert len ( upscalers ) > 0 , f " could not find upscaler named { opts . upscaler_for_img2img } "
upscaler = upscalers [ 0 ]
im = upscaler . scaler . upscale ( im , scale , upscaler . data_path )
if im . width != w or im . height != h :
im = im . resize ( ( w , h ) , resample = LANCZOS )
2022-09-30 15:38:48 +08:00
2022-09-30 19:23:41 +08:00
return im
2022-09-23 22:37:47 +08:00
2022-09-03 17:08:45 +08:00
if resize_mode == 0 :
2022-09-23 22:37:47 +08:00
res = resize ( im , width , height )
2022-09-03 17:08:45 +08:00
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
2022-09-23 22:37:47 +08:00
resized = resize ( im , src_w , src_h )
2022-09-03 17:08:45 +08:00
res = Image . new ( " RGB " , ( width , height ) )
res . paste ( resized , box = ( width / / 2 - src_w / / 2 , height / / 2 - src_h / / 2 ) )
2022-09-23 22:37:47 +08:00
2022-09-03 17:08:45 +08:00
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
2022-09-23 22:37:47 +08:00
resized = resize ( im , src_w , src_h )
2022-09-03 17:08:45 +08:00
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 ) )
2022-09-30 16:42:40 +08:00
res . paste ( resized . resize ( ( width , fill_height ) , box = ( 0 , resized . height , width , resized . height ) ) , box = ( 0 , fill_height + src_h ) )
2022-09-03 17:08:45 +08:00
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 ) )
2022-09-30 16:42:40 +08:00
res . paste ( resized . resize ( ( fill_width , height ) , box = ( resized . width , 0 , resized . width , height ) ) , box = ( fill_width + src_w , 0 ) )
2022-09-03 17:08:45 +08:00
return res
invalid_filename_chars = ' <>: " / \\ |?* \n '
2022-09-20 14:01:32 +08:00
invalid_filename_prefix = ' '
invalid_filename_postfix = ' . '
2022-09-30 06:46:23 +08:00
re_nonletters = re . compile ( r ' [ \ s ' + string . punctuation + ' ]+ ' )
2022-09-20 14:01:32 +08:00
max_filename_part_length = 128
2022-09-03 17:08:45 +08:00
2022-09-12 20:41:30 +08:00
def sanitize_filename_part ( text , replace_spaces = True ) :
if replace_spaces :
text = text . replace ( ' ' , ' _ ' )
2022-09-03 17:08:45 +08:00
2022-09-20 14:01:32 +08:00
text = text . translate ( { ord ( x ) : ' _ ' for x in invalid_filename_chars } )
text = text . lstrip ( invalid_filename_prefix ) [ : max_filename_part_length ]
text = text . rstrip ( invalid_filename_postfix )
return text
2022-09-03 17:08:45 +08:00
2022-09-12 20:41:30 +08:00
2022-09-13 05:44:08 +08:00
def apply_filename_pattern ( x , p , seed , prompt ) :
2022-09-23 00:47:43 +08:00
max_prompt_words = opts . directories_max_prompt_words
2022-09-13 05:44:08 +08:00
if seed is not None :
x = x . replace ( " [seed] " , str ( seed ) )
2022-09-20 14:01:32 +08:00
2022-09-13 05:44:08 +08:00
if prompt is not None :
2022-09-20 14:01:32 +08:00
x = x . replace ( " [prompt] " , sanitize_filename_part ( prompt ) )
2022-09-30 11:01:32 +08:00
if " [prompt_no_styles] " in x :
prompt_no_style = prompt
for style in shared . prompt_styles . get_style_prompts ( p . styles ) :
2022-09-30 15:37:18 +08:00
if len ( style ) > 0 :
style_parts = [ y for y in style . split ( " {prompt} " ) ]
for part in style_parts :
prompt_no_style = prompt_no_style . replace ( part , " " ) . replace ( " , , " , " , " ) . strip ( ) . strip ( ' , ' )
prompt_no_style = prompt_no_style . replace ( style , " " ) . strip ( ) . strip ( ' , ' ) . strip ( )
2022-09-30 11:01:32 +08:00
x = x . replace ( " [prompt_no_styles] " , sanitize_filename_part ( prompt_no_style , replace_spaces = False ) )
2022-09-20 14:01:32 +08:00
x = x . replace ( " [prompt_spaces] " , sanitize_filename_part ( prompt , replace_spaces = False ) )
2022-09-13 05:44:08 +08:00
if " [prompt_words] " in x :
words = [ x for x in re_nonletters . split ( prompt or " " ) if len ( x ) > 0 ]
if len ( words ) == 0 :
words = [ " empty " ]
2022-09-30 16:42:40 +08:00
x = x . replace ( " [prompt_words] " , sanitize_filename_part ( " " . join ( words [ 0 : max_prompt_words ] ) , replace_spaces = False ) )
2022-09-13 05:44:08 +08:00
if p is not None :
x = x . replace ( " [steps] " , str ( p . steps ) )
x = x . replace ( " [cfg] " , str ( p . cfg_scale ) )
x = x . replace ( " [width] " , str ( p . width ) )
x = x . replace ( " [height] " , str ( p . height ) )
2022-09-30 11:01:32 +08:00
x = x . replace ( " [styles] " , sanitize_filename_part ( " , " . join ( [ x for x in p . styles if not x == " None " ] ) , replace_spaces = False ) )
2022-09-30 16:42:40 +08:00
x = x . replace ( " [sampler] " , sanitize_filename_part ( sd_samplers . samplers [ p . sampler_index ] . name , replace_spaces = False ) )
2022-09-13 05:44:08 +08:00
2022-09-17 17:05:04 +08:00
x = x . replace ( " [model_hash] " , shared . sd_model . sd_model_hash )
2022-09-13 05:44:08 +08:00
x = x . replace ( " [date] " , datetime . date . today ( ) . isoformat ( ) )
2022-09-26 14:44:52 +08:00
x = x . replace ( " [datetime] " , datetime . datetime . now ( ) . strftime ( " % Y % m %d % H % M % S " ) )
2022-09-25 20:45:20 +08:00
x = x . replace ( " [job_timestamp] " , shared . state . job_timestamp )
2022-09-13 05:44:08 +08:00
2022-09-15 09:05:00 +08:00
if cmd_opts . hide_ui_dir_config :
x = re . sub ( r ' ^[ \\ /]+| \ . { 2,}[ \\ /]+|[ \\ /]+ \ . { 2,} ' , ' ' , x )
2022-09-13 05:44:08 +08:00
return x
2022-09-30 06:46:23 +08:00
2022-09-14 20:40:16 +08:00
def get_next_sequence_number ( path , basename ) :
2022-09-13 22:43:08 +08:00
"""
Determines and returns the next sequence number to use when saving an image in the specified directory .
The sequence starts at 0.
"""
result = - 1
2022-09-14 20:40:16 +08:00
if basename != ' ' :
basename = basename + " - "
prefix_length = len ( basename )
2022-09-13 22:43:08 +08:00
for p in os . listdir ( path ) :
2022-09-14 20:40:16 +08:00
if p . startswith ( basename ) :
2022-09-30 16:42:40 +08:00
l = os . path . splitext ( p [ prefix_length : ] ) [ 0 ] . split ( ' - ' ) # splits the filename (removing the basename first if one is defined, so the sequence number is always the first element)
2022-09-14 20:40:16 +08:00
try :
2022-09-13 23:46:05 +08:00
result = max ( int ( l [ 0 ] ) , result )
2022-09-14 20:40:16 +08:00
except ValueError :
2022-09-13 23:46:05 +08:00
pass
2022-09-13 22:43:08 +08:00
return result + 1
2022-09-13 05:44:08 +08:00
2022-09-30 06:46:23 +08:00
2022-09-30 16:42:40 +08:00
def save_image ( image , path , basename , seed = None , prompt = None , extension = ' png ' , info = None , short_filename = False , no_prompt = False , grid = False , pnginfo_section_name = ' parameters ' , p = None , existing_info = None , forced_filename = None , suffix = " " ) :
2022-09-03 17:08:45 +08:00
if short_filename or prompt is None or seed is None :
file_decoration = " "
elif opts . save_to_dirs :
2022-09-13 05:44:08 +08:00
file_decoration = opts . samples_filename_pattern or " [seed] "
2022-09-03 17:08:45 +08:00
else :
2022-09-13 05:44:08 +08:00
file_decoration = opts . samples_filename_pattern or " [seed]-[prompt_spaces] "
2022-09-11 23:17:13 +08:00
2022-09-12 23:59:53 +08:00
if file_decoration != " " :
file_decoration = " - " + file_decoration . lower ( )
2022-09-22 18:54:50 +08:00
file_decoration = apply_filename_pattern ( file_decoration , p , seed , prompt ) + suffix
2022-09-13 01:47:46 +08:00
2022-09-03 17:08:45 +08:00
if extension == ' png ' and opts . enable_pnginfo and info is not None :
pnginfo = PngImagePlugin . PngInfo ( )
2022-09-12 23:59:53 +08:00
if existing_info is not None :
for k , v in existing_info . items ( ) :
2022-09-13 12:34:35 +08:00
pnginfo . add_text ( k , str ( v ) )
2022-09-12 23:59:53 +08:00
2022-09-11 16:31:16 +08:00
pnginfo . add_text ( pnginfo_section_name , info )
2022-09-03 17:08:45 +08:00
else :
pnginfo = None
2022-09-14 04:28:03 +08:00
save_to_dirs = ( grid and opts . grid_save_to_dirs ) or ( not grid and opts . save_to_dirs and not no_prompt )
2022-09-10 18:36:16 +08:00
2022-09-13 05:44:08 +08:00
if save_to_dirs :
dirname = apply_filename_pattern ( opts . directories_filename_pattern or " [prompt_words] " , p , seed , prompt )
2022-09-03 17:08:45 +08:00
path = os . path . join ( path , dirname )
os . makedirs ( path , exist_ok = True )
2022-09-20 06:13:12 +08:00
if forced_filename is None :
basecount = get_next_sequence_number ( path , basename )
fullfn = " a.png "
fullfn_without_extension = " a "
for i in range ( 500 ) :
2022-09-30 06:46:23 +08:00
fn = f " { basecount + i : 05 } " if basename == ' ' else f " { basename } - { basecount + i : 04 } "
2022-09-20 06:13:12 +08:00
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
else :
fullfn = os . path . join ( path , f " { forced_filename } . { extension } " )
fullfn_without_extension = os . path . join ( path , forced_filename )
2022-09-03 17:08:45 +08:00
2022-09-17 13:32:15 +08:00
def exif_bytes ( ) :
2022-09-15 19:54:29 +08:00
return piexif . dump ( {
2022-09-14 00:23:55 +08:00
" Exif " : {
2022-09-15 19:54:29 +08:00
piexif . ExifIFD . UserComment : piexif . helper . UserComment . dump ( info or " " , encoding = " unicode " )
2022-09-14 18:38:40 +08:00
} ,
2022-09-14 00:23:55 +08:00
} )
2022-09-15 19:54:29 +08:00
if extension . lower ( ) in ( " jpg " , " jpeg " , " webp " ) :
2022-09-17 13:32:15 +08:00
image . save ( fullfn , quality = opts . jpeg_quality )
if opts . enable_pnginfo and info is not None :
piexif . insert ( exif_bytes ( ) , fullfn )
2022-09-15 19:54:29 +08:00
else :
image . save ( fullfn , quality = opts . jpeg_quality , pnginfo = pnginfo )
2022-09-03 17:08:45 +08:00
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 )
2022-09-17 13:32:15 +08:00
image . save ( fullfn_without_extension + " .jpg " , quality = opts . jpeg_quality )
if opts . enable_pnginfo and info is not None :
2022-09-17 20:39:20 +08:00
piexif . insert ( exif_bytes ( ) , fullfn_without_extension + " .jpg " )
2022-09-03 17:08:45 +08:00
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 " )
2022-09-04 23:54:12 +08:00