model

llava_inference.py

@@ -0,0 +1,25 @@
+from llava.model.builder import load_pretrained_model
+from llava.mm_utils import process_images, tokenizer_image_token
+from transformers import AutoTokenizer
+import torch
+
+class LLaVAHelper:
+    def __init__(self, model_name="llava-hf/llava-1.5-7b-hf"):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model, self.image_processor, _ = load_pretrained_model(model_name, None)
+        self.model.eval()
+
+    def generate_answer(self, image, question):
+        # Preprocess
+        image_tensor = process_images([image], self.image_processor, self.model.config)[0].unsqueeze(0).to("cuda" if torch.cuda.is_available() else "cpu")
+        prompt = f"###Human: <image>\n{question}\n###Assistant:"
+        input_ids = tokenizer_image_token(prompt, self.tokenizer, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
+
+        with torch.no_grad():
+            output_ids = self.model.generate(
+                input_ids=input_ids.input_ids,
+                images=image_tensor,
+                max_new_tokens=512
+            )
+        output = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        return output.split("###Assistant:")[-1].strip()