← back to Pattern Vault

whimsical-compare/scripts/gen_ours.py

174 lines

#!/usr/bin/env python3
"""
Generate OUR original version of each top-100 whimsical bestseller.

Pipeline per design:
  1. qwen2.5vl (local Ollama, $0) looks at the ORIGINAL thumbnail and returns a
     structured brief: style, motifs, palette, whimsy notes, and whether any
     settlement-banned element (bird/butterfly/banana/grape) is visible.
  2. We compose an ORIGINAL, settlement-SAFE SDXL prompt — inspired-by the lane
     (whimsical + palette + motif family), never a copy, with banned elements
     hard-excluded in the prompt AND the negative prompt.
  3. Replicate SDXL renders a seamless tile (~1 cent/img).
  4. Post-gen: qwen2.5vl vision-checks the OUTPUT for banned elements. Any hit
     -> settlement_verdict=BLOCK and the file is discarded (no copy in place).
  5. Store our_brief / our_prompt / our_design_path / settlement_verdict in PG.

$0 for the vision steps (local). Replicate SDXL is the only paid part.
Usage: python3 gen_ours.py [--limit N] [--only <sf_design_id>] [--redo]
"""
import os, re, json, time, sys, base64, urllib.request, urllib.error, subprocess

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OURS = f"{ROOT}/public/ours"
ORIG = f"{ROOT}/public/originals"
os.makedirs(OURS, exist_ok=True)
PG = "postgresql:///dw_unified?host=/tmp"
OLLAMA = os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434")
VMODEL = "qwen2.5vl:7b"

args = sys.argv[1:]
def argv(flag, default=None):
    return args[args.index(flag)+1] if flag in args else default
LIMIT = int(argv("--limit", "1000"))
ONLY = argv("--only")
REDO = "--redo" in args

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

UA = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36"
NEG = ("bird, birds, butterfly, butterflies, banana, bananas, banana leaves, banana pods, "
       "grape, grapes, grape clusters, text, watermark, signature, logo, blurry, low quality, "
       "jpeg artifacts, seams, mismatched edges, frame, border, vignette")

def http_json(url, payload=None, headers=None, timeout=180):
    hdr = {"User-Agent": UA, "Content-Type": "application/json"}
    if headers: hdr.update(headers)
    data = json.dumps(payload).encode() if payload is not None else None
    req = urllib.request.Request(url, data=data, headers=hdr, method="POST" if data else "GET")
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.load(r)

def ollama_vision(img_path, prompt):
    b64 = base64.b64encode(open(img_path, "rb").read()).decode()
    out = http_json(f"{OLLAMA}/api/generate", {
        "model": VMODEL, "prompt": prompt, "images": [b64],
        "stream": False, "options": {"temperature": 0.4}
    }, timeout=180)
    return out.get("response", "")

def describe(img_path):
    p = ("You are a textile design analyst. Look at this repeating wallpaper/fabric pattern and "
         "reply ONLY with compact JSON: {\"style\":\"...\",\"motifs\":\"...\",\"palette\":\"...\","
         "\"whimsy\":\"...\",\"has_bird\":bool,\"has_butterfly\":bool,\"has_banana\":bool,\"has_grape\":bool}. "
         "style=aesthetic era (e.g. cottagecore, mid-century, art-nouveau, folk, watercolor). "
         "motifs=the main repeated shapes. palette=3-5 dominant colors in words. "
         "whimsy=what makes it playful/whimsical. Flags true only if that element is clearly present.")
    raw = ollama_vision(img_path, p)
    m = re.search(r"\{.*\}", raw, re.S)
    try:
        d = json.loads(m.group(0)) if m else {}
    except Exception:
        d = {}
    # coerce any list/None fields to clean strings
    def s(v, default):
        if isinstance(v, list): return ", ".join(str(x) for x in v)
        return str(v) if v not in (None, "") else default
    return {
        "style": s(d.get("style"), "whimsical"),
        "motifs": s(d.get("motifs"), "scattered playful motifs"),
        "palette": s(d.get("palette"), "soft multicolor"),
        "whimsy": s(d.get("whimsy"), "playful"),
        "has_bird": bool(d.get("has_bird")), "has_butterfly": bool(d.get("has_butterfly")),
        "has_banana": bool(d.get("has_banana")), "has_grape": bool(d.get("has_grape")),
    }

def check_output(img_path):
    p = ("Look at this wallpaper image. Reply ONLY JSON: "
         "{\"bird\":bool,\"butterfly\":bool,\"banana\":bool,\"grape\":bool}. "
         "true only if that specific element is clearly depicted.")
    raw = ollama_vision(img_path, p)
    m = re.search(r"\{.*\}", raw, re.S)
    try: return json.loads(m.group(0)) if m else {}
    except Exception: return {}

def our_prompt(d):
    style = d.get("style", "whimsical")
    motifs = d.get("motifs", "playful scattered motifs")
    palette = d.get("palette", "soft multicolor")
    # Rewrite motifs to strip any banned element the analyst saw — we go our own way.
    safe_motifs = re.sub(r"\b(bird|birds|butterfly|butterflies|banana[s]?|grape[s]?)\b", "botanical sprig", motifs, flags=re.I)
    base = (f"Original hand-illustrated {style} wallpaper: {safe_motifs}, whimsical and charming, "
            f"gentle hand-drawn linework, our own artistic interpretation, not a copy of any existing design. "
            f"seamless repeating wallpaper pattern, tileable, flat lay, even lighting, high-end wallcovering. "
            f"Colorway: {palette}.")
    return base

def sdxl(prompt, seed):
    ver = http_json("https://api.replicate.com/v1/models/stability-ai/sdxl",
                    headers={"Authorization": f"Bearer {tok}"})["latest_version"]["id"]
    pred = http_json("https://api.replicate.com/v1/predictions", {
        "version": ver, "input": {"prompt": prompt, "negative_prompt": NEG,
        "width": 1024, "height": 1024, "num_inference_steps": 30,
        "guidance_scale": 7.5, "seed": seed, "refine": "no_refiner"}},
        headers={"Authorization": f"Bearer {tok}"})
    for _ in range(90):
        if pred.get("status") in ("succeeded","failed","canceled"): break
        time.sleep(2)
        pred = http_json(pred["urls"]["get"], headers={"Authorization": f"Bearer {tok}"})
    return pred

def rows():
    q = "SELECT rank, sf_design_id FROM whimsical_top100"
    if ONLY: q += f" WHERE sf_design_id={int(ONLY)}"
    elif not REDO: q += " WHERE our_design_path IS NULL"
    q += " ORDER BY rank LIMIT %d" % LIMIT
    out = subprocess.run(["psql", PG, "-t", "-A", "-F", "\t", "-c", q], capture_output=True, text=True).stdout
    return [l.split("\t") for l in out.strip().splitlines() if l.strip()]

def sqlset(did, brief, prompt, path, verdict):
    b = brief.replace("'", "''")
    p = prompt.replace("'", "''")
    path_sql = "NULL" if not path else "'" + path.replace("'", "''") + "'"
    sql = ("UPDATE whimsical_top100 SET our_brief='%s', our_prompt='%s', our_design_path=%s, "
           "settlement_verdict='%s', updated_at=now() WHERE sf_design_id=%d;"
           % (b, p, path_sql, verdict, int(did)))
    subprocess.run(["psql", PG, "-c", sql], capture_output=True, text=True)

total_secs=0.0; made=0; blocked=0
todo=rows()
print(f"generating {len(todo)} designs | model=stability-ai/sdxl | vision={VMODEL} ($0 local)", flush=True)
for rank, did in todo:
    op = f"{ORIG}/{did}.jpg"
    if not os.path.exists(op): print(f"[{did}] no original, skip", flush=True); continue
    d = describe(op)
    brief = f"{d.get('style','whimsical')} · {d.get('motifs','')[:60]} · {d.get('palette','')[:40]}"
    prompt = our_prompt(d)
    pred = sdxl(prompt, seed=int(did) % 100000)
    secs = (pred.get("metrics") or {}).get("predict_time", 0) or 0
    total_secs += secs
    if pred.get("status") != "succeeded":
        print(f"[{did}] gen {pred.get('status')}: {pred.get('error')}", flush=True)
        sqlset(did, brief, prompt, None, "GEN_FAIL"); continue
    out = pred["output"]; imgurl = out[0] if isinstance(out, list) else out
    dest = f"{OURS}/{did}.png"
    urllib.request.urlretrieve(imgurl, dest)
    chk = check_output(dest)
    if any(chk.get(k) for k in ("bird","butterfly","banana","grape")):
        os.remove(dest); blocked += 1
        print(f"[{did}] SETTLEMENT BLOCK (post-gen vision saw {chk}) — discarded", flush=True)
        sqlset(did, brief, prompt, None, "BLOCK"); continue
    made += 1
    sqlset(did, brief, prompt, f"/ours/{did}.png", "OK")
    print(f"[{did}] rank {rank} OK {secs:.1f}s -> {dest}", flush=True)
    time.sleep(0.8)

RATE=0.000725
cost=total_secs*RATE
print(f"\n=== made={made} blocked={blocked} | predict {total_secs:.1f}s | est ${cost:.4f} (SDXL A40 ${RATE}/s) ===", flush=True)