Spaces:
Runtime error
Runtime error
| import os | |
| import random | |
| import inspect | |
| import pathlib | |
| import re | |
| import time | |
| import glob | |
| from random import choices | |
| import yaml | |
| import modules.scripts as scripts | |
| import modules.images as images | |
| import gradio as gr | |
| from modules.processing import Processed, process_images | |
| from modules.shared import opts, cmd_opts, state | |
| from modules import scripts, script_callbacks, shared | |
| from modules.styles import StyleDatabase | |
| import modules.textual_inversion.textual_inversion | |
| from modules.sd_samplers import samplers, samplers_for_img2img | |
| ALL_KEY = 'all yaml files' | |
| def get_index(items, item): | |
| try: | |
| return items.index(item) | |
| except Exception: | |
| return None | |
| def parse_tag(tag): | |
| return tag.replace("__", "").replace('<', '').replace('>', '').strip() | |
| def read_file_lines(file): | |
| f_lines = file.read().splitlines() | |
| lines = [] | |
| for line in f_lines: | |
| line = line.strip() | |
| # Check if we have a line that is not empty or starts with an # | |
| if line and not line.startswith('#'): | |
| lines.append(line) | |
| return lines | |
| # Wildcards | |
| class TagLoader: | |
| files = [] | |
| wildcard_location = os.path.join( | |
| pathlib.Path(inspect.getfile(lambda: None)).parent.parent, "wildcards") | |
| loaded_tags = {} | |
| missing_tags = set() | |
| def load_tags(self, file_path, verbose=False, cache_files=True): | |
| if cache_files and self.loaded_tags.get(file_path): | |
| return self.loaded_tags.get(file_path) | |
| txt_file_path = os.path.join(self.wildcard_location, f'{file_path}.txt') | |
| yaml_file_path = os.path.join(self.wildcard_location, | |
| f'{file_path}.yaml') | |
| if (file_path == ALL_KEY): | |
| key = ALL_KEY | |
| else: | |
| key = file_path.lower() | |
| if self.wildcard_location and os.path.isfile(txt_file_path): | |
| with open(txt_file_path, encoding="utf8") as file: | |
| self.files.append(f"{file_path}.txt") | |
| self.loaded_tags[key] = read_file_lines(file) | |
| if key is ALL_KEY and self.wildcard_location: | |
| files = glob.glob(os.path.join(self.wildcard_location, '**/*.yaml'), recursive=True) | |
| output = {} | |
| for file in files: | |
| with open(file, encoding="utf8") as file: | |
| self.files.append(f"{file_path}.yaml") | |
| try: | |
| data = yaml.safe_load(file) | |
| for item in data: | |
| if (hasattr(output, item) and verbose): | |
| print(f"Duplicate key {item} in {file}") | |
| output[item] = { | |
| x.lower().strip() | |
| for i, x in enumerate(data[item]['Tags']) | |
| } | |
| except yaml.YAMLError as exc: | |
| print(exc) | |
| self.loaded_tags[key] = output | |
| if self.wildcard_location and os.path.isfile(yaml_file_path): | |
| with open(yaml_file_path, encoding="utf8") as file: | |
| self.files.append(f"{file_path}.yaml") | |
| try: | |
| data = yaml.safe_load(file) | |
| output = {} | |
| for item in data: | |
| output[item] = { | |
| x.lower().strip() | |
| for i, x in enumerate(data[item]['Tags']) | |
| } | |
| self.loaded_tags[key] = output | |
| except yaml.YAMLError as exc: | |
| print(exc) | |
| if not os.path.isfile(yaml_file_path) and not os.path.isfile( | |
| txt_file_path): | |
| self.missing_tags.add(file_path) | |
| return self.loaded_tags.get(key) if self.loaded_tags.get( | |
| key) else [] | |
| # <yaml:[tag]> notation | |
| class TagSelector: | |
| def __init__(self, tag_loader, options): | |
| self.tag_loader = tag_loader | |
| self.previously_selected_tags = {} | |
| self.selected_options = dict(options).get('selected_options', {}) | |
| self.verbose = dict(options).get('verbose', False) | |
| self.cache_files = dict(options).get('cache_files', True) | |
| def get_tag_choice(self, parsed_tag, tags): | |
| if self.selected_options.get(parsed_tag.lower()) is not None: | |
| return tags[self.selected_options.get(parsed_tag.lower())] | |
| return choices(tags)[0] if len(tags) > 0 else "" | |
| def get_tag_group_choice(self, parsed_tag, groups, tags): | |
| #print('selected_options', self.selected_options) | |
| #print('groups', groups) | |
| #print('parsed_tag', parsed_tag) | |
| neg_groups = [x.strip().lower() for x in groups if x.startswith('--')] | |
| neg_groups_set = {x.replace('--', '') for x in neg_groups} | |
| any_groups = [{y.strip() | |
| for i, y in enumerate(x.lower().split('|'))} | |
| for x in groups if '|' in x] | |
| pos_groups = [ | |
| x.strip().lower() for x in groups | |
| if not x.startswith('--') and '|' not in x | |
| ] | |
| pos_groups_set = {x for x in pos_groups} | |
| # print('pos_groups', pos_groups_set) | |
| # print('negative_groups', neg_groups_set) | |
| # print('any_groups', any_groups) | |
| candidates = [] | |
| for tag in tags: | |
| tag_set = tags[tag] | |
| if len(list(pos_groups_set & tag_set)) != len(pos_groups_set): | |
| continue | |
| if len(list(neg_groups_set & tag_set)) > 0: | |
| continue | |
| if len(any_groups) > 0: | |
| any_groups_found = 0 | |
| for any_group in any_groups: | |
| if len(list(any_group & tag_set)) == 0: | |
| break | |
| any_groups_found += 1 | |
| if len(any_groups) != any_groups_found: | |
| continue | |
| candidates.append(tag) | |
| if len(candidates) > 0: | |
| if self.verbose: | |
| print( | |
| f'UmiAI: Found {len(candidates)} candidates for "{parsed_tag}" with tags: {groups}, first 10: {candidates[:10]}' | |
| ) | |
| return choices(candidates)[0] | |
| print( | |
| f'UmiAI: No tag candidates found for: "{parsed_tag}" with tags: {groups}' | |
| ) | |
| return "" | |
| def select(self, tag, groups=None): | |
| self.previously_selected_tags.setdefault(tag, 0) | |
| if tag.count(':')==2: | |
| return False | |
| if self.previously_selected_tags.get(tag) < 100: | |
| self.previously_selected_tags[tag] += 1 | |
| parsed_tag = parse_tag(tag) | |
| tags = self.tag_loader.load_tags(parsed_tag, self.verbose, self.cache_files) | |
| if groups and len(groups) > 0: | |
| return self.get_tag_group_choice(parsed_tag, groups, tags) | |
| if len(tags) > 0: | |
| return self.get_tag_choice(parsed_tag, tags) | |
| else: | |
| print( | |
| f'UmiAI: No tags found in wildcard file "{parsed_tag}" or file does not exist' | |
| ) | |
| return False | |
| print(f'loaded tag more than 100 times {parsed_tag}') | |
| return False | |
| class TagReplacer: | |
| def __init__(self, tag_selector, options): | |
| self.tag_selector = tag_selector | |
| self.options = options | |
| self.wildcard_regex = re.compile('((__|<)(.*?)(__|>))') | |
| self.opts_regexp = re.compile('(?<=\[)(.*?)(?=\])') | |
| def replace_wildcard(self, matches): | |
| if matches is None or len(matches.groups()) == 0: | |
| return "" | |
| match = matches.groups()[2] | |
| match_and_opts = match.split(':') | |
| if (len(match_and_opts) == 2): | |
| selected_tags = self.tag_selector.select( | |
| match_and_opts[0], self.opts_regexp.findall(match_and_opts[1])) | |
| else: | |
| global_opts = self.opts_regexp.findall(match) | |
| if len(global_opts) > 0: | |
| selected_tags = self.tag_selector.select(ALL_KEY, global_opts) | |
| else: | |
| selected_tags = self.tag_selector.select(match) | |
| if selected_tags: | |
| return selected_tags | |
| return matches[0] | |
| def replace_wildcard_recursive(self, prompt): | |
| p = self.wildcard_regex.sub(self.replace_wildcard, prompt) | |
| while p != prompt: | |
| prompt = p | |
| p = self.wildcard_regex.sub(self.replace_wildcard, prompt) | |
| return p | |
| def replace(self, prompt): | |
| return self.replace_wildcard_recursive(prompt) | |
| # handle {1$$this | that} notation | |
| class DynamicPromptReplacer: | |
| def __init__(self): | |
| self.re_combinations = re.compile(r"\{([^{}]*)}") | |
| def get_variant_weight(self, variant): | |
| split_variant = variant.split("%") | |
| if len(split_variant) == 2: | |
| num = split_variant[0] | |
| try: | |
| return int(num) | |
| except ValueError: | |
| print(f'{num} is not a number') | |
| return 0 | |
| def get_variant(self, variant): | |
| split_variant = variant.split("%") | |
| if len(split_variant) == 2: | |
| return split_variant[1] | |
| return variant | |
| def parse_range(self, range_str, num_variants): | |
| if range_str is None: | |
| return None | |
| parts = range_str.split("-") | |
| if len(parts) == 1: | |
| low = high = min(int(parts[0]), num_variants) | |
| elif len(parts) == 2: | |
| low = int(parts[0]) if parts[0] else 0 | |
| high = min(int(parts[1]), | |
| num_variants) if parts[1] else num_variants | |
| else: | |
| raise Exception(f"Unexpected range {range_str}") | |
| return min(low, high), max(low, high) | |
| def replace_combinations(self, match): | |
| if match is None or len(match.groups()) == 0: | |
| return "" | |
| combinations_str = match.groups()[0] | |
| variants = [s.strip() for s in combinations_str.split("|")] | |
| weights = [self.get_variant_weight(var) for var in variants] | |
| variants = [self.get_variant(var) for var in variants] | |
| splits = variants[0].split("$$") | |
| quantity = splits.pop(0) if len(splits) > 1 else str(1) | |
| variants[0] = splits[0] | |
| low_range, high_range = self.parse_range(quantity, len(variants)) | |
| quantity = random.randint(low_range, high_range) | |
| summed = sum(weights) | |
| zero_weights = weights.count(0) | |
| weights = list( | |
| map(lambda x: (100 - summed) / zero_weights | |
| if x == 0 else x, weights)) | |
| try: | |
| #print(f"choosing {quantity} tag from:\n{' , '.join(variants)}") | |
| picked = [] | |
| for x in range(quantity): | |
| choice = random.choices(variants, weights)[0] | |
| picked.append(choice) | |
| index = variants.index(choice) | |
| variants.pop(index) | |
| weights.pop(index) | |
| #print(f"Picked:\n{' , '.join(picked)}\n") | |
| return " , ".join(picked) | |
| except ValueError as e: | |
| return "" | |
| def replace(self, template): | |
| if template is None: | |
| return None | |
| return self.re_combinations.sub(self.replace_combinations, template) | |
| class OptionGenerator: | |
| def __init__(self, tag_loader): | |
| self.tag_loader = tag_loader | |
| def get_configurable_options(self): | |
| return self.tag_loader.load_tags('configuration') | |
| def get_option_choices(self, tag): | |
| return self.tag_loader.load_tags(parse_tag(tag)) | |
| def parse_options(self, options): | |
| tag_presets = {} | |
| for i, tag in enumerate(self.get_configurable_options()): | |
| parsed_tag = parse_tag(tag) | |
| location = get_index(self.tag_loader.load_tags(parsed_tag), | |
| options[i]) | |
| if location is not None: | |
| tag_presets[parsed_tag.lower()] = location | |
| return tag_presets | |
| class PromptGenerator: | |
| def __init__(self, options): | |
| self.tag_loader = TagLoader() | |
| self.tag_selector = TagSelector(self.tag_loader, options) | |
| self.negative_tag_generator = NegativePromptGenerator() | |
| self.settings_generator = SettingsGenerator() | |
| self.replacers = [ | |
| self.settings_generator, | |
| TagReplacer(self.tag_selector, options), | |
| DynamicPromptReplacer(), | |
| self.negative_tag_generator | |
| ] | |
| self.verbose = dict(options).get('verbose', False) | |
| def use_replacers(self, prompt): | |
| for replacer in self.replacers: | |
| prompt = replacer.replace(prompt) | |
| return prompt | |
| def generate_single_prompt(self, original_prompt): | |
| previous_prompt = original_prompt | |
| start = time.time() | |
| prompt = self.use_replacers(original_prompt) | |
| while previous_prompt != prompt: | |
| previous_prompt = prompt | |
| prompt = self.use_replacers(prompt) | |
| end = time.time() | |
| if self.verbose: | |
| print(f"Prompt generated in {end - start} seconds") | |
| return prompt | |
| def get_negative_tags(self): | |
| return self.negative_tag_generator.get_negative_tags() | |
| def get_setting_overrides(self): | |
| return self.settings_generator.get_setting_overrides() | |
| class NegativePromptGenerator: | |
| def __init__(self): | |
| self.negative_tag = set() | |
| def strip_negative_tags(self, tags): | |
| matches = re.findall('\*\*.*?\*\*', tags) | |
| if matches: | |
| for match in matches: | |
| self.negative_tag.add(match.replace("**", "")) | |
| tags = tags.replace(match, "") | |
| return tags | |
| def replace(self, prompt): | |
| return self.strip_negative_tags(prompt) | |
| def get_negative_tags(self): | |
| return " ".join(self.negative_tag) | |
| # @@settings@@ notation | |
| class SettingsGenerator: | |
| def __init__(self): | |
| self.re_setting_tags = re.compile(r"@@(.*?)@@") | |
| self.setting_overrides = {} | |
| self.type_mapping = { | |
| 'cfg_scale': float, | |
| 'sampler': str, | |
| 'steps': int, | |
| } | |
| def strip_setting_tags(self, prompt): | |
| matches = self.re_setting_tags.findall(prompt) | |
| if matches: | |
| for match in matches: | |
| for assignment in match.split("|"): | |
| key_raw, value = assignment.split("=") | |
| if not value: | |
| print( | |
| f"Invalid setting {assignment}, settings should assign a value" | |
| ) | |
| continue | |
| key_found = False | |
| for key in self.type_mapping.keys(): | |
| if key.startswith(key_raw): | |
| self.setting_overrides[key] = self.type_mapping[ | |
| key](value) | |
| key_found = True | |
| break | |
| if not key_found: | |
| print( | |
| f"Unknown setting {key_raw}, setting should be the starting part of: {', '.join(self.type_mapping.keys())}" | |
| ) | |
| prompt = prompt.replace('@@' + match + '@@', "") | |
| return prompt | |
| def replace(self, prompt): | |
| return self.strip_setting_tags(prompt) | |
| def get_setting_overrides(self): | |
| return self.setting_overrides | |
| class Script(scripts.Script): | |
| is_txt2img = False | |
| embedding_db = modules.textual_inversion.textual_inversion.EmbeddingDatabase() | |
| def __init__(self): | |
| pass | |
| embedding_dir = os.path.join( | |
| pathlib.Path(inspect.getfile(lambda: None)).parent.parent, "embeddings") | |
| self.embedding_db.add_embedding_dir(embedding_dir) | |
| self.embedding_db.load_textual_inversion_embeddings(force_reload=True) | |
| def title(self): | |
| return "Prompt generator" | |
| def show(self, is_img2img): | |
| return scripts.AlwaysVisible | |
| def ui(self, is_img2img): | |
| self.is_txt2img = is_img2img == False | |
| with gr.Group(): | |
| with gr.Row(): | |
| enabled = gr.Checkbox(label="UmiAI enabled", value=True) | |
| verbose = gr.Checkbox(label="Verbose logging", value=False) | |
| cache_files = gr.Checkbox(label="Cache files", value=True) | |
| same_seed = gr.Checkbox(label='Same prompt in batch', | |
| value=False) | |
| negative_prompt = gr.Checkbox(label='**negative keywords**', | |
| value=True) | |
| shared_seed = gr.Checkbox(label="Static wildcards", | |
| value=False) | |
| option_generator = OptionGenerator(TagLoader()) | |
| options = [ | |
| gr.Dropdown(label=opt, | |
| choices=["RANDOM"] + | |
| option_generator.get_option_choices(opt), | |
| value="RANDOM") | |
| for opt in option_generator.get_configurable_options() | |
| ] | |
| return [enabled, verbose, cache_files, same_seed, negative_prompt, shared_seed | |
| ] + options | |
| def process(self, p, enabled, verbose, cache_files, same_seed, negative_prompt, | |
| shared_seed, *args): | |
| if not enabled: | |
| return | |
| debug = False | |
| if debug: print(f'\nModel: {p.sampler_name}, Seed: {int(p.seed)}, Batch Count: {p.n_iter}, Batch Size: {p.batch_size}, CFG: {p.cfg_scale}, Steps: {p.steps}\nOriginal Prompt: "{p.prompt}"\nOriginal Negatives: "{p.negative_prompt}"\n') | |
| original_prompt = p.all_prompts[0] | |
| if hasattr(p, "all_negative_prompts"): # hasattr to fix crash on old webui versions | |
| original_negative = p.all_negative_prompts[0] | |
| else: | |
| original_negative = "" | |
| TagLoader.files.clear() | |
| original_prompt = p.all_prompts[0] | |
| option_generator = OptionGenerator(TagLoader()) | |
| options = { | |
| 'selected_options': option_generator.parse_options(args), | |
| 'verbose': verbose, | |
| 'cache_files': cache_files, | |
| } | |
| prompt_generator = PromptGenerator(options) | |
| for cur_count in range(p.n_iter): #Batch count | |
| for cur_batch in range(p.batch_size): #Batch Size | |
| index = p.batch_size * cur_count + cur_batch | |
| # pick same wildcard for a given seed | |
| if (shared_seed): | |
| random.seed(p.all_seeds[p.batch_size *cur_count if same_seed else index]) | |
| else: | |
| random.seed(time.time()) | |
| if debug: print(f'{"Batch #"+str(cur_count) if same_seed else "Prompt #"+str(index):=^30}') | |
| prompt_generator.negative_tag_generator.negative_tag = set() | |
| prompt = prompt_generator.generate_single_prompt(original_prompt) | |
| p.all_prompts[index] = prompt | |
| if debug: print(f'Prompt: "{prompt}"') | |
| negative = original_negative | |
| if negative_prompt and hasattr(p, "all_negative_prompts"): # hasattr to fix crash on old webui versions | |
| negative += prompt_generator.get_negative_tags() | |
| p.all_negative_prompts[index] = negative | |
| if debug: print(f'Negative: "{negative}\n"') | |
| # same prompt per batch | |
| if (same_seed): | |
| for index in range(index, index + p.batch_size): | |
| p.all_prompts[index] = prompt | |
| break | |
| def find_sampler_index(sampler_list, value): | |
| for index, elem in enumerate(sampler_list): | |
| if elem[0] == value or value in elem[2]: | |
| return index | |
| att_override = prompt_generator.get_setting_overrides() | |
| #print(att_override) | |
| for att in att_override.keys(): | |
| if not att.startswith("__"): | |
| if att == 'sampler': | |
| sampler_name = att_override[att] | |
| if self.is_txt2img: | |
| sampler_index = find_sampler_index( | |
| samplers, sampler_name) | |
| else: | |
| sampler_index = find_sampler_index( | |
| samplers_for_img2img, sampler_name) | |
| if (sampler_index != None): | |
| setattr(p, 'sampler_index', sampler_index) | |
| else: | |
| print( | |
| f"Sampler {sampler_name} not found in prompt {p.all_prompts[0]}" | |
| ) | |
| continue | |
| setattr(p, att, att_override[att]) | |
| if original_prompt != p.all_prompts[0]: | |
| p.extra_generation_params["Wildcard prompt"] = original_prompt | |
| if verbose: | |
| p.extra_generation_params["File includes"] = "|".join( | |
| TagLoader.files) | |
| from modules import sd_hijack | |
| path = os.path.join(scripts.basedir(), "embeddings") | |
| sd_hijack.model_hijack.embedding_db.add_embedding_dir(path) |