gig-line/Tools/comfyui_generate.py

#!/usr/bin/env python3
"""Generate image assets from the local ComfyUI instance.

The project needs textures it cannot photograph: skin detail, hair strands, fabric
and surface grime. ComfyUI runs locally at `127.0.0.1:8188`, so this is a
**human-triggered local tool**, not a network service — nothing in CI calls it, and
no credential is involved.

    # what is installed
    python3 Tools/comfyui_generate.py models

    # one image
    python3 Tools/comfyui_generate.py txt2img \
        --prompt "seamless photoreal human skin texture, fine pores, neutral light" \
        --negative "text, watermark, hair, clothing" \
        --model juggernautXL_ragnarok.safetensors \
        --size 1024x1024 --seed 7 \
        --out Art/textures/generated/T_Character_SkinDetail.png

Outputs land in the given path and a JSON sidecar records the exact prompt, model,
seed and workflow, so an asset can be regenerated or rejected later without
guessing how it was made.

Licence: the output is a generated image from a locally hosted model. `THIRD_PARTY
NOTICES.md` must record the model and its licence before the texture ships; the
sidecar carries what is needed to do that.
"""

from __future__ import annotations

import argparse
import json
import os
import random
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

HOST = os.environ.get("COMFYUI_HOST", "http://127.0.0.1:8188")
DEFAULT_MODEL = "juggernautXL_ragnarok.safetensors"
# A second checkpoint by default: one model's texture is a lottery ticket, two is a comparison. Both are
# generated from the identical prompt and seed, then scored, and only the winner takes the output path.
DEFAULT_ALT_MODEL = "sd_xl_base_1.0.safetensors"


def score_image(path):
    """A deterministic, offline 'how much detail did this texture actually get' score.

    Mean absolute Laplacian: high-frequency energy. Generated textures that are blurry or flat score low,
    which is the failure this comparison exists to catch. Pillow is used only here, so a host without it
    still generates and simply skips scoring.
    """
    try:
        from PIL import Image, ImageFilter
    except ImportError:  # pragma: no cover - host without Pillow
        return None
    with Image.open(path) as image:
        grey = image.convert("L").filter(ImageFilter.FIND_EDGES)
        pixels = list(grey.getdata())
    return sum(pixels) / max(1, len(pixels))


def api(path: str) -> dict:
    with urllib.request.urlopen(f"{HOST}{path}", timeout=30) as response:
        return json.loads(response.read().decode())


def post(path: str, body: dict) -> dict:
    request = urllib.request.Request(
        f"{HOST}{path}", data=json.dumps(body).encode(), method="POST",
        headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(request, timeout=60) as response:
        raw = response.read().decode()
        return json.loads(raw) if raw.strip() else {}


def checkpoints() -> list[str]:
    info = api("/object_info/CheckpointLoaderSimple")
    return info["CheckpointLoaderSimple"]["input"]["required"]["ckpt_name"][0]


def workflow(prompt: str, negative: str, model: str, width: int, height: int,
             seed: int, steps: int, cfg: float, sampler: str, scheduler: str) -> dict:
    """A minimal SDXL/SD1.5 txt2img graph in ComfyUI's API format.

    Node ids are strings and references are `[node_id, output_index]`, which is the
    format `/prompt` expects — the UI's workflow JSON is a different shape and cannot
    be posted directly.
    """
    return {
        "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": model}},
        "2": {"class_type": "CLIPTextEncode",
              "inputs": {"text": prompt, "clip": ["1", 1]}},
        "3": {"class_type": "CLIPTextEncode",
              "inputs": {"text": negative, "clip": ["1", 1]}},
        "4": {"class_type": "EmptyLatentImage",
              "inputs": {"width": width, "height": height, "batch_size": 1}},
        "5": {"class_type": "KSampler",
              "inputs": {"seed": seed, "steps": steps, "cfg": cfg,
                         "sampler_name": sampler, "scheduler": scheduler,
                         "denoise": 1.0, "model": ["1", 0],
                         "positive": ["2", 0], "negative": ["3", 0],
                         "latent_image": ["4", 0]}},
        "6": {"class_type": "VAEDecode", "inputs": {"samples": ["5", 0], "vae": ["1", 2]}},
        "7": {"class_type": "SaveImage",
              "inputs": {"images": ["6", 0], "filename_prefix": "gigline"}},
    }


def wait_for_image(prompt_id: str, timeout: float = 300.0, poll: float = 1.5):
    """Poll /history until the job has an image, then fetch its bytes."""
    deadline = time.time() + timeout
    while time.time() < deadline:
        history = api(f"/history/{prompt_id}")
        entry = history.get(prompt_id)
        if entry and entry.get("outputs"):
            images = []
            for node in entry["outputs"].values():
                images.extend(node.get("images", []))
            if not images:
                return None
            image = images[0]
            query = urllib.parse.urlencode({
                "filename": image["filename"],
                "subfolder": image.get("subfolder", ""),
                "type": image.get("type", "output")})
            with urllib.request.urlopen(f"{HOST}/view?{query}", timeout=60) as response:
                return response.read()
        time.sleep(poll)
    return None


def generate(args) -> int:
    width, height = (int(v) for v in args.size.lower().split("x"))
    seed = args.seed if args.seed is not None else random.randrange(0, 2**31)
    graph = workflow(args.prompt, args.negative, args.model, width, height, seed,
                     args.steps, args.cfg, args.sampler, args.scheduler)
    queued = post("/prompt", {"prompt": graph})
    prompt_id = queued.get("prompt_id")
    if not prompt_id:
        print(f"error: ComfyUI rejected the prompt: {queued}", file=sys.stderr)
        return 1
    print(f"queued {prompt_id} ({args.model}, {width}x{height}, seed {seed})")

    data = wait_for_image(prompt_id, timeout=args.timeout, poll=args.poll)
    if data is None:
        print("error: no image came back (timeout or empty job)", file=sys.stderr)
        return 1
    os.makedirs(os.path.dirname(os.path.abspath(args.out)) or ".", exist_ok=True)
    with open(args.out, "wb") as handle:
        handle.write(data)
    sidecar = os.path.splitext(args.out)[0] + ".json"
    with open(sidecar, "w") as handle:
        json.dump({
            "tool": "Tools/comfyui_generate.py",
            "generated_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            "prompt": args.prompt,
            "negative": args.negative,
            "model": args.model,
            "size": [width, height],
            "seed": seed,
            "steps": args.steps,
            "cfg": args.cfg,
            "sampler": args.sampler,
            "scheduler": args.scheduler,
            "workflow": graph,
            "comfyui_prompt_id": prompt_id,
            "note": "generated locally; record the model licence in "
                    "THIRD_PARTY_NOTICES.md before shipping",
        }, handle, indent=2)
    print(f"wrote {args.out} (+ sidecar)")
    return 0




def generate_pair(args) -> int:
    """Generate the same prompt with two checkpoints and keep the better one at `--out`.

    Both are written beside the output as `<out stem>__<model>.png` with their own sidecars, so the losing
    candidate is kept for inspection instead of thrown away. The winner is chosen by `score_image()`; when
    Pillow is missing or the scores tie, the primary model keeps the output and the run says so.
    """
    import dataclasses

    stem, ext = os.path.splitext(args.out)
    candidates = []
    for label, model in (("primary", args.model), ("alternate", args.alt_model)):
        out = f"{stem}__{model.replace('/', '_')[:40]}{ext}"
        sub_args = dataclasses.replace(args, model=model, out=out)
        if generate(sub_args) != 0:
            print(f"{label} model {model} failed; continuing with what worked")
            continue
        candidates.append((model, out, score_image(out)))
    if not candidates:
        return 1
    scored = [c for c in candidates if c[2] is not None]
    if not scored:
        print(f"no scoring available; keeping {candidates[0][1]} as the output")
        winner = candidates[0]
    else:
        winner = max(scored, key=lambda c: c[2])
    for model, path, score in candidates:
        print(f"  {model}: {score if score is None else round(score, 2)}  {path}")
    os.replace(winner[1], args.out)
    sidecar = winner[1] + ".json"
    if os.path.exists(sidecar):
        os.replace(sidecar, args.out + ".json")
    print(f"chose {winner[0]} -> {args.out}")
    return 0


def main(argv=None) -> int:
    argv = sys.argv[1:] if argv is None else argv
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)
    sub.add_parser("models", help="list installed checkpoints")
    t = sub.add_parser("txt2img", help="generate one image")
    t.add_argument("--prompt", required=True)
    t.add_argument("--negative", default="text, watermark, signature, blurry, low quality")
    t.add_argument("--model", default=DEFAULT_MODEL)
    t.add_argument("--alt-model", default=DEFAULT_ALT_MODEL,
                   help=f"second checkpoint generated from the same prompt and seed, then compared "
                        f"(default {DEFAULT_ALT_MODEL}); pass an empty string to generate one image only")
    t.add_argument("--size", default="1024x1024")
    t.add_argument("--seed", type=int, default=None)
    t.add_argument("--steps", type=int, default=28)
    t.add_argument("--cfg", type=float, default=5.5)
    t.add_argument("--sampler", default="dpmpp_2m")
    t.add_argument("--scheduler", default="karras")
    t.add_argument("--out", required=True)
    t.add_argument("--timeout", type=float, default=300.0)
    t.add_argument("--poll", type=float, default=1.5)
    args = ap.parse_args(argv)

    if args.cmd == "models":
        for name in checkpoints():
            print(name)
        return 0
    if args.alt_model and args.alt_model != args.model:
        return generate_pair(args)
    return generate(args)


if __name__ == "__main__":
    raise SystemExit(main())