from transformers import LlavaForConditionalGeneration, LlavaProcessor
from config import LLAVA_MODEL_PATH

class LlavaWrapper:
    def __init__(self, device="cuda"):
        self.model = LlavaForConditionalGeneration.from_pretrained(LLAVA_MODEL_PATH).to(device)
        self.processor = LlavaProcessor.from_pretrained(LLAVA_MODEL_PATH)
        self.device = device
