Update vision_model.py

vision_model.py

@@ -1,20 +1,17 @@
-from transformers import AutoProcessor, LlavaForConditionalGeneration
+from transformers import BlipProcessor, BlipForConditionalGeneration
 from PIL import Image
 import torch
 import os
 
 os.environ["HF_HOME"] = "/app/.cache"
 
-model_id = "llava-hf/llava-1.5-7b-hf"
-processor = AutoProcessor.from_pretrained(model_id)
+model_id = "Salesforce/blip-image-captioning-large"
+processor = BlipProcessor.from_pretrained(model_id)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = LlavaForConditionalGeneration.from_pretrained(
-    model_id,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-).to(device)
+model = BlipForConditionalGeneration.from_pretrained(model_id).to(device)
 
-def describe_image(image_path, prompt="Describe this sculpt and any missing regions."):
+def describe_image(image_path, prompt="Describe this image."):
     image = Image.open(image_path).convert("RGB")
-    inputs = processor(prompt, image, return_tensors="pt").to(device, model.dtype)
-    output = model.generate(**inputs, max_new_tokens=150)
-    return processor.decode(output[0], skip_special_tokens=True)
+    inputs = processor(image, prompt, return_tensors="pt").to(device)
+    output = model.generate(**inputs, max_new_tokens=100)
+    return processor.decode(output[0], skip_special_tokens=True)