2023-07-03 18:10:20 +08:00
from __future__ import annotations
2022-09-13 05:44:08 +08:00
import datetime
2022-10-24 19:03:58 +08:00
2022-10-22 20:11:15 +08:00
import pytz
2022-10-14 23:15:03 +08:00
import io
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
2023-06-28 11:16:44 +08:00
from PIL import Image , ImageFont , ImageDraw , ImageColor , PngImagePlugin
2022-09-10 13:45:16 +08:00
import string
2022-11-27 21:28:32 +08:00
import json
2023-02-05 14:48:51 +08:00
import hashlib
2022-09-03 17:08:45 +08:00
2023-02-19 15:12:45 +08:00
from modules import sd_samplers , shared , script_callbacks , errors
2023-05-11 15:07:01 +08:00
from modules . paths_internal import roboto_ttf_file
2023-05-10 13:43:42 +08:00
from modules . shared import opts
2022-09-03 17:08:45 +08:00
LANCZOS = ( Image . Resampling . LANCZOS if hasattr ( Image , ' Resampling ' ) else Image . LANCZOS )
2023-05-11 15:06:19 +08:00
def get_font ( fontsize : int ) :
try :
2023-05-11 15:07:01 +08:00
return ImageFont . truetype ( opts . font or roboto_ttf_file , fontsize )
2023-05-11 15:06:19 +08:00
except Exception :
2023-05-11 15:07:01 +08:00
return ImageFont . truetype ( roboto_ttf_file , fontsize )
2023-05-11 15:06:19 +08:00
2022-09-03 17:08:45 +08: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 23:07:24 +08:00
elif opts . grid_prevent_empty_spots :
2022-10-14 20:19:39 +08:00
rows = math . floor ( math . sqrt ( len ( imgs ) ) )
while len ( imgs ) % rows != 0 :
rows - = 1
2022-10-14 23:07:24 +08:00
else :
rows = math . sqrt ( len ( imgs ) )
rows = round ( rows )
2023-01-29 08:25:52 +08:00
if rows > len ( imgs ) :
rows = len ( imgs )
2022-09-03 17:08:45 +08:00
cols = math . ceil ( len ( imgs ) / rows )
2023-01-01 23:37:37 +08:00
params = script_callbacks . ImageGridLoopParams ( imgs , cols , rows )
script_callbacks . image_grid_callback ( params )
2022-09-03 17:08:45 +08:00
w , h = imgs [ 0 ] . size
2023-01-01 23:37:37 +08:00
grid = Image . new ( ' RGB ' , size = ( params . cols * w , params . rows * h ) , color = ' black ' )
2022-09-03 17:08:45 +08:00
2023-01-01 23:37:37 +08:00
for i , img in enumerate ( params . imgs ) :
grid . paste ( img , box = ( i % params . cols * w , i / / params . cols * h ) )
2022-09-03 17:08:45 +08: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-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
2023-02-04 18:28:53 +08:00
def draw_grid_annotations ( im , width , height , hor_texts , ver_texts , margin = 0 ) :
2023-06-28 18:57:34 +08:00
2023-06-28 18:55:58 +08:00
color_active = ImageColor . getcolor ( opts . grid_text_active_color , ' RGB ' )
color_inactive = ImageColor . getcolor ( opts . grid_text_inactive_color , ' RGB ' )
color_background = ImageColor . getcolor ( opts . grid_background_color , ' RGB ' )
2022-09-03 17:08:45 +08: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-17 13:45:43 +08:00
def draw_texts ( drawing , draw_x , draw_y , lines , initial_fnt , initial_fontsize ) :
2023-05-10 16:37:18 +08:00
for line in lines :
2022-12-17 13:45:43 +08: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 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
2022-12-17 13:45:43 +08:00
fnt = get_font ( fontsize )
2022-09-13 00:17:02 +08:00
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 } '
2023-06-28 18:55:58 +08:00
calc_img = Image . new ( " RGB " , ( 1 , 1 ) , color_background )
2022-09-03 17:08:45 +08:00
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-17 13:45:43 +08:00
line . allowed_width = allowed_width
2022-09-03 17:08:45 +08:00
hor_text_heights = [ sum ( [ line . size [ 1 ] + line_spacing for line in lines ] ) - line_spacing for lines in hor_texts ]
2023-02-04 18:28:53 +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
2023-01-24 15:24:32 +08:00
pad_top = 0 if sum ( hor_text_heights ) == 0 else max ( hor_text_heights ) + line_spacing * 2
2022-09-03 17:08:45 +08:00
2023-06-28 18:55:58 +08:00
result = Image . new ( " RGB " , ( im . width + pad_left + margin * ( cols - 1 ) , im . height + pad_top + margin * ( rows - 1 ) ) , color_background )
2023-02-04 18:28:53 +08: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 17:08:45 +08:00
d = ImageDraw . Draw ( result )
for col in range ( cols ) :
2023-02-04 18:28:53 +08:00
x = pad_left + ( width + margin ) * col + width / 2
2022-09-03 17:08:45 +08:00
y = pad_top / 2 - hor_text_heights [ col ] / 2
2022-12-17 13:45:43 +08:00
draw_texts ( d , x , y , hor_texts [ col ] , fnt , fontsize )
2022-09-03 17:08:45 +08:00
for row in range ( rows ) :
x = pad_left / 2
2023-02-04 18:28:53 +08:00
y = pad_top + ( height + margin ) * row + height / 2 - ver_text_heights [ row ] / 2
2022-09-03 17:08:45 +08:00
2022-12-17 13:45:43 +08:00
draw_texts ( d , x , y , ver_texts [ row ] , fnt , fontsize )
2022-09-03 17:08:45 +08:00
return result
2023-02-05 16:44:56 +08:00
def draw_prompt_matrix ( im , width , height , all_prompts , margin = 0 ) :
2022-09-03 17:08:45 +08:00
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
2023-02-05 16:44:56 +08:00
return draw_grid_annotations ( im , width , height , hor_texts , ver_texts , margin )
2022-09-03 17:08:45 +08:00
2023-01-03 00:42:10 +08: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 22:37:47 +08:00
def resize ( im , w , h ) :
2023-01-03 00:42:10 +08:00
if upscaler_name is None or upscaler_name == " 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 :
2023-01-03 00:42:10 +08:00
upscalers = [ x for x in shared . sd_upscalers if x . name == upscaler_name ]
2023-03-26 16:01:32 +08: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 19:23:41 +08: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 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
2023-07-12 16:51:50 +08:00
if fill_height > 0 :
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 ) )
2022-09-03 17:08:45 +08:00
elif ratio > src_ratio :
fill_width = width / / 2 - src_w / / 2
2023-07-12 16:51:50 +08:00
if fill_width > 0 :
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 ) )
2022-09-03 17:08:45 +08:00
return res
2023-08-05 10:59:47 +08:00
invalid_filename_chars = ' <>: " / \\ |?* \n \r \t '
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-10-25 05:21:31 +08:00
re_pattern = re . compile ( r " (.*?)(?: \ [([^ \ [ \ ]]+) \ ]|$) " )
2022-10-24 19:03:58 +08:00
re_pattern_arg = re . compile ( r " (.*)<([^>]*)>$ " )
2022-09-20 14:01:32 +08:00
max_filename_part_length = 128
2023-04-29 21:48:43 +08:00
NOTHING_AND_SKIP_PREVIOUS_TEXT = object ( )
2022-09-03 17:08:45 +08:00
2022-09-12 20:41:30 +08:00
def sanitize_filename_part ( text , replace_spaces = True ) :
2022-10-24 19:03:58 +08:00
if text is None :
return None
2022-09-12 20:41:30 +08:00
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-10-24 19:03:58 +08:00
class FilenameGenerator :
replacements = {
' seed ' : lambda self : self . seed if self . seed is not None else ' ' ,
2023-05-24 23:17:02 +08:00
' seed_first ' : lambda self : self . seed if self . p . batch_size == 1 else self . p . all_seeds [ 0 ] ,
' seed_last ' : lambda self : NOTHING_AND_SKIP_PREVIOUS_TEXT if self . p . batch_size == 1 else self . p . all_seeds [ - 1 ] ,
2022-10-24 19:03:58 +08:00
' steps ' : lambda self : self . p and self . p . steps ,
' cfg ' : lambda self : self . p and self . p . cfg_scale ,
2022-10-29 07:07:01 +08:00
' width ' : lambda self : self . image . width ,
' height ' : lambda self : self . image . height ,
2022-10-24 19:03:58 +08: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 17:01:51 +08:00
' sampler ' : lambda self : self . p and sanitize_filename_part ( self . p . sampler_name , replace_spaces = False ) ,
2022-10-24 19:03:58 +08:00
' model_hash ' : lambda self : getattr ( self . p , " sd_model_hash " , shared . sd_model . sd_model_hash ) ,
2023-07-19 20:49:31 +08:00
' model_name ' : lambda self : sanitize_filename_part ( shared . sd_model . sd_checkpoint_info . name_for_extra , replace_spaces = False ) ,
2022-10-24 19:03:58 +08: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-08-18 14:07:40 +08:00
' prompt_hash ' : lambda self , * args : self . string_hash ( self . prompt , * args ) ,
' negative_prompt_hash ' : lambda self , * args : self . string_hash ( self . p . negative_prompt , * args ) ,
' full_prompt_hash ' : lambda self , * args : self . string_hash ( f " { self . p . prompt } { self . p . negative_prompt } " , * args ) , # a space in between to create a unique string
2022-10-24 19:03:58 +08: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-05-24 23:17:02 +08:00
' batch_number ' : lambda self : NOTHING_AND_SKIP_PREVIOUS_TEXT if self . p . batch_size == 1 or self . zip else self . p . batch_index + 1 ,
' batch_size ' : lambda self : self . p . batch_size ,
' generation_number ' : lambda self : NOTHING_AND_SKIP_PREVIOUS_TEXT if ( self . p . n_iter == 1 and self . p . batch_size == 1 ) or self . zip else self . p . iteration * self . p . batch_size + self . p . batch_index + 1 ,
2023-04-29 21:48:43 +08:00
' hasprompt ' : lambda self , * args : self . hasprompt ( * args ) , # accepts formats:[hasprompt<prompt1|default><prompt2>..]
2023-04-16 16:14:11 +08:00
' clip_skip ' : lambda self : opts . data [ " CLIP_stop_at_last_layers " ] ,
2023-05-06 01:20:33 +08:00
' denoising ' : lambda self : self . p . denoising_strength if self . p and self . p . denoising_strength else NOTHING_AND_SKIP_PREVIOUS_TEXT ,
2023-06-15 23:53:16 +08:00
' user ' : lambda self : self . p . user ,
2023-05-23 14:56:08 +08:00
' vae_filename ' : lambda self : self . get_vae_filename ( ) ,
2023-08-18 13:04:46 +08:00
' none ' : lambda self : ' ' , # Overrides the default, so you can get just the sequence number
' image_hash ' : lambda self , * args : self . image_hash ( * args ) # accepts formats: [image_hash<length>] default full hash
2022-10-24 19:03:58 +08:00
}
default_time_format = ' % Y % m %d % H % M % S '
2023-05-24 23:17:02 +08:00
def __init__ ( self , p , seed , prompt , image , zip = False ) :
2022-10-24 19:03:58 +08:00
self . p = p
self . seed = seed
self . prompt = prompt
2022-10-29 07:07:01 +08:00
self . image = image
2023-05-24 23:17:02 +08:00
self . zip = zip
2023-05-11 23:28:15 +08:00
2023-08-09 15:25:35 +08:00
def get_vae_filename ( self ) :
""" Get the name of the VAE file. """
import modules . sd_vae as sd_vae
if sd_vae . loaded_vae_file is None :
return " NoneType "
file_name = os . path . basename ( sd_vae . loaded_vae_file )
split_file_name = file_name . split ( ' . ' )
if len ( split_file_name ) > 1 and split_file_name [ 0 ] == ' ' :
return split_file_name [ 1 ] # if the first character of the filename is "." then [1] is obtained.
else :
return split_file_name [ 0 ]
2023-04-15 23:20:08 +08: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 19:03:58 +08: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 ) :
2023-06-02 19:58:10 +08:00
if style :
2022-10-24 19:03:58 +08:00
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 ) :
2023-06-02 19:58:10 +08:00
words = [ x for x in re_nonletters . split ( self . prompt or " " ) if x ]
2022-10-24 19:03:58 +08:00
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 ( )
2023-06-02 19:58:10 +08:00
time_format = args [ 0 ] if ( args and args [ 0 ] != " " ) else self . default_time_format
2022-10-24 20:13:36 +08:00
try :
time_zone = pytz . timezone ( args [ 1 ] ) if len ( args ) > 1 else None
2023-05-10 12:52:45 +08:00
except pytz . exceptions . UnknownTimeZoneError :
2022-10-24 20:13:36 +08:00
time_zone = None
2022-10-24 19:03:58 +08:00
time_zone_time = time_datetime . astimezone ( time_zone )
try :
formatted_time = time_zone_time . strftime ( time_format )
2023-05-10 12:52:45 +08:00
except ( ValueError , TypeError ) :
2022-10-24 19:03:58 +08:00
formatted_time = time_zone_time . strftime ( self . default_time_format )
return sanitize_filename_part ( formatted_time , replace_spaces = False )
2023-08-18 13:04:46 +08:00
def image_hash ( self , * args ) :
length = int ( args [ 0 ] ) if ( args and args [ 0 ] != " " ) else None
return hashlib . sha256 ( self . image . tobytes ( ) ) . hexdigest ( ) [ 0 : length ]
2023-08-18 14:07:40 +08:00
def string_hash ( self , text , * args ) :
length = int ( args [ 0 ] ) if ( args and args [ 0 ] != " " ) else 8
return hashlib . sha256 ( text . encode ( ) ) . hexdigest ( ) [ 0 : length ]
2022-10-24 19:03:58 +08:00
def apply ( self , x ) :
res = ' '
for m in re_pattern . finditer ( x ) :
text , pattern = m . groups ( )
2023-04-07 20:04:46 +08:00
2022-10-24 19:03:58 +08:00
if pattern is None :
2023-04-29 21:48:43 +08:00
res + = text
2022-10-24 19:03:58 +08:00
continue
2022-10-04 19:16:52 +08:00
2022-10-24 19:03:58 +08: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
2023-06-01 00:56:37 +08:00
errors . report ( f " Error adding [ { pattern } ] to filename " , exc_info = True )
2022-10-24 19:03:58 +08:00
2023-04-29 21:48:43 +08:00
if replacement == NOTHING_AND_SKIP_PREVIOUS_TEXT :
continue
elif replacement is not None :
res + = text + str ( replacement )
2022-10-25 05:21:31 +08:00
continue
2022-09-13 05:44:08 +08:00
2023-04-29 21:48:43 +08:00
res + = f ' { text } [ { pattern } ] '
2022-09-15 09:05:00 +08:00
2022-10-24 19:03:58 +08:00
return res
2022-09-13 05:44:08 +08:00
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 != ' ' :
2023-05-10 03:17:58 +08:00
basename = f " { basename } - "
2022-09-14 20:40:16 +08:00
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 ) :
2023-05-10 13:25:25 +08: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 20:40:16 +08:00
try :
2023-05-10 13:25:25 +08:00
result = max ( int ( parts [ 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
2023-07-08 18:43:42 +08:00
def save_image_with_geninfo ( image , geninfo , filename , extension = None , existing_pnginfo = None , pnginfo_section_name = ' parameters ' ) :
"""
Saves image to filename , including geninfo as text information for generation info .
For PNG images , geninfo is added to existing pnginfo dictionary using the pnginfo_section_name argument as key .
For JPG images , there ' s no dictionary and geninfo just replaces the EXIF description.
"""
2023-05-18 02:03:41 +08:00
if extension is None :
extension = os . path . splitext ( filename ) [ 1 ]
image_format = Image . registered_extensions ( ) [ extension ]
if extension . lower ( ) == ' .png ' :
2023-07-08 18:43:42 +08:00
existing_pnginfo = existing_pnginfo or { }
if opts . enable_pnginfo :
existing_pnginfo [ pnginfo_section_name ] = geninfo
2023-06-01 00:20:19 +08:00
if opts . enable_pnginfo :
pnginfo_data = PngImagePlugin . PngInfo ( )
for k , v in ( existing_pnginfo or { } ) . items ( ) :
pnginfo_data . add_text ( k , str ( v ) )
else :
pnginfo_data = None
2023-05-18 02:03:41 +08:00
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-05 00:19:50 +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 = " " , save_to_dirs = None ) :
2022-10-24 19:03:58 +08:00
""" Save an image.
2022-10-09 13:01:10 +08: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-15 10:01:32 +08:00
seed , prompt , short_filename ,
2022-10-09 13:01:10 +08: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 19:03:58 +08:00
"""
2022-10-29 07:07:01 +08:00
namegen = FilenameGenerator ( p , seed , prompt , image )
2022-09-13 01:47:46 +08:00
2023-08-19 15:50:43 +08:00
# WebP and JPG formats have maximum dimension limits of 16383 and 65535 respectively. switch to PNG which has a much higher limit
if ( image . height > 65535 or image . width > 65535 ) and extension . lower ( ) in ( " jpg " , " jpeg " ) or ( image . height > 16383 or image . width > 16383 ) and extension . lower ( ) == " webp " :
print ( ' Image dimensions too large; saving as PNG ' )
extension = " .png "
2022-10-05 00:19:50 +08: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 18:36:16 +08:00
2022-09-13 05:44:08 +08:00
if save_to_dirs :
2022-10-24 19:03:58 +08:00
dirname = namegen . apply ( opts . directories_filename_pattern or " [prompt_words] " ) . lstrip ( ' ' ) . rstrip ( ' \\ / ' )
2023-01-04 23:20:38 +08:00
path = os . path . join ( path , dirname )
2022-09-03 17:08:45 +08:00
2022-10-20 05:23:48 +08:00
os . makedirs ( path , exist_ok = True )
2022-09-03 17:08:45 +08:00
2022-09-20 06:13:12 +08:00
if forced_filename is None :
2022-10-24 19:03:58 +08:00
if short_filename or seed is None :
2022-10-22 16:16:55 +08:00
file_decoration = " "
2022-10-24 23:46:28 +08:00
elif opts . save_to_dirs :
2022-10-24 19:03:58 +08:00
file_decoration = opts . samples_filename_pattern or " [seed] "
2022-10-24 23:46:28 +08:00
else :
file_decoration = opts . samples_filename_pattern or " [seed]-[prompt_spaces] "
2022-10-24 19:03:58 +08:00
2023-07-15 09:52:00 +08:00
file_decoration = namegen . apply ( file_decoration ) + suffix
2022-10-24 19:03:58 +08:00
add_number = opts . save_images_add_number or file_decoration == ' '
2022-10-22 16:16:55 +08:00
2022-10-24 19:03:58 +08:00
if file_decoration != " " and add_number :
2023-05-10 03:17:58 +08:00
file_decoration = f " - { file_decoration } "
2022-10-22 16:16:55 +08:00
2022-10-24 19:03:58 +08:00
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 23:20:38 +08:00
fullfn = os . path . join ( path , f " { fn } { file_decoration } . { extension } " )
2022-10-24 19:03:58 +08:00
if not os . path . exists ( fullfn ) :
break
else :
2023-01-04 23:20:38 +08:00
fullfn = os . path . join ( path , f " { file_decoration } . { extension } " )
2022-09-20 06:13:12 +08:00
else :
2023-01-04 23:20:38 +08:00
fullfn = os . path . join ( path , f " { forced_filename } . { extension } " )
2022-10-26 18:12:44 +08: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 17:08:45 +08:00
2022-11-26 23:07:51 +08:00
def _atomically_save_image ( image_to_save , filename_without_extension , extension ) :
2023-05-18 02:03:41 +08:00
"""
save image with . tmp extension to avoid race condition when another process detects new image in the directory
"""
2023-05-10 03:17:58 +08:00
temp_file_path = f " { filename_without_extension } .tmp "
2022-10-26 18:12:44 +08:00
2023-07-08 18:43:42 +08:00
save_image_with_geninfo ( image_to_save , info , temp_file_path , extension , existing_pnginfo = params . pnginfo , pnginfo_section_name = pnginfo_section_name )
2022-11-26 23:07:51 +08:00
os . replace ( temp_file_path , filename_without_extension + extension )
fullfn_without_extension , extension = os . path . splitext ( params . filename )
2023-03-13 03:30:31 +08:00
if hasattr ( os , ' statvfs ' ) :
max_name_len = os . statvfs ( path ) . f_namemax
2023-03-13 03:33:29 +08:00
fullfn_without_extension = fullfn_without_extension [ : max_name_len - max ( 4 , len ( extension ) ) ]
2023-03-13 03:30:31 +08:00
params . filename = fullfn_without_extension + extension
2023-03-13 03:36:04 +08:00
fullfn = params . filename
2022-11-26 23:07:51 +08:00
_atomically_save_image ( image , fullfn_without_extension , extension )
2022-09-03 17:08:45 +08:00
2022-11-27 16:52:53 +08:00
image . already_saved_as = fullfn
2023-02-19 15:12:45 +08: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 17:08:45 +08:00
ratio = image . width / image . height
2023-07-05 20:39:04 +08:00
resize_to = None
2022-09-03 17:08:45 +08:00
if oversize and ratio > 1 :
2023-07-05 20:39:04 +08:00
resize_to = round ( opts . target_side_length ) , round ( image . height * opts . target_side_length / image . width )
2022-09-03 17:08:45 +08:00
elif oversize :
2023-07-05 20:39:04 +08:00
resize_to = round ( image . width * opts . target_side_length / image . height ) , round ( opts . target_side_length )
2022-09-03 17:08:45 +08:00
2023-07-05 20:39:04 +08:00
if resize_to is not None :
try :
# Resizing image with LANCZOS could throw an exception if e.g. image mode is I;16
image = image . resize ( resize_to , LANCZOS )
2023-07-08 19:21:40 +08:00
except Exception :
2023-07-05 20:39:04 +08:00
image = image . resize ( resize_to )
2023-02-19 15:12:45 +08: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 17:08:45 +08:00
if opts . save_txt and info is not None :
2022-10-09 13:01:10 +08:00
txt_fullfn = f " { fullfn_without_extension } .txt "
with open ( txt_fullfn , " w " , encoding = " utf8 " ) as file :
2023-05-10 03:17:58 +08:00
file . write ( f " { info } \n " )
2022-10-09 13:01:10 +08:00
else :
txt_fullfn = None
2022-09-03 17:08:45 +08:00
2022-10-26 18:12:44 +08:00
script_callbacks . image_saved_callback ( params )
2022-10-09 13:01:10 +08:00
return fullfn , txt_fullfn
2022-10-13 07:17:26 +08:00
2023-07-03 18:09:37 +08:00
IGNORED_INFO_KEYS = {
' jfif ' , ' jfif_version ' , ' jfif_unit ' , ' jfif_density ' , ' dpi ' , ' exif ' ,
' loop ' , ' background ' , ' timestamp ' , ' duration ' , ' progressive ' , ' progression ' ,
' icc_profile ' , ' chromaticity ' , ' photoshop ' ,
}
2023-07-03 18:10:20 +08:00
def read_info_from_image ( image : Image . Image ) - > tuple [ str | None , dict ] :
2023-07-03 18:10:42 +08:00
items = ( image . info or { } ) . copy ( )
2022-11-27 21:28:32 +08:00
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-18 04:54:23 +08:00
if exif_comment :
items [ ' exif comment ' ] = exif_comment
geninfo = exif_comment
2022-11-27 21:28:32 +08:00
2023-07-03 18:09:37 +08:00
for field in IGNORED_INFO_KEYS :
2023-05-25 20:33:40 +08:00
items . pop ( field , None )
2022-11-27 21:28:32 +08:00
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 :
2023-06-01 00:56:37 +08:00
errors . report ( " Error parsing NovelAI image generation parameters " , exc_info = True )
2022-11-27 21:28:32 +08:00
return geninfo , items
2022-10-14 23:15:03 +08:00
def image_data ( data ) :
2023-03-25 12:29:51 +08:00
import gradio as gr
2022-10-14 23:15:03 +08:00
try :
image = Image . open ( io . BytesIO ( data ) )
2022-11-27 21:28:32 +08:00
textinfo , _ = read_info_from_image ( image )
2022-10-14 23:15:03 +08:00
return textinfo , None
except Exception :
pass
try :
text = data . decode ( ' utf8 ' )
assert len ( text ) < 10000
return text , None
except Exception :
pass
2023-03-25 12:29:51 +08:00
return gr . update ( ) , None
2022-12-24 14:46:35 +08: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 ' )