Update app.py

app.py

@@ -3,16 +3,23 @@ from PIL import Image
 import requests
 from io import BytesIO
 import re
-
 import torch
-from transformers import BlipForConditionalGeneration, BlipProcessor
+from transformers import BlipForConditionalGeneration, BlipProcessor, T5ForConditionalGeneration, T5Tokenizer
+
+# Device: CPU-only
+device = torch.device("cpu")
 
-# Load model (free, publicly available) once at startup
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# Load caption model (small, CPU-friendly)
+# HF model: Salesforce/blip-image-captioning-base
 processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)
+model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+model.to(device)
+
+# Lightweight rewriter model (T5-small) to improve fluency/detail without big cost
+rewriter_tokenizer = T5Tokenizer.from_pretrained("t5-small")
+rewriter = T5ForConditionalGeneration.from_pretrained("t5-small").to(device)
 
-# Simple content-safety checker (avoid explicit sexual/graphic violent text)
+# Simple safety patterns (non-bypass)
 SENSITIVE_PATTERNS = [
     r"\b(nude|naked|porn|sex|sexual|explicit|hardcore)\b",
     r"\b(blood|gore|mutilat|disembowel|organs)\b",
@@ -33,29 +40,48 @@ def is_caption_allowed(text: str):
         return False
     return SENSITIVE_RE.search(text) is None
 
-def describe_image(url: str, max_length: int = 32):
+def generate_caption(img: Image.Image, max_len:int=30):
+    # prepare inputs
+    inputs = processor(images=img, return_tensors="pt").to(device)
+    # generate with small beams to balance speed/quality
+    out = model.generate(**inputs, max_length=max_len, num_beams=3, early_stopping=True)
+    caption = processor.decode(out[0], skip_special_tokens=True).strip()
+    return caption
+
+def rewrite_caption(caption: str, max_len:int=64):
+    # Use T5-small to expand/clean the caption. Prefix "paraphrase:" helps instruct T5.
+    input_text = "paraphrase: " + caption
+    tok = rewriter_tokenizer(input_text, return_tensors="pt", truncation=True).to(device)
+    out = rewriter.generate(**tok, max_length=max_len, num_beams=2, early_stopping=True)
+    rewritten = rewriter_tokenizer.decode(out[0], skip_special_tokens=True).strip()
+    return rewritten
+
+def describe_image(url: str, max_caption_len: int = 30, expand: bool = True):
     img, err = load_image_from_url(url)
     if err:
         return None, err
-    # prepare and generate
-    inputs = processor(images=img, return_tensors="pt").to(device)
-    out = model.generate(**inputs, max_length=max_length, num_beams=5, early_stopping=True)
-    caption = processor.decode(out[0], skip_special_tokens=True).strip()
-    # apply simple safety filter
+    caption = generate_caption(img, max_len=max_caption_len)
     if not is_caption_allowed(caption):
-        # return a sanitized, non-graphic alternative
-        return img, "Caption omitted for safety. The image may contain explicit content."
+        return img, "Caption omitted for safety."
+    if expand:
+        try:
+            caption = rewrite_caption(caption, max_len=64)
+            if not is_caption_allowed(caption):
+                return img, "Caption omitted for safety."
+        except Exception:
+            pass
     return img, caption
 
 with gr.Blocks() as demo:
-    gr.Markdown("## Image URL → Display + Caption (open-source, free model)")
+    gr.Markdown("## CPU-optimized Image Captioning (BLIP-base + T5-small rewriter)")
     with gr.Row():
         url_in = gr.Textbox(label="Image URL", placeholder="https://example.com/photo.jpg")
-        max_len = gr.Slider(minimum=10, maximum=80, value=32, label="Max caption length")
+        max_len = gr.Slider(minimum=10, maximum=60, value=30, label="Max caption token length")
+        expand_chk = gr.Checkbox(label="Expand/rewite caption (slower, more natural)", value=True)
     img_out = gr.Image(type="pil", label="Image")
     caption_out = gr.Textbox(label="Caption")
     go = gr.Button("Load & Describe")
-    go.click(fn=describe_image, inputs=[url_in, max_len], outputs=[img_out, caption_out])
+    go.click(fn=describe_image, inputs=[url_in, max_len, expand_chk], outputs=[img_out, caption_out])
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)