stable-diffusion-webui-prompt-parser2/scripts/prompt-refactor.py

# -*- coding: utf-8 -*-
# Prompt Parser Analyzer for A1111 WebUI
# Полноценная замена "Prompt Refactor": анализ, расписания, ошибки, конструкторы грамматики
#
# Требования: modules.prompt_parser = "второй файл"
#   (schedule_parser, get_schedule, at_step, at_step_from_schedule, CollectSteps, ScheduleTransformer и т.д.)
#
# Кладите этот файл в папку scripts/. В UI появится секция "Prompt Parser Analyzer".

import os
import re
import traceback
import importlib
from typing import Any, List, Tuple, Optional

import gradio as gr
from modules import script_callbacks, scripts, shared
from modules.shared import opts

# Импортируем парсер как модуль, чтобы уметь reload(...) после смены ENV-флагов
from modules import prompt_parser as pp  # должен быть «второй файл»

# -----------------------------
# Утилиты для ошибок и форматирования
# -----------------------------

def _highlight_at(text: str, line: int, column: int, context: int = 2) -> str:
    """Возвращает компактный подсвет с кареткой и близлежащими строками."""
    lines = text.splitlines()
    i = max(0, line - 1)
    start = max(0, i - context)
    end = min(len(lines), i + context + 1)

    out = []
    for idx in range(start, end):
        prefix = ">> " if idx == i else "   "
        out.append(f"{prefix}{idx+1:>4}: {lines[idx]}")
        if idx == i:
            caret = " " * (column + 7 + len(str(idx+1))) + "^"
            out.append(caret)
    return "\n".join(out)

def _guess_suggestion(src: str) -> List[str]:
    """Простые эвристики подсказок по частым ошибкам."""
    tips = []
    if src.count("(") != src.count(")"):
        tips.append("Скобки () несбалансированы — проверьте пары.")
    if src.count("[") != src.count("]"):
        tips.append("Скобки [] несбалансированы — проверьте пары.")
    if src.count("{") != src.count("}"):
        tips.append("Скобки {} несбалансированы — проверьте пары.")
    if "::" in src and not any(term in src for term in ["!", ";", "!!"]):
        tips.append("Похоже, вы используете sequence (::), но не поставили терминатор (! или ; / !! для top-level).")
    if re.search(r":[A-Za-z]", src):
        tips.append("После ':' обычно ожидается число для веса (например, :1.2) — проверьте места с двоеточием.")
    if re.search(r"\[[^]]+\]!\s*!", src):
        tips.append("alternate_distinct записывается как [a|b]! — одно '!' сразу после ']'.")
    return tips

def format_lark_error(prompt: str, err: Exception) -> str:
    """Превращаем Lark-исключение в понятный отчёт с подсветкой и советами."""
    title = f"**Ошибка парсинга:** `{type(err).__name__}`: {getattr(err, 'message', str(err))}"
    try:
        line = getattr(err, 'line', None)
        column = getattr(err, 'column', None)
        pos = getattr(err, 'pos_in_stream', None)
    except Exception:
        line = column = pos = None

    block = ""
    if line is not None and column is not None:
        block = "```\n" + _highlight_at(prompt, line, column) + "\n```"
    elif pos is not None:
        prefix = prompt[:pos]
        l = prefix.count("\n") + 1
        c = len(prefix.split("\n")[-1]) + 1
        block = "```\n" + _highlight_at(prompt, l, c) + "\n```"

    tips = _guess_suggestion(prompt)
    tips_md = "\n".join([f"- {t}" for t in tips]) if tips else "- Проверьте синтаксис согласно подсказкам во вкладке «Справка и примеры»."
    return f"{title}\n\n{block}\n\n**Как исправить:**\n{tips_md}\n"

# -----------------------------
# Работа с ENV и reload parser
# -----------------------------

class EnvState:
    def __init__(self):
        self.allow_empty_alt = os.environ.get("ALLOW_EMPTY_ALTERNATE", "0")
        self.expand_alt_per_step = os.environ.get("EXPAND_ALTERNATE_PER_STEP", "1")
        self.group_combo_limit = os.environ.get("GROUP_COMBO_LIMIT", "100")

    def apply(self, allow_empty_alt: bool, expand_alt: bool, combo_limit: int) -> bool:
        """Вернёт True, если переменные изменились и модуль нужно перезагрузить."""
        changed = False
        new_allow = "1" if allow_empty_alt else "0"
        new_expand = "1" if expand_alt else "0"
        new_limit = str(int(combo_limit))

        if new_allow != self.allow_empty_alt:
            os.environ["ALLOW_EMPTY_ALTERNATE"] = new_allow
            self.allow_empty_alt = new_allow
            changed = True

        if new_expand != self.expand_alt_per_step:
            os.environ["EXPAND_ALTERNATE_PER_STEP"] = new_expand
            self.expand_alt_per_step = new_expand
            changed = True

        if new_limit != self.group_combo_limit:
            os.environ["GROUP_COMBO_LIMIT"] = new_limit
            self.group_combo_limit = new_limit
            changed = True

        return changed

ENV_STATE = EnvState()

def maybe_reload_prompt_parser(allow_empty_alt: bool, expand_alt: bool, combo_limit: int):
    """Если флаги изменились, перезагружаем modules.prompt_parser, чтобы обновить грамматику."""
    changed = ENV_STATE.apply(allow_empty_alt, expand_alt, combo_limit)
    if changed:
        global pp
        from modules import prompt_parser as _pp
        pp = importlib.reload(_pp)

# -----------------------------
# Бизнес-логика анализа
# -----------------------------

def _safe_parse_tree(prompt: str):
    """Пробует распарсить и вернуть (tree_str, error_md)."""
    try:
        tree = pp.schedule_parser.parse(prompt)
        tree_str = tree.pretty() if hasattr(tree, "pretty") else str(tree)
        return tree_str, ""
    except Exception as e:
        return "", format_lark_error(prompt, e)

def _make_schedule(prompt: str, steps: int, seed: Optional[int], use_visitor: bool) -> Tuple[List[List[Any]], str]:
    """Возвращает (schedule, error_md). schedule: [[end_step, text], ...]"""
    try:
        schedule = pp.get_schedule(prompt, steps, True, seed, use_visitor=use_visitor)
        if not schedule:
            return [], "**Пустое расписание.** Проверьте синтаксис."
        return schedule, ""
    except Exception as e:
        return [], f"**Ошибка при построении расписания:** {e}"

def _schedule_table_md(schedule: List[List[Any]]) -> str:
    """Markdown-таблица из расписания."""
    rows = ["| end_step | text |", "|---:|---|"]
    for end_step, text in schedule:
        safe_text = text.replace("|", "\\|")
        rows.append(f"| {end_step} | {safe_text} |")
    return "\n".join(rows)

def _per_step_timeline_md(schedule: List[List[Any]], steps: int) -> str:
    """Показывает активный текст на каждом шаге 1..steps."""
    rows = ["| step | active text |", "|---:|---|"]
    for s in range(1, steps + 1):
        txt = pp.at_step_from_schedule(s, schedule)
        safe_txt = txt.replace("|", "\\|")
        rows.append(f"| {s} | {safe_txt} |")
    return "\n".join(rows)

def _compare_vis_vs_trans(prompt: str, steps: int, seed: Optional[int]) -> str:
    """Сравнение use_visitor=True/False."""
    vis, e1 = _make_schedule(prompt, steps, seed, True)
    trn, e2 = _make_schedule(prompt, steps, seed, False)
    if e1 or e2:
        return (e1 or "") + (e2 or "")
    md = ["**Visitor (сбор в визиторе)**", _schedule_table_md(vis),
          "", "**Transformer (пересбор текста на end_step)**", _schedule_table_md(trn)]
    return "\n\n".join(md)

def analyze(prompt: str,
            steps: int,
            seed: Optional[int],
            use_visitor: bool,
            allow_empty_alt: bool,
            expand_alt_per_step: bool,
            group_combo_limit: int,
            view_mode: str,
            at_step_n: int) -> Tuple[str, str, str, str, str]:
    """
    Главная точка: возвращаем 5 выходов:
      1) status_md — статус/ошибки/подсказки
      2) schedule_md — таблица расписания
      3) timeline_md — поминутка по шагам
      4) text_at_step_md — текст на указанном шаге
      5) tree_md — дерево/сравнение, в зависимости от режима
    """
    # 1) Перезагружаем парсер при смене флагов
    try:
        maybe_reload_prompt_parser(allow_empty_alt, expand_alt_per_step, group_combo_limit)
    except Exception as e:
        err = f"**Не удалось применить ENV-флаги:** {e}"
        return err, "", "", "", ""

    status_md = ""
    tree_md = ""
    schedule_md = ""
    timeline_md = ""
    text_at_step_md = ""

    # 2) Парс-дерево и ошибки
    tree_str, err_md = _safe_parse_tree(prompt)
    if err_md:
        status_md = err_md
    else:
        status_md = "**OK:** Парсинг прошёл успешно."
        if view_mode == "Parse tree":
            tree_md = "```\n" + tree_str + "\n```"

    # 3) Строим расписание
    schedule, sched_err = _make_schedule(prompt, steps, seed, use_visitor)
    if sched_err:
        status_md += "\n\n" + sched_err
        return status_md, "", "", "", tree_md

    # 4) Режимы вывода
    if view_mode in ("Schedule", "All"):
        schedule_md = _schedule_table_md(schedule)

    if view_mode in ("Timeline (per step)", "All"):
        timeline_md = _per_step_timeline_md(schedule, steps)

    # 5) Текст на конкретном шаге
    try:
        t = pp.at_step_from_schedule(at_step_n, schedule)
        text_at_step_md = f"**Шаг {at_step_n}:** `{t}`"
    except Exception as e:
        text_at_step_md = f"Ошибка получения текста на шаге {at_step_n}: {e}"

    # 6) Доп. режим сравнения
    if view_mode == "Visitor vs Transformer":
        tree_md = _compare_vis_vs_trans(prompt, steps, seed)

    return status_md, schedule_md, timeline_md, text_at_step_md, tree_md

# -----------------------------
# Конструкторы грамматики (helpers)
# -----------------------------

def build_group(items_csv: str) -> str:
    """'a, b, c' -> '{ a, b, c }'"""
    items = [s.strip() for s in items_csv.split(",") if s.strip()]
    return "{ " + ", ".join(items) + " }" if items else "{}"

def build_alternate(items_bar: str, distinct: bool) -> str:
    """'a|b|c' -> '[ a|b|c ]' или '[ a|b|c ]!'"""
    core = items_bar.strip()
    if not core:
        return "[]"
    return f"[ {core} ]!" if distinct else f"[ {core} ]"

def build_numbered(n: int, options_bar_or_group: str, distinct: bool) -> str:
    """N и 'a|b|c' или '{...}' -> 'N! {...}' / 'N {...}'"""
    n = max(1, int(n))
    body = options_bar_or_group.strip()
    mark = "!" if distinct else ""
    return f"{n}{mark} {body}"

def build_sequence(owner: str, pairs_text: str, terminator: str) -> str:
    """
    owner='character'
    pairs_text='outfit: leather jacket; mood: brooding'
    terminator = '!' или ';'
    """
    owner = owner.strip()
    pairs = []
    for chunk in pairs_text.split(";"):
        chunk = chunk.strip()
        if not chunk:
            continue
        if ":" not in chunk:
            pairs.append(chunk)
        else:
            k, v = chunk.split(":", 1)
            pairs.append(f"{k.strip()} :: {v.strip()}")
    term = "!" if terminator == "!" else ";"
    return f"{owner} :: " + ", ".join(pairs) + f" {term}"

def build_top_level_sequence(owner: str, pairs_text: str, trailing_plain: str) -> str:
    """top-level: owner ::: ... !! , trailing"""
    seq = build_sequence(owner, pairs_text, "!")
    inner = seq.replace(f"{owner} :: ", "").rstrip("!;").strip()
    tail = (", " + trailing_plain.strip()) if trailing_plain.strip() else ""
    return f"{owner.strip()} ::: {inner} !!{tail}"

def helpers_demo_md() -> str:
    return """
### Быстрые шпаргалки (конструкторы)

- **Group:** `{ a, b, c }`  
  Смысл: держим фразы как один блок (можно затем оборачивать в `1! { ... }`, `N! {...}`).

- **Alternates:** `[ a | b | c ]` / `[ a | b | c ]!`  
  Карусель по шагам (если EXPAND_ALT_PER_STEP=1) или один выбор на прогон (`]!`).

- **Numbered:** `3! { a | b | c | d }`  
  Возьми 3 уникальных варианта без повторов (или `3 { ... }` с повторами).

- **Sequence:**  
  `character :: outfit :: leather jacket, mood :: brooding !`

- **Top-level sequence + хвост:**  
  `portrait ::: hair :: auburn, eyes :: blue !!, ultra-detailed skin texture`

- **Scheduled (особый случай с одним элементом):**  
  `[ "add rim light" ] : 25` — до 25-го шага пусто, после — фраза включается.
"""

# -----------------------------
# Gradio UI
# -----------------------------

class PromptParserAnalyzerScript(scripts.Script):
    def title(self):
        return "Prompt Parser Analyzer"

    def show(self, is_txt2img):
        return scripts.AlwaysVisible

    def ui(self, is_txt2img):
        with gr.Group():
            with gr.Accordion("Prompt Parser Analyzer", open=False):
                with gr.Tabs():
                    with gr.Tab("Analyzer"):
                        prompt_in = gr.Textbox(label="Source Prompt", lines=6, placeholder="Вставьте промпт…")
                        with gr.Row():
                            steps = gr.Slider(1, 300, value=50, step=1, label="Steps")
                            seed = gr.Number(value=42, precision=0, label="Seed (для случайных конструкций)")
                            at_step_n = gr.Slider(1, 300, value=25, step=1, label="Показать текст на шаге N")
                        with gr.Row():
                            use_visitor = gr.Checkbox(value=True, label="use_visitor (быстрее; alternate2 может отличаться)")
                            view_mode = gr.Dropdown(
                                choices=["All", "Schedule", "Timeline (per step)", "Parse tree", "Visitor vs Transformer"],
                                value="All",
                                label="Режим вывода",
                            )
                        with gr.Row():
                            allow_empty_alt = gr.Checkbox(value=False, label="ALLOW_EMPTY_ALTERNATE (разрешать пустые [a||b])")
                            expand_alt = gr.Checkbox(value=True, label="EXPAND_ALTERNATE_PER_STEP (чередовать по шагам)")
                            group_limit = gr.Slider(1, 5000, value=100, step=1, label="GROUP_COMBO_LIMIT")

                        analyze_btn = gr.Button("Analyze", variant="primary")

                        status_md = gr.Markdown()
                        schedule_md = gr.Markdown()
                        timeline_md = gr.Markdown()
                        text_at_step_md = gr.Markdown()
                        tree_md = gr.Markdown()

                    with gr.Tab("Builders"):
                        gr.Markdown("### Конструкторы Lark-грамматики (собирают корректный синтаксис)")
                        with gr.Row():
                            grp_items = gr.Textbox(label="Group items (через запятую)", placeholder="a, b, c")
                            grp_out = gr.Textbox(label="Результат", lines=2)
                        grp_btn = gr.Button("Build {group}")
                        with gr.Row():
                            alt_items = gr.Textbox(label="Alternates (через |)", placeholder="a|b|c")
                            alt_distinct = gr.Checkbox(value=False, label="distinct (!)")
                            alt_out = gr.Textbox(label="Результат", lines=2)
                        alt_btn = gr.Button("Build [alternate]")
                        with gr.Row():
                            num_n = gr.Number(value=3, precision=0, label="N (сколько взять)")
                            num_body = gr.Textbox(label="Опции: '{...}' или 'a|b|c'", placeholder="{ a | b | c } или a|b|c")
                            num_distinct = gr.Checkbox(value=True, label="distinct (!)")
                            num_out = gr.Textbox(label="Результат", lines=2)
                        num_btn = gr.Button("Build numbered")
                        gr.Markdown("—")
                        with gr.Row():
                            seq_owner = gr.Textbox(label="Sequence owner", value="character")
                            seq_pairs = gr.Textbox(label="Sequence pairs (k:v; k:v)", value="outfit: leather jacket; mood: brooding")
                            seq_term = gr.Dropdown(choices=["!", ";"], value="!", label="Terminator")
                            seq_out = gr.Textbox(label="Результат", lines=2)
                        seq_btn = gr.Button("Build sequence")
                        with gr.Row():
                            top_owner = gr.Textbox(label="Top-level owner", value="portrait")
                            top_pairs = gr.Textbox(label="Pairs (k:v; k:v)", value="hair: auburn; eyes: blue")
                            top_tail = gr.Textbox(label="Trailing plain (опционально)", value="ultra-detailed skin texture")
                            top_out = gr.Textbox(label="Результат", lines=2)
                        top_btn = gr.Button("Build top-level sequence")

                analyze_btn.click(
                    fn=analyze,
                    inputs=[prompt_in, steps, seed, use_visitor, allow_empty_alt, expand_alt, group_limit, view_mode, at_step_n],
                    outputs=[status_md, schedule_md, timeline_md, text_at_step_md, tree_md],
                )

                grp_btn.click(fn=build_group, inputs=[grp_items], outputs=[grp_out])
                alt_btn.click(fn=build_alternate, inputs=[alt_items, alt_distinct], outputs=[alt_out])
                num_btn.click(fn=build_numbered, inputs=[num_n, num_body, num_distinct], outputs=[num_out])
                seq_btn.click(fn=build_sequence, inputs=[seq_owner, seq_pairs, seq_term], outputs=[seq_out])
                top_btn.click(fn=build_top_level_sequence, inputs=[top_owner, top_pairs, top_tail], outputs=[top_out])

        return [prompt_in, steps, seed, at_step_n, use_visitor]

def on_ui_settings():
    section = ('prompt-parser-analyzer', "Prompt Parser Analyzer")
    # Пример добавления глобальных опций (если понадобится):
    # shared.opts.add_option("ppa_default_steps", shared.OptionInfo(50, "Default steps", section=section))

script_callbacks.on_ui_settings(on_ui_settings)