playground/prompter/prompter.py

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "bottle",
# ]
# ///
"""Browse images one-by-one and assign text prompts via a local web UI.

Saves results to a JSONL file: {"filename": "...", "prompt": "..."}
Images navigated past without a prompt produce no output line.
"""

import argparse
import json
import socket
import sys
import threading
import time
import webbrowser
from pathlib import Path

import bottle

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

CANNED_PROMPTS: list[str] = [
    "photorealistic. skin must be flawless",
]
"""Pre-populated prompt history shown in the UI on first launch."""

BUILDING_BLOCKS: list[str] = [
    "make her prettier",
    "age her by 5 years and make her look like a 25 year old",
    "keep the dark skin",
    "wet skin",
    "skin covered with baby oil, shiny skin",
    "soaking wet skin and hair, wet clothes clinging to her body",
    "long, flowy hair",
    "japanese",
    "greek",
    "no extra or missing fingers. each hand must have 5 fingers with the same hand pose as the original image",
    # expression
    "direct her gaze at the camera",
    "give her a serious and seductive look",
    "give her a playful expression",
    "slightly parted lips",
    "french kiss, eyes closed, tongue out",
    "closed eyes",
    "lower her eyelids and slightly part her lips in a seductive fashion",
    "exaggerated expression",
    "soft dramatic lighting",
    # outfit
    "replace the outfit with glossy latex",
    "remove the tan lines, she doesn't wear white bra",
]
"""Palette of text snippets the user can append to any prompt."""

VALID_EXTENSIONS: frozenset[str] = frozenset({".webp", ".jpg", ".jpeg", ".png"})

HTML: str = (Path(__file__).parent / "prompter.html").read_text(encoding="utf-8")

# ---------------------------------------------------------------------------
# App
# ---------------------------------------------------------------------------


def find_free_port() -> int:
    """Bind to port 0 and return the OS-assigned port number."""
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("127.0.0.1", 0))
        return s.getsockname()[1]


class PrompterApp:
    def __init__(self, images: list[str], output_path: Path) -> None:
        self.images = images
        self.output_path = output_path
        self.app = bottle.Bottle()
        self.app.route("/")(self.index)
        self.app.route("/api/images")(self.api_images)
        self.app.route("/image/<idx:int>")(self.serve_image)
        self.app.route("/api/save", method="POST")(self.api_save)

    def index(self) -> str:
        return HTML

    def api_images(self) -> str:
        """Return image list, canned prompts, building blocks, and existing prompts."""
        stem_to_path: dict[str, str] = {Path(p).stem: p for p in self.images}
        existing_by_stem = self._load_existing_prompts()
        existing_prompts = {
            stem_to_path[stem]: prompt
            for stem, prompt in existing_by_stem.items()
            if stem in stem_to_path
        }
        bottle.response.content_type = "application/json"
        return json.dumps(
            {
                "images": self.images,
                "canned_prompts": CANNED_PROMPTS,
                "building_blocks": BUILDING_BLOCKS,
                "existing_prompts": existing_prompts,
            }
        )

    def serve_image(self, idx: int) -> bottle.HTTPResponse:
        if idx < 0 or idx >= len(self.images):
            bottle.abort(404, "Image not found")
        path = Path(self.images[idx])
        return bottle.static_file(path.name, root=str(path.parent))

    def api_save(self) -> str:
        data: dict = bottle.request.json or {}
        file_path: str = data.get("file_path", "")
        prompt: str = data.get("prompt", "")
        if file_path and prompt:
            self._save_entry(file_path, prompt)
        bottle.response.content_type = "application/json"
        return json.dumps({"ok": True})

    def _load_existing_prompts(self) -> dict[str, str]:
        """Read the output JSONL and return a mapping of stem → prompt.

        When a filename appears multiple times the last entry wins, so re-running
        the tool and overwriting a previous prompt works naturally.
        """
        if not self.output_path.exists():
            return {}
        prompts: dict[str, str] = {}
        with open(self.output_path, encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    entry = json.loads(line)
                    if "filename" in entry and "prompt" in entry:
                        prompts[entry["filename"]] = entry["prompt"]
                except json.JSONDecodeError:
                    pass
        return prompts

    def _save_entry(self, file_path: str, prompt: str) -> None:
        """Append one {filename, prompt} record to the JSONL output file."""
        with open(self.output_path, "a", encoding="utf-8") as f:
            f.write(
                json.dumps({"filename": Path(file_path).stem, "prompt": prompt}) + "\n"
            )

    def run(self, port: int) -> None:
        bottle.run(self.app, host="localhost", port=port, quiet=True)


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Browse images and assign text prompts via a web UI.",
    )
    parser.add_argument("images", nargs="+", help="Image paths to browse")
    parser.add_argument(
        "-o",
        "--output-path",
        default="prompts.jsonl",
        metavar="FILE",
        help="Output JSONL file path (default: prompts.jsonl)",
    )
    args = parser.parse_args()

    images = [
        str(Path(p).resolve())
        for p in args.images
        if Path(p).suffix.lower() in VALID_EXTENSIONS
    ]

    if not images:
        exts = ", ".join(sorted(VALID_EXTENSIONS))
        print(
            f"Error: no valid images found. Expected extensions: {exts}",
            file=sys.stderr,
        )
        sys.exit(1)

    print(f"Loaded {len(images)} image(s).")

    port = find_free_port()
    url = f"http://localhost:{port}"
    print(f"Starting server at {url}")

    threading.Thread(
        target=lambda: (time.sleep(0.8), webbrowser.open(url)),
        daemon=True,
    ).start()

    PrompterApp(images, Path(args.output_path)).run(port)


if __name__ == "__main__":
    main()