2022-09-12 15:44:08 -06:00
import datetime
2022-10-24 05:03:58 -06:00
import sys
import traceback
2022-10-22 06:11:15 -06:00
import pytz
2022-10-14 09:15:03 -06:00
import io
2022-09-03 03:08:45 -06:00
import math
import os
from collections import namedtuple
import re
import numpy as np
2022-09-12 05:40:02 -06:00
import piexif
import piexif . helper
2022-09-03 03:08:45 -06:00
from PIL import Image , ImageFont , ImageDraw , PngImagePlugin
2022-09-09 23:45:16 -06:00
import string
2022-11-27 06:28:32 -07:00
import json
2023-02-04 23:48:51 -07:00
import hashlib
2022-09-03 03:08:45 -06:00
2023-02-19 00:12:45 -07:00
from modules import sd_samplers , shared , script_callbacks , errors
2023-05-11 01:07:01 -06:00
from modules . paths_internal import roboto_ttf_file
2023-05-09 23:43:42 -06:00
from modules . shared import opts
2022-09-03 03:08:45 -06:00
LANCZOS = ( Image . Resampling . LANCZOS if hasattr ( Image , ' Resampling ' ) else Image . LANCZOS )
2023-05-11 01:06:19 -06:00
def get_font ( fontsize : int ) :
try :
2023-05-11 01:07:01 -06:00
return ImageFont . truetype ( opts . font or roboto_ttf_file , fontsize )
2023-05-11 01:06:19 -06:00
except Exception :
2023-05-11 01:07:01 -06:00
return ImageFont . truetype ( roboto_ttf_file , fontsize )
2023-05-11 01:06:19 -06:00
2022-09-03 03:08:45 -06:00
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
2022-10-14 09:07:24 -06:00
elif opts . grid_prevent_empty_spots :
2022-10-14 06:19:39 -06:00
rows = math . floor ( math . sqrt ( len ( imgs ) ) )
while len ( imgs ) % rows != 0 :
rows - = 1
2022-10-14 09:07:24 -06:00
else :
rows = math . sqrt ( len ( imgs ) )
rows = round ( rows )
2023-01-28 17:25:52 -07:00
if rows > len ( imgs ) :
rows = len ( imgs )
2022-09-03 03:08:45 -06:00
cols = math . ceil ( len ( imgs ) / rows )
2023-01-01 08:37:37 -07:00
params = script_callbacks . ImageGridLoopParams ( imgs , cols , rows )
script_callbacks . image_grid_callback ( params )
2022-09-03 03:08:45 -06:00
w , h = imgs [ 0 ] . size
2023-01-01 08:37:37 -07:00
grid = Image . new ( ' RGB ' , size = ( params . cols * w , params . rows * h ) , color = ' black ' )
2022-09-03 03:08:45 -06:00
2023-01-01 08:37:37 -07:00
for i , img in enumerate ( params . imgs ) :
grid . paste ( img , box = ( i % params . cols * w , i / / params . cols * h ) )
2022-09-03 03:08:45 -06:00
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-03 16:29:43 -06:00
non_overlap_width = tile_w - overlap
non_overlap_height = tile_h - overlap
2022-09-03 03:08:45 -06:00
2022-09-03 16:29:43 -06:00
cols = math . ceil ( ( w - overlap ) / non_overlap_width )
rows = math . ceil ( ( h - overlap ) / non_overlap_height )
2022-09-29 16:46:23 -06: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 03:08:45 -06:00
grid = Grid ( [ ] , tile_w , tile_h , w , h , overlap )
for row in range ( rows ) :
row_images = [ ]
2022-09-04 09:54:12 -06:00
y = int ( row * dy )
2022-09-03 03:08:45 -06:00
if y + tile_h > = h :
y = h - tile_h
for col in range ( cols ) :
2022-09-04 09:54:12 -06:00
x = int ( col * dx )
2022-09-03 03:08:45 -06:00
2022-09-29 16:46:23 -06:00
if x + tile_w > = w :
2022-09-03 03:08:45 -06: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 02:42:40 -06: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 03:08:45 -06: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
2023-02-04 03:28:53 -07:00
def draw_grid_annotations ( im , width , height , hor_texts , ver_texts , margin = 0 ) :
2022-09-03 03:08:45 -06:00
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
2022-12-16 22:45:43 -07:00
def draw_texts ( drawing , draw_x , draw_y , lines , initial_fnt , initial_fontsize ) :
2023-05-10 02:37:18 -06:00
for line in lines :
2022-12-16 22:45:43 -07:00
fnt = initial_fnt
fontsize = initial_fontsize
while drawing . multiline_textsize ( line . text , font = fnt ) [ 0 ] > line . allowed_width and fontsize > 0 :
fontsize - = 1
fnt = get_font ( fontsize )
2022-09-30 02:42:40 -06: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 03:08:45 -06:00
if not line . is_active :
2022-09-30 02:42:40 -06: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 03:08:45 -06:00
draw_y + = line . size [ 1 ] + line_spacing
fontsize = ( width + height ) / / 25
line_spacing = fontsize / / 2
2022-09-12 10:17:02 -06:00
2022-12-16 22:45:43 -07:00
fnt = get_font ( fontsize )
2022-09-12 10:17:02 -06:00
2022-09-03 03:08:45 -06:00
color_active = ( 0 , 0 , 0 )
color_inactive = ( 153 , 153 , 153 )
2022-09-09 08:54:04 -06: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 03:08:45 -06: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 ] )
2022-12-16 22:45:43 -07:00
line . allowed_width = allowed_width
2022-09-03 03:08:45 -06:00
hor_text_heights = [ sum ( [ line . size [ 1 ] + line_spacing for line in lines ] ) - line_spacing for lines in hor_texts ]
2023-02-04 03:28:53 -07: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 03:08:45 -06:00
2023-01-24 00:24:32 -07:00
pad_top = 0 if sum ( hor_text_heights ) == 0 else max ( hor_text_heights ) + line_spacing * 2
2022-09-03 03:08:45 -06:00
2023-02-05 01:44:56 -07:00
result = Image . new ( " RGB " , ( im . width + pad_left + margin * ( cols - 1 ) , im . height + pad_top + margin * ( rows - 1 ) ) , " white " )
2023-02-04 03:28:53 -07:00
for row in range ( rows ) :
for col in range ( cols ) :
cell = im . crop ( ( width * col , height * row , width * ( col + 1 ) , height * ( row + 1 ) ) )
result . paste ( cell , ( pad_left + ( width + margin ) * col , pad_top + ( height + margin ) * row ) )
2022-09-03 03:08:45 -06:00
d = ImageDraw . Draw ( result )
for col in range ( cols ) :
2023-02-04 03:28:53 -07:00
x = pad_left + ( width + margin ) * col + width / 2
2022-09-03 03:08:45 -06:00
y = pad_top / 2 - hor_text_heights [ col ] / 2
2022-12-16 22:45:43 -07:00
draw_texts ( d , x , y , hor_texts [ col ] , fnt , fontsize )
2022-09-03 03:08:45 -06:00
for row in range ( rows ) :
x = pad_left / 2
2023-02-04 03:28:53 -07:00
y = pad_top + ( height + margin ) * row + height / 2 - ver_text_heights [ row ] / 2
2022-09-03 03:08:45 -06:00
2022-12-16 22:45:43 -07:00
draw_texts ( d , x , y , ver_texts [ row ] , fnt , fontsize )
2022-09-03 03:08:45 -06:00
return result
2023-02-05 01:44:56 -07:00
def draw_prompt_matrix ( im , width , height , all_prompts , margin = 0 ) :
2022-09-03 03:08:45 -06:00
prompts = all_prompts [ 1 : ]
boundary = math . ceil ( len ( prompts ) / 2 )
prompts_horiz = prompts [ : boundary ]
prompts_vert = prompts [ boundary : ]
2022-09-30 02:42:40 -06: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 03:08:45 -06:00
2023-02-05 01:44:56 -07:00
return draw_grid_annotations ( im , width , height , hor_texts , ver_texts , margin )
2022-09-03 03:08:45 -06:00
2023-01-02 09:42:10 -07:00
def resize_image ( resize_mode , im , width , height , upscaler_name = None ) :
"""
Resizes an image with the specified resize_mode , width , and height .
Args :
resize_mode : The mode to use when resizing the image .
0 : Resize the image to the specified width and height .
1 : Resize the image to fill the specified width and height , maintaining the aspect ratio , and then center the image within the dimensions , cropping the excess .
2 : Resize the image to fit within the specified width and height , maintaining the aspect ratio , and then center the image within the dimensions , filling empty with data from image .
im : The image to resize .
width : The width to resize the image to .
height : The height to resize the image to .
upscaler_name : The name of the upscaler to use . If not provided , defaults to opts . upscaler_for_img2img .
"""
upscaler_name = upscaler_name or opts . upscaler_for_img2img
2022-09-23 08:37:47 -06:00
def resize ( im , w , h ) :
2023-01-02 09:42:10 -07:00
if upscaler_name is None or upscaler_name == " None " or im . mode == ' L ' :
2022-09-23 08:37:47 -06:00
return im . resize ( ( w , h ) , resample = LANCZOS )
2022-09-30 01:38:48 -06:00
scale = max ( w / im . width , h / im . height )
2022-09-30 05:23:41 -06:00
if scale > 1.0 :
2023-01-02 09:42:10 -07:00
upscalers = [ x for x in shared . sd_upscalers if x . name == upscaler_name ]
2023-03-26 02:01:32 -06:00
if len ( upscalers ) == 0 :
upscaler = shared . sd_upscalers [ 0 ]
print ( f " could not find upscaler named { upscaler_name or ' <empty string> ' } , using { upscaler . name } as a fallback " )
else :
upscaler = upscalers [ 0 ]
2022-09-30 05:23:41 -06:00
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 01:38:48 -06:00
2022-09-30 05:23:41 -06:00
return im
2022-09-23 08:37:47 -06:00
2022-09-03 03:08:45 -06:00
if resize_mode == 0 :
2022-09-23 08:37:47 -06:00
res = resize ( im , width , height )
2022-09-03 03:08:45 -06: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 08:37:47 -06:00
resized = resize ( im , src_w , src_h )
2022-09-03 03:08:45 -06:00
res = Image . new ( " RGB " , ( width , height ) )
res . paste ( resized , box = ( width / / 2 - src_w / / 2 , height / / 2 - src_h / / 2 ) )
2022-09-23 08:37:47 -06:00
2022-09-03 03:08:45 -06: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 08:37:47 -06:00
resized = resize ( im , src_w , src_h )
2022-09-03 03:08:45 -06: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 02:42:40 -06:00
res . paste ( resized . resize ( ( width , fill_height ) , box = ( 0 , resized . height , width , resized . height ) ) , box = ( 0 , fill_height + src_h ) )
2022-09-03 03:08:45 -06: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 02:42:40 -06:00
res . paste ( resized . resize ( ( fill_width , height ) , box = ( resized . width , 0 , resized . width , height ) ) , box = ( fill_width + src_w , 0 ) )
2022-09-03 03:08:45 -06:00
return res
invalid_filename_chars = ' <>: " / \\ |?* \n '
2022-09-20 00:01:32 -06:00
invalid_filename_prefix = ' '
invalid_filename_postfix = ' . '
2022-09-29 16:46:23 -06:00
re_nonletters = re . compile ( r ' [ \ s ' + string . punctuation + ' ]+ ' )
2022-10-24 15:21:31 -06:00
re_pattern = re . compile ( r " (.*?)(?: \ [([^ \ [ \ ]]+) \ ]|$) " )
2022-10-24 05:03:58 -06:00
re_pattern_arg = re . compile ( r " (.*)<([^>]*)>$ " )
2022-09-20 00:01:32 -06:00
max_filename_part_length = 128
2023-04-29 07:48:43 -06:00
NOTHING_AND_SKIP_PREVIOUS_TEXT = object ( )
2022-09-03 03:08:45 -06:00
2022-09-12 06:41:30 -06:00
def sanitize_filename_part ( text , replace_spaces = True ) :
2022-10-24 05:03:58 -06:00
if text is None :
return None
2022-09-12 06:41:30 -06:00
if replace_spaces :
text = text . replace ( ' ' , ' _ ' )
2022-09-03 03:08:45 -06:00
2022-09-20 00:01:32 -06: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 03:08:45 -06:00
2022-09-12 06:41:30 -06:00
2022-10-24 05:03:58 -06:00
class FilenameGenerator :
replacements = {
' seed ' : lambda self : self . seed if self . seed is not None else ' ' ,
' steps ' : lambda self : self . p and self . p . steps ,
' cfg ' : lambda self : self . p and self . p . cfg_scale ,
2022-10-28 17:07:01 -06:00
' width ' : lambda self : self . image . width ,
' height ' : lambda self : self . image . height ,
2022-10-24 05:03:58 -06:00
' styles ' : lambda self : self . p and sanitize_filename_part ( " , " . join ( [ style for style in self . p . styles if not style == " None " ] ) or " None " , replace_spaces = False ) ,
2022-11-19 02:01:51 -07:00
' sampler ' : lambda self : self . p and sanitize_filename_part ( self . p . sampler_name , replace_spaces = False ) ,
2022-10-24 05:03:58 -06:00
' model_hash ' : lambda self : getattr ( self . p , " sd_model_hash " , shared . sd_model . sd_model_hash ) ,
2022-11-23 07:44:20 -07:00
' model_name ' : lambda self : sanitize_filename_part ( shared . sd_model . sd_checkpoint_info . model_name , replace_spaces = False ) ,
2022-10-24 05:03:58 -06:00
' date ' : lambda self : datetime . datetime . now ( ) . strftime ( ' % Y- % m- %d ' ) ,
' datetime ' : lambda self , * args : self . datetime ( * args ) , # accepts formats: [datetime], [datetime<Format>], [datetime<Format><Time Zone>]
' job_timestamp ' : lambda self : getattr ( self . p , " job_timestamp " , shared . state . job_timestamp ) ,
2023-02-04 23:48:51 -07:00
' prompt_hash ' : lambda self : hashlib . sha256 ( self . prompt . encode ( ) ) . hexdigest ( ) [ 0 : 8 ] ,
2022-10-24 05:03:58 -06:00
' prompt ' : lambda self : sanitize_filename_part ( self . prompt ) ,
' prompt_no_styles ' : lambda self : self . prompt_no_style ( ) ,
' prompt_spaces ' : lambda self : sanitize_filename_part ( self . prompt , replace_spaces = False ) ,
' prompt_words ' : lambda self : self . prompt_words ( ) ,
2023-04-29 07:48:43 -06:00
' batch_number ' : lambda self : NOTHING_AND_SKIP_PREVIOUS_TEXT if self . p . batch_size == 1 else self . p . batch_index + 1 ,
' generation_number ' : lambda self : NOTHING_AND_SKIP_PREVIOUS_TEXT if self . p . n_iter == 1 and self . p . batch_size == 1 else self . p . iteration * self . p . batch_size + self . p . batch_index + 1 ,
' hasprompt ' : lambda self , * args : self . hasprompt ( * args ) , # accepts formats:[hasprompt<prompt1|default><prompt2>..]
2023-04-16 02:14:11 -06:00
' clip_skip ' : lambda self : opts . data [ " CLIP_stop_at_last_layers " ] ,
2023-05-05 11:20:33 -06:00
' denoising ' : lambda self : self . p . denoising_strength if self . p and self . p . denoising_strength else NOTHING_AND_SKIP_PREVIOUS_TEXT ,
2022-10-24 05:03:58 -06:00
}
default_time_format = ' % Y % m %d % H % M % S '
2022-10-28 17:07:01 -06:00
def __init__ ( self , p , seed , prompt , image ) :
2022-10-24 05:03:58 -06:00
self . p = p
self . seed = seed
self . prompt = prompt
2022-10-28 17:07:01 -06:00
self . image = image
2023-05-11 09:28:15 -06:00
2023-04-15 09:20:08 -06:00
def hasprompt ( self , * args ) :
lower = self . prompt . lower ( )
if self . p is None or self . prompt is None :
return None
outres = " "
for arg in args :
if arg != " " :
division = arg . split ( " | " )
expected = division [ 0 ] . lower ( )
default = division [ 1 ] if len ( division ) > 1 else " "
if lower . find ( expected ) > = 0 :
outres = f ' { outres } { expected } '
else :
outres = outres if default == " " else f ' { outres } { default } '
return sanitize_filename_part ( outres )
2022-10-24 05:03:58 -06:00
def prompt_no_style ( self ) :
if self . p is None or self . prompt is None :
return None
prompt_no_style = self . prompt
for style in shared . prompt_styles . get_style_prompts ( self . p . styles ) :
if len ( style ) > 0 :
for part in style . split ( " {prompt} " ) :
prompt_no_style = prompt_no_style . replace ( part , " " ) . replace ( " , , " , " , " ) . strip ( ) . strip ( ' , ' )
prompt_no_style = prompt_no_style . replace ( style , " " ) . strip ( ) . strip ( ' , ' ) . strip ( )
return sanitize_filename_part ( prompt_no_style , replace_spaces = False )
def prompt_words ( self ) :
words = [ x for x in re_nonletters . split ( self . prompt or " " ) if len ( x ) > 0 ]
if len ( words ) == 0 :
words = [ " empty " ]
return sanitize_filename_part ( " " . join ( words [ 0 : opts . directories_max_prompt_words ] ) , replace_spaces = False )
def datetime ( self , * args ) :
time_datetime = datetime . datetime . now ( )
2022-10-25 08:39:21 -06:00
time_format = args [ 0 ] if len ( args ) > 0 and args [ 0 ] != " " else self . default_time_format
2022-10-24 06:13:36 -06:00
try :
time_zone = pytz . timezone ( args [ 1 ] ) if len ( args ) > 1 else None
2023-05-09 22:52:45 -06:00
except pytz . exceptions . UnknownTimeZoneError :
2022-10-24 06:13:36 -06:00
time_zone = None
2022-10-24 05:03:58 -06:00
time_zone_time = time_datetime . astimezone ( time_zone )
try :
formatted_time = time_zone_time . strftime ( time_format )
2023-05-09 22:52:45 -06:00
except ( ValueError , TypeError ) :
2022-10-24 05:03:58 -06:00
formatted_time = time_zone_time . strftime ( self . default_time_format )
return sanitize_filename_part ( formatted_time , replace_spaces = False )
def apply ( self , x ) :
res = ' '
for m in re_pattern . finditer ( x ) :
text , pattern = m . groups ( )
2023-04-07 06:04:46 -06:00
2022-10-24 05:03:58 -06:00
if pattern is None :
2023-04-29 07:48:43 -06:00
res + = text
2022-10-24 05:03:58 -06:00
continue
2022-10-04 05:16:52 -06:00
2022-10-24 05:03:58 -06:00
pattern_args = [ ]
while True :
m = re_pattern_arg . match ( pattern )
if m is None :
break
pattern , arg = m . groups ( )
pattern_args . insert ( 0 , arg )
fun = self . replacements . get ( pattern . lower ( ) )
if fun is not None :
try :
replacement = fun ( self , * pattern_args )
except Exception :
replacement = None
print ( f " Error adding [ { pattern } ] to filename " , file = sys . stderr )
print ( traceback . format_exc ( ) , file = sys . stderr )
2023-04-29 07:48:43 -06:00
if replacement == NOTHING_AND_SKIP_PREVIOUS_TEXT :
continue
elif replacement is not None :
res + = text + str ( replacement )
2022-10-24 15:21:31 -06:00
continue
2022-09-12 15:44:08 -06:00
2023-04-29 07:48:43 -06:00
res + = f ' { text } [ { pattern } ] '
2022-09-14 19:05:00 -06:00
2022-10-24 05:03:58 -06:00
return res
2022-09-12 15:44:08 -06:00
2022-09-29 16:46:23 -06:00
2022-09-14 06:40:16 -06:00
def get_next_sequence_number ( path , basename ) :
2022-09-13 08:43:08 -06: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 06:40:16 -06:00
if basename != ' ' :
2023-05-09 13:17:58 -06:00
basename = f " { basename } - "
2022-09-14 06:40:16 -06:00
prefix_length = len ( basename )
2022-09-13 08:43:08 -06:00
for p in os . listdir ( path ) :
2022-09-14 06:40:16 -06:00
if p . startswith ( basename ) :
2023-05-09 23:25:25 -06:00
parts = 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 06:40:16 -06:00
try :
2023-05-09 23:25:25 -06:00
result = max ( int ( parts [ 0 ] ) , result )
2022-09-14 06:40:16 -06:00
except ValueError :
2022-09-13 09:46:05 -06:00
pass
2022-09-13 08:43:08 -06:00
return result + 1
2022-09-12 15:44:08 -06:00
2022-09-29 16:46:23 -06:00
2023-05-17 12:03:41 -06:00
def save_image_with_geninfo ( image , geninfo , filename , extension = None , existing_pnginfo = None ) :
if extension is None :
extension = os . path . splitext ( filename ) [ 1 ]
image_format = Image . registered_extensions ( ) [ extension ]
existing_pnginfo = existing_pnginfo or { }
if opts . enable_pnginfo :
existing_pnginfo [ ' parameters ' ] = geninfo
if extension . lower ( ) == ' .png ' :
pnginfo_data = PngImagePlugin . PngInfo ( )
for k , v in ( existing_pnginfo or { } ) . items ( ) :
pnginfo_data . add_text ( k , str ( v ) )
image . save ( filename , format = image_format , quality = opts . jpeg_quality , pnginfo = pnginfo_data )
elif extension . lower ( ) in ( " .jpg " , " .jpeg " , " .webp " ) :
if image . mode == ' RGBA ' :
image = image . convert ( " RGB " )
elif image . mode == ' I;16 ' :
image = image . point ( lambda p : p * 0.0038910505836576 ) . convert ( " RGB " if extension . lower ( ) == " .webp " else " L " )
image . save ( filename , format = image_format , quality = opts . jpeg_quality , lossless = opts . webp_lossless )
if opts . enable_pnginfo and geninfo is not None :
exif_bytes = piexif . dump ( {
" Exif " : {
piexif . ExifIFD . UserComment : piexif . helper . UserComment . dump ( geninfo or " " , encoding = " unicode " )
} ,
} )
piexif . insert ( exif_bytes , filename )
else :
image . save ( filename , format = image_format , quality = opts . jpeg_quality )
2022-10-04 10:19:50 -06: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 = " " , save_to_dirs = None ) :
2022-10-24 05:03:58 -06:00
""" Save an image.
2022-10-08 23:01:10 -06:00
Args :
image ( ` PIL . Image ` ) :
The image to be saved .
path ( ` str ` ) :
The directory to save the image . Note , the option ` save_to_dirs ` will make the image to be saved into a sub directory .
basename ( ` str ` ) :
The base filename which will be applied to ` filename pattern ` .
2022-12-14 19:01:32 -07:00
seed , prompt , short_filename ,
2022-10-08 23:01:10 -06:00
extension ( ` str ` ) :
Image file extension , default is ` png ` .
pngsectionname ( ` str ` ) :
Specify the name of the section which ` info ` will be saved in .
info ( ` str ` or ` PngImagePlugin . iTXt ` ) :
PNG info chunks .
existing_info ( ` dict ` ) :
Additional PNG info . ` existing_info == { pngsectionname : info , . . . } `
no_prompt :
TODO I don ' t know its meaning.
p ( ` StableDiffusionProcessing ` )
forced_filename ( ` str ` ) :
If specified , ` basename ` and filename pattern will be ignored .
save_to_dirs ( bool ) :
If true , the image will be saved into a subdirectory of ` path ` .
Returns : ( fullfn , txt_fullfn )
fullfn ( ` str ` ) :
The full path of the saved imaged .
txt_fullfn ( ` str ` or None ) :
If a text file is saved for this image , this will be its full path . Otherwise None .
2022-10-24 05:03:58 -06:00
"""
2022-10-28 17:07:01 -06:00
namegen = FilenameGenerator ( p , seed , prompt , image )
2022-09-12 11:47:46 -06:00
2022-10-04 10:19:50 -06:00
if save_to_dirs is None :
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 04:36:16 -06:00
2022-09-12 15:44:08 -06:00
if save_to_dirs :
2022-10-24 05:03:58 -06:00
dirname = namegen . apply ( opts . directories_filename_pattern or " [prompt_words] " ) . lstrip ( ' ' ) . rstrip ( ' \\ / ' )
2023-01-04 08:20:38 -07:00
path = os . path . join ( path , dirname )
2022-09-03 03:08:45 -06:00
2022-10-19 15:23:48 -06:00
os . makedirs ( path , exist_ok = True )
2022-09-03 03:08:45 -06:00
2022-09-19 16:13:12 -06:00
if forced_filename is None :
2022-10-24 05:03:58 -06:00
if short_filename or seed is None :
2022-10-22 02:16:55 -06:00
file_decoration = " "
2022-10-24 09:46:28 -06:00
elif opts . save_to_dirs :
2022-10-24 05:03:58 -06:00
file_decoration = opts . samples_filename_pattern or " [seed] "
2022-10-24 09:46:28 -06:00
else :
file_decoration = opts . samples_filename_pattern or " [seed]-[prompt_spaces] "
2022-10-24 05:03:58 -06:00
add_number = opts . save_images_add_number or file_decoration == ' '
2022-10-22 02:16:55 -06:00
2022-10-24 05:03:58 -06:00
if file_decoration != " " and add_number :
2023-05-09 13:17:58 -06:00
file_decoration = f " - { file_decoration } "
2022-10-22 02:16:55 -06:00
2022-10-24 05:03:58 -06:00
file_decoration = namegen . apply ( file_decoration ) + suffix
if add_number :
basecount = get_next_sequence_number ( path , basename )
fullfn = None
for i in range ( 500 ) :
fn = f " { basecount + i : 05 } " if basename == ' ' else f " { basename } - { basecount + i : 04 } "
2023-01-04 08:20:38 -07:00
fullfn = os . path . join ( path , f " { fn } { file_decoration } . { extension } " )
2022-10-24 05:03:58 -06:00
if not os . path . exists ( fullfn ) :
break
else :
2023-01-04 08:20:38 -07:00
fullfn = os . path . join ( path , f " { file_decoration } . { extension } " )
2022-09-19 16:13:12 -06:00
else :
2023-01-04 08:20:38 -07:00
fullfn = os . path . join ( path , f " { forced_filename } . { extension } " )
2022-10-26 04:12:44 -06:00
pnginfo = existing_info or { }
if info is not None :
pnginfo [ pnginfo_section_name ] = info
params = script_callbacks . ImageSaveParams ( image , p , fullfn , pnginfo )
script_callbacks . before_image_saved_callback ( params )
image = params . image
fullfn = params . filename
info = params . pnginfo . get ( pnginfo_section_name , None )
2022-09-03 03:08:45 -06:00
2022-11-26 08:07:51 -07:00
def _atomically_save_image ( image_to_save , filename_without_extension , extension ) :
2023-05-17 12:03:41 -06:00
"""
save image with . tmp extension to avoid race condition when another process detects new image in the directory
"""
2023-05-09 13:17:58 -06:00
temp_file_path = f " { filename_without_extension } .tmp "
2022-10-26 04:12:44 -06:00
2023-05-17 12:03:41 -06:00
save_image_with_geninfo ( image_to_save , info , temp_file_path , extension , params . pnginfo )
2022-11-26 08:07:51 -07:00
os . replace ( temp_file_path , filename_without_extension + extension )
fullfn_without_extension , extension = os . path . splitext ( params . filename )
2023-03-12 13:30:31 -06:00
if hasattr ( os , ' statvfs ' ) :
max_name_len = os . statvfs ( path ) . f_namemax
2023-03-12 13:33:29 -06:00
fullfn_without_extension = fullfn_without_extension [ : max_name_len - max ( 4 , len ( extension ) ) ]
2023-03-12 13:30:31 -06:00
params . filename = fullfn_without_extension + extension
2023-03-12 13:36:04 -06:00
fullfn = params . filename
2022-11-26 08:07:51 -07:00
_atomically_save_image ( image , fullfn_without_extension , extension )
2022-09-03 03:08:45 -06:00
2022-11-27 01:52:53 -07:00
image . already_saved_as = fullfn
2023-02-19 00:12:45 -07:00
oversize = image . width > opts . target_side_length or image . height > opts . target_side_length
if opts . export_for_4chan and ( oversize or os . stat ( fullfn ) . st_size > opts . img_downscale_threshold * 1024 * 1024 ) :
2022-09-03 03:08:45 -06:00
ratio = image . width / image . height
if oversize and ratio > 1 :
2023-02-20 01:30:13 -07:00
image = image . resize ( ( round ( opts . target_side_length ) , round ( image . height * opts . target_side_length / image . width ) ) , LANCZOS )
2022-09-03 03:08:45 -06:00
elif oversize :
2023-02-20 01:30:13 -07:00
image = image . resize ( ( round ( image . width * opts . target_side_length / image . height ) , round ( opts . target_side_length ) ) , LANCZOS )
2022-09-03 03:08:45 -06:00
2023-02-19 00:12:45 -07:00
try :
_atomically_save_image ( image , fullfn_without_extension , " .jpg " )
except Exception as e :
errors . display ( e , " saving image as downscaled JPG " )
2022-09-03 03:08:45 -06:00
if opts . save_txt and info is not None :
2022-10-08 23:01:10 -06:00
txt_fullfn = f " { fullfn_without_extension } .txt "
with open ( txt_fullfn , " w " , encoding = " utf8 " ) as file :
2023-05-09 13:17:58 -06:00
file . write ( f " { info } \n " )
2022-10-08 23:01:10 -06:00
else :
txt_fullfn = None
2022-09-03 03:08:45 -06:00
2022-10-26 04:12:44 -06:00
script_callbacks . image_saved_callback ( params )
2022-10-08 23:01:10 -06:00
return fullfn , txt_fullfn
2022-10-12 17:17:26 -06:00
2022-11-27 06:28:32 -07:00
def read_info_from_image ( image ) :
items = image . info or { }
geninfo = items . pop ( ' parameters ' , None )
if " exif " in items :
exif = piexif . load ( items [ " exif " ] )
exif_comment = ( exif or { } ) . get ( " Exif " , { } ) . get ( piexif . ExifIFD . UserComment , b ' ' )
try :
exif_comment = piexif . helper . UserComment . load ( exif_comment )
except ValueError :
exif_comment = exif_comment . decode ( ' utf8 ' , errors = " ignore " )
2023-01-17 13:54:23 -07:00
if exif_comment :
items [ ' exif comment ' ] = exif_comment
geninfo = exif_comment
2022-11-27 06:28:32 -07:00
for field in [ ' jfif ' , ' jfif_version ' , ' jfif_unit ' , ' jfif_density ' , ' dpi ' , ' exif ' ,
' loop ' , ' background ' , ' timestamp ' , ' duration ' ] :
items . pop ( field , None )
if items . get ( " Software " , None ) == " NovelAI " :
try :
json_info = json . loads ( items [ " Comment " ] )
sampler = sd_samplers . samplers_map . get ( json_info [ " sampler " ] , " Euler a " )
geninfo = f """ { items [ " Description " ] }
Negative prompt : { json_info [ " uc " ] }
Steps : { json_info [ " steps " ] } , Sampler : { sampler } , CFG scale : { json_info [ " scale " ] } , Seed : { json_info [ " seed " ] } , Size : { image . width } x { image . height } , Clip skip : 2 , ENSD : 31337 """
except Exception :
2022-12-24 12:35:29 -07:00
print ( " Error parsing NovelAI image generation parameters: " , file = sys . stderr )
2022-11-27 06:28:32 -07:00
print ( traceback . format_exc ( ) , file = sys . stderr )
return geninfo , items
2022-10-14 09:15:03 -06:00
def image_data ( data ) :
2023-03-24 22:29:51 -06:00
import gradio as gr
2022-10-14 09:15:03 -06:00
try :
image = Image . open ( io . BytesIO ( data ) )
2022-11-27 06:28:32 -07:00
textinfo , _ = read_info_from_image ( image )
2022-10-14 09:15:03 -06:00
return textinfo , None
except Exception :
pass
try :
text = data . decode ( ' utf8 ' )
assert len ( text ) < 10000
return text , None
except Exception :
pass
2023-03-24 22:29:51 -06:00
return gr . update ( ) , None
2022-12-23 23:46:35 -07:00
def flatten ( img , bgcolor ) :
""" replaces transparency with bgcolor (example: " #ffffff " ), returning an RGB mode image with no transparency """
if img . mode == " RGBA " :
background = Image . new ( ' RGBA ' , img . size , bgcolor )
background . paste ( img , mask = img )
img = background
return img . convert ( ' RGB ' )