Update app.py

app.py

@@ -66,33 +66,8 @@ def search(text_query, image_query, k=5):
         return [], "Build the index first."
 
     with torch.no_grad():
-        if text_query and text_query.strip():
-            inputs = processor(text=[text_query.strip()], return_tensors="pt")
-            q = model.get_text_features(**inputs)           # [1, 512]
-        elif image_query is not None:
-            pil = image_query.convert("RGB")
-            inputs = processor(images=pil, return_tensors="pt")
-            q = model.get_image_features(**inputs)          # [1, 512]
-        else:
-            return [], "Enter text or upload an image."
-
-        q = F.normalize(q, p=2, dim=-1)[0]                  # [512]
-        sims = (INDEX["feats"] @ q).cpu()                   # [N]
-        topk = torch.topk(sims, k=min(int(k), sims.shape[0]))
-
-        items = []
-        for idx in topk.indices.tolist():
-            cap = f"id: {INDEX['ids'][idx]}  score: {float(sims[idx]):.3f}  band: {INDEX['bands'][idx]}"
-            items.append((INDEX["thumbs"][idx], cap))
-        return items, f"Returned {len(items)} results."
-
-    def search(text_query, image_query, k=5):
-    if INDEX["feats"] is None:
-        return [], "Build the index first."
-
-    with torch.no_grad():
-        if text_query and text_query.strip():
-            inputs = processor(text=[text_query.strip()], return_tensors="pt")
+        if text_query and str(text_query).strip():
+            inputs = processor(text=[str(text_query).strip()], return_tensors="pt")
             q = model.get_text_features(**inputs)          # [1, 512]
         elif image_query is not None:
             pil = image_query.convert("RGB")
@@ -103,13 +78,16 @@ def search(text_query, image_query, k=5):
 
         q = F.normalize(q, p=2, dim=-1)[0]                 # [512]
         sims = (INDEX["feats"] @ q).cpu()                  # [N]
-        topk = torch.topk(sims, k=min(int(k), sims.shape[0]))
+        k = min(int(k), sims.shape[0])
+        topk = torch.topk(sims, k=k)
 
         items = []
         for idx in topk.indices.tolist():
             cap = f"id: {INDEX['ids'][idx]}  score: {float(sims[idx]):.3f}  band: {INDEX['bands'][idx]}"
             items.append((INDEX["thumbs"][idx], cap))
-        return items, f"Returned {len(items)} results."
+
+        return items, f"Returned {k} results."
+
 
 # ---------- UI ----------
 with gr.Blocks() as demo: