Update app.py

app.py

@@ -2,25 +2,28 @@ import gradio as gr
 from PIL import Image
 import requests
 from io import BytesIO
-import re
 import torch
-from transformers import BlipForConditionalGeneration, BlipProcessor, T5ForConditionalGeneration, T5Tokenizer
+from transformers import (
+    VisionEncoderDecoderModel,
+    ViTImageProcessor,
+    AutoTokenizer,
+    T5ForConditionalGeneration,
+    T5Tokenizer,
+)
+import re
 
 device = torch.device("cpu")
 
-# Load models (CPU-friendly)
-processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)
+# Image captioning model (nlpconnect/vit-gpt2-image-captioning)
+processor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning").to(device)
+model.eval()
 
+# Rewriter (T5)
 rewriter_tokenizer = T5Tokenizer.from_pretrained("t5-small")
 rewriter = T5ForConditionalGeneration.from_pretrained("t5-small").to(device)
-
-# Safety patterns (simple filter)
-SENSITIVE_PATTERNS = [
-    r"\b(nude|naked|porn|sex|sexual|explicit|hardcore)\b",
-    r"\b(blood|gore|mutilat|disembowel|organs)\b",
-]
-SENSITIVE_RE = re.compile("|".join(SENSITIVE_PATTERNS), flags=re.IGNORECASE)
+rewriter.eval()
 
 def load_image_from_url(url: str, timeout=10):
     try:
@@ -31,18 +34,14 @@ def load_image_from_url(url: str, timeout=10):
     except Exception as e:
         return None, f"Error loading image: {e}"
 
-def is_caption_allowed(text: str):
-    if not text:
-        return False
-    return SENSITIVE_RE.search(text) is None
-
-def generate_caption(img: Image.Image, max_len:int=30):
-    inputs = processor(images=img, return_tensors="pt").to(device)
-    out = model.generate(**inputs, max_length=max_len, num_beams=3, early_stopping=True)
-    caption = processor.decode(out[0], skip_special_tokens=True).strip()
+def generate_caption(img: Image.Image, max_len: int = 30):
+    inputs = processor(images=img, return_tensors="pt")
+    pixel_values = inputs.pixel_values.to(device)
+    out = model.generate(pixel_values, max_length=max_len, num_beams=2, early_stopping=True)
+    caption = tokenizer.decode(out[0], skip_special_tokens=True).strip()
     return caption
 
-def rewrite_caption(caption: str, max_len:int=64):
+def rewrite_caption(caption: str, max_len: int = 64):
     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)
@@ -54,29 +53,18 @@ def describe_image(url: str, max_caption_len: int = 30, expand: bool = True):
     if err:
         return None, f"Error: {err}"
     caption = generate_caption(img, max_len=max_caption_len)
-    if not is_caption_allowed(caption):
-        # Provide a neutral, respectful safety message with next steps
-        safety_msg = ("A descriptive caption was not provided because the image may contain explicit or graphic content. "
-                      "If this is unexpected, try a different image or upload a cropped/edited version that removes sensitive content.")
-        return img, safety_msg
     if expand:
         try:
             caption = rewrite_caption(caption, max_len=64)
-            if not is_caption_allowed(caption):
-                safety_msg = ("A descriptive caption was not provided because the generated text may describe explicit or graphic content. "
-                              "Try a different image or disable the expansion option.")
-                return img, safety_msg
         except Exception:
             pass
-    # Make caption more descriptive by appending structural cues (objects, colors, setting)
-    # Quick heuristic: if short, expand with a simple template
     if len(caption.split()) < 6:
         caption = f"{caption}. The scene appears to contain: {caption.lower()}."
     return img, caption
 
 # Gradio UI: image left, caption right
 with gr.Blocks() as demo:
-    gr.Markdown("## Image captioning — image on the left, descriptive caption on the right (CPU-optimized)")
+    gr.Markdown("## Image captioning — image on the left, descriptive caption on the right (CPU-optimized, uncensored)")
     with gr.Row():
         with gr.Column(scale=1):
             url_in = gr.Textbox(label="Image URL", placeholder="https://example.com/photo.jpg")