Update app.py

app.py

@@ -1,35 +1,30 @@
 import gradio as gr
-from transformers import Blip2Processor, Blip2ForConditionalGeneration, BitsAndBytesConfig
+from transformers import Blip2Processor, Blip2ForConditionalGeneration
 from gtts import gTTS
 from tempfile import NamedTemporaryFile
 from PIL import Image
 import torch
 import whisper
 
-# Set device
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load BLIP-2 model
+# Load BLIP-2 (smaller model)
 processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
-model = Blip2ForConditionalGeneration.from_pretrained(
-    "Salesforce/blip2-opt-2.7b", device_map="auto"
-).to(device)
+model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b").to(device)
 
-# Load Whisper model
-whisper_model = whisper.load_model("small")
+# Load Whisper
+whisper_model = whisper.load_model("base")
 
-# Transcribe function
 def transcribe(audio_path):
     result = whisper_model.transcribe(audio_path)
     return result["text"]
 
-# Main function
 def ask_image(image, audio):
     question = transcribe(audio)
     inputs = processor(images=image, text=question, return_tensors="pt").to(device)
     generated_ids = model.generate(**inputs)
     answer = processor.decode(generated_ids[0], skip_special_tokens=True)
-    
+
     tts = gTTS(answer)
     with NamedTemporaryFile(delete=False, suffix=".mp3") as f:
         tts.save(f.name)
@@ -37,7 +32,6 @@ def ask_image(image, audio):
 
     return answer, audio_out
 
-# Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("## 🎤🖼️ Ask-the-Image: Ask questions about an image using your voice")
 
@@ -50,5 +44,4 @@ with gr.Blocks() as demo:
     btn = gr.Button("Ask")
     btn.click(fn=ask_image, inputs=[image_input, audio_input], outputs=[text_output, audio_output])
 
-
 demo.launch()