Create app.py

app.py

@@ -0,0 +1,159 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+import torch
+import open_clip
+from mobileclip.modules.common.mobileone import reparameterize_model
+from PIL import Image
+import requests
+from io import BytesIO
+import logging
+import numpy as np
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+app = FastAPI(
+    title="MobileCLIP API",
+    description="API for MobileCLIP image and text embeddings",
+    version="1.0.0"
+)
+
+# Global variables for model
+model = None
+preprocess = None
+tokenizer = None
+
+class TextRequest(BaseModel):
+    text: str
+
+class ImageRequest(BaseModel):
+    image_url: str
+
+class SimilarityRequest(BaseModel):
+    image_url: str
+    text: str
+
+class EmbeddingResponse(BaseModel):
+    embedding: list
+
+class SimilarityResponse(BaseModel):
+    similarity: float
+
+def load_model():
+    """Load and initialize the MobileCLIP model"""
+    global model, preprocess, tokenizer
+    
+    try:
+        logger.info("Loading MobileCLIP model...")
+        model, _, preprocess = open_clip.create_model_and_transforms('MobileCLIP-S2', pretrained='datacompdr')
+        tokenizer = open_clip.get_tokenizer('MobileCLIP-S2')
+        
+        # Reparameterize for inference
+        model.eval()
+        model = reparameterize_model(model)
+        logger.info("Model loaded successfully!")
+        
+    except Exception as e:
+        logger.error(f"Failed to load model: {str(e)}")
+        raise e
+
+def download_image(url: str) -> Image.Image:
+    """Download image from URL"""
+    try:
+        response = requests.get(url, timeout=10)
+        response.raise_for_status()
+        image = Image.open(BytesIO(response.content))
+        return image.convert('RGB')
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Failed to download image: {str(e)}")
+
+def get_image_embedding(image: Image.Image):
+    """Get embedding for an image"""
+    try:
+        image_tensor = preprocess(image).unsqueeze(0)
+        with torch.no_grad():
+            image_features = model.encode_image(image_tensor)
+            # Normalize the embedding
+            image_features = image_features / image_features.norm(dim=-1, keepdim=True)
+        return image_features.squeeze().cpu().numpy()
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to process image: {str(e)}")
+
+def get_text_embedding(text: str):
+    """Get embedding for text"""
+    try:
+        text_tokens = tokenizer([text])
+        with torch.no_grad():
+            text_features = model.encode_text(text_tokens)
+            # Normalize the embedding
+            text_features = text_features / text_features.norm(dim=-1, keepdim=True)
+        return text_features.squeeze().cpu().numpy()
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to process text: {str(e)}")
+
+def calculate_similarity(embedding1: np.ndarray, embedding2: np.ndarray) -> float:
+    """Calculate cosine similarity between two embeddings"""
+    return float(np.dot(embedding1, embedding2))
+
+@app.on_event("startup")
+async def startup_event():
+    """Load model on startup"""
+    load_model()
+
+@app.get("/")
+async def root():
+    """Health check endpoint"""
+    return {"message": "MobileCLIP API is running!", "status": "healthy"}
+
+@app.post("/image-embedding", response_model=EmbeddingResponse)
+async def image_embedding(request: ImageRequest):
+    """Get embedding for an image given its URL"""
+    if model is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+    
+    try:
+        image = download_image(request.image_url)
+        embedding = get_image_embedding(image)
+        return EmbeddingResponse(embedding=embedding.tolist())
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Error in image_embedding: {str(e)}")
+        raise HTTPException(status_code=500, detail="Internal server error")
+
+@app.post("/text-embedding", response_model=EmbeddingResponse)
+async def text_embedding(request: TextRequest):
+    """Get embedding for text"""
+    if model is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+    
+    try:
+        embedding = get_text_embedding(request.text)
+        return EmbeddingResponse(embedding=embedding.tolist())
+    except Exception as e:
+        logger.error(f"Error in text_embedding: {str(e)}")
+        raise HTTPException(status_code=500, detail="Internal server error")
+
+@app.post("/similarity", response_model=SimilarityResponse)
+async def similarity(request: SimilarityRequest):
+    """Calculate similarity between image and text"""
+    if model is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+    
+    try:
+        image = download_image(request.image_url)
+        image_embedding = get_image_embedding(image)
+        text_embedding = get_text_embedding(request.text)
+        
+        similarity_score = calculate_similarity(image_embedding, text_embedding)
+        return SimilarityResponse(similarity=similarity_score)
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Error in similarity: {str(e)}")
+        raise HTTPException(status_code=500, detail="Internal server error")
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)