diff --git a/modules/detailer/__init__.py b/modules/detailer/__init__.py index 2b6901196..a318fcd48 100644 --- a/modules/detailer/__init__.py +++ b/modules/detailer/__init__.py @@ -1,5 +1,5 @@ from .models import detailer_models -from .helper import detailer_opt, DetailerResult, list_models +from .helper import detailer_opt, DetailerResult, list_models, assign_prompts from .detailer import Detailer diff --git a/modules/detailer/detailer.py b/modules/detailer/detailer.py index 7470e00e1..29d07139d 100644 --- a/modules/detailer/detailer.py +++ b/modules/detailer/detailer.py @@ -5,7 +5,7 @@ import gradio as gr from PIL import Image, ImageDraw from modules.logger import log from modules import shared, processing, devices, processing_class, ui_common, ui_components, ui_symbols, images, extra_networks, sd_models -from modules.detailer import DetailerResult, detailer_opt +from modules.detailer import DetailerResult, detailer_opt, assign_prompts class Detailer(): @@ -175,8 +175,6 @@ class Detailer(): else: negative = negative.replace('[PROMPT]', orig_negative) negative = negative.replace('[prompt]', orig_negative) - prompt_lines = 99 * [p.strip() for p in prompt.split('\n')] - negative_lines = 99 * [n.strip() for n in negative.split('\n')] args = { 'detailer': True, @@ -234,13 +232,15 @@ class Detailer(): if detailer_opt(p, 'detailer_include_detections', 'detailer_save'): annotated = self.draw_masks(annotated, items, p=p) + resolved_prompts = assign_prompts(prompt, items) + resolved_negatives = assign_prompts(negative, items) for j, item in enumerate(items): if item.mask is None: continue pc.keep_prompts = True shared.sd_model.fail_on_switch_error = True - pc.prompt = prompt_lines[i*len(items)+j] - pc.negative_prompt = negative_lines[i*len(items)+j] + pc.prompt = resolved_prompts[j] + pc.negative_prompt = resolved_negatives[j] pc.prompts = [pc.prompt] pc.negative_prompts = [pc.negative_prompt] pc.prompts, pc.network_data = extra_networks.parse_prompts(pc.prompts, pc.network_data) diff --git a/modules/detailer/helper.py b/modules/detailer/helper.py index 70bfae236..5875d5a1a 100644 --- a/modules/detailer/helper.py +++ b/modules/detailer/helper.py @@ -1,8 +1,12 @@ import os +import re from PIL import Image from modules.logger import log +class_tag_re = re.compile(r'^\[class\s*=\s*([^\]]+)\]\s*(.*)$', re.IGNORECASE) + + def list_models(self): from modules.detailer import detailer_models from modules import shared @@ -34,6 +38,51 @@ def detailer_opt(p, attr, opts_attr=None): return getattr(shared.opts, opts_attr or attr, None) +def parse_prompt_lines(text: str): + """Split a detailer prompt into class-tagged templates and positional fallback lines. + + A line starting with '[CLASS=name]' or '[CLASS=name1,name2]' assigns its text to every + detection whose label matches one of the given class names (case-insensitive). All + other non-empty lines are kept, in order, as the legacy positional fallback used for + detections that don't match any class tag. + """ + class_map: dict[str, str] = {} + fallback: list[str] = [] + for line in (text or '').split('\n'): + line = line.strip() + m = class_tag_re.match(line) + if m: + names = [n.strip().lower() for n in m.group(1).split(',') if n.strip()] + for name in names: + class_map[name] = m.group(2).strip() + else: + fallback.append(line) + return class_map, fallback + + +def assign_prompts(text: str, items: list) -> list[str]: + """Resolve a detailer prompt/negative-prompt string into one entry per detection. + + Detections whose YOLO label matches a '[CLASS=name]' tag get that tag's text. + Remaining detections fall back to the untagged lines, applied positionally in + detection order and cycling if there are more detections than fallback lines + (matching prior behavior when no class tags are used). + """ + class_map, fallback = parse_prompt_lines(text) + if len(fallback) == 0: + fallback = [''] + resolved = [] + fallback_idx = 0 + for item in items: + label = (getattr(item, 'label', None) or '').strip().lower() + if label in class_map: + resolved.append(class_map[label]) + else: + resolved.append(fallback[fallback_idx % len(fallback)]) + fallback_idx += 1 + return resolved + + class DetailerResult: def __init__(self, cls: int, label: str, score: float, box: list[int], mask: Image.Image = None, item: Image.Image = None, width = 0, height = 0, args = None): if args is None: