Update app.py

app.py

@@ -67,37 +67,74 @@
 # # demo.launch(share=True)
 
 import gradio as gr
-from transformers import BlipProcessor, BlipForConditionalGeneration, pipeline
+from transformers import (
+    BlipProcessor, 
+    BlipForConditionalGeneration, 
+    BlipForQuestionAnswering, 
+    pipeline
+)
 from PIL import Image
 import torch
 from gtts import gTTS
 import tempfile
-import os
 
 # ----------------------
-# Load BLIP (Large) for Captioning
+# Device setup
 # ----------------------
-caption_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
-caption_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # ----------------------
-# Translation pipelines
+# Load Models Once
 # ----------------------
+print("🔄 Loading models...")
+
+# Captioning
+caption_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
+caption_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large").to(device)
+
+# VQA
+vqa_processor = BlipProcessor.from_pretrained("Salesforce/blip-vqa-base")
+vqa_model = BlipForQuestionAnswering.from_pretrained("Salesforce/blip-vqa-base").to(device)
+
+# Translation
 translation_models = {
     "Hindi": pipeline("translation", model="Helsinki-NLP/opus-mt-en-hi"),
     "French": pipeline("translation", model="Helsinki-NLP/opus-mt-en-fr"),
     "Spanish": pipeline("translation", model="Helsinki-NLP/opus-mt-en-es"),
 }
 
+# Safety Moderation Pipeline
+moderation_model = pipeline("text-classification", model="unitary/toxic-bert")
+
+print("✅ All models loaded!")
+
+# ----------------------
+# Safety Filter Function
+# ----------------------
+def is_caption_safe(caption):
+    result = moderation_model(caption)[0]
+    label = result["label"]
+    score = result["score"]
+
+    # toxic-bert gives "toxic" or "non-toxic"
+    if label.lower() == "toxic" and score > 0.7:
+        return False
+    return True
+
 # ----------------------
-# Caption + Translate + Speak Function
+# Caption + Translate + Speak
 # ----------------------
 def generate_caption_translate_speak(image, target_lang):
     # Step 1: Caption
-    inputs = caption_processor(images=image, return_tensors="pt")
-    out = caption_model.generate(**inputs, max_new_tokens=50)
+    inputs = caption_processor(images=image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        out = caption_model.generate(**inputs, max_new_tokens=50)
     english_caption = caption_processor.decode(out[0], skip_special_tokens=True)
 
+    # Step 1.5: Safety Check
+    if not is_caption_safe(english_caption):
+        return "⚠️ Warning: Unsafe or inappropriate content detected!", "", None
+
     # Step 2: Translate
     if target_lang in translation_models:
         translated = translation_models[target_lang](english_caption)[0]['translation_text']
@@ -108,37 +145,37 @@ def generate_caption_translate_speak(image, target_lang):
     tts = gTTS(english_caption, lang="en")
     tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
     tts.save(tmp_file.name)
-    audio_file = tmp_file.name
 
-    return english_caption, translated, audio_file
+    return english_caption, translated, tmp_file.name
 
 # ----------------------
-# VQA Function (using BLIP VQA)
+# VQA
 # ----------------------
-from transformers import BlipProcessor, BlipForQuestionAnswering
-
-vqa_processor = BlipProcessor.from_pretrained("Salesforce/blip-vqa-base")
-vqa_model = BlipForQuestionAnswering.from_pretrained("Salesforce/blip-vqa-base").to("cuda" if torch.cuda.is_available() else "cpu")
-
 def vqa_answer(image, question):
-    inputs = vqa_processor(image, question, return_tensors="pt").to(vqa_model.device)
-    out = vqa_model.generate(**inputs, max_new_tokens=50)
+    inputs = vqa_processor(image, question, return_tensors="pt").to(device)
+    with torch.no_grad():
+        out = vqa_model.generate(**inputs, max_new_tokens=50)
     answer = vqa_processor.decode(out[0], skip_special_tokens=True)
+
+    # Run safety filter on answers too
+    if not is_caption_safe(answer):
+        return "⚠️ Warning: Unsafe or inappropriate content detected!"
+
     return answer
 
 # ----------------------
 # Gradio UI
 # ----------------------
 with gr.Blocks(title="BLIP Vision App") as demo:
-    gr.Markdown("## 🖼️ BLIP: Image Captioning + Translation + Speech + VQA")
+    gr.Markdown("## 🖼️ BLIP: Image Captioning + Translation + Speech + VQA (with Safety Filter)")
 
     with gr.Tab("Caption + Translate + Speak"):
         with gr.Row():
             img_in = gr.Image(type="pil", label="Upload Image")
-            lang_in = gr.Dropdown(["Hindi", "French", "Spanish"], label="Translate To")
+            lang_in = gr.Dropdown(["Hindi", "French", "Spanish"], label="Translate To", value="Hindi")
         eng_out = gr.Textbox(label="English Caption")
         trans_out = gr.Textbox(label="Translated Caption")
-        audio_out = gr.Audio(label="Spoken Caption")
+        audio_out = gr.Audio(label="Spoken Caption", type="filepath")
         btn1 = gr.Button("Generate Caption, Translate & Speak")
         btn1.click(generate_caption_translate_speak, inputs=[img_in, lang_in], outputs=[eng_out, trans_out, audio_out])