`scripts/gen_briefs.py`

113 lines
#!/usr/bin/env python3
"""Generate settlement-cleared brief variants via Replicate SDXL. Curator-mode, $0-local-first fallback."""
import os, re, json, time, sys, urllib.request, urllib.error

ROOT = "/Users/macstudio3/Projects/pattern-vault"
OUT = f"{ROOT}/assets/generated"
os.makedirs(OUT, exist_ok=True)

# token
tok = None
for p in [os.path.expanduser("~/Projects/secrets-manager/.env")]:
    with open(p) as f:
        m = re.search(r"REPLICATE_API_TOKEN=([A-Za-z0-9_-]+)", f.read())
        if m: tok = m.group(1)
assert tok, "no replicate token"

briefs = json.load(open(f"{ROOT}/data/trend-board.json"))["briefs"]

NEG = ("birds, bird, butterfly, butterflies, bananas, banana leaves, grapes, grape clusters, "
       "tropical monstera palm fronds, text, watermark, signature, logo, blurry, low quality, "
       "jpeg artifacts, seams, mismatched edges, frame, border, vignette")

# Learned 2026-07-02: SDXL defaults grounds to muddy-dark and ignores a color named late in the
# prompt. To force a ground color, front-load it as the FIRST token, negate dark backgrounds, and
# bump guidance so the color binds. Add a "groundColor" field to a brief to trigger this mode.
DARK_NEG = "black background, dark background, navy background, charcoal background, gradient background"

def sdxl_prompt(b):
    base = b["prompt"]
    if b["id"] == "BRIEF-2026-mural-01":
        tile = "fine-art wall mural panorama, painterly, cohesive scene"
        w, h = 1216, 832
    else:
        tile = "seamless repeating wallpaper pattern, tileable, flat lay, even lighting, high-end wallcovering"
        w, h = 1024, 1024
    ground = b.get("groundColor")
    if ground:
        # front-load the flat ground color; bump guidance; add dark-bg negative
        prompt = (f"flat solid {ground} background color filling the entire canvas, "
                  f"{base} {tile}. Colorway: {b['colorway']}.")
        return prompt, w, h, 9.5, DARK_NEG
    return f"{base} {tile}. Colorway: {b['colorway']}.", w, h, 7.5, ""

UA = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36"
def post(url, payload, extra_headers=None):
    hdr = {"Authorization": f"Bearer {tok}", "Content-Type": "application/json", "User-Agent": UA}
    if extra_headers: hdr.update(extra_headers)
    req = urllib.request.Request(url, data=json.dumps(payload).encode(), headers=hdr, method="POST")
    with urllib.request.urlopen(req, timeout=120) as r:
        return json.load(r)

def get(url):
    req = urllib.request.Request(url, headers={"Authorization": f"Bearer {tok}", "User-Agent": UA})
    with urllib.request.urlopen(req, timeout=60) as r:
        return json.load(r)

MODEL = "stability-ai/sdxl"
VERSION = get(f"https://api.replicate.com/v1/models/{MODEL}")["latest_version"]["id"]
print(f"using {MODEL}@{VERSION[:12]}", flush=True)
results = []
total_secs = 0.0
seeds = [101, 202]  # 2 variants/brief

for b in briefs:
    prompt, w, h, guidance, extra_neg = sdxl_prompt(b)
    neg = NEG + (", " + extra_neg if extra_neg else "")
    for i, seed in enumerate(seeds, 1):
        tag = f"{b['id']}_v{i}"
        payload = {"version": VERSION, "input": {
            "prompt": prompt, "negative_prompt": neg,
            "width": w, "height": h, "num_inference_steps": 30,
            "guidance_scale": guidance, "seed": seed, "refine": "no_refiner"
        }}
        ok = False
        for attempt in range(4):
            try:
                pred = post("https://api.replicate.com/v1/predictions", payload); ok = True; break
            except urllib.error.HTTPError as e:
                body = e.read().decode()[:150]
                if e.code == 429:
                    print(f"[{tag}] 429 throttled, backoff {6*(attempt+1)}s", flush=True)
                    time.sleep(6*(attempt+1)); continue
                print(f"[{tag}] HTTP {e.code}: {body}", flush=True); break
        if not ok:
            results.append({"tag": tag, "status": "error"}); continue
        # poll if not done
        for _ in range(60):
            if pred.get("status") in ("succeeded", "failed", "canceled"): break
            time.sleep(2); pred = get(pred["urls"]["get"])
        st = pred.get("status")
        secs = (pred.get("metrics") or {}).get("predict_time", 0) or 0
        total_secs += secs
        if st == "succeeded":
            out = pred["output"]
            imgurl = out[0] if isinstance(out, list) else out
            dest = f"{OUT}/{tag}.png"
            urllib.request.urlretrieve(imgurl, dest)
            sz = os.path.getsize(dest)
            print(f"[{tag}] OK {secs:.1f}s -> {dest} ({sz//1024}KB)", flush=True)
            results.append({"tag": tag, "brief": b["id"], "path": dest, "status": "ok",
                            "predict_time": secs, "prompt": prompt})
        else:
            print(f"[{tag}] {st}: {pred.get('error')}", flush=True)
            results.append({"tag": tag, "brief": b["id"], "status": st})
        time.sleep(1.5)

# Replicate SDXL on A40 ~ $0.000725/sec (Nvidia A40 Large). Report both.
RATE = 0.000725
cost = total_secs * RATE
print(f"\n=== TOTAL predict_time {total_secs:.1f}s  est cost ${cost:.4f} (rate ${RATE}/s A40) ===", flush=True)
json.dump({"results": results, "total_predict_secs": total_secs, "est_cost_usd": round(cost,4)},
          open(f"{OUT}/_gen_manifest.json", "w"), indent=2)