Model Deployed2

app.py

@@ -1,30 +1,31 @@
 import gradio as gr
 from PIL import Image
 import torch
-from transformers import BlipProcessor, BlipForConditionalGeneration
+from transformers import Blip2ForConditionalGeneration, AutoProcessor
 
-# Load model and processor
-processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+# Load your fine-tuned model and processor from local directories
+processor = AutoProcessor.from_pretrained("./processor")
+model = Blip2ForConditionalGeneration.from_pretrained("./model", device_map="auto", torch_dtype=torch.float16)
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model = model.to(device)
-
-# Define the function to generate caption
-def generate_caption(image):
-    inputs = processor(images=image, return_tensors="pt").to(device)
+# Inference function
+def generate_caption(image: Image.Image) -> str:
+    # Convert image to RGB and process
+    image = image.convert("RGB")
+    inputs = processor(images=image, return_tensors="pt").to(model.device, torch.float16)
+    
+    # Generate caption
     generated_ids = model.generate(pixel_values=inputs.pixel_values, max_length=25)
     caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
     return caption
 
-# Create Gradio interface
+# Gradio UI
 iface = gr.Interface(
     fn=generate_caption,
     inputs=gr.Image(type="pil"),
     outputs="text",
-    title="Image Caption Generator",
-    description="Upload an image to generate a caption."
+    title="🖼️ Image Captioning with Fine-Tuned BLIP2",
+    description="Upload an image to generate a caption using your custom fine-tuned BLIP2 model.",
 )
 
-# Launch
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()