`scripts/controlnet-anchored-reroll.py`

152 lines
#!/usr/bin/env python3
"""controlnet-anchored-reroll.py — GENERALIZED seam-fix re-roll generator for the
seam-wave scale-up. Accepts arbitrary --ids (any animal-forward FAIL root), not
just the two hard-scoped glassware roots of controlnet-anchored-reroll-2.py.

It is a thin generalization of controlnet-anchored-reroll-2.py — it reuses that
script's PROVEN image-level retry machinery byte-for-byte (reroll_one):

  • ControlNet canny edge-anchor on a CROPPED single focal motif (find_focal_crop
    + build_canny_hint + gen_controlnet_seamless from controlnet-anchored-redo) —
    the canny path is the one empirically shown to produce clean seamless tiles
    (reroll-2's own note: text-led produced busy edge-FAILING tiles 0/12; canny
    cleared PASS-both-axes reliably). Subject-loss — the only canny failure mode —
    is handled by the subject-recovery vision gate below.
  • Circular-pad native-seamless txt2img (SeamlessTile + MakeCircularVAE, denoise
    1.0) → the tile wraps by construction, ZERO post-hoc smear.
  • IMAGE-LEVEL SEAM-RETRY LOOP: re-roll the seed until the raw motif tile clears
    cross-seam PASS on BOTH axes (in-memory verify-edge-seamless) AND coverage is
    in band AND the subject-recovery gate (free local llava) confirms the animal
    is present. Cap MAX_RETRIES. INSERT the DB row ONLY for the final passing
    candidate — exactly one new id per root max (all_designs is INSERT-only; a
    failed attempt must never touch the DB).
  • Composite onto a REAL seamless natural-fibre ground, re-verify the written
    composite, then insert.

DEFAULT config (DEFAULT_CFG) is the animal-forward canny profile: a moderate
focal crop, weak-canny strength so layout stays open, coverage band 8-40%, and
canny ENABLED (no_canny=False). Per-root overrides can still be supplied via
ROOT_CFG (imported from reroll-2 so 54076/54266 keep their tuned profiles if ever
re-run through here). The ground rotates across roots so Steve sees range.

SACRED ROOTS: never overwrites a root. Final candidate → NEW file + NEW id,
is_published=FALSE, user_removed=FALSE, parent_design_id=<root>. Never published
live. Settlement gate untouched (children of already-passed roots, no new prompt
class). DATABASE_URL self-resolves, never echoed. Does NOT use make_seamless.py /
force-edge-seamless.py.

Usage (called by scripts/seam-wave.sh):
  python3 scripts/controlnet-anchored-reroll.py --ids 43199,54385,55272
  python3 scripts/controlnet-anchored-reroll.py --ids 43199 --max-retries 8
"""
import argparse, datetime, json, os, sys
from pathlib import Path

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
QUEUE = os.path.join(ROOT, 'data', 'seam-fix-queue.jsonl')

sys.path.insert(0, os.path.join(ROOT, 'scripts'))
from importlib import import_module
cnar = import_module('controlnet-anchored-redo')        # GROUNDS, machinery
reroll2 = import_module('controlnet-anchored-reroll-2')  # proven reroll_one + ROOT_CFG
edgeverify = import_module('verify-edge-seamless')

# Reuse reroll-2's per-root special profiles where they exist (e.g. the two
# glassware roots if ever routed through here). Everything else uses DEFAULT_CFG.
ROOT_CFG = dict(reroll2.ROOT_CFG)

# DEFAULT animal-forward canny profile. The animal is the dominant structure (we
# pre-filter for that), so the canny crop anchors the RIGHT subject and a weak
# strength keeps the surround open. Subject-recovery gate still guards against the
# rare roll where the crop loses the figure.
#   ground is filled per-root from cnar.GROUNDS at call time (range, not one note).
DEFAULT_CFG = {
    'no_canny': False,
    'cov_lo': 8.0, 'cov_hi': 40.0,   # reject empty (too low) / over-busy (too high)
    'crop_frac': 0.30, 'motif_px': 170, 'cn_strength': 0.50, 'cn_end': 0.55,
}


def cfg_for(root, ground):
    cfg = dict(ROOT_CFG.get(root, DEFAULT_CFG))
    cfg.setdefault('ground', ground)
    # a per-root override may pin its own ground; otherwise use the rotated one
    if root not in ROOT_CFG:
        cfg['ground'] = ground
    return cfg


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--ids', required=True,
                    help='comma-separated root ids (any animal-forward FAIL root)')
    ap.add_argument('--max-retries', type=int, default=4)  # throttled 8→4 (load relief, 2026-06-11)
    ap.add_argument('--pass-tol', type=float, default=edgeverify.PASS_TOL)
    ap.add_argument('--warn-tol', type=float, default=edgeverify.WARN_TOL)
    ap.add_argument('--ground', help='force one ground for all (else rotates)')
    args = ap.parse_args()

    ids = [int(x) for x in args.ids.split(',') if x.strip()]
    if not ids:
        print(json.dumps({'error': 'no ids'}), file=sys.stderr)
        sys.exit(2)

    # Idempotency: roots that already have a PASS candidate in the queue are DONE.
    # Re-processing them would INSERT a duplicate (all_designs is INSERT-only — dupes
    # can't be cleaned up). Skip them. This makes the wave runner safe to re-invoke.
    done_roots = set()
    try:
        with open(QUEUE) as _q:
            for _line in _q:
                try:
                    _row = json.loads(_line)
                    if _row.get('ok') and _row.get('edge_verdict') == 'PASS':
                        done_roots.add(int(_row.get('root_id')))
                except Exception:
                    pass
    except FileNotFoundError:
        pass

    results = []
    for i, rid in enumerate(ids):
        if rid in done_roots:
            results.append({'root_id': rid, 'ok': True, 'skipped': 'already has PASS candidate',
                            'edge_verdict': 'PASS', 'generator': 'controlnet-anchored-reroll'})
            print(f"=> root {rid}: SKIP (already has PASS candidate)", file=sys.stderr)
            continue
        ground = args.ground or cnar.GROUNDS[i % len(cnar.GROUNDS)]
        cfg = cfg_for(rid, ground)
        try:
            # DEFAULT_CFG roots are canny (no_canny=False); reroll_one honors that.
            # force_canny stays False — the cfg's no_canny flag already drives the path.
            r = reroll2.reroll_one(rid, cfg, args.max_retries,
                                   args.pass_tol, args.warn_tol, force_canny=False)
        except Exception as e:
            r = {'root_id': rid, 'ok': False, 'error': str(e), 'ground': cfg.get('ground')}
        # stamp the generator name so the ledger/queue is honest about which script ran
        if isinstance(r, dict):
            r.setdefault('generator', 'controlnet-anchored-reroll')
            r['generator'] = 'controlnet-anchored-reroll'
        results.append(r)
        if r.get('ok'):
            with open(QUEUE, 'a') as q:
                q.write(json.dumps(r) + '\n')
        v = r.get('edge_verdict', r.get('error', '?'))
        print(f"=> root {rid}: {'OK' if r.get('ok') else 'FAIL'} edge={v} "
              f"new_id={r.get('new_id')} attempts={r.get('attempts_used')} "
              f"ground={r.get('ground')}", file=sys.stderr)

    print(json.dumps(results, indent=2))
    # Exit 0 when the wave RAN — a root that "couldn't clear" after retries is a
    # normal terminal outcome (it won't clear on a re-run), so the wave runner MUST
    # advance its manifest index past it. Exit non-zero ONLY if every root threw an
    # execution exception (e.g. ComfyUI down) — that's a real infra failure where the
    # runner should NOT advance (it retries; the loop's zero-PASS-streak stop catches
    # a sustained outage). Per-root non-clears never block manifest progress.
    all_errored = bool(results) and all('error' in r for r in results)
    sys.exit(1 if all_errored else 0)


if __name__ == '__main__':
    main()