Trying to faster model

app.py

@@ -93,6 +93,7 @@ from encoder import FashionCLIPEncoder
 REQUESTS_HEADERS = {
     'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
 }
+BATCH_SIZE = 30  # Define batch size for processing
 
 # Initialize encoder
 encoder = FashionCLIPEncoder()
@@ -116,8 +117,9 @@ def batch_process_images(image_urls: str):
     if not urls:
         return {"error": "No valid image URLs provided."}
 
-    embeddings = []
     results = []
+    batch_urls, batch_images = [], []
+
     for url in urls:
         try:
             # Download image
@@ -125,10 +127,32 @@ def batch_process_images(image_urls: str):
             if not image:
                 results.append({"image_url": url, "error": "Failed to download image"})
                 continue
-            
-            # Generate embedding
-            embedding = encoder.encode_images([image])[0]
-            
+
+            batch_urls.append(url)
+            batch_images.append(image)
+
+            # Process batch when reaching batch size
+            if len(batch_images) == BATCH_SIZE:
+                process_batch(batch_urls, batch_images, results)
+                batch_urls, batch_images = [], []
+
+        except Exception as e:
+            results.append({"image_url": url, "error": str(e)})
+
+    # Process remaining images in the last batch
+    if batch_images:
+        process_batch(batch_urls, batch_images, results)
+
+    return results
+
+
+# Helper function to process a batch
+def process_batch(batch_urls, batch_images, results):
+    try:
+        # Generate embeddings
+        embeddings = encoder.encode_images(batch_images)
+        
+        for url, embedding in zip(batch_urls, embeddings):
             # Normalize embedding
             embedding_normalized = embedding / np.linalg.norm(embedding)
             
@@ -138,9 +162,9 @@ def batch_process_images(image_urls: str):
                 "embedding_preview": embedding_normalized[:5].tolist(),  # First 5 values for preview
                 "success": True
             })
-        except Exception as e:
+    except Exception as e:
+        for url in batch_urls:
             results.append({"image_url": url, "error": str(e)})
-    return results
 
 
 # Gradio Interface

encoder.py

@@ -1,4 +1,5 @@
 from typing import List, Dict, Optional
+
 import torch
 from PIL.Image import Image
 from torch.utils.data import DataLoader
@@ -15,15 +16,14 @@ class FashionCLIPEncoder:
         )
         self.model = AutoModel.from_pretrained(MODEL_NAME, trust_remote_code=True)
 
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.model.to(self.device)
         self.model.eval()
+        self.device = self.model.device
 
     def encode_images(
         self, images: List[Image], batch_size: Optional[int] = None
     ) -> List[List[float]]:
         if batch_size is None:
-            batch_size = min(len(images), 32)  # Default to a safe batch size
+            batch_size = len(images)
 
         def transform_fn(el: Dict):
             return self.processor(
@@ -39,12 +39,9 @@ class FashionCLIPEncoder:
 
         with torch.no_grad():
             for batch in dataloader:
-                try:
-                    batch = {k: v.to(self.device) for k, v in batch.items()}
-                    embeddings = self._encode_images(batch)
-                    image_embeddings.extend(embeddings)
-                except Exception as e:
-                    print(f"Error encoding image batch: {e}")
+                batch = {k: v.to(self.device) for k, v in batch.items()}
+                embeddings = self._encode_images(batch)
+                image_embeddings.extend(embeddings)
 
         return image_embeddings
 
@@ -52,7 +49,7 @@ class FashionCLIPEncoder:
         self, text: List[str], batch_size: Optional[int] = None
     ) -> List[List[float]]:
         if batch_size is None:
-            batch_size = min(len(text), 32)  # Default to a safe batch size
+            batch_size = len(text)
 
         def transform_fn(el: Dict):
             kwargs = {
@@ -68,17 +65,13 @@ class FashionCLIPEncoder:
         )
         dataset.set_format("torch")
         dataloader = DataLoader(dataset, batch_size=batch_size)
-
         text_embeddings = []
 
         with torch.no_grad():
             for batch in dataloader:
-                try:
-                    batch = {k: v.to(self.device) for k, v in batch.items()}
-                    embeddings = self._encode_text(batch)
-                    text_embeddings.extend(embeddings)
-                except Exception as e:
-                    print(f"Error encoding text batch: {e}")
+                batch = {k: v.to(self.device) for k, v in batch.items()}
+                embeddings = self._encode_text(batch)
+                text_embeddings.extend(embeddings)
 
         return text_embeddings
 
@@ -86,4 +79,4 @@ class FashionCLIPEncoder:
         return self.model.get_image_features(**batch).detach().cpu().numpy().tolist()
 
     def _encode_text(self, batch: Dict) -> List:
-        return self.model.get_text_features(**batch).detach().cpu().numpy().tolist()
+        return self.model.get_text_features(**batch).detach().cpu().numpy().tolist()