Update app.py

app.py

@@ -1,78 +1,150 @@
+import os
 import base64
 import torch
+import numpy as np
+import faiss
+import json
 
 from fastapi import FastAPI
 from pydantic import BaseModel
+from contextlib import asynccontextmanager
+from huggingface_hub import snapshot_download
+from sentence_transformers import SentenceTransformer
 from PIL import Image
 from io import BytesIO
 
-from transformers import AutoImageProcessor, AutoModelForImageClassification
+from transformers import AutoProcessor, AutoModelForVision2Seq
 
 # ─────────────────────────────
-# MODEL (REPLACE WITH YOUR CLASSIFIER)
+# CONFIG
 # ─────────────────────────────
-MODEL_REPO = "Rady10/Plant-Disease-Qwen3VL-2B"  # or your plant model
+MODEL_REPO = "Rady10/Plant-Disease-Qwen3VL-2B"
+RAG_REPO = "Rady10/Plant-Disease-AgroRAG-Index"
 
-app = FastAPI(title="🌿 Plant Disease Classifier")
+DEVICE = "cpu"
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 # ─────────────────────────────
-# LOAD MODEL
+# GLOBALS
 # ─────────────────────────────
-print("Loading classifier...")
-
-processor = AutoImageProcessor.from_pretrained(MODEL_REPO)
-model = AutoModelForImageClassification.from_pretrained(MODEL_REPO)
+model = None
+processor = None
 
-model.eval()
-
-print("Model loaded")
+faiss_index = None
+rag_chunks = None
+embedder = None
 
 # ─────────────────────────────
-# REQUEST FORMAT
+# FASTAPI APP
 # ─────────────────────────────
-class Request(BaseModel):
-    image: str  # base64 string
+app = FastAPI(title="🌿 Plant Disease Vision API")
 
 # ─────────────────────────────
-# DECODE IMAGE
+# LOAD MODELS ONCE
 # ─────────────────────────────
-def decode_image(b64):
-    return Image.open(
-        BytesIO(base64.b64decode(b64))
-    ).convert("RGB")
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+
+    global model, processor, faiss_index, rag_chunks, embedder
+
+    print("Loading vision model...")
+
+    processor = AutoProcessor.from_pretrained(
+        MODEL_REPO,
+        trust_remote_code=True
+    )
+
+    model = AutoModelForVision2Seq.from_pretrained(
+        MODEL_REPO,
+        torch_dtype=torch.float32,
+        device_map="cpu",
+        trust_remote_code=True
+    )
+
+    model.eval()
+
+    # ───── RAG (optional but included) ─────
+    print("Loading RAG...")
+
+    rag_dir = snapshot_download(
+        repo_id=RAG_REPO,
+        repo_type="dataset",
+        local_dir="./rag"
+    )
+
+    faiss_index = faiss.read_index(
+        os.path.join(rag_dir, "agro.index")
+    )
+
+    with open(os.path.join(rag_dir, "chunks.json"), "r", encoding="utf-8") as f:
+        rag_chunks = json.load(f)
+
+    embedder = SentenceTransformer(
+        "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
+    )
+
+    print("ALL LOADED")
+
+    yield
+
+app = FastAPI(lifespan=lifespan)
 
 # ─────────────────────────────
-# PREDICTION FUNCTION
+# REQUEST MODEL
 # ─────────────────────────────
-def predict(image):
+class VisionRequest(BaseModel):
+    image: str   # base64
+    text: str = ""
 
-    inputs = processor(images=image, return_tensors="pt")
+# ─────────────────────────────
+# IMAGE DECODER
+# ─────────────────────────────
+def decode_image(base64_str):
+    img_data = base64.b64decode(base64_str)
+    return Image.open(BytesIO(img_data)).convert("RGB")
 
-    with torch.no_grad():
-        outputs = model(**inputs)
+# ─────────────────────────────
+# GENERATION
+# ─────────────────────────────
+def generate(image, text):
 
-    logits = outputs.logits
+    if text.strip() == "":
+        text = "What disease is shown in this plant image?"
 
-    pred_id = logits.argmax(-1).item()
+    inputs = processor(
+        text=text,
+        images=image,
+        return_tensors="pt"
+    )
 
-    label = model.config.id2label[pred_id]
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
 
-    return label
+    with torch.no_grad():
+        output = model.generate(
+            **inputs,
+            max_new_tokens=256,
+            temperature=0.7,
+            top_p=0.9
+        )
+
+    return processor.batch_decode(
+        output,
+        skip_special_tokens=True
+    )[0]
 
 # ─────────────────────────────
-# API ENDPOINT
+# API ROUTES
 # ──────��──────────────────────
-@app.post("/predict")
-def predict_api(req: Request):
+@app.get("/")
+def root():
+    return {"status": "vision api running"}
 
-    image = decode_image(req.image)
+@app.post("/analyze")
+def analyze(req: VisionRequest):
 
-    label = predict(image)
+    image = decode_image(req.image)
 
-    return {
-        "prediction": label
-    }
+    result = generate(image, req.text)
 
-@app.get("/")
-def home():
-    return {"status": "classifier running"}
+    return {"response": result}