olivv-cs-ComfyUI-FunPack/modules/training/lora/rated.py

"""Results the person rated, found in their saved projects.

A render's rating lives on its scene (and on each earlier take); this reads the project files and
lists every rated result that still has a file, newest project first. The wording of a rating is
the rating feature's (`app/ui/features/rating/choices.js`): a number from 6 up is liked, a lower
one disliked, and "Disliked: bad image" / "Disliked: bad composition" say which way it went.
"""

from __future__ import annotations

import json
import re

from .store import kind_of

FORGET = "-Just forget it-"
LABEL_TITLES = {"liked": "Liked results", "disliked": "Disliked results", "bad_composition": "Bad composition results",
                "bad_image": "Bad image results"}


def label_of(rating) -> str | None:
    v = str(rating or "").replace("|loved", "").strip()
    if not v or v == FORGET:
        return None
    if v == "Disliked: bad image":
        return "bad_image"
    if v == "Disliked: bad composition":
        return "bad_composition"
    if v.isdigit():
        return "liked" if int(v) >= 6 else "disliked"
    return "disliked" if v.startswith("Disliked") else None


def _media(raw):
    if not isinstance(raw, dict):
        return None
    name = raw.get("filename")
    if not isinstance(name, str) or not kind_of(name):
        return None
    return {"filename": name, "subfolder": str(raw.get("subfolder") or ""), "type": raw.get("type") if raw.get("type") in ("output", "temp", "input") else "output"}


INCLUDE = ("prefix", "main", "postfix", "variables")


def prompt_caption(text, doc=None, include=("main",), expand=None, resolve=None) -> str:
    """A rated result's prompt as a caption, from the parts asked for.

    `include` names the parts: prefix and postfix are the project's anchor and postfix, main is the scene's own words,
    variables fills in `$name` (left as written when not asked for). `expand(text)` is the app's shortcut
    expansion; `resolve(text, variables)` fills the variables in.
    """
    doc = doc if isinstance(doc, dict) else {}
    want = set(include or ())
    parts = []
    if "prefix" in want:
        parts.append(doc.get("anchor"))
    if "main" in want:
        parts.append(text)
    if "postfix" in want and doc.get("postfix_enabled") is not False:
        parts.append(doc.get("postfix"))
    joined = " ".join(" ".join(str(p or "").split()) for p in parts).strip()
    if expand and joined:
        joined = " ".join(str(expand(joined)).split())
    if "variables" in want and resolve:
        joined = resolve(joined, doc.get("variables"))
    return " ".join(joined.split())


def _graph_of(path):
    raw = None
    try:
        if str(path).lower().endswith((".png", ".webp", ".jpg", ".jpeg")):
            from PIL import Image
            with Image.open(path) as img:
                raw = (img.info or {}).get("prompt")
        else:
            import av
            with av.open(str(path)) as container:
                raw = (container.metadata or {}).get("prompt")
        graph = json.loads(raw) if raw else None
    except Exception:                                      # noqa: BLE001 -- a file that cannot be read says nothing
        return None
    return graph if isinstance(graph, dict) else None


def embedded_prompt(path) -> str:
    """The prompt the file was made from, read from its own queued graph: the text typed into the prompt node (before any rewrite), else a text encoder's. "" when it says nothing."""
    graph = _graph_of(path)
    if graph is None:
        return ""
    nodes = [(i, n) for i, n in graph.items() if isinstance(n, dict) and isinstance((n.get("inputs") or {}).get("text"), str)]
    for wanted in ("FunUIEnhancePrompt",):
        for _, n in nodes:
            if n.get("class_type") == wanted and n["inputs"]["text"].strip():
                return " ".join(n["inputs"]["text"].split())
    encoders = [n["inputs"]["text"] for i, n in nodes if n.get("class_type") == "CLIPTextEncode" and i != "negative" and n["inputs"]["text"].strip()]
    return " ".join(max(encoders, key=len).split()) if encoders else ""


def embedded_text(path) -> str | None:
    """Every piece of text the file's own queued graph carried (ComfyUI keeps it in the video or picture), or None when the file says nothing."""
    graph = _graph_of(path)
    if graph is None:
        return None
    texts = []
    for node in graph.values():
        for value in ((node or {}).get("inputs") or {}).values() if isinstance(node, dict) else ():
            if isinstance(value, str):
                texts.append(value)
    return " ".join(texts)


def _words(text) -> set[str]:
    return set(re.findall(r"[a-z0-9']{3,}", str(text or "").lower()))


def same_video(prompt_text, path, minimum: float = 0.5) -> bool:
    """False only when the file proves it was made from another prompt: file names repeat between ComfyUI instances, so a name alone does not say which video was rated.

    Judged on the words of the prompt the project kept for it; a file that carries no prompt of its own, or a rating with no words, is given the benefit of the doubt.
    """
    wanted = _words(prompt_text)
    if not wanted:
        return True
    found = embedded_text(path)
    if found is None:
        return True
    return len(wanted & _words(found)) / len(wanted) >= minimum


def is_this_video(result, path, slack_ms: int = 120_000) -> bool:
    """Whether the file under this name is the video the rating was given to. A rating is stamped with the time it was given, and a file is
    never made after the rating of it; one that is newer is another video that took the name (names repeat between ComfyUI instances),
    unless its own embedded prompt still says it is the same."""
    at = result.get("rated_at")
    try:
        newer = at is None or path.stat().st_mtime * 1000 > at + slack_ms      # a rating with no time on it cannot say: the prompt has to
    except OSError:
        return False
    return not newer or same_video(result.get("prompt_text"), path)


def scan(projects, expand=None, resolve=None, include=("main",)) -> list[dict]:
    """[{project, project_name, label, rating, media, caption}] for every rated result in these project documents.

    The caption is the prompt that made the result, built by `prompt_caption` from the parts in `include`.
    """
    out, seen = [], set()

    def take(doc, rating, media, caption, at=None):
        label, media = label_of(rating), _media(media)
        if label is None or media is None:
            return
        key = (media["type"], media["subfolder"], media["filename"])
        if key in seen:
            return
        seen.add(key)
        out.append({"project": doc.get("id", ""), "project_name": doc.get("name", ""), "label": label, "rating": str(rating),
                    "media": media, "caption": prompt_caption(caption, doc, include, expand, resolve), "prompt_text": " ".join(str(caption or "").split()),
                    "rated_at": at if isinstance(at, (int, float)) else None})

    for doc in projects:
        if not isinstance(doc, dict):
            continue
        scenes = {s.get("id"): s for s in doc.get("scenes") or [] if isinstance(s, dict)}
        for sid, render in (doc.get("scene_renders") or {}).items():
            scene, render = scenes.get(sid) or {}, render if isinstance(render, dict) else {}
            rated = scene.get("rated") if isinstance(scene.get("rated"), dict) else {}
            take(doc, scene.get("rating"), render.get("media"), rated.get("text") or render.get("text") or scene.get("text"), rated.get("at"))
        for sid, takes in (doc.get("scene_variants") or {}).items():
            for t in takes if isinstance(takes, list) else []:
                if isinstance(t, dict):
                    rated = t.get("rated") if isinstance(t.get("rated"), dict) else {}
                    take(doc, t.get("rating"), t.get("media"), rated.get("text") or t.get("text"), rated.get("at"))
    return out