`visual-search/train/mac_deploy_embed.py`

103 lines
#!/usr/bin/env python3
# LOCAL deploy embedder: embed the whole catalog with the AUGMENTED model (dw_clip_ft_aug.pt) on the
# Mac's MPS, writing a Postgres COPY file (text format) for Kamatera to ingest LOCALLY. No prod-DB
# connection from the Mac — it only produces a data file. Reuses the 90k training image cache; dead
# URLs get a NULL-embedding marker (so the index self-terminates, same as the daemon). Resumable.
import os, io, sys, csv, time, threading, queue, urllib.request, urllib.parse
os.environ.setdefault('PYTORCH_ENABLE_MPS_FALLBACK', '1')
import torch, torch.nn.functional as F, open_clip
from PIL import Image
csv.field_size_limit(10_000_000); Image.MAX_IMAGE_PIXELS = None

HERE = os.path.dirname(os.path.abspath(__file__)); CACHE = os.path.join(HERE, 'cache')
MANI = os.path.join(HERE, 'deploy_manifest.tsv'); OUT = os.path.join(HERE, 'deploy_embeddings.copy')
DONE = os.path.join(HERE, 'deploy_done.txt')
UA = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/17.0 Safari/605.1.15"
WORKERS = 20; EMB_BATCH = 128; CHUNK = 1500

model, _, pp = open_clip.create_model_and_transforms("ViT-B-32", pretrained="laion2b_s34b_b79k")
model.load_state_dict(torch.load(os.path.join(HERE, 'ckpt', 'dw_clip_ft_aug.pt'), map_location='cpu'), strict=False)
dev = 'mps' if torch.backends.mps.is_available() else 'cpu'
model = model.to(dev).float().eval(); DIM = model.visual.output_dim
print(f"AUG model loaded · device {dev} · dim {DIM}", flush=True)

def small(u):
    if 'shopify' in u:
        pr = urllib.parse.urlparse(u); q = urllib.parse.parse_qs(pr.query); q['width'] = ['512']
        return urllib.parse.urlunparse(pr._replace(query=urllib.parse.urlencode(q, doseq=True)))
    return u
def esc(s): return (s or '').replace('\\', '\\\\').replace('\t', ' ').replace('\n', ' ').replace('\r', ' ')

def load_img(row):
    # returns (row, tensor)=ok · (row,'DEAD')=genuine 404/410 → NULL marker · (row,None)=transient → retry later
    idv, dw, mfr, vc, pat, url = row
    cp = os.path.join(CACHE, idv + '.jpg')
    if os.path.exists(cp) and os.path.getsize(cp) > 200:
        try: return (row, pp(Image.open(cp).convert('RGB')))
        except Exception: pass
    dead = False
    for attempt in range(3):
        try:
            data = urllib.request.urlopen(urllib.request.Request(small(url), headers={'User-Agent': UA}), timeout=25).read()
            im = Image.open(io.BytesIO(data)).convert('RGB')
            try: im.copy().save(cp, 'JPEG', quality=88)   # populate cache for resumes
            except Exception: pass
            return (row, pp(im))
        except urllib.error.HTTPError as e:
            if e.code in (404, 410, 403 if False else 0): dead = True; break   # genuinely gone
            if e.code in (429, 500, 502, 503):            # rate-limit / transient → back off + retry
                time.sleep(1.5 * (attempt + 1)); continue
            dead = True; break
        except Exception:
            time.sleep(1.0 * (attempt + 1)); continue     # network hiccup → retry
    return (row, 'DEAD' if dead else None)

def main():
    done = set()
    if os.path.exists(DONE):
        with open(DONE) as f: done = set(l.strip() for l in f if l.strip())
    rows = []
    with open(MANI) as f:
        for r in csv.reader(f, delimiter='\t'):
            if len(r) >= 6 and r[0] and r[0] not in done and r[5].startswith('http'): rows.append(r)
    print(f"to embed: {len(rows)} (already done {len(done)})", flush=True)
    fout = open(OUT, 'a'); fdone = open(DONE, 'a')
    t0 = time.time(); total = 0
    for i in range(0, len(rows), CHUNK):
        batch = rows[i:i+CHUNK]
        got = [None]*len(batch); q = queue.Queue()
        for j, r in enumerate(batch): q.put((j, r))
        lock = threading.Lock()
        def worker():
            while True:
                try: j, r = q.get_nowait()
                except queue.Empty: return
                got[j] = load_img(r); q.task_done()
        ts = [threading.Thread(target=worker, daemon=True) for _ in range(WORKERS)]
        for t in ts: t.start()
        for t in ts: t.join()
        ok = [g for g in got if g and not isinstance(g[1], str) and g[1] is not None]
        bad = [g for g in got if g and g[1] == 'DEAD']   # genuine 404 → NULL marker
        # g[1] is None → transient failure: NOT written, NOT marked done → retried on the next resume pass
        # embed the good ones in sub-batches
        for k in range(0, len(ok), EMB_BATCH):
            sub = ok[k:k+EMB_BATCH]
            with torch.no_grad():
                v = model.encode_image(torch.stack([g[1] for g in sub]).to(dev))
                v = (v / v.norm(dim=-1, keepdim=True)).cpu().numpy().astype('float32')
            for g, emb in zip(sub, v):
                r = g[0]; hx = emb.tobytes().hex()
                fout.write(f"{r[0]}\t{esc(r[1])}\t{esc(r[2])}\t{esc(r[3])}\t{esc(r[4])}\t{esc(r[5])}\t{DIM}\t\\\\x{hx}\n")
                fdone.write(r[0] + "\n")
        for g in bad:
            r = g[0]
            fout.write(f"{r[0]}\t{esc(r[1])}\t{esc(r[2])}\t{esc(r[3])}\t{esc(r[4])}\t{esc(r[5])}\t\\N\t\\N\n")
            fdone.write(r[0] + "\n")
        fout.flush(); fdone.flush(); total += len(batch)
        print(f"{total}/{len(rows)} embedded  {total/max(1e-6,time.time()-t0):.0f}/s", flush=True)
    fout.close(); fdone.close()
    print(f"DONE — {total} rows → {OUT}", flush=True)

if __name__ == '__main__':
    main()