updaded files

src/api.py

@@ -0,0 +1,17 @@
+from fastapi import FastAPI
+from .model import load_model
+from .schemas import Query
+
+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)
+    return {"answer": predicted_answer}
+
+@app.get("/")
+async def root():
+    return {"message": "OmniParser API is running. Use /predict endpoint for predictions."}

src/model.py

@@ -0,0 +1,6 @@
+from transformers import AutoProcessor, AutoModelForVisualQuestionAnswering
+
+def load_model():
+    processor = AutoProcessor.from_pretrained("microsoft/OmniParser")
+    model = AutoModelForVisualQuestionAnswering.from_pretrained("microsoft/OmniParser")
+    return processor, model

src/schemas.py

@@ -0,0 +1,5 @@
+from pydantic import BaseModel
+
+class Query(BaseModel):
+    image_url: str
+    question: str