Update app.py

app.py

@@ -1,12 +1,12 @@
-from fastapi import FastAPI, UploadFile, File
+from fastapi import FastAPI, UploadFile, File, HTTPException
 from transformers import CLIPProcessor, CLIPModel
-from PIL import Image
+from PIL import Image, UnidentifiedImageError
 import torch
 import io
 
 app = FastAPI()
 
-# Load a compatible CLIP model and processor
+# Load the CLIP model and processor
 model = CLIPModel.from_pretrained("laion/CLIP-ViT-B-32-laion2B-s34B-b79K")
 processor = CLIPProcessor.from_pretrained("laion/CLIP-ViT-B-32-laion2B-s34B-b79K")
 
@@ -16,17 +16,21 @@ def home():
 
 @app.post("/search/")
 async def search_image(file: UploadFile = File(...), query: str = "a photo"):
-    # Load the uploaded image
-    contents = await file.read()
-    image = Image.open(io.BytesIO(contents)).convert("RGB")
+    try:
+        # Read and decode image
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents)).convert("RGB")
+    except UnidentifiedImageError:
+        raise HTTPException(status_code=400, detail="Invalid image file format.")
 
     # Preprocess image and text
     inputs = processor(text=[query], images=image, return_tensors="pt", padding=True)
 
-    # Forward pass
-    outputs = model(**inputs)
-    logits_per_image = outputs.logits_per_image
-    probs = logits_per_image.softmax(dim=1)
+    # Forward pass through the model
+    with torch.no_grad():
+        outputs = model(**inputs)
+        logits_per_image = outputs.logits_per_image
+        probs = logits_per_image.softmax(dim=1)
 
     return {
         "query": query,