Update vision_model.py

vision_model.py

@@ -1,16 +1,20 @@
-from transformers import BlipProcessor, Blip2ForConditionalGeneration
+from transformers import AutoProcessor, LlavaForConditionalGeneration
 from PIL import Image
 import torch
 import os
+
 os.environ["HF_HOME"] = "/app/.cache"
 
-model_id = "Salesforce/blip2-opt-2.7b"
-processor = BlipProcessor.from_pretrained(model_id)
+model_id = "llava-hf/llava-1.5-7b-hf"
+processor = AutoProcessor.from_pretrained(model_id)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = Blip2ForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32).to(device)
+model = LlavaForConditionalGeneration.from_pretrained(
+    model_id,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+).to(device)
 
 def describe_image(image_path, prompt="Describe this sculpt and any missing regions."):
     image = Image.open(image_path).convert("RGB")
-    inputs = processor(image, prompt, return_tensors="pt").to(device, model.dtype)
+    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)