2023-01-23 09:24:43 +03:00
import json
import html
import os
import platform
import sys
import gradio as gr
2023-01-23 15:50:32 +03:00
import subprocess as sp
2023-01-23 09:24:43 +03:00
from modules import call_queue , shared
from modules . generation_parameters_copypaste import image_from_url_text
import modules . images
2023-05-31 22:40:09 +03:00
from modules . ui_components import ToolButton
2023-08-09 10:25:35 +03:00
import modules . generation_parameters_copypaste as parameters_copypaste
2023-01-23 09:24:43 +03:00
folder_symbol = ' \U0001f4c2 ' # 📂
2023-05-31 22:40:09 +03:00
refresh_symbol = ' \U0001f504 ' # 🔄
2023-01-23 09:24:43 +03:00
def update_generation_info ( generation_info , html_info , img_index ) :
try :
generation_info = json . loads ( generation_info )
if img_index < 0 or img_index > = len ( generation_info [ " infotexts " ] ) :
return html_info , gr . update ( )
return plaintext_to_html ( generation_info [ " infotexts " ] [ img_index ] ) , gr . update ( )
except Exception :
pass
# if the json parse or anything else fails, just return the old html_info
return html_info , gr . update ( )
2023-07-15 08:07:25 +03:00
def plaintext_to_html ( text , classname = None ) :
content = " <br> \n " . join ( html . escape ( x ) for x in text . split ( ' \n ' ) )
return f " <p class= ' { classname } ' > { content } </p> " if classname else f " <p> { content } </p> "
2023-01-23 09:24:43 +03:00
def save_files ( js_data , images , do_make_zip , index ) :
import csv
filenames = [ ]
fullfns = [ ]
#quick dictionary to class object conversion. Its necessary due apply_filename_pattern requiring it
class MyObject :
def __init__ ( self , d = None ) :
if d is not None :
for key , value in d . items ( ) :
setattr ( self , key , value )
data = json . loads ( js_data )
p = MyObject ( data )
path = shared . opts . outdir_save
save_to_dirs = shared . opts . use_save_to_dirs_for_ui
extension : str = shared . opts . samples_format
start_index = 0
2023-05-24 18:17:02 +03:00
only_one = False
2023-01-23 09:24:43 +03:00
if index > - 1 and shared . opts . save_selected_only and ( index > = data [ " index_of_first_image " ] ) : # ensures we are looking at a specific non-grid picture, and we have save_selected_only
2023-05-24 18:17:02 +03:00
only_one = True
2023-01-23 09:24:43 +03:00
images = [ images [ index ] ]
start_index = index
os . makedirs ( shared . opts . outdir_save , exist_ok = True )
with open ( os . path . join ( shared . opts . outdir_save , " log.csv " ) , " a " , encoding = " utf8 " , newline = ' ' ) as file :
at_start = file . tell ( ) == 0
writer = csv . writer ( file )
if at_start :
writer . writerow ( [ " prompt " , " seed " , " width " , " height " , " sampler " , " cfgs " , " steps " , " filename " , " negative_prompt " ] )
for image_index , filedata in enumerate ( images , start_index ) :
image = image_from_url_text ( filedata )
is_grid = image_index < p . index_of_first_image
i = 0 if is_grid else ( image_index - p . index_of_first_image )
2023-05-24 18:17:02 +03:00
p . batch_index = image_index - 1
2023-01-23 09:24:43 +03:00
fullfn , txt_fullfn = modules . images . save_image ( image , path , " " , seed = p . all_seeds [ i ] , prompt = p . all_prompts [ i ] , extension = extension , info = p . infotexts [ image_index ] , grid = is_grid , p = p , save_to_dirs = save_to_dirs )
filename = os . path . relpath ( fullfn , path )
filenames . append ( filename )
fullfns . append ( fullfn )
if txt_fullfn :
filenames . append ( os . path . basename ( txt_fullfn ) )
fullfns . append ( txt_fullfn )
writer . writerow ( [ data [ " prompt " ] , data [ " seed " ] , data [ " width " ] , data [ " height " ] , data [ " sampler_name " ] , data [ " cfg_scale " ] , data [ " steps " ] , filenames [ 0 ] , data [ " negative_prompt " ] ] )
# Make Zip
if do_make_zip :
2023-05-24 18:17:02 +03:00
zip_fileseed = p . all_seeds [ index - 1 ] if only_one else p . all_seeds [ 0 ]
namegen = modules . images . FilenameGenerator ( p , zip_fileseed , p . all_prompts [ 0 ] , image , True )
zip_filename = namegen . apply ( shared . opts . grid_zip_filename_pattern or " [datetime]_[[model_name]]_[seed]-[seed_last] " )
zip_filepath = os . path . join ( path , f " { zip_filename } .zip " )
2023-01-23 09:24:43 +03:00
from zipfile import ZipFile
with ZipFile ( zip_filepath , " w " ) as zip_file :
for i in range ( len ( fullfns ) ) :
with open ( fullfns [ i ] , mode = " rb " ) as f :
zip_file . writestr ( filenames [ i ] , f . read ( ) )
fullfns . insert ( 0 , zip_filepath )
return gr . File . update ( value = fullfns , visible = True ) , plaintext_to_html ( f " Saved: { filenames [ 0 ] } " )
def create_output_panel ( tabname , outdir ) :
def open_folder ( f ) :
if not os . path . exists ( f ) :
print ( f ' Folder " { f } " does not exist. After you create an image, the folder will be created. ' )
return
elif not os . path . isdir ( f ) :
print ( f """
WARNING
An open_folder request was made with an argument that is not a folder .
This could be an error or a malicious attempt to run code on your computer .
Requested path was : { f }
""" , file=sys.stderr)
return
if not shared . cmd_opts . hide_ui_dir_config :
path = os . path . normpath ( f )
if platform . system ( ) == " Windows " :
os . startfile ( path )
elif platform . system ( ) == " Darwin " :
sp . Popen ( [ " open " , path ] )
elif " microsoft-standard-WSL2 " in platform . uname ( ) . release :
sp . Popen ( [ " wsl-open " , path ] )
else :
sp . Popen ( [ " xdg-open " , path ] )
with gr . Column ( variant = ' panel ' , elem_id = f " { tabname } _results " ) :
with gr . Group ( elem_id = f " { tabname } _gallery_container " ) :
2023-08-18 12:15:30 +03:00
result_gallery = gr . Gallery ( label = ' Output ' , show_label = False , elem_id = f " { tabname } _gallery " , columns = 4 , preview = True , height = shared . opts . gallery_height or None )
2023-01-23 09:24:43 +03:00
generation_info = None
with gr . Column ( ) :
2023-03-21 08:18:14 +03:00
with gr . Row ( elem_id = f " image_buttons_ { tabname } " , elem_classes = " image-buttons " ) :
2023-08-12 19:03:33 +03:00
open_folder_button = ToolButton ( folder_symbol , elem_id = f ' { tabname } _open_folder ' , visible = not shared . cmd_opts . hide_ui_dir_config , tooltip = " Open images output directory. " )
2023-01-23 09:24:43 +03:00
if tabname != " extras " :
2023-08-12 19:03:33 +03:00
save = ToolButton ( ' 💾 ' , elem_id = f ' save_ { tabname } ' , tooltip = f " Save the image to a dedicated directory ( { shared . opts . outdir_save } ). " )
save_zip = ToolButton ( ' 🗃️ ' , elem_id = f ' save_zip_ { tabname } ' , tooltip = f " Save zip archive with images to a dedicated directory ( { shared . opts . outdir_save } ) " )
buttons = {
' img2img ' : ToolButton ( ' 🖼️ ' , elem_id = f ' { tabname } _send_to_img2img ' , tooltip = " Send image and generation parameters to img2img tab. " ) ,
' inpaint ' : ToolButton ( ' 🎨️ ' , elem_id = f ' { tabname } _send_to_inpaint ' , tooltip = " Send image and generation parameters to img2img inpaint tab. " ) ,
' extras ' : ToolButton ( ' 📐 ' , elem_id = f ' { tabname } _send_to_extras ' , tooltip = " Send image and generation parameters to extras tab. " )
}
2023-01-23 09:24:43 +03:00
open_folder_button . click (
fn = lambda : open_folder ( shared . opts . outdir_samples or outdir ) ,
inputs = [ ] ,
outputs = [ ] ,
)
if tabname != " extras " :
2023-03-27 10:20:01 +03:00
download_files = gr . File ( None , file_count = " multiple " , interactive = False , show_label = False , visible = False , elem_id = f ' download_files_ { tabname } ' )
2023-01-23 09:24:43 +03:00
with gr . Group ( ) :
2023-03-25 12:43:14 +03:00
html_info = gr . HTML ( elem_id = f ' html_info_ { tabname } ' , elem_classes = " infotext " )
2023-07-15 08:07:25 +03:00
html_log = gr . HTML ( elem_id = f ' html_log_ { tabname } ' , elem_classes = " html-log " )
2023-01-23 09:24:43 +03:00
generation_info = gr . Textbox ( visible = False , elem_id = f ' generation_info_ { tabname } ' )
if tabname == ' txt2img ' or tabname == ' img2img ' :
generation_info_button = gr . Button ( visible = False , elem_id = f " { tabname } _generation_info_button " )
generation_info_button . click (
fn = update_generation_info ,
_js = " function(x, y, z) { return [x, y, selected_gallery_index()] } " ,
inputs = [ generation_info , html_info , html_info ] ,
outputs = [ html_info , html_info ] ,
2023-03-20 16:09:36 +03:00
show_progress = False ,
2023-01-23 09:24:43 +03:00
)
save . click (
fn = call_queue . wrap_gradio_call ( save_files ) ,
_js = " (x, y, z, w) => [x, y, false, selected_gallery_index()] " ,
inputs = [
generation_info ,
result_gallery ,
html_info ,
html_info ,
] ,
outputs = [
download_files ,
html_log ,
] ,
show_progress = False ,
)
save_zip . click (
fn = call_queue . wrap_gradio_call ( save_files ) ,
_js = " (x, y, z, w) => [x, y, true, selected_gallery_index()] " ,
inputs = [
generation_info ,
result_gallery ,
html_info ,
html_info ,
] ,
outputs = [
download_files ,
html_log ,
]
)
else :
html_info_x = gr . HTML ( elem_id = f ' html_info_x_ { tabname } ' )
2023-03-25 12:43:14 +03:00
html_info = gr . HTML ( elem_id = f ' html_info_ { tabname } ' , elem_classes = " infotext " )
2023-01-23 09:24:43 +03:00
html_log = gr . HTML ( elem_id = f ' html_log_ { tabname } ' )
2023-02-14 14:55:42 +03:00
paste_field_names = [ ]
if tabname == " txt2img " :
paste_field_names = modules . scripts . scripts_txt2img . paste_field_names
elif tabname == " img2img " :
paste_field_names = modules . scripts . scripts_img2img . paste_field_names
2023-01-30 00:25:30 +03:00
for paste_tabname , paste_button in buttons . items ( ) :
parameters_copypaste . register_paste_params_button ( parameters_copypaste . ParamBinding (
2023-02-14 14:55:42 +03:00
paste_button = paste_button , tabname = paste_tabname , source_tabname = " txt2img " if tabname == " txt2img " else None , source_image_component = result_gallery ,
paste_field_names = paste_field_names
2023-01-30 00:25:30 +03:00
) )
2023-01-23 09:24:43 +03:00
return result_gallery , generation_info if tabname != " extras " else html_info_x , html_info , html_log
2023-05-31 22:40:09 +03:00
def create_refresh_button ( refresh_component , refresh_method , refreshed_args , elem_id ) :
2023-08-03 22:46:57 +03:00
refresh_components = refresh_component if isinstance ( refresh_component , list ) else [ refresh_component ]
label = None
for comp in refresh_components :
label = getattr ( comp , ' label ' , None )
if label is not None :
break
2023-05-31 22:40:09 +03:00
def refresh ( ) :
refresh_method ( )
args = refreshed_args ( ) if callable ( refreshed_args ) else refreshed_args
for k , v in args . items ( ) :
2023-08-03 22:46:57 +03:00
for comp in refresh_components :
setattr ( comp , k , v )
2023-05-31 22:40:09 +03:00
2023-08-05 09:17:36 +03:00
return [ gr . update ( * * ( args or { } ) ) for _ in refresh_components ] if len ( refresh_components ) > 1 else gr . update ( * * ( args or { } ) )
2023-05-31 22:40:09 +03:00
2023-08-03 22:46:57 +03:00
refresh_button = ToolButton ( value = refresh_symbol , elem_id = elem_id , tooltip = f " { label } : refresh " if label else " Refresh " )
2023-05-31 22:40:09 +03:00
refresh_button . click (
fn = refresh ,
inputs = [ ] ,
2023-08-04 14:56:27 +03:00
outputs = refresh_components
2023-05-31 22:40:09 +03:00
)
return refresh_button
2023-08-03 22:46:57 +03:00
def setup_dialog ( button_show , dialog , * , button_close = None ) :
""" Sets up the UI so that the dialog (gr.Box) is invisible, and is only shown when buttons_show is clicked, in a fullscreen modal window. """
dialog . visible = False
button_show . click (
fn = lambda : gr . update ( visible = True ) ,
inputs = [ ] ,
outputs = [ dialog ] ,
2023-08-27 09:39:37 +03:00
) . then ( fn = None , _js = " function() { popupId( ' " + dialog . elem_id + " ' ); } " )
2023-08-03 22:46:57 +03:00
if button_close :
button_close . click ( fn = None , _js = " closePopup " )