| import os |
| import io |
| import threading |
| import torch |
|
|
| from PIL import Image |
| from fastapi import FastAPI, UploadFile, File, Form |
| from fastapi.responses import JSONResponse |
|
|
| from unsloth import FastVisionModel |
| from transformers import TextIteratorStreamer |
| from groq import Groq |
|
|
| MODEL_REPO = "Rady10/Plant-Disease-Qwen3VL-2B" |
|
|
| MAX_SEQ_LENGTH = 512 |
|
|
| VISION_MAX_NEW_TOKENS = 256 |
|
|
| SUMMARY_MAX_TOKENS = 220 |
|
|
| GROQ_API_KEY = os.environ["GROQ_API_KEY"] |
|
|
| GROQ_MODEL = "llama-3.3-70b-versatile" |
|
|
| print("Loading vision model...") |
|
|
| model, tokenizer = FastVisionModel.from_pretrained( |
| MODEL_REPO, |
| load_in_4bit=True, |
| max_seq_length=MAX_SEQ_LENGTH, |
| use_gradient_checkpointing="unsloth", |
| ) |
|
|
| FastVisionModel.for_inference(model) |
|
|
| print("Vision model loaded.") |
|
|
| groq_client = Groq(api_key=GROQ_API_KEY) |
|
|
| VISION_SYSTEM_PROMPT = """ |
| You are a professional plant disease expert. |
| |
| You identify diseases from plant images and answer agriculture questions. |
| """ |
|
|
| SUMMARY_SYSTEM_PROMPT = """ |
| You are an agriculture assistant. |
| |
| Rewrite the diagnosis in simple Egyptian Arabic. |
| |
| Rules: |
| - 3-6 short sentences. |
| - Mention disease. |
| - Mention likely cause. |
| - Mention treatment. |
| - Do not invent information. |
| """ |
|
|
| def run_vision_model(image, question): |
|
|
| messages = [ |
| { |
| "role": "system", |
| "content": [ |
| { |
| "type":"text", |
| "text":VISION_SYSTEM_PROMPT |
| } |
| ] |
| }, |
| { |
| "role":"user", |
| "content":[ |
| { |
| "type":"image", |
| "image":image |
| }, |
| { |
| "type":"text", |
| "text":question or "Diagnose this plant." |
| } |
| ] |
| } |
| ] |
|
|
| prompt = tokenizer.apply_chat_template( |
| messages, |
| add_generation_prompt=True |
| ) |
|
|
| inputs = tokenizer( |
| image, |
| prompt, |
| add_special_tokens=False, |
| return_tensors="pt" |
| ) |
|
|
| if torch.cuda.is_available(): |
| inputs = inputs.to("cuda") |
|
|
| streamer = TextIteratorStreamer( |
| tokenizer, |
| skip_prompt=True, |
| skip_special_tokens=True, |
| ) |
|
|
| kwargs = dict( |
| **inputs, |
| streamer=streamer, |
| max_new_tokens=VISION_MAX_NEW_TOKENS, |
| temperature=0.7, |
| top_p=0.9, |
| repetition_penalty=1.15, |
| ) |
|
|
| thread = threading.Thread( |
| target=model.generate, |
| kwargs=kwargs, |
| ) |
|
|
| thread.start() |
|
|
| text = "" |
|
|
| for piece in streamer: |
| text += piece |
|
|
| return text.strip() |
|
|
| def summarize_with_groq(raw, question): |
|
|
| response = groq_client.chat.completions.create( |
| model=GROQ_MODEL, |
| messages=[ |
| { |
| "role":"system", |
| "content":SUMMARY_SYSTEM_PROMPT |
| }, |
| { |
| "role":"user", |
| "content":f""" |
| Question: |
| {question} |
| |
| Diagnosis: |
| |
| {raw} |
| """ |
| } |
| ], |
| max_tokens=SUMMARY_MAX_TOKENS, |
| temperature=0.6, |
| ) |
|
|
| return response.choices[0].message.content |
|
|
| app = FastAPI(title="Plant Disease API") |
|
|
| @app.post("/analyze") |
| async def analyze( |
| image: UploadFile = File(...), |
| question: str = Form("") |
| ): |
|
|
| image_bytes = await image.read() |
|
|
| image = Image.open(io.BytesIO(image_bytes)).convert("RGB") |
|
|
| raw = run_vision_model(image, question) |
|
|
| summary = summarize_with_groq(raw, question) |
|
|
| return JSONResponse( |
| { |
| "question": question, |
| "summary": summary, |
| "technical_diagnosis": raw |
| } |
| ) |
