scripts/sync-interiors.py

#!/usr/bin/env python3
"""
Two-way sync between this Space's interiors and vai-avatar's interior library.

WHY THIS CANNOT JUST COPY FILES
-------------------------------
The field searches interiors with its own CLIP matrix
(`image-embeddings_25k-512-32bit.bin`, 512 dims). vai-avatar's library uses a
different embedding model. A new image therefore cannot be copied across: it
has to be encoded with the field's own model first, or it lands in a 512-dim
space that is not the one the field searches.

This was measured rather than assumed. Cosine between the two datasets on the
377 images they share is 0.6347; cosine between two DIFFERENT images is 0.6336.
Same dimensionality, no signal -- the spaces are not interchangeable, so a file
copy would produce confidently wrong search results.

WHAT IT DOES
------------
  added    live on the API, absent from the field -> download, embed, upload,
           append to both index files
  removed  present in the field, gone from the API  -> drop from the index and
           the vector blob, delete the bitmap

The source of truth is the API's own `/latest`, not the dataset's
`metadata.json`. That file is a history, not a manifest: it holds 1,752 rows
for 1,122 live images, so deleted rows linger in it and a file-driven sync
would resurrect whatever the user removed on the other side.

RUN
---
    pip install -r scripts/requirements-sync.txt
    python scripts/sync-interiors.py --dry-run    # report, change nothing
    python scripts/sync-interiors.py              # apply
    python scripts/sync-interiors.py --limit 5    # apply a few, for testing
"""

from __future__ import annotations

import argparse
import base64
import io
import json
import os
import struct
import sys
import time
import urllib.error
import urllib.request

from PIL import Image

# --------------------------------------------------------------------------
# Where things live
# --------------------------------------------------------------------------

API = os.environ.get("INTERIOR_API_URL", "https://bep40-vai-avatar2-interior-api.hf.space")
DATASET = os.environ.get("FIELD_DATASET", "bep40/interior-image-field-assets")
HUB = os.environ.get("HF_HUB") or os.environ.get("HF_TOKEN")

INDEX_NAME = "image-embeddings.json"
BLOB_NAME = "image-embeddings_25k-512-32bit.bin"

RESOLVE = f"https://huggingface.co/datasets/{DATASET}/resolve/main/"

# The field's image model, matching what produced the existing vectors.
IMAGE_MODEL = "sentence-transformers/clip-ViT-B-32-multilingual-v1"
DIMS = 512

BRANDS = ("malloca", "grob", "hafele", "eurogold", "canzy", "konox", "demax")


def log(message: str) -> None:
    print(message, flush=True)


def get(url: str, timeout: int = 120) -> bytes:
    request = urllib.request.Request(url, headers={"User-Agent": "interior-sync/1.0"})
    with urllib.request.urlopen(request, timeout=timeout) as response:
        return response.read()


def post_json(url: str, payload: dict, timeout: int = 120) -> dict:
    request = urllib.request.Request(
        url,
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json", "User-Agent": "interior-sync/1.0"},
    )
    with urllib.request.urlopen(request, timeout=timeout) as response:
        return json.loads(response.read().decode())


# --------------------------------------------------------------------------
# The field's vector format
#
# `image-embeddings_25k-512-32bit.bin` is headerless little-endian float32,
# 512 floats per row, row order matching `image-embeddings.json`. The "25k" in
# the filename is decorative: the file holds exactly as many rows as the index
# has entries.
# --------------------------------------------------------------------------

def decode_blob(raw: bytes) -> list[list[float]]:
    count, remainder = divmod(len(raw), 4 * DIMS)
    if remainder:
        raise ValueError(
            f"blob is {len(raw)} bytes, not a whole number of {DIMS}-float rows "
            f"(remainder {remainder}) -- the index and the vectors disagree"
        )
    flat = struct.unpack(f"<{count * DIMS}f", raw[: count * 4 * DIMS])
    return [list(flat[i * DIMS : (i + 1) * DIMS]) for i in range(count)]


def encode_blob(rows: list[list[float]]) -> bytes:
    for index, row in enumerate(rows):
        if len(row) != DIMS:
            raise ValueError(f"row {index} has {len(row)} values, expected {DIMS}")
    return struct.pack(f"<{len(rows) * DIMS}f", *(v for row in rows for v in row))


# --------------------------------------------------------------------------
# Reading the current state
# --------------------------------------------------------------------------

def fetch_live_images() -> dict[str, dict]:
    """Every image the API currently serves, keyed by image_id.

    The paging parameter is `top_k`, not `limit`; sending `limit` silently
    returns the default page of 24 no matter what is asked for.
    """
    first = post_json(f"{API}/latest", {"top_k": 1, "offset": 0})
    total = int(first.get("total", 0))
    live: dict[str, dict] = {}
    offset = 0
    while offset < total:
        page = post_json(f"{API}/latest", {"top_k": 250, "offset": offset})
        rows = page.get("results", [])
        if not rows:
            break
        for row in rows:
            live[row["image_id"]] = row
        offset += len(rows)
        time.sleep(0.4)
    log(f"  API reports {total} live images; collected {len(live)}")
    if len(live) != total:
        log(f"  WARNING: expected {total} ids but collected {len(live)}")
    return live


def load_field_state() -> tuple[list[dict], list[list[float]]]:
    index = json.loads(get(RESOLVE + INDEX_NAME).decode())
    blob = decode_blob(get(RESOLVE + BLOB_NAME))
    if len(index) != len(blob):
        raise ValueError(
            f"index has {len(index)} entries but the blob holds {len(blob)} vectors; "
            "refusing to write a file the field cannot load"
        )
    return index, blob


# --------------------------------------------------------------------------
# Record conventions, mirrored from the field's own scraper output
# --------------------------------------------------------------------------

def aspect_ratio(image: Image.Image) -> float:
    width, height = image.size
    return round(width / height, 4) if height else 1.0


def blur_placeholder(image: Image.Image) -> str:
    """
    A tiny base64 JPEG, like the placeholders the scraper writes.

    The field paints `blur` behind each tile, so an image without one shows an
    empty box for the frames before its bitmap arrives. This costs ~60 bytes and
    removes that flash.
    """
    buffer = io.BytesIO()
    image.convert("RGB").resize((4, 3)).save(buffer, format="JPEG", quality=40)
    return base64.b64encode(buffer.getvalue()).decode()


def title_from(entry: dict) -> str:
    for key in ("desc", "title"):
        value = entry.get(key)
        if isinstance(value, str) and value.strip():
            return " ".join(value.split())[:120]
    for key in ("labels", "keywords"):
        value = entry.get(key)
        if isinstance(value, list) and value:
            return str(value[0]).strip()[:120]
    return ""


def categories_from(entry: dict) -> list[str]:
    out: list[str] = []
    for key in ("categories", "room_hints"):
        for value in entry.get(key) or []:
            text = str(value).strip()
            if text and text not in out:
                out.append(text)
    return out[:4] or ["Uncategorized"]


def build_record(entry: dict, image: Image.Image) -> dict:
    """One interior row, in exactly the shape image-embeddings.json uses."""
    image_id = entry["image_id"]
    return {
        "id": image_id,
        # The field resolves an interior's bitmap relative to its own dataset and
        # the scraper stored a bare filename, so keep that convention.
        "url": f"{image_id}.jpg",
        "ar": aspect_ratio(image),
        "blur": blur_placeholder(image),
        "title": title_from(entry),
        "categories": categories_from(entry),
        "source": str(entry.get("source") or ""),
    }


# --------------------------------------------------------------------------
# Embedding
# --------------------------------------------------------------------------

class Embedder:
    """The field's own image model, loaded lazily so a no-op sync is cheap."""

    def __init__(self) -> None:
        self._model = None

    def _load(self):
        if self._model is None:
            from sentence_transformers import SentenceTransformer

            log(f"  loading {IMAGE_MODEL} (first run downloads ~600 MB)")
            self._model = SentenceTransformer(IMAGE_MODEL)
        return self._model

    def image(self, image: Image.Image) -> list[float]:
        vector = self._load().encode([image], convert_to_numpy=True)[0]
        norm = float((vector**2).sum() ** 0.5)
        if norm > 0:
            vector = vector / norm
        if len(vector) != DIMS:
            raise ValueError(f"model returned {len(vector)} dims, expected {DIMS}")
        return [float(v) for v in vector]


# --------------------------------------------------------------------------
# The sync itself
# --------------------------------------------------------------------------

def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dry-run", action="store_true", help="report changes, write nothing")
    parser.add_argument("--no-upload", action="store_true", help="build everything but do not push")
    parser.add_argument("--limit", type=int, default=0, help="cap how many images to add")
    args = parser.parse_args()

    if not HUB and not args.dry_run and not args.no_upload:
        log("HF_TOKEN is not set. Export it, or pass --no-upload to build without pushing.")
        return 2

    log("Reading the field...")
    index, vectors = load_field_state()
    log(f"  field holds {len(index)} interiors")

    log("Reading the API...")
    live = fetch_live_images()

    have = {str(row.get("id")) for row in index}
    to_add = sorted(set(live) - have)
    to_drop = sorted(have - set(live))

    log("")
    log(f"  to add   : {len(to_add)}")
    log(f"  to remove: {len(to_drop)}")
    if to_drop:
        log(f"    {', '.join(to_drop[:6])}{' ...' if len(to_drop) > 6 else ''}")
    if to_add:
        log(f"    {', '.join(to_add[:6])}{' ...' if len(to_add) > 6 else ''}")

    if args.limit:
        to_add = to_add[: args.limit]
        log(f"  limited to {len(to_add)} additions for this run")

    if not to_add and not to_drop:
        log("\nNothing to do -- the two sides already agree.")
        return 0
    if args.dry_run:
        log("\nDry run: nothing was written.")
        return 0
    if not to_add:
        log("\nOnly removals pending, which the same run applies below.")
    if not to_add and not to_drop:
        return 0

    added_rows: list[dict] = []
    added_vectors: list[list[float]] = []
    bitmaps: list[tuple[str, bytes]] = []

    if to_add:
        log(f"\nEmbedding {len(to_add)} new interiors...")
        embedder = Embedder()
        for n, image_id in enumerate(to_add, 1):
            entry = live[image_id]
            try:
                raw = get(f"{API}/image/{image_id}.jpg")
                image = Image.open(io.BytesIO(raw))
                image.load()
            except (urllib.error.URLError, OSError) as error:
                log(f"  [{n}/{len(to_add)}] {image_id}: fetch failed ({error}) -- skipped")
                continue
            try:
                vector = embedder.image(image)
            except Exception as error:  # noqa: BLE001 - report and keep going
                log(f"  [{n}/{len(to_add)}] {image_id}: embed failed ({error}) -- skipped")
                continue
            added_rows.append(build_record(entry, image))
            added_vectors.append(vector)
            bitmaps.append((f"images/{image_id}.jpg", raw))
            log(f"  [{n}/{len(to_add)}] {image_id}  {image.size[0]}x{image.size[1]}  {title_from(entry)[:44]}")

    new_index = index + added_rows
    new_vectors = vectors + added_vectors

    if to_drop:
        # Rebuild both files from the surviving rows, so the index and the blob
        # cannot drift out of alignment: the field refuses to load them otherwise.
        drop = set(to_drop)
        keep = [i for i, row in enumerate(new_index) if str(row.get("id")) not in drop]
        new_index = [new_index[i] for i in keep]
        new_vectors = [new_vectors[i] for i in keep]

    log("")
    log(f"Final: {len(new_index)} interiors "
        f"(+{len(added_rows)} added, -{len(to_drop)} removed)")

    if not added_rows and not to_drop:
        log("Nothing changed.")
        return 0
    if args.no_upload:
        log("Skipping push (--no-upload).")
        return 0

    from huggingface_hub import CommitOperationAdd, CommitOperationDelete, HfApi

    client = HfApi(token=HUB) if HUB else HfApi()
    operations = [CommitOperationAdd(path_in_repo=name, path_or_fileobj=data)
                  for name, data in bitmaps]
    operations.append(CommitOperationAdd(
        path_in_repo=INDEX_NAME,
        path_or_fileobj=json.dumps(new_index, ensure_ascii=False).encode(),
    ))
    operations.append(CommitOperationAdd(
        path_in_repo=BLOB_NAME, path_or_fileobj=encode_blob(new_vectors)))
    if to_drop:
        # Only delete what is actually on the Hub, or the whole commit 404s.
        present = set(client.list_repo_files(DATASET, repo_type="dataset"))
        removed = 0
        for image_id in to_drop:
            path = f"images/{image_id}.jpg"
            if path in present:
                operations.append(CommitOperationDelete(path_in_repo=path))
                removed += 1
        log(f"  deleting {removed} of {len(to_drop)} bitmaps (the rest were already gone)")

    commit = client.create_commit(
        repo_id=DATASET,
        repo_type="dataset",
        operations=operations,
        commit_message=(
            f"sync interiors with vai-avatar: +{len(added_rows)} -{len(to_drop)}\n\n"
            "Driven by the API's live image list rather than metadata.json: that file\n"
            "retains deleted rows (1,752 rows for 1,122 live images), so a\n"
            "file-driven sync resurrects whatever the user removed on the other side.\n\n"
            "New images are embedded with the field's own CLIP model. The two sides\n"
            "use different embedding models, so copying a file would place it in an\n"
            "unrelated vector space and search would return noise.\n\n"
            "Generated by scripts/sync-interiors.py."
        ),
    )
    log(f"Pushed {commit.oid} to {DATASET}")
    log("The field polls the dataset revision, so these appear within a minute.")
    return 0


if __name__ == "__main__":
    sys.exit(main())