Updated FastAPI app and requirements

app.py

@@ -7,7 +7,6 @@ import base64
 import os
 import logging
 from huggingface_hub import login
-from torchvision import transforms
 
 # Enable logging
 logging.basicConfig(level=logging.INFO)
@@ -22,16 +21,22 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 os.environ["HF_HOME"] = "/tmp/huggingface"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
 
-# Login to Hugging Face
+# Ensure Hugging Face Token is Set
 HF_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
-if HF_TOKEN:
+if not HF_TOKEN:
+    logger.error("Hugging Face token not found! Set HUGGINGFACE_TOKEN in the environment.")
+    raise RuntimeError("Hugging Face token missing. Set it in your environment.")
+
+# Login to Hugging Face
+try:
     login(HF_TOKEN)
-else:
-    logger.warning("Hugging Face token not found. Set HUGGINGFACE_TOKEN.")
+except Exception as e:
+    logger.error(f"Failed to authenticate Hugging Face token: {e}")
+    raise RuntimeError("Authentication with Hugging Face failed.")
 
 # Configure Quantization
 quantization_config = BitsAndBytesConfig(
-    load_in_4bit=True,
+    load_in_4bit=True,  # Change to load_in_8bit=True if 4-bit fails
     bnb_4bit_compute_dtype=torch.float16,
     bnb_4bit_use_double_quant=True,
 )
@@ -57,13 +62,6 @@ except Exception as e:
 # Allowed Formats
 ALLOWED_FORMATS = {"jpeg", "jpg", "png", "bmp", "tiff"}
 
-def preprocess_image(image: Image.Image):
-    transform = transforms.Compose([
-        transforms.Resize((512, 512)),
-        transforms.ToTensor(),
-    ])
-    return transform(image).unsqueeze(0)
-
 @app.post("/predict/")
 async def predict(file: UploadFile = File(...)):
     try:
@@ -74,7 +72,6 @@ async def predict(file: UploadFile = File(...)):
         # Read Image
         image_bytes = await file.read()
         image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-        image_tensor = preprocess_image(image)
 
         # Convert Image to Base64
         buffered = io.BytesIO()
@@ -84,7 +81,7 @@ async def predict(file: UploadFile = File(...)):
         # Validation Step
         validation_prompt = "Is this a medical X-ray or CT scan? Answer only 'yes' or 'no'."
         validation_inputs = processor(
-            text=validation_prompt, images=image_tensor, return_tensors="pt"
+            text=validation_prompt, images=image, return_tensors="pt"
         ).to(DEVICE)
 
         with torch.no_grad():
@@ -109,7 +106,7 @@ async def predict(file: UploadFile = File(...)):
         Recommendations:
         • [Follow-up Actions]
         """
-        analysis_inputs = processor(text=analysis_prompt, images=image_tensor, return_tensors="pt").to(DEVICE)
+        analysis_inputs = processor(text=analysis_prompt, images=image, return_tensors="pt").to(DEVICE)
 
         with torch.no_grad():
             analysis_output = model.generate(