Create app.py

app.py

@@ -0,0 +1,120 @@
+# 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.")