File size: 4,938 Bytes
508ae61 |
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 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
# app.py
"""
Streamlit BLIP-2 Image Captioning demo
- Uses HuggingFace transformers' Blip2Processor + Blip2ForConditionalGeneration
- Caches the model & processor with st.cache_resource so they load once per Space/session.
- Designed for deployment on Hugging Face Spaces (use Docker SDK / Streamlit template).
"""
import streamlit as st
from PIL import Image
import io
import torch
from transformers import Blip2Processor, Blip2ForConditionalGeneration
st.set_page_config(
page_title="BLIP-2 Image Captioning",
layout="wide",
initial_sidebar_state="expanded",
)
# --- Sidebar / Info ---
st.sidebar.title("BLIP-2 Caption Demo")
st.sidebar.markdown(
"""
Upload an image and BLIP-2 will generate a caption.
- Model choices: choose a BLIP-2 model (large models may need GPU / won’t fit on CPU).
- For Spaces deployment, prefer smaller/flan-xl variants or use inference API.
"""
)
# Recommended default model (change if you want)
DEFAULT_MODEL = "Salesforce/blip2-opt-2.7b"
@st.cache_resource(show_spinner=False)
def load_model_and_processor(model_name: str):
"""Load and cache the BLIP-2 processor and model."""
# Note: large models will require a GPU; smaller variants or hosted inference endpoints recommended for CPU-only Spaces.
processor = Blip2Processor.from_pretrained(model_name)
model = Blip2ForConditionalGeneration.from_pretrained(model_name)
# move to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
return processor, model, device
def generate_caption(processor, model, device, pil_image: Image.Image, max_new_tokens=50, num_beams=4):
"""Generate caption text for a PIL image using BLIP-2."""
if pil_image.mode != "RGB":
pil_image = pil_image.convert("RGB")
inputs = processor(images=pil_image, return_tensors="pt").to(device)
# Generate - tune generation args as needed
generated_ids = model.generate(**inputs, max_new_tokens=max_new_tokens, num_beams=num_beams)
# decode using the tokenizer in the processor
caption = processor.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
return caption
# --- UI layout ---
col1, col2 = st.columns([1, 1.2])
with col1:
st.header("Upload image")
uploaded = st.file_uploader("Choose an image", type=["png", "jpg", "jpeg"], accept_multiple_files=False)
st.markdown("**Model selection**")
model_name = st.selectbox(
"Pick BLIP-2 model (large models may not run on CPU)",
options=[
"Salesforce/blip2-flan-t5-xl",
"Salesforce/blip2-opt-2.7b",
"Salesforce/blip2-flan-t5-xxl",
],
index=1 if DEFAULT_MODEL.endswith("2.7b") else 0,
help="Large models require GPU or HF Inference API; choose smaller if you have no GPU.",
)
max_tokens = st.slider("Max caption length (tokens)", min_value=10, max_value=200, value=50)
num_beams = st.slider("Beam search width (num_beams)", min_value=1, max_value=8, value=4)
st.write("---")
st.markdown("Tips:")
st.markdown(
"- If deploying on CPU-only Spaces, use a smaller/flan model or use the Hugging Face Inference API.\n"
"- Model loading is cached to speed up subsequent requests."
)
with col2:
st.header("Preview & Caption")
if uploaded is None:
st.info("Upload an image on the left to generate a caption.")
st.empty()
else:
# display image
image_bytes = uploaded.read()
pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
st.image(pil_image, use_column_width=True)
# Load model & processor (cached)
with st.spinner("Loading model (cached after first load)..."):
processor, model, device = load_model_and_processor(model_name)
# Generate caption
if st.button("Generate caption"):
with st.spinner("Generating caption..."):
try:
caption = generate_caption(processor, model, device, pil_image, max_new_tokens=max_tokens, num_beams=num_beams)
st.success("Caption generated")
st.markdown(f"**Caption:** {caption}")
# Provide a copy button and simple download
st.download_button("Download caption (.txt)", caption, file_name="caption.txt")
except Exception as e:
st.error(f"Error during generation: {e}")
st.info("If model is too large or out-of-memory, try a smaller model or use GPU.")
# --- Footer / Resources ---
st.markdown("---")
st.markdown(
"Built with BLIP-2 + Transformers. For production or public Spaces hosting, consider using Hugging Face Inference API or a smaller model variant to avoid OOM on CPU-only hosts."
)
st.caption("Docs: BLIP-2 (Transformers), Hugging Face Spaces (Streamlit), Streamlit caching & uploader.")
|