app.py

# ============================================================
# BLIP Image Captioning — Spaces Minimal (CPU, Gradio v5)
# ------------------------------------------------------------
# • Designed ONLY for Hugging Face Spaces (lean, fast UI)
# • Loads model from the Hub (public by default, token optional)
# • Gallery-only output (no tables, no downloads, no file I/O)
# • CPU-friendly defaults; avoids torchvision by using slow processor
# ============================================================

import os
import re
from typing import Dict, List, Optional

import torch
from PIL import Image, ImageOps
import gradio as gr
from transformers import (
    BlipForConditionalGeneration,
    BlipProcessor,
    AutoTokenizer,
    GenerationConfig,
    __version__ as TF_VER,
)

# ---------------------------
# Environment / Config
# ---------------------------

os.environ["TOKENIZERS_PARALLELISM"] = "false"

# Point to your Hub model repo (public).
# You can also set this in the Space's "Variables" as MODEL_ID_OR_PATH.
MODEL_ID = os.getenv("MODEL_ID_OR_PATH", "YaekobB/blip-caption-model")

# Optional: set HF_TOKEN in Space "Secrets" if the model is private.
HF_TOKEN = os.getenv("HF_TOKEN", None)

device = torch.device("cpu")
# Keep threads modest on Spaces CPU machines
torch.set_num_threads(max(1, (os.cpu_count() or 2) - 1))

print("🔧 torch:", torch.__version__, "| transformers:", TF_VER)
print("🔄 Loading model from:", MODEL_ID)

# ---------------------------
# Load model + generation config
# ---------------------------

model = BlipForConditionalGeneration.from_pretrained(MODEL_ID, token=HF_TOKEN)
try:
    model.generation_config = GenerationConfig.from_pretrained(MODEL_ID, token=HF_TOKEN)
except Exception:
    pass

# Use slow processor (no torchvision dependency on Spaces)
# If your repo has preprocessor_config.json, this will align correctly.
try:
    processor = BlipProcessor.from_pretrained(MODEL_ID, use_fast=False, token=HF_TOKEN)
    print("ℹ️ Processor:", MODEL_ID)
except Exception as e:
    # Final fallback to base (still slow/CPU-friendly). Should rarely be needed.
    print("⚠️ Processor load from model repo failed; reason:", e)
    processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base", use_fast=False)
    print("ℹ️ Processor fallback:", "Salesforce/blip-image-captioning-base")

# Match your training/inference preprocessing (Kaggle style: 224 + center-crop)
processor.image_processor.size = {"height": 224, "width": 224}
if hasattr(processor.image_processor, "do_center_crop"):
    processor.image_processor.do_center_crop = True
if hasattr(processor.image_processor, "crop_size"):
    processor.image_processor.crop_size = {"height": 224, "width": 224}

# ---------------------------
# Tokenizer alignment with LM head
# ---------------------------

def _lm_head_rows(m: BlipForConditionalGeneration) -> int:
    try:
        return m.text_decoder.cls.predictions.decoder.weight.shape[0]
    except Exception:
        return m.get_input_embeddings().weight.shape[0]

lm_rows = _lm_head_rows(model)

def _load_tokenizer(model_head_rows: int):
    # Try model repo first (likely contains your fine-tuned tokenizer)
    try:
        tok = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
        print("ℹ️ Tokenizer:", MODEL_ID)
    except Exception as e:
        print("⚠️ Tokenizer load from model repo failed; reason:", e)
        tok = AutoTokenizer.from_pretrained("Salesforce/blip-image-captioning-base")

    extra_tokens: List[str] = []
    if len(tok) < model_head_rows:
        need = model_head_rows - len(tok)
        extra_tokens = [f"<extra_tok_{i}>" for i in range(need)]
        tok.add_tokens(extra_tokens)
        tok.add_special_tokens({"additional_special_tokens": extra_tokens})
        print(f"ℹ️ Added {len(extra_tokens)} dummy *special* tokens to match LM head.")
    return tok, extra_tokens

tokenizer, EXTRA_TOKENS = _load_tokenizer(lm_rows)

# Pad/eos & left padding (decoder-only-friendly)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left"
model.config.pad_token_id = tokenizer.pad_token_id
model.generation_config.pad_token_id = tokenizer.pad_token_id
if tokenizer.eos_token_id is not None:
    model.generation_config.eos_token_id = tokenizer.eos_token_id

processor.tokenizer = tokenizer

# If extra tokens were added, ban them during generation
BAD_WORDS_IDS: Optional[List[List[int]]] = None
if EXTRA_TOKENS:
    bad = [tokenizer.convert_tokens_to_ids(t) for t in EXTRA_TOKENS]
    BAD_WORDS_IDS = [[i] for i in bad if i is not None]

def _strip_extra_tokens(text: str) -> str:
    if not EXTRA_TOKENS:
        return text
    return re.sub(r"\s*<extra_tok_\d+>\s*", " ", text).strip()

model.to(device).eval()
print("✅ Ready. Device:", device)
print("📏 Inference image size:", getattr(processor.image_processor, "size", None))
print(f"🧪 vocab check -> lm_head rows: {lm_rows} | tokenizer size: {len(tokenizer)}")

# ---------------------------
# Decoding presets (CPU-friendly defaults)
# ---------------------------

BASE_ARGS = dict(
    min_length=5,
    no_repeat_ngram_size=2,
    early_stopping=True,
    do_sample=False,
)

PRESETS: Dict[str, Dict] = {
    "Fast (CPU)": dict(num_beams=1, max_length=22, length_penalty=1.0, **BASE_ARGS),
    "Balanced":   dict(num_beams=2, max_length=28, length_penalty=1.0, **BASE_ARGS),
    "Quality":    dict(num_beams=4, max_length=32, length_penalty=1.05, **BASE_ARGS),
}
DEFAULT_PRESET = "Fast (CPU)"   # Spaces default

# ---------------------------
# Tiny speed win: fast open
# ---------------------------

def _open_rgb_fast(path: str) -> Image.Image:
    img = Image.open(path)
    img = ImageOps.exif_transpose(img)          # fix orientation cheaply
    img.thumbnail((768, 768), Image.BILINEAR)    # downscale hint before processor
    return img.convert("RGB")

# ---------------------------
# Captioning
# ---------------------------

@torch.no_grad()
def caption_image(path: str, preset: str, beams: int, maxlen: int, lenpen: float, progress=gr.Progress(track_tqdm=True)) -> str:
    # Resolve preset defaults
    p = PRESETS[preset]
    beams  = int(beams or p["num_beams"])
    maxlen = int(maxlen or p["max_length"])
    lenpen = float(lenpen or p["length_penalty"])

    img = _open_rgb_fast(path)
    batch = processor(images=img, return_tensors="pt").to(device)

    gen_kwargs = dict(
        num_beams=beams,
        max_length=maxlen,
        length_penalty=lenpen,
        **BASE_ARGS,
    )
    if BAD_WORDS_IDS is not None:
        gen_kwargs["bad_words_ids"] = BAD_WORDS_IDS

    ids = model.generate(pixel_values=batch["pixel_values"], **gen_kwargs)
    text = tokenizer.decode(ids[0], skip_special_tokens=True)
    return _strip_extra_tokens(text)

def run_batch(paths: List[str], preset: str, beams: int, maxlen: int, lenpen: float, progress=gr.Progress(track_tqdm=True)):
    progress(0, desc="Preparing images")
    results = []
    gallery = []
    n = len(paths or [])
    for i, p in enumerate(paths or []):
        cap = caption_image(p, preset, beams, maxlen, lenpen)
        results.append((p, cap))
        gallery.append((p, cap))
        progress((i + 1) / max(1, n), desc=f"Captioned {i+1}/{n}")
    return gallery

# ---------------------------
# Gradio UI (minimal, gallery-only)
# ---------------------------

from gradio.themes.base import Base
from gradio.themes.utils import colors

THEME = Base(primary_hue=colors.green, secondary_hue=colors.gray)

CUSTOM_CSS = """
h1, h2, h3 { color: #228B22 !important; }  /* forest green headings */
#gallery .grid-wrap .label {
  background: rgba(255,255,255,0.92);
  border-radius: 10px; padding: 6px 10px;
  font-size: 0.95rem; line-height: 1.25rem;
}
"""

with gr.Blocks(title="Multimodal Image Captioning with BLIP", theme=THEME, css=CUSTOM_CSS, fill_height=True) as demo:
    gr.Markdown(
        """
# 🖼️ Multimodal Image Captioning with BLIP
Upload one or many images and get captions from a fine-tuned BLIP model.
        """.strip()
    )

    with gr.Row():
        with gr.Column(scale=6):
            uploader = gr.File(
                label="Upload image(s)",
                file_types=[".jpg", ".jpeg", ".png", ".bmp", ".webp"],
                file_count="multiple",
                type="filepath",
            )
            gr.Markdown("_Supported: JPG, PNG, WebP, BMP. Inference at 224×224 (center-crop)._")
        with gr.Column(scale=4):
            preset = gr.Radio(choices=list(PRESETS.keys()), value=DEFAULT_PRESET, label="Preset")
            with gr.Accordion("Advanced (optional)", open=False):
                beams  = gr.Slider(1, 6,  value=PRESETS[DEFAULT_PRESET]["num_beams"], step=1, label="num_beams")
                maxlen = gr.Slider(16, 48, value=PRESETS[DEFAULT_PRESET]["max_length"], step=1, label="max_length")
                lenpen = gr.Slider(0.8, 1.3, value=PRESETS[DEFAULT_PRESET]["length_penalty"], step=0.05, label="length_penalty")
            run = gr.Button("🚀 Generate Captions", variant="primary")

    gallery = gr.Gallery(label="Results (image + caption below)", elem_id="gallery", columns=2, height="auto")

    def _dispatch(files, preset, beams, maxlen, lenpen):
        return run_batch(files, preset, beams, maxlen, lenpen)

    run.click(_dispatch, inputs=[uploader, preset, beams, maxlen, lenpen], outputs=[gallery])

# Keep queue small for CPU Spaces
# new (Gradio 5+)
demo.queue(default_concurrency_limit=1, max_size=16).launch()

if __name__ == "__main__":
    demo.launch()