Back to BLIP

app.py

@@ -1,37 +1,25 @@
 import streamlit as st
 from PIL import Image
-
 import torch
-from transformers import MllamaForConditionalGeneration, AutoProcessor
 
-model_id = "meta-llama/Llama-3.2-11B-Vision-Instruct"
+from transformers import Blip2Processor, Blip2ForConditionalGeneration
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-model = MllamaForConditionalGeneration.from_pretrained(
-    model_id,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
+processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
+model = Blip2ForConditionalGeneration.from_pretrained(
+    "Salesforce/blip2-opt-2.7b", load_in_8bit=True, device_map={"": 0}, torch_dtype=torch.float16
 )
-processor = AutoProcessor.from_pretrained(model_id)
+
 
 enable = st.checkbox("Enable camera")
 picture = st.camera_input("Take a picture", disabled=not enable)
 
 if picture:
 	image = Image.open(picture)
+	prompt = "Question: At what location is this person most likely attending this online meeting? Answer:"
+	inputs = processor(images=image, text=prompt, return_tensors="pt").to(device="cuda", dtype=torch.float16)
 
-	messages = [
-		{"role": "user", "content": [
-			{"type": "image"},
-			{"type": "text", "text": "Provide your best guess as to where this person is holding his online meeting. Just state your guess of location in your response."}
-		]}
-	]
-	input_text = processor.apply_chat_template(messages, add_generation_prompt=True)
-	inputs = processor(
-		image,
-		input_text,
-		add_special_tokens=False,
-		return_tensors="pt"
-	).to(model.device)
-
-	output = model.generate(**inputs, max_new_tokens=30)
-	print(processor.decode(output[0]))
+	generated_ids = model.generate(**inputs)
+	generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+	st.write(generated_text)