HitmanLoges-comfyui-glitch-lab-sampler/tools/render_samples.py

"""
Drive ComfyUI's HTTP API to render the README's gallery.

Renders six images at 512x512 with a single photorealistic prompt:
- one UNPATCHED reference (no glitch-lab nodes; standard KSampler)
- five preset renders (baseline, soft-bleed, veiled, schedule-burn, slag)

A consistent prompt + an unpatched baseline lets the viewer actually see what
each parameter changes, instead of the previous gallery's per-preset prompts
which (per blepping's review) let the abstract subject hide the degradation.

Run with the ComfyUI venv's python:

    & "C:\\Users\\Loghan\\Documents\\ComfyUI\\.venv\\Scripts\\python.exe" tools\\render_samples.py

Assumes:
- ComfyUI is running on http://127.0.0.1:8000 (Desktop installer default).
- The glitch-lab custom node package is loaded (three nodes registered).
- v1-5-pruned-emaonly.safetensors is in models/checkpoints/.

stdlib only. No third-party deps.
"""

from __future__ import annotations

import json
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid
from pathlib import Path


COMFY_URL = "http://127.0.0.1:8000"
CHECKPOINT = "v1-5-pruned-emaonly.safetensors"
WIDTH = 512
HEIGHT = 512
CFG = 7.0
CLIENT_ID = str(uuid.uuid4())

# A single photorealistic prompt used for every render. Picked because:
# - Strong geometric structure (the box) so degradation is obvious
# - Inanimate subject (no risk of distorted figures)
# - SD 1.5 handles it well at the unpatched baseline
PROMPT = "a red telephone box on a cobblestone street, london, late afternoon, photograph, sharp focus, 35mm"
NEGATIVE = "blurry, distorted, cartoon, painting"

REPO_ROOT = Path(__file__).resolve().parent.parent
LIBRARY_PATH = REPO_ROOT / "presets" / "library.json"
SAMPLES_DIR = REPO_ROOT / "samples"

# Which presets to render. All five use the same PROMPT above so the only
# variable across the gallery is the preset's parameters.
PRESET_IDS: list[str] = ["baseline", "soft-bleed", "veiled", "schedule-burn", "slag"]

REQUIRED_NODES = {
    "GlitchSigmas",
    "AttentionDropoutPatcher",
    "LatentFrameInterleave",
}


def http_get(path: str) -> dict:
    url = f"{COMFY_URL}{path}"
    with urllib.request.urlopen(url, timeout=30) as r:
        return json.loads(r.read().decode("utf-8"))


def http_get_bytes(path: str) -> bytes:
    url = f"{COMFY_URL}{path}"
    with urllib.request.urlopen(url, timeout=120) as r:
        return r.read()


def http_post(path: str, body: dict) -> dict:
    url = f"{COMFY_URL}{path}"
    data = json.dumps(body).encode("utf-8")
    req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"}, method="POST")
    with urllib.request.urlopen(req, timeout=30) as r:
        return json.loads(r.read().decode("utf-8"))


def verify_environment() -> None:
    print(f"checking ComfyUI at {COMFY_URL} ...")
    try:
        info = http_get("/object_info")
    except urllib.error.URLError as e:
        print(f"FAIL: cannot reach ComfyUI: {e}")
        sys.exit(2)

    missing = REQUIRED_NODES - set(info.keys())
    if missing:
        print(f"FAIL: glitch-lab nodes missing from ComfyUI: {sorted(missing)}")
        print("Hint: restart ComfyUI Desktop so it re-scans custom_nodes/")
        sys.exit(2)
    print(f"  three glitch-lab nodes registered (GlitchSampler removed — using built-in)")

    if "SamplerEulerAncestral" not in info:
        print("FAIL: ComfyUI build doesn't expose SamplerEulerAncestral")
        sys.exit(2)
    print("  SamplerEulerAncestral available")

    ckpts_field = info.get("CheckpointLoaderSimple", {}).get("input", {}).get("required", {}).get("ckpt_name")
    available: list[str] = []
    if isinstance(ckpts_field, list) and ckpts_field and isinstance(ckpts_field[0], list):
        available = ckpts_field[0]
    if CHECKPOINT not in available:
        print(f"FAIL: checkpoint not in ComfyUI's checkpoint list: {CHECKPOINT}")
        print(f"  available: {available}")
        sys.exit(2)
    print(f"  checkpoint present: {CHECKPOINT}")


def load_preset(token_id: str) -> dict:
    with LIBRARY_PATH.open(encoding="utf-8") as f:
        library = json.load(f)
    for t in library:
        if t["id"] == token_id:
            return t
    raise SystemExit(f"preset not found in library.json: {token_id}")


def build_unpatched_workflow(seed: int = 0) -> dict:
    """Vanilla KSampler with the same prompt + checkpoint, no glitch-lab nodes."""
    return {
        "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": CHECKPOINT}},
        "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": {
                "model": ["1", 0],
                "positive": ["2", 0],
                "negative": ["3", 0],
                "latent_image": ["4", 0],
                "seed": seed,
                "steps": 20,
                "cfg": CFG,
                "sampler_name": "euler_ancestral",
                "scheduler": "karras",
                "denoise": 1.0,
            },
        },
        "6": {"class_type": "VAEDecode", "inputs": {"samples": ["5", 0], "vae": ["1", 2]}},
        "7": {"class_type": "SaveImage", "inputs": {"images": ["6", 0], "filename_prefix": "glitch-lab/unpatched"}},
    }


def build_preset_workflow(preset: dict) -> dict:
    """Workflow using the three custom nodes + built-in SamplerEulerAncestral."""
    p = preset["parameters"]
    return {
        "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": CHECKPOINT}},
        "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": "GlitchSigmas",
            "inputs": {
                "model": ["1", 0],
                "scheduler": p["scheduler"],
                "steps": p["steps"],
                "denoise": 1.0,
                "jitter": p["jitter"],
                "seed": p["seed"],
            },
        },
        "6": {
            "class_type": "AttentionDropoutPatcher",
            "inputs": {
                "model": ["1", 0],
                "self_attn_dropout": p["self_attn_dropout"],
                "cross_attn_dropout": p["cross_attn_dropout"],
                "seed": p["seed"],
            },
        },
        "7": {
            "class_type": "SamplerEulerAncestral",
            "inputs": {"eta": 1.0, "s_noise": p["s_noise"]},
        },
        "8": {"class_type": "RandomNoise", "inputs": {"noise_seed": p["seed"]}},
        "9": {
            "class_type": "CFGGuider",
            "inputs": {"model": ["6", 0], "positive": ["2", 0], "negative": ["3", 0], "cfg": CFG},
        },
        "10": {
            "class_type": "SamplerCustomAdvanced",
            "inputs": {
                "noise": ["8", 0],
                "guider": ["9", 0],
                "sampler": ["7", 0],
                "sigmas": ["5", 0],
                "latent_image": ["4", 0],
            },
        },
        "11": {
            "class_type": "LatentFrameInterleave",
            "inputs": {"latent": ["10", 0], "intensity": p["frame_interleave_intensity"], "seed": p["seed"]},
        },
        "12": {"class_type": "VAEDecode", "inputs": {"samples": ["11", 0], "vae": ["1", 2]}},
        "13": {"class_type": "SaveImage", "inputs": {"images": ["12", 0], "filename_prefix": f"glitch-lab/{preset['id']}"}},
    }


def submit_and_wait(workflow: dict, label: str, timeout_s: float = 300.0) -> list[dict]:
    body = {"prompt": workflow, "client_id": CLIENT_ID}
    res = http_post("/prompt", body)
    prompt_id = res.get("prompt_id")
    if not prompt_id:
        raise SystemExit(f"[{label}] no prompt_id in response: {res}")
    print(f"  submitted, prompt_id={prompt_id[:8]}")

    deadline = time.time() + timeout_s
    while time.time() < deadline:
        hist = http_get(f"/history/{prompt_id}")
        if hist and prompt_id in hist:
            entry = hist[prompt_id]
            status = entry.get("status", {})
            if status.get("completed"):
                outputs = entry.get("outputs", {}) or {}
                images: list[dict] = []
                for _node_id, out in outputs.items():
                    images.extend(out.get("images", []) or [])
                return images
            if status.get("status_str") == "error":
                msgs = status.get("messages", [])
                raise SystemExit(f"[{label}] execution error: {msgs}")
        time.sleep(2.0)
    raise SystemExit(f"[{label}] timed out after {timeout_s}s")


def download_image(image_info: dict, dest: Path) -> None:
    qs = urllib.parse.urlencode(
        {
            "filename": image_info["filename"],
            "subfolder": image_info.get("subfolder", ""),
            "type": image_info.get("type", "output"),
        }
    )
    data = http_get_bytes(f"/view?{qs}")
    dest.parent.mkdir(parents=True, exist_ok=True)
    dest.write_bytes(data)
    print(f"  saved -> {dest} ({len(data)} bytes)")


def main() -> int:
    verify_environment()
    SAMPLES_DIR.mkdir(parents=True, exist_ok=True)

    print(f"\n[unpatched] (reference — no glitch-lab nodes) prompt: {PROMPT!r}")
    images = submit_and_wait(build_unpatched_workflow(seed=0), "unpatched")
    if images:
        for i, img in enumerate(images):
            suffix = "" if i == 0 else f"_{i}"
            download_image(img, SAMPLES_DIR / f"unpatched{suffix}.png")

    for token_id in PRESET_IDS:
        print(f"\n[{token_id}] prompt: {PROMPT!r}")
        preset = load_preset(token_id)
        workflow = build_preset_workflow(preset)
        images = submit_and_wait(workflow, token_id)
        if not images:
            print(f"  WARN: no images in output")
            continue
        for i, img in enumerate(images):
            suffix = "" if i == 0 else f"_{i}"
            download_image(img, SAMPLES_DIR / f"{token_id}{suffix}.png")

    print("\ndone.")
    return 0


if __name__ == "__main__":
    sys.exit(main())