mirror of
https://github.com/vladmandic/automatic
synced 2026-08-30 00:50:59 +02:00
566b6629bd
Styles in nested subfolders that shared a base JSON name were being overwritten because load_style keyed self.styles by the raw style["name"] while the Style object stored the prefixed name. The extra networks page iterates dict keys, so only one card per name surfaced. Key the dict by the same prefixed name, and warn when a true duplicate (same prefix and same name) is still detected.
522 lines
22 KiB
Python
522 lines
22 KiB
Python
from __future__ import annotations
|
|
from typing import TYPE_CHECKING
|
|
import re
|
|
import os
|
|
import csv
|
|
import json
|
|
import time
|
|
import random
|
|
from modules import files_cache, shared, infotext, sd_models, sd_vae
|
|
from modules.logger import log
|
|
|
|
if TYPE_CHECKING:
|
|
from modules.processing_class import StableDiffusionProcessing
|
|
|
|
debug_enabled = os.environ.get('SD_STYLES_DEBUG', None) is not None
|
|
|
|
|
|
class Style:
|
|
def __init__(self, name: str, desc: str = "", prompt: str = "", negative_prompt: str = "", extra: str = "", wildcards: str = "", filename: str = "", preview: str = "", mtime: float = 0):
|
|
self.name = name
|
|
self.description = desc
|
|
self.prompt = prompt
|
|
self.negative_prompt = negative_prompt
|
|
self.extra = extra
|
|
self.wildcards = wildcards
|
|
self.filename = filename
|
|
self.preview = preview
|
|
self.mtime = mtime
|
|
|
|
|
|
def merge_prompts(style_prompt: str, prompt: str) -> str:
|
|
if "{prompt}" in style_prompt:
|
|
res = style_prompt.replace("{prompt}", prompt)
|
|
else:
|
|
original_prompt = prompt.strip()
|
|
style_prompt = style_prompt.strip()
|
|
parts = filter(None, (original_prompt, style_prompt))
|
|
if original_prompt.endswith(","):
|
|
res = " ".join(parts)
|
|
else:
|
|
res = ", ".join(parts)
|
|
return res
|
|
|
|
|
|
def apply_styles_to_prompt(prompt, styles):
|
|
for style in styles:
|
|
prompt = merge_prompts(style, prompt)
|
|
return prompt
|
|
|
|
|
|
def select_from_weighted_list(inner: str) -> str:
|
|
def _split_weight(p: str):
|
|
"""Split 'name:weight' where the colon separator must not be inside brackets. Returns (name, wstr) or None."""
|
|
depth = 0
|
|
last_colon = -1
|
|
for i, c in enumerate(p):
|
|
if c in '<([{':
|
|
depth += 1
|
|
elif c in '>)]}':
|
|
if depth > 0:
|
|
depth -= 1
|
|
elif c == ':' and depth == 0:
|
|
last_colon = i
|
|
if last_colon < 0:
|
|
return p, 1.0
|
|
return p[:last_colon].strip(), p[last_colon + 1:].strip()
|
|
|
|
if not inner or len(inner.strip()) == 0:
|
|
return ''
|
|
|
|
parts = [p.strip() for p in inner.split('|')]
|
|
weighted: dict[str, float] = {}
|
|
|
|
for p in parts:
|
|
name, weight = _split_weight(p)
|
|
try:
|
|
w = float(weight)
|
|
except Exception:
|
|
w = 1.0
|
|
weighted[name] = weighted.get(name, 0.0) + max(0.0, w)
|
|
|
|
W = sum(weighted.values())
|
|
if len(weighted) == 0 or W <= 0.0:
|
|
return inner
|
|
weighted = {k: v / W for k, v in weighted.items()} # normalize to sum=1
|
|
names, weights = zip(*weighted.items(), strict=False)
|
|
choice = random.choices(names, weights=weights, k=1)[0]
|
|
return choice
|
|
|
|
|
|
def apply_curly_braces_to_prompt(prompt, seed=-1):
|
|
# unweighted: woman with {white|green|{purple|yellow}} highlights and {red|blue} dress
|
|
# weighted: woman with {white:0.6|green:0.2|{purple|yellow}} highlights and {red:.6|blue:.4} dress
|
|
if not isinstance(prompt, str) or len(prompt) == 0:
|
|
return prompt
|
|
old_state = None
|
|
if seed > 0:
|
|
old_state = random.getstate()
|
|
random.seed(seed)
|
|
try:
|
|
pattern = re.compile(r'\{([^{}]*)\}', re.DOTALL) # innermost braces
|
|
while True:
|
|
m = pattern.search(prompt)
|
|
if not m:
|
|
break
|
|
inner = m.group(1)
|
|
choice = select_from_weighted_list(inner)
|
|
prompt = prompt[:m.start()] + choice + prompt[m.end():] # replace this specific span (slice-based) to avoid accidental other replacements
|
|
finally:
|
|
if old_state is not None:
|
|
random.setstate(old_state)
|
|
return prompt
|
|
|
|
|
|
def apply_file_wildcards(prompt, replaced = None, not_found = None, recursion=0, seed=-1, p: StableDiffusionProcessing | None = None):
|
|
if not_found is None:
|
|
not_found = []
|
|
if replaced is None:
|
|
replaced = []
|
|
def check_wildcard_files(prompt, wildcard, files, file_only=True):
|
|
trimmed = wildcard.replace('\\', os.path.sep).replace('/', os.path.sep).strip().lower()
|
|
for file in files:
|
|
if file_only:
|
|
paths = [os.path.splitext(file)[0].lower(), os.path.splitext(os.path.basename(file).lower())[0]] # fullname and basename
|
|
else:
|
|
paths = [os.path.splitext(p.lower())[0] for p in os.path.normpath(file).split(os.path.sep)] # every path component
|
|
paths.insert(0, os.path.splitext(file)[0].lower())
|
|
if (trimmed in paths) or (os.path.sep in trimmed and trimmed in paths[0]):
|
|
try:
|
|
with open(file, encoding='utf-8') as f:
|
|
lines = f.readlines()
|
|
lines = [line.split('#')[0].strip('\n').strip() for line in lines]
|
|
lines = [line for line in lines if len(line) > 0]
|
|
if len(lines) > 0:
|
|
choice = random.choice(lines)
|
|
if '|' in choice:
|
|
choice = random.choice(choice.split('|')).strip(' []{}\n')
|
|
prompt = prompt.replace(f"__{wildcard}__", choice, 1)
|
|
log.debug(f'Apply wildcard: select="{wildcard}" choice="{choice}" file="{file}" choices={len(lines)}')
|
|
replaced.append(wildcard)
|
|
if p is not None:
|
|
p.extra_generation_params['Wildcards'] = p.extra_generation_params.get('Wildcards', []) + [trimmed]
|
|
return prompt, True
|
|
except Exception as e:
|
|
log.error(f'Wildcards: wildcard={wildcard} file={file} {e}')
|
|
if not file_only:
|
|
return prompt, False
|
|
return check_wildcard_files(prompt, wildcard, files, file_only=False)
|
|
|
|
def get_wildcards(prompt):
|
|
matches = re.findall(r'__(.*?)__', prompt, re.DOTALL)
|
|
matches = [m for m in matches if m not in not_found]
|
|
# matches = [m for m in matches if m not in replaced]
|
|
return matches
|
|
|
|
recursion += 1
|
|
if not shared.opts.wildcards_enabled or recursion >= 10 or not isinstance(prompt, str) or len(prompt) == 0:
|
|
return prompt, replaced, not_found
|
|
wildcards = get_wildcards(prompt)
|
|
if len(wildcards) == 0:
|
|
return prompt, replaced, not_found
|
|
files = list(files_cache.list_files(shared.opts.wildcards_dir, ext_filter=[".txt"], recursive=True))
|
|
if len(files) == 0:
|
|
return prompt, replaced, not_found
|
|
for wildcard in wildcards:
|
|
prompt, found = check_wildcard_files(prompt, wildcard, files)
|
|
if found and wildcard in not_found:
|
|
not_found.remove(wildcard)
|
|
elif not found and wildcard not in not_found:
|
|
not_found.append(wildcard)
|
|
prompt, replaced, not_found = apply_file_wildcards(prompt, replaced, not_found, recursion, seed) # recursive until we get early return
|
|
return prompt, replaced, not_found
|
|
|
|
|
|
def apply_wildcards_to_prompt(prompt, all_wildcards, seed=-1, silent=False, p: StableDiffusionProcessing | None = None):
|
|
if prompt is None or len(prompt) == 0:
|
|
return prompt
|
|
old_state = None
|
|
if seed > 0 and len(all_wildcards) > 0:
|
|
old_state = random.getstate()
|
|
random.seed(seed)
|
|
replaced = {}
|
|
t0 = time.time()
|
|
for style_wildcards in all_wildcards:
|
|
wildcards = [x.strip() for x in style_wildcards.replace('\n', ' ').split(";") if len(x.strip()) > 0]
|
|
for wildcard in wildcards:
|
|
try:
|
|
what, words = wildcard.split("=", 1)
|
|
if what in prompt:
|
|
words = [x.strip() for x in words.split(",") if len(x.strip()) > 0]
|
|
word = random.choice(words)
|
|
prompt = prompt.replace(what, word)
|
|
replaced[what] = word
|
|
except Exception as e:
|
|
log.error(f'Wildcards: wildcard="{wildcard}" error={e}')
|
|
t1 = time.time()
|
|
prompt, replaced_file, not_found = apply_file_wildcards(prompt, [], [], recursion=0, seed=seed, p=p)
|
|
t2 = time.time()
|
|
if replaced and not silent:
|
|
log.debug(f'Apply wildcards: {replaced} path="{shared.opts.wildcards_dir}" type=style time={t1-t0:.2f}')
|
|
if (len(replaced_file) > 0 or len(not_found) > 0) and not silent:
|
|
log.debug(f'Apply wildcards: found={replaced_file} missing={not_found} path="{shared.opts.wildcards_dir}" type=file seed={seed} time={t2-t2:.2f}')
|
|
if old_state is not None:
|
|
random.setstate(old_state)
|
|
return prompt
|
|
|
|
|
|
def get_reference_style():
|
|
if getattr(shared.sd_model, 'sd_checkpoint_info', None) is None:
|
|
return None
|
|
name = shared.sd_model.sd_checkpoint_info.name
|
|
name = name.replace('\\', '/').replace('Diffusers/', '')
|
|
for k, v in shared.reference_models.items():
|
|
model_file = os.path.splitext(v.get('path', '').split('@')[0])[0].replace('huggingface/', '')
|
|
if k == name or model_file == name:
|
|
return v.get('extras', None)
|
|
return None
|
|
|
|
|
|
def apply_styles_to_extra(p, style: Style):
|
|
if style is None:
|
|
return
|
|
name_map = {
|
|
'sampler': 'sampler_name',
|
|
'size-1': 'width',
|
|
'size-2': 'height',
|
|
'model': 'sd_model_checkpoint',
|
|
'vae': 'sd_vae',
|
|
'unet': 'sd_unet',
|
|
'te': 'sd_text_encoder',
|
|
'refine': 'enable_hr',
|
|
'hires': 'hr_force',
|
|
}
|
|
name_exclude = [
|
|
'size',
|
|
]
|
|
reference_style = get_reference_style()
|
|
extra = infotext.parse(reference_style) if shared.opts.extra_network_reference_values else {}
|
|
style_extra = apply_wildcards_to_prompt(style.extra, [style.wildcards], silent=True, p=p)
|
|
style_extra = ' ' + style_extra.lower()
|
|
extra.update(infotext.parse(style_extra))
|
|
extra.pop('Prompt', None)
|
|
extra.pop('Negative prompt', None)
|
|
params = []
|
|
settings = []
|
|
skipped = []
|
|
|
|
for k, v in extra.items():
|
|
k = k.lower().replace(' ', '_')
|
|
if k in name_map: # rename some fields
|
|
k = name_map[k]
|
|
if k in name_exclude: # exclude some fields
|
|
continue
|
|
if hasattr(p, k):
|
|
orig = getattr(p, k)
|
|
if (type(orig) != type(v)) and (orig is not None):
|
|
if not (type(orig) == int and type(v) == float): # dont convert float to int
|
|
v = type(orig)(v)
|
|
setattr(p, k, v)
|
|
if debug_enabled:
|
|
log.trace(f'Apply style param: {k}={v}')
|
|
params.append(f'{k}={v}')
|
|
elif shared.opts.data_labels.get(k, None) is not None:
|
|
if debug_enabled:
|
|
log.trace(f'Apply style setting: {k}={v}')
|
|
shared.opts.data[k] = v
|
|
if k == 'sd_model_checkpoint':
|
|
sd_models.reload_model_weights()
|
|
if k == 'sd_vae':
|
|
sd_vae.reload_vae_weights()
|
|
settings.append(f'{k}={v}')
|
|
else:
|
|
if debug_enabled:
|
|
log.trace(f'Apply style skip: {k}={v}')
|
|
skipped.append(f'{k}={v}')
|
|
log.debug(f'Apply style: name="{style.name}" params={params} settings={settings} unknown={skipped} reference={True if reference_style else False}')
|
|
|
|
|
|
class StyleDatabase:
|
|
def __init__(self, opts):
|
|
from modules import paths
|
|
|
|
self.no_style = Style("None")
|
|
self.styles = {}
|
|
self.path = opts.styles_dir
|
|
self.built_in = opts.extra_networks_styles
|
|
if os.path.isfile(opts.styles_dir) or opts.styles_dir.endswith(".csv"):
|
|
legacy_file = opts.styles_dir
|
|
self.load_csv(legacy_file)
|
|
opts.styles_dir = os.path.join(paths.models_path, "styles")
|
|
self.path = opts.styles_dir
|
|
try:
|
|
os.makedirs(opts.styles_dir, exist_ok=True)
|
|
self.save_styles(opts.styles_dir, verbose=True)
|
|
log.debug(f'Migrated styles: file="{legacy_file}" folder="{opts.styles_dir}"')
|
|
self.reload()
|
|
except Exception as e:
|
|
log.error(f'styles failed to migrate: file="{legacy_file}" error={e}')
|
|
if not os.path.isdir(opts.styles_dir):
|
|
opts.styles_dir = os.path.join(paths.models_path, "styles")
|
|
self.path = opts.styles_dir
|
|
try:
|
|
os.makedirs(opts.styles_dir, exist_ok=True)
|
|
except Exception:
|
|
pass
|
|
|
|
def load_style(self, fn, prefix=None):
|
|
with open(fn, encoding='utf-8') as f:
|
|
new_style = None
|
|
try:
|
|
all_styles = json.load(f)
|
|
if type(all_styles) is dict:
|
|
all_styles = [all_styles]
|
|
for style in all_styles:
|
|
if type(style) is not dict or "name" not in style:
|
|
raise ValueError('cannot parse style')
|
|
basename = os.path.splitext(os.path.basename(fn))[0]
|
|
name = re.sub(r'[\t\r\n]', '', style.get("name", basename)).strip()
|
|
if prefix is not None:
|
|
name = os.path.join(prefix, name)
|
|
else:
|
|
name = os.path.join(os.path.dirname(os.path.relpath(fn, self.path)), name)
|
|
new_style = Style(
|
|
name=name,
|
|
desc=style.get('description', name),
|
|
prompt=style.get("prompt", ""),
|
|
negative_prompt=style.get("negative", ""),
|
|
extra=style.get("extra", ""),
|
|
wildcards=style.get("wildcards", ""),
|
|
preview=style.get("preview", None),
|
|
filename=fn,
|
|
mtime=os.path.getmtime(fn),
|
|
)
|
|
# key by the prefixed name so styles with the same base name in different
|
|
# subfolders do not collide and overwrite each other
|
|
if name in self.styles:
|
|
log.warning(f'Style duplicate name: name="{name}" file="{fn}" existing="{self.styles[name].filename}"')
|
|
self.styles[name] = new_style
|
|
except Exception as e:
|
|
log.error(f'Failed to load style: file="{fn}" error={e}')
|
|
return new_style
|
|
|
|
def reload(self):
|
|
t0 = time.time()
|
|
self.styles.clear()
|
|
|
|
def list_folder(folder):
|
|
import concurrent.futures
|
|
future_items = {}
|
|
style_files = list(files_cache.list_files(folder, ext_filter=['.json'], recursive=files_cache.not_hidden))
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=shared.max_workers) as executor:
|
|
for fn in style_files:
|
|
future_items[executor.submit(self.load_style, fn, None)] = fn
|
|
if self.built_in:
|
|
fn = os.path.join('data', 'art-styles.json')
|
|
future_items[executor.submit(self.load_style, fn, 'Reference')] = fn
|
|
for future in concurrent.futures.as_completed(future_items):
|
|
future.result()
|
|
|
|
self.built_in = shared.opts.extra_networks_styles
|
|
list_folder(self.path)
|
|
self.styles = dict(sorted(self.styles.items(), key=lambda style: style[1].filename))
|
|
t1 = time.time()
|
|
log.info(f'Available Styles: path="{self.path}" items={len(self.styles.keys())} time={t1-t0:.2f}')
|
|
|
|
def find_style(self, name):
|
|
found = [style for style in self.styles.values() if style.name == name]
|
|
return found[0] if len(found) > 0 else self.no_style
|
|
|
|
def get_style_prompts(self, styles):
|
|
if styles is None:
|
|
return []
|
|
if not isinstance(styles, list):
|
|
log.error(f'Styles invalid: {styles}')
|
|
return []
|
|
return [self.find_style(x).prompt for x in styles]
|
|
|
|
def get_negative_style_prompts(self, styles):
|
|
if styles is None:
|
|
return []
|
|
if not isinstance(styles, list):
|
|
log.error(f'Styles invalid: {styles}')
|
|
return []
|
|
return [self.find_style(x).negative_prompt for x in styles]
|
|
|
|
def apply_styles_to_prompts(self, prompts, negatives, styles, seeds, p: StableDiffusionProcessing | None = None):
|
|
if styles is None:
|
|
return prompts, negatives
|
|
if not isinstance(styles, list):
|
|
log.error(f'Styles invalid styles: {styles}')
|
|
return prompts, negatives
|
|
if prompts is None or not isinstance(prompts, list):
|
|
log.error(f'Styles invalid prompts: {prompts}')
|
|
return prompts, negatives
|
|
if seeds is None or not isinstance(seeds, list):
|
|
log.error(f'Styles invalid seeds: {seeds}')
|
|
return prompts, negatives
|
|
if p is not None:
|
|
if p.all_templates is None or len(p.all_templates) < len(prompts):
|
|
p.all_templates = prompts
|
|
if p.all_negative_templates is None or len(p.all_negative_templates) < len(negatives):
|
|
p.all_negative_templates = negatives
|
|
|
|
jobid = shared.state.begin('Styles')
|
|
parsed_positive = []
|
|
parsed_negative = []
|
|
random_state = random.getstate()
|
|
|
|
for i in range(len(prompts)):
|
|
if seeds[i]> 0:
|
|
random.seed(seeds[i])
|
|
prompt = prompts[i]
|
|
prompt = apply_curly_braces_to_prompt(prompt, seeds[i])
|
|
prompt = apply_styles_to_prompt(prompt, [self.find_style(x).prompt for x in styles])
|
|
prompt = apply_wildcards_to_prompt(prompt, [self.find_style(x).wildcards for x in styles], seeds[i], p=p)
|
|
parsed_positive.append(prompt)
|
|
for i in range(len(negatives)):
|
|
if seeds[i]> 0:
|
|
random.seed(seeds[i])
|
|
prompt = negatives[i]
|
|
prompt = apply_curly_braces_to_prompt(prompt, seeds[i])
|
|
prompt = apply_styles_to_prompt(prompt, [self.find_style(x).negative_prompt for x in styles])
|
|
prompt = apply_wildcards_to_prompt(prompt, [self.find_style(x).wildcards for x in styles], seeds[i], p=p)
|
|
parsed_negative.append(prompt)
|
|
|
|
random.setstate(random_state)
|
|
shared.state.end(jobid)
|
|
return parsed_positive, parsed_negative
|
|
|
|
def apply_styles_to_prompt(self, prompt, styles, wildcards:bool=True):
|
|
if styles is None:
|
|
return prompt
|
|
if not isinstance(styles, list):
|
|
log.error(f'Styles invalid: {styles}')
|
|
return prompt
|
|
prompt = apply_styles_to_prompt(prompt, [self.find_style(x).prompt for x in styles])
|
|
if wildcards:
|
|
prompt = apply_wildcards_to_prompt(prompt, [self.find_style(x).wildcards for x in styles])
|
|
return prompt
|
|
|
|
def apply_negative_styles_to_prompt(self, prompt, styles, wildcards:bool=True):
|
|
if styles is None:
|
|
return prompt
|
|
if not isinstance(styles, list):
|
|
log.error(f'Styles invalid: {styles}')
|
|
return prompt
|
|
prompt = apply_styles_to_prompt(prompt, [self.find_style(x).negative_prompt for x in styles])
|
|
if wildcards:
|
|
prompt = apply_wildcards_to_prompt(prompt, [self.find_style(x).wildcards for x in styles])
|
|
return prompt
|
|
|
|
def apply_styles_to_extra(self, p):
|
|
if len(getattr(p, 'original_prompt', '')) == 0:
|
|
p.original_prompt = p.prompt
|
|
if len(getattr(p, 'original_negative', '')) == 0:
|
|
p.original_negative = p.negative_prompt
|
|
|
|
if p.styles is None:
|
|
return
|
|
if p.styles is None or not isinstance(p.styles, list):
|
|
log.error(f'Styles invalid: {p.styles}')
|
|
return
|
|
|
|
for style in p.styles:
|
|
if style is None or style == '':
|
|
continue
|
|
s = self.find_style(style)
|
|
if s == self.no_style:
|
|
log.warning(f'Apply style: name="{style}" not found')
|
|
continue
|
|
apply_styles_to_extra(p, s)
|
|
|
|
def extract_comments(self, p):
|
|
if not isinstance(p.prompt, str):
|
|
return
|
|
match = re.search(r'/\*.*?\*/', p.prompt, flags=re.DOTALL)
|
|
if match:
|
|
comment = match.group()
|
|
p.prompt = p.prompt.replace(comment, '')
|
|
p.extra_generation_params['Comment'] = comment.replace('/*', '').replace('*/', '')
|
|
|
|
def save_styles(self, path, verbose=False):
|
|
for name in list(self.styles):
|
|
style = {
|
|
"name": name,
|
|
"prompt": self.styles[name].prompt,
|
|
"negative": self.styles[name].negative_prompt,
|
|
"extra": "",
|
|
"preview": "",
|
|
}
|
|
keepcharacters = (' ','.','_')
|
|
fn = "".join(c for c in name if c.isalnum() or c in keepcharacters).strip()
|
|
fn = os.path.join(path, fn + ".json")
|
|
try:
|
|
with open(fn, 'w', encoding='utf-8') as f:
|
|
json.dump(style, f, indent=2)
|
|
if verbose:
|
|
log.debug(f'Saved style: name={name} file="{fn}"')
|
|
except Exception as e:
|
|
log.error(f'Failed to save style: name={name} file="{path}" error={e}')
|
|
count = len(list(self.styles))
|
|
if count > 0:
|
|
log.debug(f'Saved styles: folder="{path}" items={count}')
|
|
|
|
def load_csv(self, legacy_file):
|
|
if not os.path.isfile(legacy_file):
|
|
return
|
|
with open(legacy_file, encoding="utf-8-sig", newline='') as file:
|
|
reader = csv.DictReader(file, skipinitialspace=True)
|
|
num = 0
|
|
for row in reader:
|
|
try:
|
|
name = row["name"]
|
|
prompt = row["prompt"] if "prompt" in row else row["text"]
|
|
negative = row.get("negative_prompt", "") if "negative_prompt" in row else row.get("negative", "")
|
|
self.styles[name] = Style(name, desc=name, prompt=prompt, negative_prompt=negative)
|
|
log.debug(f'Migrated style: {self.styles[name].__dict__}')
|
|
num += 1
|
|
except Exception:
|
|
log.error(f'Styles error: file="{legacy_file}" row={row}')
|
|
log.info(f'Load legacy styles: file="{legacy_file}" loaded={num} created={len(list(self.styles))}')
|