prompt_sort.py

"""Auto-sort: a prompt reorganised into the Prompt Frame's boxes, by Ollama.

Your ask (2026-08-17): the auto prompt already leans on the local Ollama
server; the same engine can read what is written across Style, Subject,
Surroundings and Light & colour, and put every phrase into the box it belongs
in - one button that tidies a prompt somebody typed as a lump.

Contract:
  - the model may MOVE and lightly TIDY phrases, never invent content and
    never drop it: every idea in goes to some box out
  - Style and Lighting may name a preset from the frame's own dropdown lists
    ONLY when the text clearly says so; otherwise they stay untouched
  - strict JSON out; anything else is a soft failure and the frame stays as
    it was, with one console line saying why
"""
import json

try:
    from . import autoprompt as _ap
    from .style_library import CHOICES as STYLE_CHOICES
    from .prompt_lists import LIGHTING_CHOICES
    from .prompt_tools import SWAP_MODES, STYLE_MODES, ACT_MODES, convert_text
except ImportError:  # loaded as a plain file (tests)
    import autoprompt as _ap
    from style_library import CHOICES as STYLE_CHOICES
    from prompt_lists import LIGHTING_CHOICES
    from prompt_tools import SWAP_MODES, STYLE_MODES, ACT_MODES, convert_text

FIELDS = ("subject", "surroundings", "style_extra", "light_and_colour", "placement")
# the boxes read in: the five above, and the Anything else box, which is unsorted
# text the sorter empties into them
IN_FIELDS = FIELDS + ("extra",)

SYSTEM = """You reorganise an image prompt into labelled boxes. You are a sorter, not a writer.

Boxes:
- subject: who or what the picture is of, and how they look (person, clothes, expression, pose, objects that are the point).
- surroundings: where it is - the place, the environment, background elements, weather.
- style_extra: how the picture is MADE - medium, rendering, camera/lens words, film stock, art style wording.
- light_and_colour: lighting mood, time of day light, colour palette, tonal mood.
- placement: where the subject stands IN the scene, as one short phrase (e.g. "standing at the water's edge"), or empty.
An "extra" line in the input is unsorted text: file every phrase of it into one of the boxes above.

Rules:
1. Every phrase from the input must land in exactly one box. Do not invent anything. Do not drop anything.
2. Keep the writer's wording; you may split, merge and lightly punctuate, nothing more.
3. If the input clearly names a style from this list, put its exact name in "style" (else ""): %s
4. If the input clearly names a lighting from this list, put its exact name in "lighting" (else ""): %s
5. Answer with ONLY a JSON object with keys: subject, surroundings, style_extra, light_and_colour, placement, style, lighting. Strings only. No commentary, no markdown fences."""


# REWRITE: the same boxes, the same model, a writer this time. Every fact stays;
# the wording gets concrete, and a style tag says what kind of picture it is for.
REWRITE_STYLES = ("keep", "photoreal", "cinematic", "illustration", "variation")

_STYLE_NOTES = {
    "keep": "Keep the medium and style exactly as written.",
    "photoreal": "Write it as a photograph: name the camera feel, lens, film or sensor "
                 "look and the light where they fit, in plain photographic terms.",
    "cinematic": "Write it as a film still: framing, lens, colour grade and the light, "
                 "in the terms a cinematographer uses.",
    "illustration": "Write it as an illustration: the medium, line, brush and palette, "
                    "in the terms an illustrator uses.",
    "variation": "Keep the medium and style exactly as written.",
}

REWRITE_SYSTEM = """You rewrite an image prompt so it reads well for a text to image model, box by box. You are a writer, not an inventor.

Boxes: subject, surroundings, style_extra, light_and_colour, placement (the same boxes in, the same boxes out).
An "extra" line in the input is unsorted text: fold every fact of it into the right box, so nothing is left out.

Rules:
1. Every fact in the input stays: who, what, where, what they wear, the light, the colours. Add nothing that changes the picture.
2. Improve the wording: concrete nouns, specific materials and light, one clear sentence or two per box, no filler, no lists of adjectives, no words like "stunning" or "masterpiece".
3. %s
4. An empty box stays empty unless the style note above asks for style words in style_extra.
5. Answer with ONLY a JSON object with keys: subject, surroundings, style_extra, light_and_colour, placement. Strings only. No commentary, no markdown fences."""


# VARIATION: the user's own instruction (2026-09-18). Same boxes in and out, but
# the model is asked for a fresh take on the same situation rather than a
# faithful rewrite: pose, gesture, clothing detail, props, framing and light may
# move, the concept may not. Runs warmer than a rewrite for that reason.
VARIATION_SYSTEM = """You rewrite an image prompt box by box into a polished text-to-image prompt, while creating a fresh variation of the same overall situation.

Boxes: subject, surroundings, style_extra, light_and_colour, placement. Return the same boxes in the same structure.

An "extra" line in the input is unsorted text. Interpret every useful detail from it and fold it into the most appropriate box.

Your goal is not to reproduce the input exactly. Use the input as visual direction and create a new image that feels clearly related: the same general subject, mood, setting, activity, and visual idea, but with creatively varied details.

Rules:

1. Preserve the core concept.
   Keep the important identity of the scene: who or what the image is about, the general activity, setting, mood, clothing category, visual style, and major visual features.

2. Create a variation, not a copy.
   You may naturally change secondary details such as pose, gesture, expression, clothing design or materials, nearby objects, background details, environmental features, camera framing, composition, lighting direction, and colour balance.

3. Do not change the scene into a different concept.
   A woman relaxing on a tropical beach should remain a woman in a beach-related situation, but she could be walking beside the water, sitting beneath a parasol, leaning against a beach bar, or standing near the shoreline instead of repeating the original pose exactly.

4. Be creatively specific.
   Replace vague wording with concrete visual information. Prefer specific actions, materials, surfaces, environmental details, camera relationships, and observable lighting.

5. Improve the writing.
   Each box should read naturally for a text-to-image model. Use one or two coherent sentences per box. Avoid filler, excessive adjective stacking, keyword spam, and words such as "stunning", "beautiful", "masterpiece", or "best quality".

6. Subject controls the main person or object.
   Describe appearance, clothing, pose, expression, action, and relevant physical details here. You may reinterpret secondary details while keeping the same general character concept.

7. Surroundings controls the environment.
   Describe the location, architecture, landscape, furniture, props, weather, atmosphere, and background activity. Expand sparse environments with plausible details that support the original situation without changing its theme.

8. style_extra controls visual treatment.

%s
Follow the style note when provided. Style words belong here rather than being scattered unnecessarily through other boxes.

9. light_and_colour controls illumination and palette.
Describe light source, direction, softness or hardness, time-of-day qualities, reflections, shadows, and dominant colour relationships. Preserve the general mood of the input while allowing a visually interesting variation.

10. placement controls composition and camera relationship.
Describe where the subject appears in the frame, camera height, viewing angle, shot distance, orientation, foreground/background relationship, and important spatial arrangement. You may improve or vary the composition rather than reproducing it exactly.

11. Treat the input as inspiration rather than immutable instructions.
When several possible interpretations exist, choose the version that creates the clearest, most visually interesting image while remaining recognisably connected to the original prompt.

12. Do not introduce unrelated concepts.
Creative additions must logically belong to the existing scene. Do not add new characters, major objects, fantasy elements, text, logos, animals, vehicles, or other dominant elements unless they are already suggested by the input.

13. An empty box stays empty unless:
- information from "extra" clearly belongs there, or
- the style note specifically requires style information in style_extra.

14. Do not mention that you changed, interpreted, expanded, or rewrote anything.

15. Output ONLY a valid JSON object with exactly these keys:
"subject"
"surroundings"
"style_extra"
"light_and_colour"
"placement"

Every value must be a JSON string. No commentary, markdown, headings, explanations, or additional keys."""


def rewrite_fields(fields, model, style="keep", url=_ap.OLLAMA_URL, transport=None,
                   generate=None, note=""):
    """fields: the frame's text boxes -> the same boxes rewritten, or None.

    `note` is an instruction of your own and REPLACES the style note when given:
    the five styles are shortcuts, not the whole of what can be asked for. It
    lands in the same slot they do, inside the system prompt, so the rules about
    keeping every fact and returning JSON still hold around it. The boxes are
    still the question; nothing about the shape of the exchange changes.

    `generate` may be injected (tests); it takes (model, system, prompt) and
    returns the reply text. Returns None on any failure, having printed why."""
    style = style if style in REWRITE_STYLES else "keep"
    note = str(note or "").strip()[:2000]
    lump = "\n".join("%s: %s" % (k, str(fields.get(k) or "").strip())
                     for k in IN_FIELDS if str(fields.get(k) or "").strip())
    if not lump.strip():
        print("[RedNode Prompt Rewrite] nothing to rewrite", flush=True)
        return None
    vary = style == "variation"
    gen = generate or (lambda m, s, p: _ap.ollama_generate(
        m, s, p, url=url, options={"temperature": 0.8 if vary else 0.4, "num_predict": 900 if vary else 700},
        keep_alive=0, **({"transport": transport} if transport else {})))
    system = (VARIATION_SYSTEM if vary else REWRITE_SYSTEM) % (note or _STYLE_NOTES[style])
    if note:
        print("[RedNode Prompt Rewrite] using your own instruction rather than a "
              "style: %s" % (note[:90] + ("..." if len(note) > 90 else "")), flush=True)
    reply = gen(model, system,
                "Input boxes:\n\n" + lump + "\n\nReturn the JSON.")
    if not reply:
        print("[RedNode Prompt Rewrite] the model returned nothing", flush=True)
        return None
    try:
        data = _extract_json(reply)
    except Exception as exc:
        print("[RedNode Prompt Rewrite] could not read the reply (%s): %r"
              % (exc, str(reply)[:200]), flush=True)
        return None
    out = {}
    for k in FIELDS:
        v = data.get(k, "")
        out[k] = str(v).strip() if isinstance(v, (str, int, float)) else ""
    if not any(out[k] for k in FIELDS):
        print("[RedNode Prompt Rewrite] the reply had no text in any box; keeping "
              "the frame as it was", flush=True)
        return None
    return out


# ---------------------------------------------------------------------------
# THE FINAL PROMPT: the Editor tab's Converter page. The active rig's row as it
# will be encoded, every caption, wildcard and people merge already in, reworked
# once more at queue time: the local model's rewrite when switched on, then the
# converter's tables and rules, so a swap or a rule of your own has the last word.
# ---------------------------------------------------------------------------
FINAL_STYLES = ("keep", "photoreal", "cinematic", "illustration")

FINAL_SYSTEM = """You rewrite one finished image prompt so it reads well for a text to image model. You are a writer, not an inventor.

Rules:
1. Keep every fact: every person, name, object, colour, action, place, pose, camera and light detail stays.
2. Do not add people, objects, actions or details that are not in the prompt.
3. Keep names exactly as written.
4. %s
5. Write plain prose in one paragraph. No lists, headings, labels or quotation marks.
6. Do not mention that you changed, interpreted or rewrote anything.

Output ONLY the rewritten prompt."""

_FINAL_CACHE = {}


def parse_final(raw):
    """The Editor Converter's final-prompt block, normalised. The panel mirrors it
    in readCfg (Trap 13). The converter is on and empty by default, which changes
    nothing; the rewrite is an explicit option and starts off."""
    r = raw if isinstance(raw, dict) else {}
    return {
        "on": r.get("on", True) is not False,
        "gender": r.get("gender") if r.get("gender") in SWAP_MODES else "off",
        "style": r.get("style") if r.get("style") in STYLE_MODES else "off",
        "act": r.get("act") if r.get("act") in ACT_MODES else "off",
        "remove_cum": bool(r.get("remove_cum")),
        "shave": bool(r.get("shave")),
        "rules": str(r.get("rules") or ""),
        "llm": bool(r.get("llm")),
        "llm_style": r.get("llm_style") if r.get("llm_style") in FINAL_STYLES else "keep",
        "llm_note": str(r.get("llm_note") or "")[:2000],
        "llm_fresh": bool(r.get("llm_fresh")),
    }


def rewrite_text(text, model, style="keep", note="", url=_ap.OLLAMA_URL, keep_alive=0,
                 reuse=True, generate=None):
    """One finished prompt -> the same prompt rewritten by the local model, or "".

    `note` is an instruction of your own and replaces the style's, as it does for
    the frame rewrite. REUSE keeps the answer for the same prompt, style and note,
    so a queue that changes nothing asks nothing; a new wildcard roll is a new
    prompt and is asked again. `generate` may be injected (tests)."""
    text = str(text or "").strip()
    if not text:
        return ""
    style = style if style in FINAL_STYLES else "keep"
    note = str(note or "").strip()[:2000]
    key = json.dumps([text, style, note, str(model)])
    if reuse and key in _FINAL_CACHE:
        print("[RedNode Final Prompt] rewrite: reused", flush=True)
        return _FINAL_CACHE[key]
    gen = generate or (lambda m, s, pr: _ap.ollama_generate(
        m, s, pr, url=url, options={"temperature": 0.4, "num_predict": 900},
        keep_alive=keep_alive))
    reply = gen(model, FINAL_SYSTEM % (note or _STYLE_NOTES[style]),
                "Prompt:\n\n" + text + "\n\nWrite the rewritten prompt.")
    out = _ap._strip_think(str(reply or "")).strip().strip('"').strip()
    if not out:
        print("[RedNode Final Prompt] rewrite: the model returned nothing; the prompt "
              "is used as it was", flush=True)
        return ""
    if len(_FINAL_CACHE) >= 64:
        _FINAL_CACHE.pop(next(iter(_FINAL_CACHE)))
    _FINAL_CACHE[key] = out
    return out


def finish_prompt(text, fin, model="", url=_ap.OLLAMA_URL, keep_alive=0, generate=None):
    """The final prompt through the Editor's Converter -> (text, what was done, warning).

    The rewrite first when switched on, then the converter. Any failure keeps the
    text it had; the warning is for the run log."""
    out = str(text or "")
    did, warn = [], ""
    if not fin.get("on", True) or not out.strip():
        return out, did, warn
    if fin.get("llm"):
        if not str(model or "").strip():
            warn = ("The final prompt rewrite is on, but no Ollama model is chosen on an "
                    "Auto Prompt section, so it was skipped.")
            print("[RedNode Final Prompt] %s" % warn, flush=True)
        else:
            r = rewrite_text(out, model, fin.get("llm_style", "keep"), fin.get("llm_note", ""),
                             url=url, keep_alive=keep_alive,
                             reuse=not fin.get("llm_fresh"), generate=generate)
            if r:
                out = r
                did.append("rewritten by %s" % model)
    before = out
    out = convert_text(out, gender_swap=fin.get("gender", "off"),
                       style_convert=fin.get("style", "off"),
                       nsfw_act_swap=fin.get("act", "off"),
                       nsfw_remove_cum=bool(fin.get("remove_cum")),
                       nsfw_shave_pubic=bool(fin.get("shave")),
                       custom_rules=fin.get("rules", ""), lock_to_authority=False)
    if out != before:
        did.append("converted")
    return out, did, warn


def _system():
    styles = ", ".join(str(s) for s in STYLE_CHOICES if str(s).lower() != "none")
    lights = ", ".join(str(l) for l in LIGHTING_CHOICES if str(l).lower() != "none")
    return SYSTEM % (styles, lights)


def _extract_json(text):
    """The first {...} object in a reply; models like to wrap it in prose."""
    text = str(text or "").strip()
    if text.startswith("```"):
        text = text.strip("`")
        if text.lower().startswith("json"):
            text = text[4:]
    a = text.find("{")
    b = text.rfind("}")
    if a < 0 or b <= a:
        raise ValueError("no JSON object in the reply")
    return json.loads(text[a:b + 1])


def sort_fields(fields, model, url=_ap.OLLAMA_URL, transport=None, generate=None):
    """fields: dict of the frame's text boxes -> a dict of sorted boxes, or None.

    `generate` may be injected (tests); it takes (model, system, prompt) and
    returns the reply text. Returns None on any failure, having printed why.
    """
    lump = "\n".join("%s: %s" % (k, str(fields.get(k) or "").strip())
                     for k in IN_FIELDS if str(fields.get(k) or "").strip())
    if not lump.strip():
        print("[RedNode Prompt Sort] nothing to sort", flush=True)
        return None
    gen = generate or (lambda m, s, p: _ap.ollama_generate(
        m, s, p, url=url, options={"temperature": 0.1, "num_predict": 700},
        keep_alive=0, **({"transport": transport} if transport else {})))
    reply = gen(model, _system(), "Input boxes as currently filled (some may be "
                                  "wrong or lumped together):\n\n" + lump
                                  + "\n\nReturn the JSON.")
    if not reply:
        print("[RedNode Prompt Sort] the model returned nothing", flush=True)
        return None
    try:
        data = _extract_json(reply)
    except Exception as exc:
        print("[RedNode Prompt Sort] could not read the reply (%s): %r"
              % (exc, str(reply)[:200]), flush=True)
        return None
    out = {}
    for k in FIELDS:
        v = data.get(k, "")
        out[k] = str(v).strip() if isinstance(v, (str, int, float)) else ""
    # the presets: only an EXACT name from the lists is honoured
    st = str(data.get("style") or "").strip()
    lt = str(data.get("lighting") or "").strip()
    out["style"] = st if st in [str(x) for x in STYLE_CHOICES] else ""
    out["lighting"] = lt if lt in [str(x) for x in LIGHTING_CHOICES] else ""
    # nothing lost: if the sorter emptied every text box, refuse the result
    if not any(out[k] for k in FIELDS):
        print("[RedNode Prompt Sort] the reply had no text in any box; keeping "
              "the frame as it was", flush=True)
        return None
    return out


try:
    from server import PromptServer
    from aiohttp import web

    @PromptServer.instance.routes.post("/rednode/prompt_sort")
    async def _rn_prompt_sort(request):
        try:
            body = await request.json()
        except Exception:
            return web.json_response({"error": "bad request"}, status=400)
        model = str(body.get("model") or "").strip()
        url = _ap.OLLAMA_URL                        # the server's setting, not the caller's
        if not model:
            return web.json_response(
                {"error": "no Ollama model chosen: pick one on the Auto Prompt "
                          "section (any tab) first"}, status=400)
        import asyncio
        loop = asyncio.get_event_loop()
        result = await loop.run_in_executor(
            None, lambda: sort_fields(body.get("fields") or {}, model, url=url))
        if result is None:
            return web.json_response(
                {"error": "the sorter could not produce a result; the console "
                          "says why"}, status=502)
        return web.json_response({"fields": result})

    @PromptServer.instance.routes.post("/rednode/prompt_rewrite")
    async def _rn_prompt_rewrite(request):
        try:
            body = await request.json()
        except Exception:
            return web.json_response({"error": "bad request"}, status=400)
        model = str(body.get("model") or "").strip()
        url = _ap.OLLAMA_URL                        # the server's setting, not the caller's
        if not model:
            return web.json_response(
                {"error": "no Ollama model chosen: pick one on the Auto Prompt "
                          "section (any tab) first"}, status=400)
        import asyncio
        loop = asyncio.get_event_loop()
        result = await loop.run_in_executor(
            None, lambda: rewrite_fields(body.get("fields") or {}, model,
                                         str(body.get("style") or "keep"), url=url,
                                         note=str(body.get("note") or "")))
        if result is None:
            return web.json_response(
                {"error": "the rewrite could not produce a result; the console "
                          "says why"}, status=502)
        return web.json_response({"fields": result})
except Exception as _e:
    print("[RedNode Prompt Sort] route not registered: %s" % _e, flush=True)