File size: 3,451 Bytes
3abbd00 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
import streamlit as st
import torch
from transformers import Blip2Processor, Blip2ForConditionalGeneration
from PIL import Image
# Set up the Streamlit page configuration
st.set_page_config(
page_title="BLIP-2 Image Captioning",
page_icon="📸",
layout="wide",
)
# --- Model Loading (using caching for efficiency) ---
# The @st.cache_resource decorator ensures the model and processor are loaded only once.
# This is crucial for a performant Streamlit app on Hugging Face Spaces.
@st.cache_resource
def load_model():
"""
Loads the BLIP-2 model and processor from Hugging Face Hub.
We're using a smaller version (`blip2-opt-2.7b`) that is more suitable for
Hugging Face's free tier, though it may still require significant resources.
We load the model in 8-bit to reduce memory usage.
"""
# Check if a CUDA-enabled GPU is available. If not, use CPU.
device = "cuda" if torch.cuda.is_available() else "cpu"
try:
# Load the processor and model
processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
# We load the model in 8-bit to save memory, which is important for
# deployment on platforms like Hugging Face Spaces.
model = Blip2ForConditionalGeneration.from_pretrained(
"Salesforce/blip2-opt-2.7b",
device_map="auto",
load_in_8bit=True,
torch_dtype=torch.float16
)
return processor, model, device
except Exception as e:
st.error(f"Error loading the model: {e}")
st.info("The model is very large and may require a GPU with at least 15GB of VRAM. "
"If you're seeing this error, the free tier of Hugging Face Spaces might not be enough.")
return None, None, None
# --- Main App Interface ---
st.title("📸 BLIP-2 Image Captioning AI")
st.write(
"Upload an image, and this application will generate a descriptive caption using the powerful "
"[BLIP-2 model](https://huggingface.co/Salesforce/blip2-opt-2.7b) from Hugging Face."
)
# Load the model and processor
processor, model, device = load_model()
if model and processor:
# Create a file uploader widget
uploaded_file = st.file_uploader(
"Choose an image...",
type=["jpg", "jpeg", "png", "bmp"],
help="Upload an image file to get a caption."
)
if uploaded_file is not None:
# Display the uploaded image
st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
st.write("")
st.info("Generating caption...")
try:
# Open the uploaded image file as a PIL Image
raw_image = Image.open(uploaded_file).convert("RGB")
# Preprocess the image and generate the caption
inputs = processor(images=raw_image, return_tensors="pt").to(device, torch.float16)
out = model.generate(**inputs, max_new_tokens=50) # Increased max_new_tokens for longer captions
# Decode the generated tokens to text
caption = processor.decode(out[0], skip_special_tokens=True).strip()
# Display the generated caption
st.success(f"**Caption:** {caption}")
except Exception as e:
st.error(f"An error occurred during caption generation: {e}")
else:
st.warning("The application could not be initialized. Please check the logs for details.")
|