updaed files

requirements.txt

@@ -4,4 +4,5 @@ torch==2.0.0
 transformers
 pydantic
 python-multipart
-pillow
+pillow
+requests

src/api.py

@@ -1,17 +1,27 @@
 from fastapi import FastAPI
 from .model import load_model
 from .schemas import Query
+import requests
+from PIL import Image
+from io import BytesIO
 
 app = FastAPI()
 processor, model = load_model()
 
 @app.post("/predict")
 async def predict(query: Query):
-    inputs = processor(images=query.image_url, text=query.question, return_tensors="pt")
-    outputs = model(**inputs)
-    predicted_answer = processor.decode(outputs.logits.argmax(-1)[0], skip_special_tokens=True)
+    # Download and process the image
+    response = requests.get(query.image_url)
+    image = Image.open(BytesIO(response.content))
+    
+    # Process the image and question
+    inputs = processor(images=image, text=query.question, return_tensors="pt")
+    outputs = model.generate(**inputs)
+    
+    # Get the predicted answer
+    predicted_answer = processor.decode(outputs[0], skip_special_tokens=True)
     return {"answer": predicted_answer}
 
 @app.get("/")
 async def root():
-    return {"message": "OmniParser API is running. Use /predict endpoint for predictions."}
+    return {"message": "Visual QA API is running. Use /predict endpoint for predictions."}

src/model.py

@@ -1,6 +1,6 @@
-from transformers import AutoProcessor, AutoModelForVisualQuestionAnswering
+from transformers import BlipProcessor, BlipForQuestionAnswering
 
 def load_model():
-    processor = AutoProcessor.from_pretrained("microsoft/OmniParser")
-    model = AutoModelForVisualQuestionAnswering.from_pretrained("microsoft/OmniParser")
+    processor = BlipProcessor.from_pretrained("Salesforce/blip-vqa-base")
+    model = BlipForQuestionAnswering.from_pretrained("Salesforce/blip-vqa-base")
     return processor, model