kgilper-krea-reference/docs/deepstack-layers/generate_sweep.py

"""Single-chunk deepstack sweep - the fresh empirical determination (section 16.4).

Renders the experiment that measures what each of the 12 deepstack chunks
carries: hold one reference at fixed strength, and for each chunk L render a
gain table that is 1.0 everywhere except a spike at L. Any visual difference
between spike-L and the all-ones control is attributable to chunk L, because
the pooled channel and every other chunk are held identical.

Each spike is expressed as a V10 custom recipe (role balanced -> even base,
neutral shape/global, so only the layers array differs). Rendering therefore
REQUIRES the kg_krea_v10 nodes installed on the target ComfyUI - the single-
chunk tables are not expressible through the V9 widgets. Run
`--dry-run` anywhere to generate the recipe files and API graphs without a
server (validates construction against the real node code via the stubs).

Usage:
  python generate_sweep.py --dry-run                 # build + validate, no render
  python generate_sweep.py --server http://HOST:PORT # render on a V10 box
  python generate_sweep.py --gains 2 4 6             # multi-level (default: 4)

Outputs (under docs/deepstack-layers/sweep_out/):
  recipes/chunk-NN-gainG.yaml   one custom recipe per spike
  grids/<ref>__chunk-NN-gainG.png   rendered result (when a server is given)
  grids/<ref>__control.png          the all-ones baseline
  sweep-manifest.json               every cell: ref, chunk, gain, seed, file
"""

import argparse
import json
import sys
import time
import urllib.parse
import urllib.request
from pathlib import Path

HERE = Path(__file__).resolve().parent
REPO_ROOT = HERE.parents[1]
OUT_DIR = HERE / "sweep_out"
sys.path.insert(0, str(REPO_ROOT / "tests"))

CHUNKS = 12
MODEL = "krea2\\krea2_turbo_nvfp4.safetensors"
CLIP = "krea2\\qwen3vl_4b_fp8_scaled.safetensors"
VAE = "krea2\\qwen_image_vae.safetensors"
WIDTH, HEIGHT, STEPS, CFG, SEED = 512, 768, 8, 1.0, 424242

# References chosen so every aspect CAN respond: a clear-subject product shot
# (structure + material) and an abstract palette study (palette + finish).
REFERENCES = {
    "subject": "krea-reference-examples/slot1_content_anchor.png",
    "palette": "krea-reference-examples/slot2_style_reference.png",
}
# Neutral prompt: let the reference speak so the chunk's contribution shows.
PROMPT = "a product photograph on a plain seamless background, soft even light"
NEGATIVE = "boring, dull, blurry, low-quality, oversaturated colours"
CARD_STRENGTH = 0.7  # fixed across the whole sweep; literal feel curve for exactness


def spike_table(chunk, gain):
    table = [1.0] * CHUNKS
    table[chunk] = float(gain)
    return table


def spike_recipe(chunk, gain):
    """A V10 custom recipe isolating one chunk. Role balanced keeps shape and
    global at 1.0, so between cells only the `layers` array differs."""
    label = "sweep chunk {:02d} gain {:g}".format(chunk, gain)
    return label, {
        "label": label,
        "description": "Deepstack sweep: chunk {} spiked to {}x, all others 1.0.".format(chunk, gain),
        "role": "balanced",
        "treatment": "normal",
        "color": 1.0,
        "detail": 1.0,
        "study": "384",
        "framing": "stack",
        "subject": "allow",
        "early": 1.0,
        "late": 1.0,
        "guard": False,
        "cap": 3.0,
        "shape": 1.0,
        "global": 1.0,
        "layers": spike_table(chunk, gain),
    }


def control_recipe():
    return "sweep control", {
        "label": "sweep control",
        "description": "Deepstack sweep control: all chunks 1.0 (native).",
        "role": "balanced", "treatment": "normal", "color": 1.0, "detail": 1.0,
        "study": "384", "framing": "stack", "subject": "allow", "early": 1.0,
        "late": 1.0, "guard": False, "cap": 3.0, "shape": 1.0, "global": 1.0,
        "layers": [1.0] * CHUNKS,
    }


def validate_recipes(recipes):
    """Run every generated recipe through the REAL V10 validation + resolution."""
    from _kg_stub_env import load_module
    nodes, _ = load_module("kg_krea_v10", "kg_deepstack_sweep_validate")
    card_cls = nodes.KGKrea2ImageGuideCardV10
    reserved = set(card_cls.PURPOSE_LABELS) | set(card_cls.QUICK_RECIPES)
    problems = []
    for label, bundle in recipes.items():
        try:
            _resolved_label, resolved = nodes.custom_recipes.validate_recipe(bundle, reserved)
            if resolved["layers"] != bundle["layers"]:
                problems.append("{}: layers not preserved".format(label))
        except Exception as exc:  # noqa: BLE001
            problems.append("{}: {}".format(label, exc))
    return problems


def build_api_graph(recipe_label, recipe_bundle, reference_image):
    """API graph: a V10 guide card carrying the spike recipe -> V10 stack."""
    card_inputs = {
        "Reference image": ["load", 0],
        "How strongly this image guides": CARD_STRENGTH,
        "Use image for": recipe_label,
    }
    stack_inputs = {
        "Krea CLIP": ["clip", 0],
        "Final image prompt": PROMPT,
        "Written prompt strength": 1.0,
        "Image slider feel": "literal slider values",
        "Image detail level": "medium - balanced default (384)",
        "Image framing": "keep full image shape",
        "When images guide": "guide the whole image",
        "Early-to-final handoff": 0.4,
        "Text/logo guard prompt handling": "gentle guard - keep my prompt words",
        "Balance strong cards": "off - use my values",
        "Reuse image studies": "always re-study",
        "Reference 1 guide card": ["card", 0],
    }
    neg_inputs = dict(stack_inputs)
    neg_inputs.pop("Reference 1 guide card")
    neg_inputs["Final image prompt"] = NEGATIVE
    return {
        "unet": {"class_type": "UNETLoader", "inputs": {"unet_name": MODEL, "weight_dtype": "default"}},
        "sampling": {"class_type": "ModelSamplingAuraFlow", "inputs": {"shift": 3.14, "model": ["unet", 0]}},
        "clip": {"class_type": "CLIPLoader", "inputs": {"clip_name": CLIP, "type": "krea2", "device": "default"}},
        "vae": {"class_type": "VAELoader", "inputs": {"vae_name": VAE}},
        "load": {"class_type": "LoadImage", "inputs": {"image": reference_image, "upload": "image"}},
        "card": {"class_type": "KGKrea2ImageGuideCardV10", "inputs": card_inputs},
        "pos": {"class_type": "KGTextEncodeKreaImageReferencesV10", "inputs": stack_inputs},
        "neg": {"class_type": "KGTextEncodeKreaImageReferencesV10", "inputs": neg_inputs},
        "latent": {"class_type": "EmptyLatentImage", "inputs": {"width": WIDTH, "height": HEIGHT, "batch_size": 1}},
        "sampler": {"class_type": "KSampler", "inputs": {
            "seed": SEED, "steps": STEPS, "cfg": CFG, "sampler_name": "euler",
            "scheduler": "simple", "denoise": 1.0, "model": ["sampling", 0],
            "positive": ["pos", 0], "negative": ["neg", 0], "latent_image": ["latent", 0]}},
        "decode": {"class_type": "VAEDecode", "inputs": {"samples": ["sampler", 0], "vae": ["vae", 0]}},
        "save": {"class_type": "SaveImage", "inputs": {"images": ["decode", 0],
                 "filename_prefix": "krea_deepstack_sweep/tmp"}},
    }


def post_and_wait(server, graph, timeout=420):
    payload = json.dumps({"prompt": graph}).encode("utf-8")
    req = urllib.request.Request(server + "/prompt", data=payload, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=30) as r:
        prompt_id = json.load(r)["prompt_id"]
    start = time.time()
    while time.time() - start < timeout:
        with urllib.request.urlopen(server + "/history/" + prompt_id, timeout=30) as r:
            entry = json.load(r).get(prompt_id)
        if entry and entry.get("outputs"):
            return entry["outputs"]["save"]["images"][0]
        if entry and entry.get("status", {}).get("status_str") == "error":
            raise RuntimeError(json.dumps(entry["status"])[:800])
        time.sleep(2)
    raise TimeoutError(prompt_id)


def download(server, info, dest):
    q = urllib.parse.urlencode({"filename": info["filename"],
                                "subfolder": info.get("subfolder", ""), "type": info.get("type", "output")})
    with urllib.request.urlopen(server + "/view?" + q, timeout=60) as r:
        dest.write_bytes(r.read())


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--server", help="ComfyUI base URL with V10 nodes installed")
    ap.add_argument("--dry-run", action="store_true", help="build + validate only")
    ap.add_argument("--gains", type=float, nargs="+", default=[4.0])
    args = ap.parse_args()

    (OUT_DIR / "recipes").mkdir(parents=True, exist_ok=True)
    (OUT_DIR / "grids").mkdir(parents=True, exist_ok=True)

    recipes = dict([control_recipe()])
    for gain in args.gains:
        for chunk in range(CHUNKS):
            label, bundle = spike_recipe(chunk, gain)
            recipes[label] = bundle
    for label, bundle in recipes.items():
        slug = label.replace(" ", "-")
        (OUT_DIR / "recipes" / (slug + ".yaml")).write_text(
            "\n".join("{}: {}".format(k, json.dumps(v)) for k, v in bundle.items()) + "\n", encoding="utf-8")

    problems = validate_recipes(recipes)
    if problems:
        print("VALIDATION FAILED:")
        for p in problems:
            print("  -", p)
        return 1
    print("validated {} recipes (1 control + {} spikes) against the real V10 node code".format(
        len(recipes), len(recipes) - 1))

    cells = []
    for ref_name, ref_image in REFERENCES.items():
        for label, bundle in recipes.items():
            cells.append({"reference": ref_name, "reference_image": ref_image,
                          "recipe_label": label, "layers": bundle["layers"], "seed": SEED})

    if args.dry_run or not args.server:
        (OUT_DIR / "sweep-manifest.json").write_text(json.dumps(
            {"rendered": False, "prompt": PROMPT, "card_strength": CARD_STRENGTH,
             "seed": SEED, "gains": args.gains, "cells": cells}, indent=2) + "\n", encoding="utf-8")
        print("dry run: {} recipe files + {} planned cells written to {}".format(
            len(recipes), len(cells), OUT_DIR))
        print("to render: install kg_krea_v10 on a ComfyUI box, then rerun with --server URL")
        return 0

    server = args.server.rstrip("/")
    for cell in cells:
        img = build_api_graph(cell["recipe_label"], recipes[cell["recipe_label"]], cell["reference_image"])
        slug = cell["recipe_label"].replace(" ", "-")
        dest = OUT_DIR / "grids" / "{}__{}.png".format(cell["reference"], slug)
        print("rendering", dest.name, flush=True)
        info = post_and_wait(server, img)
        download(server, info, dest)
        cell["output"] = "grids/{}__{}.png".format(cell["reference"], slug)
    (OUT_DIR / "sweep-manifest.json").write_text(json.dumps(
        {"rendered": True, "server": server, "prompt": PROMPT, "card_strength": CARD_STRENGTH,
         "seed": SEED, "gains": args.gains, "cells": cells}, indent=2) + "\n", encoding="utf-8")
    print("rendered {} cells -> {}".format(len(cells), OUT_DIR / "grids"))
    print("next: score with SCORING.md and record verdicts in README.md")
    return 0


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