# ============================================================ # 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"" 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*\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()