Update app.py

app.py

@@ -4,13 +4,18 @@ import gradio as gr
 import tempfile
 import secrets
 from pathlib import Path
-from transformers import AutoModelForCausalLM, AutoTokenizer, BlipForConditionalGeneration, AutoProcessor
+from transformers import AutoModelForCausalLM, AutoTokenizer, BlipForConditionalGeneration, AutoProcessor, Qwen2VLForConditionalGeneration
+from qwen_vl_utils import process_vision_info
 from PIL import Image
 
 # Load Vision-Language Model
-vl_model_name = "Salesforce/blip-image-captioning-large"
-vl_model = BlipForConditionalGeneration.from_pretrained(vl_model_name)
-vl_processor = AutoProcessor.from_pretrained(vl_model_name)
+vl_model = Qwen2VLForConditionalGeneration.from_pretrained(
+    "Qwen/Qwen2-VL-2B-Instruct",
+    torch_dtype=torch.bfloat16,
+    attn_implementation="flash_attention_2",
+    device_map="auto",
+)
+vl_processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct")
 
 # Load Text Model
 model_name = "Qwen/Qwen2.5-Math-1.5B-Instruct"
@@ -31,7 +36,13 @@ def process_image(image, shouldConvert=False):
 
     # Convert the image to tensor
     inputs = vl_processor(images=image, return_tensors="pt")
-    output = vl_model.generate(**inputs)
+    generated_ids = vl_model.generate(**inputs)
+    generated_ids_trimmed = [
+        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+    ]
+    output = processor.batch_decode(
+        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )
     description = vl_processor.batch_decode(output, skip_special_tokens=True)[0]
 
     return f"Math-related content detected: {description}"