Upload ai_service.py

ai_service.py

@@ -0,0 +1,111 @@
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+from typing import Optional
+from sentence_transformers import SentenceTransformer
+from transformers import CLIPProcessor, CLIPModel
+from PIL import Image
+import torch
+import io
+
+app = FastAPI(title="AI Embedding Service")
+
+class ModelLoader:
+    def __init__(self):
+        self._text_model = None
+        self._clip_model = None
+        self._clip_processor = None
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    @property
+    def text_model(self):
+        if self._text_model is None:
+            print("Loading text model (lazy initialization)...")
+            self._text_model = SentenceTransformer("BAAI/bge-large-en")
+        return self._text_model
+
+    @property
+    def clip_model(self):
+        if self._clip_model is None:
+            print("Loading image model (lazy initialization)...")
+            # Load in fp16 to save memory, especially for Hugging Face Spaces
+            self._clip_model = CLIPModel.from_pretrained(
+                "openai/clip-vit-large-patch14",
+                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
+            ).to(self.device)
+        return self._clip_model
+
+    @property
+    def clip_processor(self):
+        if self._clip_processor is None:
+            self._clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
+        return self._clip_processor
+
+models = ModelLoader()
+
+@app.get("/health")
+async def health():
+    return {
+        "status": "healthy",
+        "text_model_loaded": models._text_model is not None,
+        "image_model_loaded": models._clip_model is not None,
+        "device": models.device
+    }
+
+@app.post("/embed")
+async def embed(
+    property_name: Optional[str] = Form(None),
+    description: Optional[str] = Form(None),
+    images: Optional[list[UploadFile]] = File(None)
+):
+    response_data = {}
+    
+    # Process Property Name
+    if property_name and property_name.strip():
+        vec_name = models.text_model.encode(property_name, normalize_embeddings=True)
+        response_data["property_name_vector"] = vec_name.tolist()
+        
+    # Process Description
+    if description and description.strip():
+        vec_desc = models.text_model.encode(description, normalize_embeddings=True)
+        response_data["description_vector"] = vec_desc.tolist()
+        
+    # Process Multiple Images
+    if images:
+        image_vectors = []
+        for image in images:
+            if not image.filename:
+                continue
+            
+            contents = await image.read()
+            img = Image.open(io.BytesIO(contents)).convert("RGB")
+            
+            inputs = models.clip_processor(images=img, return_tensors="pt").to(models.device)
+            
+            with torch.no_grad():
+                outputs = models.clip_model.get_image_features(**inputs)
+            
+            # Extract tensor depending on transformers output format
+            if isinstance(outputs, torch.Tensor):
+                features = outputs
+            elif hasattr(outputs, "image_embeds") and outputs.image_embeds is not None:
+                features = outputs.image_embeds
+            elif hasattr(outputs, "pooler_output") and outputs.pooler_output is not None:
+                features = outputs.pooler_output
+            else:
+                features = outputs[0]
+            
+            # Apply L2 Normalization for Cosine Similarity
+            normalized_features = torch.nn.functional.normalize(features, p=2, dim=-1)
+            vec_img = normalized_features.squeeze().tolist()
+            image_vectors.append(vec_img)
+        
+        if image_vectors:
+            response_data["image_vectors"] = image_vectors
+            
+    if not response_data:
+        raise HTTPException(status_code=400, detail="Must provide at least one of property_name, description, or images")
+        
+    return response_data
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)