API tools faq
paste
RiamMar

scripts/prompt_gallery.py

RiamMar
Aug 4th, 2023 (edited)
74
0
Never
Add comment
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 29.58 KB | None | 0 0
raw download clone embed print report
  1. import copy
  2. import os
  3. import random
  4. import sys
  5. import traceback
  6. import shlex
  7. from modules.paths_internal import extensions_dir
  8. import yaml
  9. import platform
  10. import subprocess as sp
  11. import shutil
  12. import json
  13. import tempfile
  14. import gradio as gr
  15. import csv
  16. import typing
  17. import base64
  18. import io
  19. from PIL import Image
  20. import mimetypes
  21.  
  22. mimetypes.init()
  23. mimetypes.add_type('application/javascript', '.js')
  24.  
  25. import modules.generation_parameters_copypaste as parameters_copypaste
  26. from modules.generation_parameters_copypaste import image_from_url_text
  27. import modules.scripts as scripts
  28. from modules.processing import Processed, process_images
  29.  
  30. from modules.shared import opts, cmd_opts, OptionInfo, hide_dirs, state
  31. import modules.shared as shared
  32.  
  33. if '__file__' in locals().keys():
  34.     root_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
  35.     root_path = os.path.abspath(os.path.join(root_path, os.path.pardir))
  36. else:
  37.     root_path = os.path.abspath(shared.script_path)
  38.  
  39. rela_path =  os.path.join('extensions')
  40.  
  41. try:
  42.     with open(os.path.join(extensions_dir, 'prompt_gallery_name.json')) as fd:
  43.         extension_name = json.load(fd)['name']
  44. except:
  45.     extension_name = "Prompt Gallery"
  46.  
  47. OUTPATH_SAMPLES = os.path.join(extensions_dir, extension_name, 'assets', 'preview')
  48. OUTPATH_GRIDS =  os.path.join(extensions_dir, extension_name, 'assets', 'grid')
  49.  
  50. BATCH_SIZE = 2
  51. N_ITER = 2
  52. STEPS = 30
  53. CFG_SCALE = 11.5
  54. WIDTH = 512
  55. HEIGHT = 768
  56. SAMPLER_INDEX = 1
  57. RESTORE_FACE = 'true'
  58. TILING = 'false'
  59. DO_NOT_SAVE_GRID = 'false'
  60.  
  61. EXCLUDED_TAGS = ['']
  62. global SKIP_EXISTS
  63. SKIP_EXISTS = True
  64.  
  65. OUTPUTS_DICT = list()
  66. OUTPUTS = {}
  67. rawDict = {}
  68. qc_dict = {}
  69. trg_img = ''
  70. current_folder = ''
  71.  
  72.  
  73. def options_section_def(section_identifier, options_dict):
  74.     ret = {}
  75.     for k, v in options_dict.items():
  76.         v.section = section_identifier
  77.         ret[k] = v.default
  78.     return ret
  79.  
  80. pg_templates = {}
  81.  
  82. pg_templates.update(options_section_def(('saving-images', "Saving images/grids"), {
  83.     "samples_save": OptionInfo(True, "Always save all generated images"),
  84.     "samples_format": OptionInfo('png', 'File format for images'),
  85.     "samples_filename_pattern": OptionInfo("[seed]", "Images filename pattern", component_args=hide_dirs),
  86.     "save_images_add_number": OptionInfo(True, "Add number to filename when saving", component_args=hide_dirs),
  87.  
  88.     "grid_save": OptionInfo(True, "Always save all generated image grids"),
  89.     "grid_format": OptionInfo('png', 'File format for grids'),
  90.     "grid_extended_filename": OptionInfo(False, "Add extended info (seed, prompt) to filename when saving grid"),
  91.     "grid_only_if_multiple": OptionInfo(True, "Do not save grids consisting of one picture"),
  92.     "grid_prevent_empty_spots": OptionInfo(False, "Prevent empty spots in grid (when set to autodetect)"),
  93.     "n_rows": OptionInfo(-1, "Grid row count; use -1 for autodetect and 0 for it to be same as batch size", gr.Slider, {"minimum": -1, "maximum": 16, "step": 1}),
  94.  
  95.     "enable_pnginfo": OptionInfo(True, "Save text information about generation parameters as chunks to png files"),
  96.     "save_txt": OptionInfo(False, "Create a text file next to every image with generation parameters."),
  97.     "jpeg_quality": OptionInfo(80, "Quality for saved jpeg images", gr.Slider, {"minimum": 1, "maximum": 100, "step": 1}),
  98.     "export_for_4chan": OptionInfo(True, "If PNG image is larger than 4MB or any dimension is larger than 4000, downscale and save copy as JPG"),
  99.  
  100.     "use_original_name_batch": OptionInfo(False, "Use original name for output filename during batch process in extras tab"),
  101.     "save_selected_only": OptionInfo(True, "When using 'Save' button, only save a single selected image"),
  102.     "do_not_add_watermark": OptionInfo(False, "Do not add watermark to images"),
  103.  
  104. }))
  105.  
  106. # pg_templates.update(options_section_def(('saving-paths', "Paths for saving"), {
  107. #     "outdir_samples": OptionInfo("", "Output directory for images; if empty, defaults to three directories below", component_args=hide_dirs),
  108. #     "outdir_txt2img_samples": OptionInfo("outputs/txt2img-images", 'Output directory for txt2img images', component_args=hide_dirs),
  109. #     "outdir_img2img_samples": OptionInfo("outputs/img2img-images", 'Output directory for img2img images', component_args=hide_dirs),
  110. #     "outdir_extras_samples": OptionInfo("outputs/extras-images", 'Output directory for images from extras tab', component_args=hide_dirs),
  111. #     "outdir_grids": OptionInfo("", "Output directory for grids; if empty, defaults to two directories below", component_args=hide_dirs),
  112. #     "outdir_txt2img_grids": OptionInfo("outputs/txt2img-grids", 'Output directory for txt2img grids', component_args=hide_dirs),
  113. #     "outdir_img2img_grids": OptionInfo("outputs/img2img-grids", 'Output directory for img2img grids', component_args=hide_dirs),
  114. #     "outdir_save": OptionInfo("log/images", "Directory for saving images using the Save button", component_args=hide_dirs),
  115. # }))
  116.  
  117.  
  118. pg_templates.update(options_section_def(('saving-to-dirs', "Saving to a directory"), {
  119.     "save_to_dirs": OptionInfo(False, "Save images to a subdirectory"),
  120.     "grid_save_to_dirs": OptionInfo(False, "Save grids to a subdirectory"),
  121.     "use_save_to_dirs_for_ui": OptionInfo(False, "When using \"Save\" button, save images to a subdirectory"),
  122.     "directories_filename_pattern": OptionInfo("", "Directory name pattern", component_args=hide_dirs),
  123.     "directories_max_prompt_words": OptionInfo(8, "Max prompt words for [prompt_words] pattern", gr.Slider, {"minimum": 1, "maximum": 20, "step": 1, **hide_dirs}),
  124. }))
  125.  
  126. pg_templates.update(options_section_def(('upscaling', "Upscaling"), {
  127.     "ESRGAN_tile": OptionInfo(192, "Tile size for ESRGAN upscalers. 0 = no tiling.", gr.Slider, {"minimum": 0, "maximum": 512, "step": 16}),
  128.     "ESRGAN_tile_overlap": OptionInfo(8, "Tile overlap, in pixels for ESRGAN upscalers. Low values = visible seam.", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}),
  129.     "realesrgan_enabled_models": OptionInfo(["R-ESRGAN 4x+", "R-ESRGAN 4x+ Anime6B"], "Select which Real-ESRGAN models to show in the web UI. (Requires restart)", gr.CheckboxGroup, lambda: {"choices": realesrgan_models_names()}),
  130.     "upscaler_for_img2img": OptionInfo(None, "Upscaler for img2img", gr.Dropdown, lambda: {"choices": [x.name for x in sd_upscalers]}),
  131.     "use_scale_latent_for_hires_fix": OptionInfo(False, "Upscale latent space image when doing hires. fix"),
  132. }))
  133.  
  134. pg_templates.update(options_section_def(('ui', "User interface"), {
  135.     "show_progressbar": OptionInfo(True, "Show progressbar"),
  136.     "show_progress_every_n_steps": OptionInfo(0, "Show image creation progress every N sampling steps. Set to 0 to disable. Set to -1 to show after completion of batch.", gr.Slider, {"minimum": -1, "maximum": 32, "step": 1}),
  137.     "show_progress_grid": OptionInfo(True, "Show previews of all images generated in a batch as a grid"),
  138.     "return_grid": OptionInfo(True, "Show grid in results for web"),
  139.     "do_not_show_images": OptionInfo(False, "Do not show any images in results for web"),
  140.     # "disable_weights_auto_swap": OptionInfo(False, "When reading generation parameters from text into UI (from PNG info or pasted text), do not change the selected model/checkpoint."),
  141.     "send_seed": OptionInfo(True, "Send seed when sending prompt or image to other interface"),
  142.     "font": OptionInfo("", "Font for image grids that have text"),
  143.     "js_modal_lightbox": OptionInfo(True, "Enable full page image viewer"),
  144.     "js_modal_lightbox_initially_zoomed": OptionInfo(True, "Show images zoomed in by default in full page image viewer"),
  145.     "show_progress_in_title": OptionInfo(True, "Show generation progress in window title."),
  146.     'localization': OptionInfo("None", "Localization (requires restart)", gr.Dropdown, lambda: {"choices": ["None"] + list(localization.localizations.keys())}, refresh=lambda: localization.list_localizations(cmd_opts.localizations_dir)),
  147. }))
  148.  
  149. map_sampler_to_idx = {
  150.     'Euler a': 0,
  151.     'Euler': 1,
  152.     'LMS': 2,
  153.     'Heun': 3,
  154.     'DPM2': 4,
  155.     'DPM2 a': 5,
  156.     'DPM fast': 6,
  157.     'DPM adaptive': 7,
  158.     'LMS Karras': 8,
  159.     'DPM2 Karras': 9,
  160.     'DPM2 a Karras': 10,
  161.     'DDIM': 11,
  162.     'PLMS': 12,
  163.     'DPM++ 2S a Karras': 13,
  164.     'DPM++ 2M Karras': 14,
  165.     'DPM++ SDE Karras': 15}
  166.  
  167. map_keys = {
  168.     "value": "prompt",
  169.     "negative": "negative_prompt"}
  170.  
  171. map_param = {
  172.     "sd_model": "sd_model",
  173.     "outpath_samples": "outpath_samples",
  174.     "outpath_grids": "outpath_grids",
  175.     "prompt_for_display": "prompt_for_display",
  176.     "styles": "styles",
  177.     "Seed": "seed",
  178.     "Variation seed strength": "subseed_strength",
  179.     "Variation seed": "subseed",
  180.     "seed_resize_from_h": "seed_resize_from_h",
  181.     "seed_resize_from_w": "seed_resize_from_w",
  182.     "Sampler": "sampler_index",
  183.     "batch_size": "batch_size",
  184.     "n_iter": "n_iter",
  185.     "Steps": "steps",
  186.     "CFG scale": "cfg_scale",
  187.     "width": "width",
  188.     "height": "height",
  189.     "restore_faces": "restore_faces",
  190.     "tiling": "tiling",
  191.     "do_not_save_samples": "do_not_save_samples",
  192.     "do_not_save_grid": "do_not_save_grid"}
  193.  
  194. def process_string_tag(tag):
  195.     return tag
  196.  
  197.  
  198. def process_int_tag(tag):
  199.     return int(tag)
  200.  
  201.  
  202. def process_float_tag(tag):
  203.     return float(tag)
  204.  
  205.  
  206. def process_boolean_tag(tag):
  207.     return True if (tag == "true") else False
  208.  
  209.  
  210. prompt_tags = {
  211.     "sd_model": None,
  212.     "outpath_samples": process_string_tag,
  213.     "outpath_grids": process_string_tag,
  214.     "prompt_for_display": process_string_tag,
  215.     "prompt": process_string_tag,
  216.     "negative_prompt": process_string_tag,
  217.     "styles": process_string_tag,
  218.     "seed": process_int_tag,
  219.     "subseed_strength": process_float_tag,
  220.     "subseed": process_int_tag,
  221.     "seed_resize_from_h": process_int_tag,
  222.     "seed_resize_from_w": process_int_tag,
  223.     "sampler_index": process_int_tag,
  224.     "batch_size": process_int_tag,
  225.     "n_iter": process_int_tag,
  226.     "steps": process_int_tag,
  227.     "cfg_scale": process_float_tag,
  228.     "width": process_int_tag,
  229.     "height": process_int_tag,
  230.     "restore_faces": process_boolean_tag,
  231.     "tiling": process_boolean_tag,
  232.     "do_not_save_samples": process_boolean_tag,
  233.     "do_not_save_grid": process_boolean_tag
  234. }
  235.  
  236. avatar_prompts = list()
  237. avatar_names = list()
  238. avatar_negatives = list()
  239. avatar_name = ""
  240.  
  241. def cmdargs(line):
  242.     args = shlex.split(line)
  243.     args[args.index('--outpath_samples')+1] = args[args.index('--outpath_samples')+1][:-1]
  244.     args[args.index('--outpath_grids')+1] = args[args.index('--outpath_grids')+1][:-1]
  245.     pos = 0
  246.     res = {}
  247.  
  248.     while pos < len(args):
  249.         arg = args[pos]
  250.  
  251.         assert arg.startswith("--"), f'must start with "--": {arg}'
  252.         tag = arg[2:]
  253.  
  254.         func = prompt_tags.get(tag, None)
  255.         assert func, f'unknown commandline option: {arg}'
  256.  
  257.         assert pos+1 < len(args), f'missing argument for command line option {arg}'
  258.  
  259.         val = args[pos+1]
  260.  
  261.         res[tag] = func(val)
  262.  
  263.         pos += 2
  264.  
  265.     return res
  266.  
  267.  
  268.  
  269. def add_param(key, value, cur_str):
  270.     cur_str += '--{key} {value} '.format(key=key, value=value)
  271.     return cur_str
  272.  
  273. def parse_size(i_width, i_height, str_size, cur_str):
  274.     i_width = str_size.split('x')[0]
  275.     i_height = str_size.split('x')[1]  
  276.  
  277. def parse_virariant_size(str_size, cur_str):
  278.     width = str_size.split('x')[0]
  279.     height = str_size.split('x')[1]
  280.     cur_str = add_param('seed_resize_from_w', width, cur_str)
  281.     cur_str = add_param('seed_resize_from_h', height, cur_str)    
  282.     return cur_str
  283.  
  284. def parse_param(param_str):
  285.     m_batch_size = BATCH_SIZE
  286.     m_n_iter = N_ITER
  287.     m_steps = STEPS
  288.     m_cfg_scale = CFG_SCALE
  289.     m_width = WIDTH
  290.     m_height = HEIGHT
  291.     m_sampler_index = SAMPLER_INDEX
  292.     # m_tiling = TILING
  293.     m_restore_faces = RESTORE_FACE
  294.     m_do_not_save_grid = DO_NOT_SAVE_GRID
  295.     # m_sd_model = sd_model
  296.     cur_line = ""
  297.     for item in param_str.split(', '):
  298.         if item == '':
  299.             continue
  300.         group = item.split(': ')
  301.         key = group[0]
  302.         value = group[1]
  303.         if key == 'Steps':
  304.             m_steps = value
  305.         elif key == "CFG scale":
  306.             m_cfg_scale = value
  307.         elif key == 'Sampler':
  308.             m_sampler_index = map_sampler_to_idx[value]
  309.         elif key == 'Size':
  310.             parse_size(m_width, m_height, value, cur_line)
  311.         elif key == 'Seed resize from':
  312.             cur_line = parse_virariant_size(value, cur_line)
  313.         elif key == 'Seed':
  314.             cur_line = add_param("seed", value, cur_line)
  315.         elif key == 'Variation seed strength':
  316.             cur_line = add_param("subseed_strength", value, cur_line)  
  317.         elif key == 'Variation seed':
  318.             cur_line = add_param("subseed", value, cur_line)  
  319.         # elif key == 'Model hash':
  320.         #     cur_line = add_param("sd_model", m_sd_model, cur_line)  
  321.     cur_line = add_param("batch_size", m_batch_size, cur_line)
  322.     cur_line = add_param("n_iter", m_n_iter, cur_line)
  323.     cur_line = add_param("steps", m_steps, cur_line)
  324.     cur_line = add_param("cfg_scale", m_cfg_scale, cur_line)
  325.     cur_line = add_param("sampler_index", m_sampler_index, cur_line)
  326.     cur_line = add_param("width", m_width, cur_line)
  327.     cur_line = add_param("height", m_height, cur_line)
  328.     cur_line = add_param("restore_faces", m_restore_faces, cur_line)
  329. #     cur_line = add_param("tiling", m_tiling, cur_line)
  330.     cur_line = add_param("do_not_save_grid", m_do_not_save_grid, cur_line)
  331.     return cur_line
  332.  
  333. def parse_yaml_dict(rawDict, tag, avatar_prompt, avatar_name, default_negative):
  334.     # depth-first-search
  335.     if 'value' in rawDict.keys() or 'negative' in rawDict.keys():
  336.         if SKIP_EXISTS:
  337.             if os.path.exists(os.path.join(OUTPATH_SAMPLES, tag, avatar_name+'.png')) or os.path.exists(os.path.join(OUTPATH_SAMPLES,tag,'Not-available.png')):
  338.                 print("Skip "+str(tag))
  339.                 return ""
  340.         cur = ""
  341.         parsed_param = False
  342.         m_positive = avatar_prompt
  343.         m_negative = default_negative
  344.         for item in rawDict.items():
  345.             key = item[0]
  346.             value = item[1]
  347.             if key == 'param':
  348.                 params = parse_param(rawDict['param'])
  349.                 parsed_param = True
  350.             elif key == 'value':
  351.                 m_positive = value + m_positive
  352.             elif key == 'negative':
  353.                 m_negative = value +','+ m_negative
  354.  
  355.         cur += "--{key} \"{value}\" ".format(key='prompt', value= m_positive)
  356.         cur += "--{key} \"{value}\" ".format(key='negative_prompt', value= m_negative)
  357.  
  358.         if parsed_param == False:
  359.             params = parse_param("")
  360.         cur += params
  361.         cur += '--outpath_samples \"{} \"'.format(str(os.path.join(OUTPATH_SAMPLES, str(tag), '')))
  362.         cur += ' --outpath_grids \"{} \"'.format(str(os.path.join(OUTPATH_GRIDS, str(tag), '')))
  363.         return cur
  364.     else:
  365.         for item in rawDict.items():
  366.             key = item[0]
  367.             ret = parse_yaml_dict(rawDict[key], tag if key=='' else key, avatar_prompt, avatar_name, default_negative)
  368.             if len(ret) != 0:
  369.                 if tag not in EXCLUDED_TAGS:
  370.                     OUTPUTS_DICT.append({'name': key,
  371.                                     'prompt': item[1]['value'] if 'value' in item[1].keys() else '',
  372.                                     'negative_prompt': item[1]['negative'] if 'negative' in item[1].keys() else ''})
  373.                 if tag in OUTPUTS.keys():
  374.                     OUTPUTS[tag].append(ret)
  375.                 else:
  376.                     OUTPUTS[tag] = [ret]
  377.         return ""
  378.  
  379.  
  380. def rename_preview(avatar_name):
  381.     if avatar_name == '':
  382.         print("Please select avatar name first.")
  383.         return
  384.     root = OUTPATH_SAMPLES
  385.     for folder in os.listdir(root):
  386.         files = os.listdir(os.path.join(root, folder))
  387.         if 'Not-available.png' in files:
  388.             print('Skip '+ folder + ' not available.')
  389.             continue
  390.         if avatar_name + '.png' in files:
  391.             continue
  392.         for each_avatar in avatar_names:
  393.             if each_avatar + '.png' in files:
  394.                 files.remove(each_avatar + '.png')
  395.         if len(files) == 1:
  396.             os.rename(os.path.join(root, folder, files[0]), os.path.join(root, folder, avatar_name + '.png'))
  397.         else:
  398.             print('There are 0 or more than 1 files in ' + folder)
  399.  
  400. def load_prompt_file(file):
  401.     if (file is None):
  402.         lines = []
  403.     else:
  404.         lines = [x.strip() for x in file.decode('utf8', errors='ignore').split("\n")]
  405.  
  406.     return None, "\n".join(lines), gr.update(lines=7)
  407.  
  408. def copy_from_prompt_app():
  409.     return []
  410.  
  411. def open_folder(f):
  412.     if not os.path.exists(f):
  413.         print(f'Folder "{f}" does not exist. After you create an image, the folder will be created.')
  414.         return
  415.     elif not os.path.isdir(f):
  416.         print(f"""
  417. WARNING
  418. An open_folder request was made with an argument that is not a folder.
  419. This could be an error or a malicious attempt to run code on your computer.
  420. Requested path was: {f}
  421. """, file=sys.stderr)
  422.         return
  423.  
  424.     if not shared.cmd_opts.hide_ui_dir_config:
  425.         path = os.path.normpath(f)
  426.         if platform.system() == "Windows":
  427.             os.startfile(path)
  428.         elif platform.system() == "Darwin":
  429.             sp.Popen(["open", path])
  430.         else:
  431.             sp.Popen(["xdg-open", path])
  432.  
  433. class PromptStyle(typing.NamedTuple):
  434.     name: str
  435.     prompt: str
  436.     negative_prompt: str
  437.  
  438. def save_styles() -> None:
  439.     if len(OUTPUTS.keys()) == 0:
  440.         return
  441.     path = os.path.join('./', 'styles.csv')
  442.     # Write to temporary file first, so we don't nuke the file if something goes wrong
  443.     fd, temp_path = tempfile.mkstemp(".csv")
  444.     with os.fdopen(fd, "a", encoding="utf-8-sig", newline='') as file:
  445.         # _fields is actually part of the public API: typing.NamedTuple is a replacement for collections.NamedTuple,
  446.         # and collections.NamedTuple has explicit documentation for accessing _fields. Same goes for _asdict()
  447.         writer = csv.DictWriter(file, fieldnames=PromptStyle._fields)
  448.         writer.writeheader()
  449.         for row in OUTPUTS_DICT:
  450.             writer.writerow({'name': row['name'], 'prompt': row['prompt'], 'negative_prompt': row['negative_prompt']})
  451.         # writer.writerows(style._asdict() for k,     style in self.styles.items())
  452.  
  453.     # Always keep a backup file around
  454.     if os.path.exists(path) and not os.path.exists(path + ".bak"):
  455.         shutil.move(path, path + ".bak")
  456.     shutil.move(temp_path, path)
  457.  
  458. def load_prompt(file, default_negative, dropdown, skip_exist):
  459.     global SKIP_EXISTS
  460.     SKIP_EXISTS = skip_exist
  461.     if dropdown == '' or file is None:
  462.         return
  463.     rawDict = yaml.load(file, Loader = yaml.BaseLoader)
  464.     default_negative = default_negative + ',' + avatar_negatives[avatar_names.index(dropdown)]
  465.     parse_yaml_dict(rawDict, "", avatar_prompts[avatar_names.index(dropdown)], dropdown, default_negative)
  466.     prompt_txt = ""
  467.     keys = list(filter(lambda x: x not in EXCLUDED_TAGS, OUTPUTS.keys()))
  468.     for style in keys:
  469.         for each_line in OUTPUTS[style]:
  470.             prompt_txt += each_line + '\n'
  471.     return [prompt_txt, gr.Row.update(visible=True)]
  472.  
  473. def load_avartar(avatar_dict, customize_tags_positive):
  474.     avatars = yaml.load(avatar_dict, yaml.BaseLoader)
  475.  
  476.     for name, prompt in avatars.items():
  477.         avatar_names.append(name)
  478.         if 'value' in prompt.keys():
  479.             avatar_prompts.append(customize_tags_positive + ', ' +  prompt['value'])
  480.         else:
  481.             avatar_prompts.append(customize_tags_positive + ', ' +  '')
  482.         if 'negative' in prompt.keys():
  483.             avatar_negatives.append(prompt['negative'])
  484.         else:
  485.             avatar_negatives.append('')
  486.     return [gr.Dropdown.update(choices=avatar_names, value=avatar_names[0]), gr.Column.update(visible=True),  gr.Group.update(visible=True)]
  487.  
  488. def scan_outputs(avatar_name):
  489.     if avatar_name is None or len(avatar_name) == 0:
  490.         print("Please select avatar name first.")
  491.         return
  492.     root = OUTPATH_SAMPLES
  493.     global qc_dict
  494.     qc_dict = {}
  495.  
  496.     if not os.path.exists(root):
  497.         os.mkdir(root)
  498.  
  499.     for folder in os.listdir(root):
  500.         if os.path.isdir(os.path.join(root, folder)) == False:
  501.             continue
  502.            
  503.         files = os.listdir(os.path.join(root, folder))
  504.         if 'Not-available.png' in files:
  505.             print('Skip '+ folder + ' not available.')
  506.             continue
  507.         if avatar_name + '.png' in files:
  508.             continue
  509.         for each_avatar in avatar_names:
  510.             if each_avatar + '.png' in files:
  511.                 files.remove(each_avatar + '.png')
  512.         if len(files) == 0:
  513.             continue
  514.         # print("Insert file:")
  515.         for file in files:
  516.             print(os.path.join(root, folder, file))
  517.         qc_dict[folder] = [os.path.join(root, folder, file) for file in files]
  518.  
  519.     if len(qc_dict.keys()) == 0:
  520.         return gr.Dropdown.update(choices=[])
  521.     return gr.Dropdown.update(choices=list(qc_dict.keys()), value=list(qc_dict.keys())[0])
  522.  
  523. def update_gallery(dropdown, avatar):
  524.     root = OUTPATH_SAMPLES
  525.     global trg_img, current_folder
  526.     current_folder = os.path.join(root, dropdown)
  527.     trg_img = os.path.join(root, dropdown, avatar + '.png')
  528.     # print("Detected folders:")
  529.     # print(qc_dict)
  530.     # print("Selected folder:")
  531.     # print(dropdown)
  532.     # print("Detected files:")
  533.     # print(qc_dict[dropdown])
  534.     return qc_dict[dropdown]
  535.  
  536. def clean_select_picture(filename):
  537.     if current_folder == '':
  538.         print("Please select qc tag.")
  539.         return
  540.     for file in os.listdir(current_folder):
  541.         is_avatar = False
  542.         for each_avatar in avatar_names:
  543.             if each_avatar + '.png' == file:
  544.                 is_avatar = True
  545.                 break
  546.         if '-' in file and file.split('.')[0] in filename:
  547.             print("rename", os.path.join(current_folder, file), trg_img)
  548.             os.rename(os.path.join(current_folder, file), trg_img)
  549.         elif is_avatar == False:
  550.             print(file, "delete", os.path.join(current_folder, file))
  551.             os.remove(os.path.join(current_folder, file))
  552.  
  553. def image_url(filedata):
  554.     if type(filedata) == list:
  555.         if len(filedata) == 0:
  556.             return None
  557.         filedata = filedata[0]
  558.         # for item in filedata:
  559.         #     if filedata["is_file"]:
  560.         #         filedata = item
  561.         #         filename = filedata["name"]
  562.         #         tempdir = os.path.normpath(tempfile.gettempdir())
  563.         #         normfn = os.path.normpath(filename)
  564.         #         assert normfn.startswith(tempdir), 'trying to open image file not in temporary directory'
  565.  
  566.         #         image = Image.open(filename)
  567.         #         clean_select_picture(os.path.basename(filename))
  568.                 # return Image.open(filename)
  569.     if type(filedata) == dict and filedata["is_file"]:
  570.         filename = filedata["name"]
  571.         tempdir = os.path.normpath(tempfile.gettempdir())
  572.         normfn = os.path.normpath(filename)
  573.         assert normfn.startswith(tempdir), 'trying to open image file not in temporary directory'
  574.  
  575.         image = Image.open(filename)
  576.         clean_select_picture(os.path.basename(filename))
  577.         return Image.open(filename)
  578.  
  579.     elif type(filedata) == dict:
  580.         print(filedata)
  581.         print("Dict is not file.")
  582.         return
  583.  
  584.  
  585.  
  586.     elif  type(filedata) != dict and filedata.startswith("data:image/png;base64,"):
  587.         filedata = filedata[len("data:image/png;base64,"):]
  588.  
  589.         filedata = base64.decodebytes(filedata.encode('utf-8'))
  590.         image = Image.open(io.BytesIO(filedata))
  591.         return image
  592.     return None
  593.        
  594.  
  595. def dropdown_change():
  596.     global OUTPUTS, OUTPUTS_DICT
  597.     default_negative = []
  598.     OUTPUTS = {}
  599.     OUTPUTS_DICT = []
  600.     return [ gr.File.update(value=None), gr.Textbox.update(value=None)]  
  601.                    
  602.  
  603. class Script(scripts.Script):
  604.     def title(self):
  605.         return "Prompt gallery"
  606.  
  607.     def ui(self, is_img2img):
  608.         with gr.Group():
  609.             with gr.Column():
  610.                 label_avatar = gr.Label("Upload avatars config")
  611.                 avatar_dict = gr.File(label="Upload avatar prompt inputs", type='bytes')
  612.        
  613.         # copy_from_app_button = gr.Button("Copy From Prompt Preview")
  614.  
  615.         with gr.Group():
  616.             with gr.Column(visible=False) as avatar_col:
  617.                 label_presets = gr.Label("Presets")
  618.                 dropdown = gr.Dropdown(label="Choose avatar", choices=[""], value="", type="value", elem_id="dropdown")
  619.                 dropdown.save_to_config = True
  620.                 with gr.Row():
  621.                     checkbox_iterate = gr.Checkbox(label="Iterate seed every line", value=False)
  622.                     skip_exist = gr.Checkbox(value=True, label="skip exist")
  623.                 default_negative = gr.Textbox(label="default_negative", lines=1)
  624.                 default_positive = gr.Textbox(label="default_positive", lines=1)
  625.                 prompt_dict = gr.File(label="Upload prompt dictionary", type='bytes')
  626.                 with gr.Row(visible = False) as save_prompts:
  627.                     open_button = gr.Button("Open outputs directory")
  628.                     export_button = gr.Button("Export to WebUI style")
  629.             prompt_display = gr.Textbox(label="List of prompt inputs", lines=1)
  630.  
  631.        
  632.         prompt_dict.change(fn=load_prompt, inputs=[prompt_dict, default_negative, dropdown, skip_exist], outputs=[prompt_display, save_prompts])
  633.         open_button.click(fn=lambda: open_folder(OUTPATH_SAMPLES), inputs=[], outputs=[])
  634.         export_button.click(fn=save_styles, inputs=[], outputs=[])
  635.  
  636.         with gr.Group(visible=False) as qc_widgets:
  637.             label_preview = gr.Label("QC preview")
  638.             with gr.Row():
  639.                 qc_refresh = gr.Button("QC scan")
  640.                 preview_dropdown = gr.Dropdown(label="Select prompts", choices=[""], value="", type="value", elem_id="dropdown")
  641.             preview_gallery = gr.Gallery(label='Output', show_label=False, elem_id=f"preview_gallery").style(grid=4)
  642.             qc_refresh.click(fn=scan_outputs, inputs=[dropdown], outputs=preview_dropdown)
  643.             with gr.Row():
  644.                 qc_show = gr.Button(f"Show pics")
  645.                 qc_select = gr.Button(f"Select")
  646.                 rename_button = gr.Button("Auto rename")
  647.             selected_img = gr.Image(label="Selected",show_label=False, source="upload", interactive=True, type="pil").style(height=480)
  648.         qc_show.click(fn=update_gallery, inputs=[preview_dropdown, dropdown], outputs=preview_gallery)
  649.         qc_select.click(
  650.                     fn=lambda x: image_url(x),
  651.                     _js="extract_image_from_gallery",
  652.                     inputs=[preview_gallery],
  653.                     outputs=[selected_img],
  654.         )
  655.             # qc_select.click(fn=select_picture, inputs=[dropdown, preview_dropdown, preview_gallery], outputs=[])
  656.         dropdown.change(fn=dropdown_change, inputs=[], outputs=[prompt_dict, prompt_display])
  657.         rename_button.click(fn=rename_preview, inputs=[dropdown], outputs=[])
  658.             # qc_select.click(fn=scan_outputs, inputs=[], outputs=[preview_dropdown])
  659.  
  660.         avatar_dict.change(fn=load_avartar, inputs=[avatar_dict, default_positive], outputs=[dropdown, avatar_col, qc_widgets])
  661.         return [checkbox_iterate, avatar_dict, prompt_dict, default_negative, default_positive, dropdown, prompt_display, rename_button, label_avatar, open_button, export_button, skip_exist, label_presets, label_preview, preview_dropdown, preview_gallery, qc_select, qc_refresh, qc_show, selected_img]
  662.  
  663.     def run(self, p, checkbox_iterate, avatar_dict, prompt_dict, default_negative, default_positive, dropdown, prompt_display, rename_button, label_avatar, open_button, export_button, skip_exist, label_presets, label_preview, preview_dropdown, preview_gallery, qc_select, qc_refresh, qc_show, selected_img):
  664.         global pg_templates
  665.         backup = copy.deepcopy(shared.opts)
  666.  
  667.         shared.opts.data.update(pg_templates)
  668.         lines = [x.strip() for x in prompt_display.splitlines()]
  669.         lines = [x for x in lines if len(x) > 0]
  670.  
  671.         p.do_not_save_grid = True
  672.  
  673.         job_count = 0
  674.         jobs = []
  675.  
  676.         for line in lines:
  677.             if "--" in line:
  678.                 try:
  679.                     args = cmdargs(line)
  680.                 except Exception:
  681.                     print(f"Error parsing line [line] as commandline:", file=sys.stderr)
  682.                     print(traceback.format_exc(), file=sys.stderr)
  683.                     args = {"prompt": line}
  684.             else:
  685.                 args = {"prompt": line}
  686.  
  687.             n_iter = args.get("n_iter", 1)
  688.             if n_iter != 1:
  689.                 job_count += n_iter
  690.             else:
  691.                 job_count += 1
  692.  
  693.             jobs.append(args)
  694.  
  695.         # print(f"Will process {len(lines)} lines in {job_count} jobs.")
  696.         if (checkbox_iterate and p.seed == -1):
  697.             p.seed = int(random.randrange(4294967294))
  698.  
  699.         state.job_count = job_count
  700.  
  701.         images = []
  702.         for n, args in enumerate(jobs):
  703.             state.job = f"{state.job_no + 1} out of {state.job_count}"
  704.  
  705.             copy_p = copy.copy(p)
  706.             for k, v in args.items():
  707.                 setattr(copy_p, k, v)
  708.  
  709.             proc = process_images(copy_p)
  710.             images += proc.images
  711.            
  712.             if (checkbox_iterate):
  713.                 p.seed = p.seed + (p.batch_size * p.n_iter)
  714.  
  715.         OUTPUTS = {}
  716.         rawDict = {}
  717.        
  718.         shared.opts = backup
  719.         return Processed(p, images, p.seed, "")
  720.  
Tags: Stable Diffusion automatic1111 Prompt Gallery
Add Comment
Please, Sign In to add comment
Public Pastes
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy.  OK, I Understand
Not a member of Pastebin yet?
Sign Up, it unlocks many cool features!