Model Deployed

app.py

@@ -1,31 +1,30 @@
 import gradio as gr
 from PIL import Image
 import torch
-from transformers import Blip2ForConditionalGeneration, AutoProcessor
+from transformers import BlipProcessor, BlipForConditionalGeneration
 
-# Load your fine-tuned model and processor from local directories
-processor = AutoProcessor.from_pretrained("./processor")
-model = Blip2ForConditionalGeneration.from_pretrained("./model")
+# Load model and processor
+processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
 
-# 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
+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)
     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
 
-# Gradio UI
+# Create Gradio interface
 iface = gr.Interface(
     fn=generate_caption,
     inputs=gr.Image(type="pil"),
     outputs="text",
-    title="🖼️ Image Captioning with Fine-Tuned BLIP2",
-    description="Upload an image to generate a caption using your custom fine-tuned BLIP2 model.",
+    title="Image Caption Generator",
+    description="Upload an image to generate a caption."
 )
 
-if __name__ == "__main__":
-    iface.launch()
+# Launch
+iface.launch()