339 lines
13 KiB
Python
339 lines
13 KiB
Python
import re
|
|
from models import db, Character
|
|
from utils import _IDENTITY_KEYS, _WARDROBE_KEYS, _BODY_GROUP_KEYS, parse_orientation
|
|
|
|
|
|
def _dedup_tags(prompt_str):
|
|
"""Remove duplicate tags from a comma-separated prompt string, preserving first-occurrence order."""
|
|
seen = set()
|
|
result = []
|
|
for tag in prompt_str.split(','):
|
|
t = tag.strip()
|
|
if t and t.lower() not in seen:
|
|
seen.add(t.lower())
|
|
result.append(t)
|
|
return ', '.join(result)
|
|
|
|
|
|
def _cross_dedup_prompts(positive, negative):
|
|
"""Remove tags shared between positive and negative prompts.
|
|
|
|
Repeatedly strips the first occurrence from each side until the tag exists
|
|
on only one side. Equal counts cancel out completely; any excess on one side
|
|
retains the remainder, allowing deliberate overrides (e.g. adding a tag twice
|
|
in the positive while it appears once in the negative leaves one copy positive).
|
|
"""
|
|
def parse_tags(s):
|
|
return [t.strip() for t in s.split(',') if t.strip()]
|
|
|
|
pos_tags = parse_tags(positive)
|
|
neg_tags = parse_tags(negative)
|
|
|
|
shared = {t.lower() for t in pos_tags} & {t.lower() for t in neg_tags}
|
|
for tag_lower in shared:
|
|
while (
|
|
any(t.lower() == tag_lower for t in pos_tags) and
|
|
any(t.lower() == tag_lower for t in neg_tags)
|
|
):
|
|
pos_tags.pop(next(i for i, t in enumerate(pos_tags) if t.lower() == tag_lower))
|
|
neg_tags.pop(next(i for i, t in enumerate(neg_tags) if t.lower() == tag_lower))
|
|
|
|
return ', '.join(pos_tags), ', '.join(neg_tags)
|
|
|
|
|
|
def _resolve_character(character_slug):
|
|
"""Resolve a character_slug string (possibly '__random__') to a Character instance."""
|
|
if character_slug == '__random__':
|
|
return Character.query.order_by(db.func.random()).first()
|
|
if character_slug:
|
|
return Character.query.filter_by(slug=character_slug).first()
|
|
return None
|
|
|
|
|
|
def _ensure_character_fields(character, selected_fields, include_wardrobe=True, include_defaults=False):
|
|
"""Mutate selected_fields in place to include essential character identity/wardrobe/name keys.
|
|
|
|
include_wardrobe — also inject active wardrobe keys (default True)
|
|
include_defaults — also inject defaults::expression and defaults::pose (for outfit/look previews)
|
|
"""
|
|
identity = character.data.get('identity', {})
|
|
for key in _IDENTITY_KEYS:
|
|
if identity.get(key):
|
|
field_key = f'identity::{key}'
|
|
if field_key not in selected_fields:
|
|
selected_fields.append(field_key)
|
|
if include_defaults:
|
|
for key in ['expression', 'pose']:
|
|
if character.data.get('defaults', {}).get(key):
|
|
field_key = f'defaults::{key}'
|
|
if field_key not in selected_fields:
|
|
selected_fields.append(field_key)
|
|
if 'special::name' not in selected_fields:
|
|
selected_fields.append('special::name')
|
|
if include_wardrobe:
|
|
wardrobe = character.get_active_wardrobe()
|
|
for key in _WARDROBE_KEYS:
|
|
if wardrobe.get(key):
|
|
field_key = f'wardrobe::{key}'
|
|
if field_key not in selected_fields:
|
|
selected_fields.append(field_key)
|
|
|
|
|
|
def _append_background(prompts, character=None):
|
|
"""Append a (color-prefixed) simple background tag to prompts['main']."""
|
|
primary_color = character.data.get('styles', {}).get('primary_color', '') if character else ''
|
|
bg = f"{primary_color} simple background" if primary_color else "simple background"
|
|
prompts['main'] = f"{prompts['main']}, {bg}"
|
|
|
|
|
|
def build_prompt(data, selected_fields=None, default_fields=None, active_outfit='default'):
|
|
def is_selected(section, key):
|
|
# Priority:
|
|
# 1. Manual selection from form (if list is not empty)
|
|
# 2. Default fields (saved per character)
|
|
# 3. Select all (fallback)
|
|
if selected_fields is not None and len(selected_fields) > 0:
|
|
return f"{section}::{key}" in selected_fields
|
|
if default_fields:
|
|
return f"{section}::{key}" in default_fields
|
|
return True
|
|
|
|
identity = data.get('identity', {})
|
|
|
|
# Get wardrobe - handle both new nested format and legacy flat format
|
|
wardrobe_data = data.get('wardrobe', {})
|
|
if 'default' in wardrobe_data and isinstance(wardrobe_data.get('default'), dict):
|
|
# New nested format - get active outfit
|
|
wardrobe = wardrobe_data.get(active_outfit or 'default', wardrobe_data.get('default', {}))
|
|
else:
|
|
# Legacy flat format
|
|
wardrobe = wardrobe_data
|
|
|
|
defaults = data.get('defaults', {})
|
|
action_data = data.get('action', {})
|
|
style_data = data.get('style', {})
|
|
participants = data.get('participants', {})
|
|
|
|
# Helper: collect selected values from identity + wardrobe for a given group key
|
|
def _group_val(key):
|
|
parts = []
|
|
id_val = identity.get(key, '')
|
|
wd_val = wardrobe.get(key, '')
|
|
if id_val and is_selected('identity', key):
|
|
val = id_val
|
|
# Filter out conflicting tags from base if participants data is present
|
|
if participants and key == 'base':
|
|
val = re.sub(r'\b(1girl|1boy|solo)\b', '', val).replace(', ,', ',').strip(', ')
|
|
if val:
|
|
parts.append(val)
|
|
if wd_val and is_selected('wardrobe', key):
|
|
parts.append(wd_val)
|
|
return ', '.join(parts)
|
|
|
|
# 1. Main Prompt
|
|
parts = []
|
|
|
|
# Handle participants logic
|
|
if participants:
|
|
if participants.get('solo_focus') == 'true':
|
|
parts.append('(solo focus:1.2)')
|
|
orientation = participants.get('orientation', '')
|
|
if orientation:
|
|
parts.extend(parse_orientation(orientation))
|
|
else:
|
|
parts.append("(solo:1.2)")
|
|
|
|
# Use character_id (underscores to spaces) for tags compatibility
|
|
char_tag = data.get('character_id', '').replace('_', ' ')
|
|
if char_tag and is_selected('special', 'name'):
|
|
parts.append(char_tag)
|
|
|
|
# Add all body groups to main prompt
|
|
for key in _BODY_GROUP_KEYS:
|
|
val = _group_val(key)
|
|
if val:
|
|
parts.append(val)
|
|
|
|
# Add defaults (expression, pose, scene)
|
|
for key in ['expression', 'pose', 'scene']:
|
|
val = defaults.get(key)
|
|
if val and is_selected('defaults', key):
|
|
parts.append(val)
|
|
|
|
# Standard character styles
|
|
char_aesthetic = data.get('styles', {}).get('aesthetic')
|
|
if char_aesthetic and is_selected('styles', 'aesthetic'):
|
|
parts.append(f"{char_aesthetic} style")
|
|
|
|
# Styles Gallery logic
|
|
if style_data.get('artist_name') and is_selected('style', 'artist_name'):
|
|
parts.append(f"by {style_data['artist_name']}")
|
|
if style_data.get('artistic_style') and is_selected('style', 'artistic_style'):
|
|
parts.append(style_data['artistic_style'])
|
|
|
|
tags = data.get('tags', [])
|
|
if tags and is_selected('special', 'tags'):
|
|
parts.extend(tags)
|
|
|
|
lora = data.get('lora', {})
|
|
if lora.get('lora_triggers') and is_selected('lora', 'lora_triggers'):
|
|
parts.append(lora.get('lora_triggers'))
|
|
|
|
# 2. Face Prompt: head group + expression + action head
|
|
face_parts = []
|
|
if char_tag and is_selected('special', 'name'):
|
|
face_parts.append(char_tag)
|
|
head_val = _group_val('head')
|
|
if head_val:
|
|
face_parts.append(head_val)
|
|
if defaults.get('expression') and is_selected('defaults', 'expression'):
|
|
face_parts.append(defaults.get('expression'))
|
|
if action_data.get('head') and is_selected('action', 'head'):
|
|
face_parts.append(action_data.get('head'))
|
|
|
|
# 3. Hand Prompt: hands group + action hands
|
|
hand_parts = []
|
|
hands_val = _group_val('hands')
|
|
if hands_val:
|
|
hand_parts.append(hands_val)
|
|
if action_data.get('hands') and is_selected('action', 'hands'):
|
|
hand_parts.append(action_data.get('hands'))
|
|
|
|
# 4. Feet Prompt: feet group + action feet
|
|
feet_parts = []
|
|
feet_val = _group_val('feet')
|
|
if feet_val:
|
|
feet_parts.append(feet_val)
|
|
if action_data.get('feet') and is_selected('action', 'feet'):
|
|
feet_parts.append(action_data.get('feet'))
|
|
|
|
return {
|
|
"main": _dedup_tags(", ".join(parts)),
|
|
"face": _dedup_tags(", ".join(face_parts)),
|
|
"hand": _dedup_tags(", ".join(hand_parts)),
|
|
"feet": _dedup_tags(", ".join(feet_parts)),
|
|
}
|
|
|
|
|
|
def build_multi_prompt(char_a, char_b, extras_prompt=''):
|
|
"""Build prompts for a two-character generation using BREAK separation.
|
|
|
|
Returns dict with combined prompts (main, face, hand) and per-character
|
|
prompts (char_a_main, char_a_face, char_b_main, char_b_face) for the
|
|
per-person/face ADetailer passes.
|
|
"""
|
|
# Build individual prompts with all fields selected
|
|
prompts_a = build_prompt(char_a.data)
|
|
prompts_b = build_prompt(char_b.data)
|
|
|
|
# Strip solo/orientation tags from individual prompts — we'll add combined ones
|
|
_solo_orientation_tags = {
|
|
'solo', '(solo:1.2)', '(solo focus:1.2)',
|
|
'1girl', '1boy', '2girls', '2boys', '3girls', '3boys',
|
|
'hetero', 'yuri', 'yaoi',
|
|
'multiple girls', 'multiple boys',
|
|
}
|
|
|
|
def _strip_tags(prompt_str, tags_to_remove):
|
|
parts = [t.strip() for t in prompt_str.split(',') if t.strip()]
|
|
return ', '.join(p for p in parts if p.lower() not in tags_to_remove)
|
|
|
|
main_a = _strip_tags(prompts_a['main'], _solo_orientation_tags)
|
|
main_b = _strip_tags(prompts_b['main'], _solo_orientation_tags)
|
|
|
|
# Compute combined orientation
|
|
orient_a = char_a.data.get('participants', {}).get('orientation', '1F')
|
|
orient_b = char_b.data.get('participants', {}).get('orientation', '1F')
|
|
|
|
# Count total M and F across both characters
|
|
combined_m = orient_a.upper().count('M') + orient_b.upper().count('M')
|
|
combined_f = orient_a.upper().count('F') + orient_b.upper().count('F')
|
|
combined_orientation = f"{combined_m}M{combined_f}F" if combined_m else f"{combined_f}F"
|
|
if combined_f == 0:
|
|
combined_orientation = f"{combined_m}M"
|
|
orientation_tags = parse_orientation(combined_orientation)
|
|
|
|
# Build combined main prompt with BREAK separation
|
|
orientation_str = ', '.join(orientation_tags)
|
|
combined_main = f"{orientation_str}, {main_a} BREAK {main_b}"
|
|
if extras_prompt:
|
|
combined_main = f"{extras_prompt}, {combined_main}"
|
|
|
|
# Merge face/hand prompts for the hand detailer (shared, not per-character)
|
|
hand_parts = [p for p in [prompts_a['hand'], prompts_b['hand']] if p]
|
|
|
|
return {
|
|
"main": _dedup_tags(combined_main),
|
|
"face": "", # not used — per-character face prompts go to SEGS detailers
|
|
"hand": _dedup_tags(', '.join(hand_parts)),
|
|
# Per-character prompts for SEGS-based ADetailer passes
|
|
"char_a_main": _dedup_tags(main_a),
|
|
"char_a_face": _dedup_tags(prompts_a['face']),
|
|
"char_b_main": _dedup_tags(main_b),
|
|
"char_b_face": _dedup_tags(prompts_b['face']),
|
|
}
|
|
|
|
|
|
def build_extras_prompt(actions, outfits, scenes, styles, detailers):
|
|
"""Combine positive prompt text from all selected category items."""
|
|
parts = []
|
|
|
|
for action in actions:
|
|
data = action.data
|
|
lora = data.get('lora', {})
|
|
if lora.get('lora_triggers'):
|
|
parts.append(lora['lora_triggers'])
|
|
parts.extend(data.get('tags', []))
|
|
for key in _BODY_GROUP_KEYS:
|
|
val = data.get('action', {}).get(key)
|
|
if val:
|
|
parts.append(val)
|
|
|
|
for outfit in outfits:
|
|
data = outfit.data
|
|
wardrobe = data.get('wardrobe', {})
|
|
for key in _BODY_GROUP_KEYS:
|
|
val = wardrobe.get(key)
|
|
if val:
|
|
parts.append(val)
|
|
lora = data.get('lora', {})
|
|
if lora.get('lora_triggers'):
|
|
parts.append(lora['lora_triggers'])
|
|
parts.extend(data.get('tags', []))
|
|
|
|
for scene in scenes:
|
|
data = scene.data
|
|
scene_fields = data.get('scene', {})
|
|
for key in ['background', 'foreground', 'lighting']:
|
|
val = scene_fields.get(key)
|
|
if val:
|
|
parts.append(val)
|
|
lora = data.get('lora', {})
|
|
if lora.get('lora_triggers'):
|
|
parts.append(lora['lora_triggers'])
|
|
parts.extend(data.get('tags', []))
|
|
|
|
for style in styles:
|
|
data = style.data
|
|
style_fields = data.get('style', {})
|
|
if style_fields.get('artist_name'):
|
|
parts.append(f"by {style_fields['artist_name']}")
|
|
if style_fields.get('artistic_style'):
|
|
parts.append(style_fields['artistic_style'])
|
|
lora = data.get('lora', {})
|
|
if lora.get('lora_triggers'):
|
|
parts.append(lora['lora_triggers'])
|
|
|
|
for detailer in detailers:
|
|
data = detailer.data
|
|
prompt = data.get('prompt', '')
|
|
if isinstance(prompt, list):
|
|
parts.extend(p for p in prompt if p)
|
|
elif prompt:
|
|
parts.append(prompt)
|
|
lora = data.get('lora', {})
|
|
if lora.get('lora_triggers'):
|
|
parts.append(lora['lora_triggers'])
|
|
|
|
return ", ".join(p for p in parts if p)
|