Updated FastAPI app and requirements

app.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI, File, UploadFile
 import torch
-from transformers import AutoProcessor, LlavaForConditionalGeneration
+from transformers import AutoProcessor, AutoModelForImageTextToText, BitsAndBytesConfig
 from PIL import Image
 import io
 import base64
@@ -8,17 +8,27 @@ import base64
 # Initialize FastAPI app
 app = FastAPI()
 
-# Load the model and processor from Hugging Face
+# Model details
 model_name = "mervinpraison/Llama-3.2-11B-Vision-Radiology-mini"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# Configure 4-bit quantization using bitsandbytes
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,   # Use 4-bit quantization (set False for 8-bit)
+    bnb_4bit_compute_dtype=torch.float16,  # Use float16 for faster computation
+    bnb_4bit_use_double_quant=True,  # Double quantization for memory efficiency
+)
+
+# Load processor and quantized model
 processor = AutoProcessor.from_pretrained(model_name)
-model = LlavaForConditionalGeneration.from_pretrained(model_name).to(device)
+model = AutoModelForImageTextToText.from_pretrained(
+    model_name, quantization_config=quantization_config
+).to(device)
 
 @app.post("/predict/")
 async def predict(file: UploadFile = File(...)):
     try:
-        # Read image
+        # Read and process the image
         image_bytes = await file.read()
         image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
 

requirement.txt

@@ -1,6 +0,0 @@
-fastapi
-uvicorn
-torch
-transformers
-pillow
-python-multipart

requirements.txt

@@ -4,3 +4,4 @@ torch
 transformers
 pillow
 python-multipart
+bitsandbytes