llava_inference.py

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()