Update app.py

app.py

@@ -30,43 +30,43 @@ 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):
-    openai_api_key: Optional[SecretStr] = Field(
-        alias="api_key",
-        default_factory=secret_from_env("OPENROUTER_API_KEY", default=None),
-    )
-    @property
-    def lc_secrets(self) -> dict[str, str]:
-        return {"openai_api_key": "OPENROUTER_API_KEY"}
+# # --- 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):
+#     openai_api_key: Optional[SecretStr] = Field(
+#         alias="api_key",
+#         default_factory=secret_from_env("OPENROUTER_API_KEY", default=None),
+#     )
+#     @property
+#     def lc_secrets(self) -> dict[str, str]:
+#         return {"openai_api_key": "OPENROUTER_API_KEY"}
  
-    def __init__(self,
-                 openai_api_key: Optional[str] = None,
-                 **kwargs):
-        openai_api_key = (
-            openai_api_key or os.environ.get("OPENROUTER_API_KEY")
-        )
-        super().__init__(
-            base_url="https://openrouter.ai/api/v1",
-            openai_api_key=openai_api_key,
-            **kwargs
-        )
+#     def __init__(self,
+#                  openai_api_key: Optional[str] = None,
+#                  **kwargs):
+#         openai_api_key = (
+#             openai_api_key or os.environ.get("OPENROUTER_API_KEY")
+#         )
+#         super().__init__(
+#             base_url="https://openrouter.ai/api/v1",
+#             openai_api_key=openai_api_key,
+#             **kwargs
+#         )
  
-llm2 = ChatOpenRouter(
-    #model_name="deepseek/deepseek-r1-0528:free",
-    #model_name="google/gemini-2.0-flash-exp:free",
-    #model_name="deepseek/deepseek-v3-base:free",
-    model_name="deepseek/deepseek-r1:free"
-)
+# llm2 = ChatOpenRouter(
+#     #model_name="deepseek/deepseek-r1-0528:free",
+#     #model_name="google/gemini-2.0-flash-exp:free",
+#     #model_name="deepseek/deepseek-v3-base:free",
+#     model_name="deepseek/deepseek-r1:free"
+# )
 
 
 def log_execution_time(func):
@@ -79,7 +79,7 @@ def log_execution_time(func):
         return result
     return wrapper
 
-global pdf_doc
+# global pdf_doc
 # ============================== #
 #     INITIALIZE CLIP EMBEDDER   #
 # ============================== #
@@ -319,43 +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_processor = AutoImageProcessor.from_pretrained(model_name, use_fast=True)
-        _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)
+# # 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_processor = AutoImageProcessor.from_pretrained(model_name, use_fast=True)
+#         _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:
@@ -921,7 +921,6 @@ def clean_base64_for_model(raw_b64, max_bytes_threshold=4000000) -> str:
     # otherwise return original with its mime prefix (ensure prefix exists)
     return f"data:{mime};base64,{clean_b64}"
 
-
 SCRATCH_OPCODES = [
     'motion_movesteps', 'motion_turnright', 'motion_turnleft', 'motion_goto',
     'motion_gotoxy', 'motion_glideto', 'motion_glidesecstoxy', 'motion_pointindirection',
@@ -3382,43 +3381,107 @@ SPRITE_DIR   / "Abby.sprite3" / "34a175600dc009a521eb46fdbbbeeb67.png"
         CODE_BLOCKS_DIR  / "script5.jpg",
         CODE_BLOCKS_DIR  / "script6.jpg"]
     folder_image_paths = [os.path.normpath(str(p)) for p in folder_image_paths]
-    # =========================================
 
-    # -----------------------------------------
-    #  Load reference embeddings from JSON
-    # -----------------------------------------
-    with open(f"{BLOCKS_DIR}/embed.json", "r") as f:
+
+    # ============================== #
+    #      EMBED SPRITE IMAGES       #
+    #      (using CLIP again)        #
+    # ============================== #
+
+    # Make sure all buffers are at start
+    for buf in sprite_images_bytes:
+        try:
+            buf.seek(0)
+        except Exception:
+            pass
+
+    # Try the fast path: embed whole list at once (many CLIP wrappers accept a list of BytesIO/PIL)
+    try:
+        sprite_matrix = clip_embd.embed_image(sprite_images_bytes, batch_size=8)
+        sprite_matrix = np.array(sprite_matrix, dtype=np.float32)
+    except Exception:
+        sprite_feats = []
+        for buf in sprite_images_bytes:
+            buf.seek(0)
+            try:
+                feats = clip_embd.embed_image([buf])[0]
+            except Exception:
+                buf.seek(0)
+                pil_img = Image.open(buf).convert("RGB")
+                try:
+                    feats = clip_embd.embed_image([pil_img])[0]
+                except Exception:
+                    pil_arr = np.array(pil_img)
+                    feats = clip_embd.embed_image([pil_arr])[0]
+            sprite_feats.append(np.asarray(feats, dtype=np.float32))
+        sprite_matrix = np.vstack(sprite_feats)  # shape (N, D)
+
+    # --- load reference embeddings (unchanged) ---
+    with open(f"{BLOCKS_DIR}/openclip_embeddings.json", "r") as f:
         embedding_json = json.load(f)
 
+    img_matrix = np.array([img["embeddings"] for img in embedding_json], dtype=np.float32)
+
+    
+    # =========================================
+
+    # # -----------------------------------------
+    # #  Load reference embeddings from JSON
+    # # -----------------------------------------
+    # with open(f"{BLOCKS_DIR}/embed.json", "r") as f:
+    #     embedding_json = json.load(f)
+
     # =========================================
     #  Decode & embed each sprite image
     # =========================================
+    # # ============================== #
+    # #      EMBED SPRITE IMAGES       #
+    # # ============================== #
     # sprite_features = []
     # for b64 in sprite_base64:
-    #     if "," in b64:
+    #     if "," in b64:  # strip data URI prefix if present
     #         b64 = b64.split(",", 1)[1]
-            
+    
     #     img_bytes = base64.b64decode(b64)
     #     pil_img = Image.open(BytesIO(img_bytes)).convert("RGB")
+    
+    #     # optional re-encode to PNG for CLIP
     #     buf = BytesIO()
     #     pil_img.save(buf, format="PNG")
     #     buf.seek(0)
-    #     feats = clip_embd.embed_image([buf])[0]
+    
+    #     feats = clip_embd.embed_image([buf])[0]  # extract CLIP embedding
     #     sprite_features.append(feats)
+    
+    # sprite_matrix = np.array(sprite_features, dtype=np.float32) 
+    # # ============================== #
+    # #      EMBED SPRITE IMAGES       #
+    # # ============================== #
+    # # ensure model is initialized (fast no-op after first call)
+    # init_dinov2()
 
-    # ============================== #
-    #      EMBED SPRITE IMAGES       #
-    # ============================== #
-    # 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)
+    # # 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)
+    # # 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)
+    
+    def l2_normalize_rows(x: np.ndarray, eps: float = 1e-10) -> np.ndarray:
+        """
+        L2-normalize each row of a 2D numpy array.
+    
+        Args:
+            x: Array of shape (N, D).
+            eps: Small constant to avoid division by zero.
+    
+        Returns:
+            Normalized array of shape (N, D) where each row has unit norm.
+        """
+        norms = np.linalg.norm(x, axis=1, keepdims=True)
+        return x / np.maximum(norms, eps)
+    
     sprite_matrix = l2_normalize_rows(sprite_matrix)
     img_matrix = l2_normalize_rows(img_matrix)