omarkashif commited on
Commit
00c6c1d
·
verified ·
1 Parent(s): 8c2a76f

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +27 -0
app.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import BlipProcessor, BlipForConditionalGeneration
3
+ from PIL import Image
4
+ import torch
5
+
6
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
7
+
8
+ # Load fine-tuned model from HF Hub
9
+ model = BlipForConditionalGeneration.from_pretrained("omarkashif/blip-finetuned-flickr30k").to(device)
10
+ processor = BlipProcessor.from_pretrained("omarkashif/blip-finetuned-flickr30k")
11
+
12
+ def generate_caption(image):
13
+ image = image.convert("RGB")
14
+ inputs = processor(images=image, return_tensors="pt").to(device)
15
+ with torch.no_grad():
16
+ output = model.generate(**inputs, max_length=40, num_beams=5)
17
+ caption = processor.decode(output[0], skip_special_tokens=True)
18
+ return caption
19
+
20
+ # Gradio UI
21
+ gr.Interface(
22
+ fn=generate_caption,
23
+ inputs=gr.Image(type="pil"),
24
+ outputs="text",
25
+ title="BLIP Fine-Tuned Caption Generator",
26
+ description="Upload an image to generate a caption using a BLIP model fine-tuned on Flickr30k."
27
+ ).launch()