Update app.py

app.py

@@ -70,42 +70,29 @@ import gradio as gr
 from transformers import Blip2Processor, Blip2ForConditionalGeneration, pipeline
 from PIL import Image
 import torch
-import streamlit as st
 
 # ----------------------
-# Cached Model Loaders
+# Load BLIP2 for Captioning
 # ----------------------
-@st.cache_resource
-def load_caption_model():
-    processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
-    model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b")
-    return processor, model
-
-@st.cache_resource
-def load_vqa_model():
-    processor = Blip2Processor.from_pretrained("Salesforce/blip2-flan-t5-xl")
-    model = Blip2ForConditionalGeneration.from_pretrained(
-        "Salesforce/blip2-flan-t5-xl", torch_dtype=torch.float16, device_map="auto"
-    )
-    return processor, model
-
-@st.cache_resource
-def load_translation_models():
-    return {
-        "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"),
-    }
+caption_processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
+caption_model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b")
 
 # ----------------------
-# Load All Models with Spinner
+# Load BLIP2 for VQA
 # ----------------------
-with st.spinner("Loading BLIP2 models... please wait ⏳"):
-    caption_processor, caption_model = load_caption_model()
-    vqa_processor, vqa_model = load_vqa_model()
-    translation_models = load_translation_models()
+vqa_processor = Blip2Processor.from_pretrained("Salesforce/blip2-flan-t5-xl")
+vqa_model = Blip2ForConditionalGeneration.from_pretrained(
+    "Salesforce/blip2-flan-t5-xl", torch_dtype=torch.float16, device_map="auto"
+)
 
-st.success("✅ Models are ready!")
+# ----------------------
+# Translation pipelines
+# ----------------------
+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"),
+}
 
 # ----------------------
 # Caption + Translate Function
@@ -115,6 +102,7 @@ def generate_caption_translate(image, target_lang):
     out = caption_model.generate(**inputs, max_new_tokens=50)
     english_caption = caption_processor.decode(out[0], skip_special_tokens=True)
 
+    # Translate if chosen
     if target_lang in translation_models:
         translated = translation_models[target_lang](english_caption)[0]['translation_text']
     else: