Update app.py

app.py

@@ -1,30 +1,43 @@
 import gradio as gr
-from transformers import BlipProcessor, BlipForConditionalGeneration
 from PIL import Image
-import requests
+import torch
+from torchvision import transforms
+from transformers import BlipForConditionalGeneration, AutoTokenizer
 
-# Load model and processor
+# Load model and tokenizer
 model_name = "hackergeek/radiology-image-captioning"
-processor = BlipProcessor.from_pretrained(model_name)
 model = BlipForConditionalGeneration.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# Manual preprocessing
+preprocess = transforms.Compose([
+    transforms.Resize((384, 384)),  # BLIP models usually expect 384x384
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225]),
+])
 
 def generate_caption(image):
     """
-    Generates a radiology caption for a given image
+    Generate radiology caption for a PIL image.
     """
-    if isinstance(image, str):  # if image is a URL
-        image = Image.open(requests.get(image, stream=True).raw).convert("RGB")
-    elif isinstance(image, Image.Image):
+    if not isinstance(image, Image.Image):
+        image = Image.open(image).convert("RGB")
+    else:
         image = image.convert("RGB")
-
-    inputs = processor(images=image, return_tensors="pt")
-    out = model.generate(**inputs)
-    caption = processor.decode(out[0], skip_special_tokens=True)
+    
+    pixel_values = preprocess(image).unsqueeze(0)  # add batch dimension
+    with torch.no_grad():
+        outputs = model.generate(pixel_values=pixel_values)
+    caption = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return caption
 
-# Create Gradio interface
+# Gradio Interface
 title = "Radiology Image Captioning"
-description = "Upload a radiology image (X-ray, CT, MRI) and get an automatic caption generated by the `hackergeek/radiology-image-captioning` model."
+description = (
+    "Upload a radiology image (X-ray, CT, MRI) and get an automatic caption "
+    "generated by the `hackergeek/radiology-image-captioning` model."
+)
 
 iface = gr.Interface(
     fn=generate_caption,