Update app.py

app.py

@@ -8,35 +8,39 @@ import os
 import spaces
 
 # Load model/processor
-model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
-processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+model = CLIPModel.from_pretrained("openai/clip-vit-base-patch14")
+processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch14")
 model.eval()
 
 DATASET_DIR = Path("dataset")
 CACHE_FILE = "cache.pkl"
 
-def preprocess_image(image: Image.Image) -> Image.Image:
-    return image.resize((224, 224)).convert("RGB")
-
 def get_embedding(image: Image.Image, device="cpu"):
-    image = preprocess_image(image)
+    # Use CLIP's built-in preprocessing instead of custom resize
     inputs = processor(images=image, return_tensors="pt").to(device)
     model_device = model.to(device)
     with torch.no_grad():
         emb = model_device.get_image_features(**inputs)
+    # L2 normalize the embeddings
     emb = emb / emb.norm(p=2, dim=-1, keepdim=True)
     return emb
-
+    
+@space.GPU
 def get_reference_embeddings():
     if os.path.exists(CACHE_FILE):
         with open(CACHE_FILE, "rb") as f:
             return pickle.load(f)
 
     embeddings = {}
+    # Use GPU for preprocessing reference images too for consistency
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
     for img_path in DATASET_DIR.glob("*.jpg"):
         img = Image.open(img_path).convert("RGB")
-        emb = get_embedding(img)
-        embeddings[img_path.name] = emb
+        emb = get_embedding(img, device=device)
+        # Store on CPU to save GPU memory
+        embeddings[img_path.name] = emb.cpu()
+    
     with open(CACHE_FILE, "wb") as f:
         pickle.dump(embeddings, f)
     return embeddings
@@ -47,24 +51,35 @@ reference_embeddings = get_reference_embeddings()
 def search_similar(query_img):
     query_emb = get_embedding(query_img, device="cuda")
     results = []
+    
     for name, ref_emb in reference_embeddings.items():
-        sim = torch.nn.functional.cosine_similarity(query_emb, ref_emb.to("cuda")).item()
+        # Move reference embedding to same device as query
+        ref_emb_gpu = ref_emb.to("cuda")
+        # Compute cosine similarity
+        sim = torch.nn.functional.cosine_similarity(query_emb, ref_emb_gpu, dim=1).item()
         results.append((name, sim))
+    
     results.sort(key=lambda x: x[1], reverse=True)
     return [(f"dataset/{name}", f"Score: {score:.4f}") for name, score in results[:5]]
 
 def add_image(name: str, image):
     path = DATASET_DIR / f"{name}.jpg"
     image.save(path)
-    emb = get_embedding(image)
-    reference_embeddings[f"{name}.jpg"] = emb
+    
+    # Use GPU for consistency if available
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    emb = get_embedding(image, device=device)
+    
+    # Store on CPU to save memory
+    reference_embeddings[f"{name}.jpg"] = emb.cpu()
+    
     with open(CACHE_FILE, "wb") as f:
         pickle.dump(reference_embeddings, f)
     return f"Image {name} added to dataset."
 
 search_interface = gr.Interface(fn=search_similar,
                                 inputs=gr.Image(type="pil", label="Query Image"),
-                                outputs=gr.Gallery(label="Top Matches"),
+                                outputs=gr.Gallery(label="Top Matches", columns=5),
                                 allow_flagging="never")
 
 add_interface = gr.Interface(fn=add_image,