Update app.py

app.py

@@ -27,6 +27,17 @@ from langchain.chat_models import ChatOpenAI
 from langchain_openai import ChatOpenAI
 from pydantic import Field, SecretStr
 from difflib import get_close_matches
+import torch
+from transformers import AutoImageProcessor, AutoModel
+
+# --- Config (tune threads as needed) ---
+DINOV2_MODEL = "facebook/dinov2-small"   # small = best CPU latency/quality tradeoff
+DEVICE = torch.device("cpu")
+torch.set_num_threads(4)  # tune for your CPU
+
+# --- Globals for single-shot model load ---
+_dinov2_processor = None
+_dinov2_model = None
 
 os.environ["OPENROUTER_API_KEY"] = os.getenv("OPENROUTER_API_KEY", "default_key_or_placeholder")
 class ChatOpenRouter(ChatOpenAI):
@@ -308,6 +319,43 @@ agent_json_resolver = create_react_agent(
     prompt=SYSTEM_PROMPT_JSON_CORRECTOR 
 )
 
+# adding the new embedding models:
+def init_dinov2(model_name: str = DINOV2_MODEL, device: torch.device = DEVICE):
+    """Lazy-initialize DINOv2 processor & model (call once before embedding)."""
+    global _dinov2_processor, _dinov2_model
+    if _dinov2_processor is None or _dinov2_model is None:
+        _dinov2_processor = AutoImageProcessor.from_pretrained(model_name)
+        _dinov2_model = AutoModel.from_pretrained(model_name)
+        _dinov2_model.eval().to(device)
+
+def embed_bytesio_list(bytesio_list, batch_size: int = 8):
+    """
+    Accepts a list of BytesIO objects (each contains an image, like your sprite_images_bytes).
+    Returns: np.ndarray shape (N, D) of L2-normalized embeddings (dtype float32).
+    """
+    if _dinov2_processor is None or _dinov2_model is None:
+        init_dinov2()
+
+    imgs = [Image.open(b).convert("RGB") for b in bytesio_list]
+    embs = []
+    for i in range(0, len(imgs), batch_size):
+        batch = imgs[i : i + batch_size]
+        inputs = _dinov2_processor(images=batch, return_tensors="pt")
+        inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
+        with torch.no_grad():
+            out = _dinov2_model(**inputs)
+            # global image embedding from CLS token
+            cls = out.last_hidden_state[:, 0, :]  # (B, D)
+            cls = torch.nn.functional.normalize(cls, p=2, dim=1)  # L2 normalize rows
+            embs.append(cls.cpu().numpy())
+    if not embs:
+        return np.zeros((0, _dinov2_model.config.hidden_size), dtype=np.float32)
+    return np.vstack(embs).astype(np.float32)
+
+def l2_normalize_rows(a: np.ndarray, eps: float = 1e-12) -> np.ndarray:
+    norm = np.linalg.norm(a, axis=1, keepdims=True)
+    return a / (norm + eps)
+
 # Helper function to load the block catalog from a JSON file
 def _load_block_catalog(block_type: str) -> Dict:
     """
@@ -2342,7 +2390,7 @@ def similarity_matching(sprites_data: str, project_folder: str) -> str:
     # -----------------------------------------
     #  Load reference embeddings from JSON
     # -----------------------------------------
-    with open(f"{BLOCKS_DIR}/embeddings.json", "r") as f:
+    with open(f"{BLOCKS_DIR}/embed.json", "r") as f:
         embedding_json = json.load(f)
 
     # =========================================
@@ -2364,10 +2412,18 @@ def similarity_matching(sprites_data: str, project_folder: str) -> str:
     # ============================== #
     #      EMBED SPRITE IMAGES       #
     # ============================== #
-    sprite_features = clip_embd.embed_image(sprite_images_bytes)
-    
-    sprite_matrix = np.vstack(sprite_features)
-    img_matrix = np.array([img["embeddings"] for img in embedding_json])
+    # ensure model is initialized (fast no-op after first call)
+    init_dinov2()
+
+    # embed the incoming sprite BytesIO images (same data structure you already use)
+    sprite_matrix = embed_bytesio_list(sprite_images_bytes, batch_size=8)  # shape (N, D)
+
+    # load reference embeddings from JSON (they must be numeric lists)
+    img_matrix = np.array([img["embeddings"] for img in embedding_json], dtype=np.float32)
+
+    # normalize both sides (important — stored embeddings may not be normalized)
+    sprite_matrix = l2_normalize_rows(sprite_matrix)
+    img_matrix = l2_normalize_rows(img_matrix)
     
     # =========================================
     #  Compute similarities & pick best match