Update handler.py

handler.py

@@ -1,75 +1,140 @@
 import torch
-from transformers import AutoProcessor, LlavaForConditionalGeneration
-from peft import PeftModel
+from transformers import BitsAndBytesConfig
 from PIL import Image
 import requests
 from io import BytesIO
 import base64
 
+# Import LLaVA's specific tools
+from llava.model.builder import load_pretrained_model
+from llava.mm_utils import get_model_name_from_path, process_images, tokenizer_image_token
+from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
+
 class EndpointHandler:
     def __init__(self, path=""):
-        # The 'path' argument will be the path to your LoRA repo on the Hub
-        base_model_id = "llava-hf/llava-1.5-7b-hf"
-        
-        print("Loading processor...")
-        self.processor = AutoProcessor.from_pretrained(base_model_id, revision="a272c74")
-        
-        print("Loading base model...")
-        self.model = LlavaForConditionalGeneration.from_pretrained(
-            base_model_id,
-            load_in_4bit=True,
-            torch_dtype=torch.float16,
+        # path is the local path to your LoRA adapter repository
+
+        # This single, official LLaVA function handles everything:
+        # 1. Loads the base model (llava-1.5-7b-hf)
+        # 2. Correctly applies your LoRA adapter from the `path`
+        # 3. Loads the correct processor and tokenizer
+        model_name = get_model_name_from_path(path)
+        self.tokenizer, self.model, self.image_processor, self.context_len = load_pretrained_model(
+            model_path=path,                         # Path to your LoRA adapter
+            model_base="llava-hf/llava-1.5-7b-hf",   # Base model to load underneath
+            model_name=model_name,
+            load_in_4bit=True,                       # Load in 4-bit for efficiency
             device_map="auto"
         )
-        
-        print(f"Loading LoRA adapters from repository path: {path}...")
-        self.model = PeftModel.from_pretrained(self.model, path)
-        print("✅ Model and adapters loaded successfully.")
+        self.model.eval()
+        print("✅ Base model and LoRA adapter loaded successfully using LLaVA's official loader.")
+
     def __call__(self, data: dict) -> dict:
+        # Get the payload from the 'inputs' key
         payload = data.pop("inputs", data)
-    
-        # Optional dynamic override of LoRA adapters (costly if done per request)
-        lora_override = payload.pop("lora_path", None)
-        if lora_override:
-            # reload adapters on the fly (might be slow)
-            self.model = PeftModel.from_pretrained(self.model.base_model, lora_override)
-    
-
-        # Now, get the prompt and image from the extracted payload
+            
         prompt_text = payload.pop("prompt", "Describe the image in detail.")
         image_url = payload.pop("image_url", None)
         image_b64 = payload.pop("image_b64", None)
         max_new_tokens = payload.pop("max_new_tokens", 200)
 
-        # Load image from either a URL or a base64 string
         if image_url:
             try:
                 response = requests.get(image_url)
-                response.raise_for_status() # Raise an exception for bad status codes
-                image = Image.open(BytesIO(response.content))
+                response.raise_for_status()
+                image = Image.open(BytesIO(response.content)).convert("RGB")
             except Exception as e:
                 return {"error": f"Failed to load image from URL: {e}"}
         elif image_b64:
             try:
                 image_bytes = base64.b64decode(image_b64)
-                image = Image.open(BytesIO(image_bytes))
+                image = Image.open(BytesIO(image_bytes)).convert("RGB")
             except Exception as e:
                 return {"error": f"Failed to decode base64 image: {e}"}
         else:
             return {"error": "No image provided. Please use 'image_url' or 'image_b64'."}
 
-        # Format the prompt for LLaVA
-        prompt = f"USER: <image>\n{prompt_text} ASSISTANT:"
+        # Process the image
+        image_tensor = process_images([image], self.image_processor, self.model.config)
+        image_tensor = image_tensor.to(self.model.device, dtype=torch.float16)
+
+        # Format the prompt correctly for LLaVA v1.5
+        # Note: The format is slightly different from the previous generic one
+        conv_mode = "llava_v1"
+        conv = Conversation(
+            system="A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.",
+            roles=("USER", "ASSISTANT"),
+            version="v1",
+            messages=(),
+            offset=0,
+            sep_style=SeparatorStyle.TWO,
+            sep=" ",
+            sep2="</s>",
+        )
+        
+        prompt = DEFAULT_IMAGE_TOKEN + '\n' + prompt_text
+        conv.append_message(conv.roles[0], prompt)
+        conv.append_message(conv.roles[1], None)
+        prompt_with_template = conv.get_prompt()
         
-        # Process inputs
-        inputs = self.processor(text=prompt, images=image, return_tensors="pt").to("cuda")
+        input_ids = tokenizer_image_token(prompt_with_template, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()
         
         # Generate a response
-        with torch.no_grad():
-            output = self.model.generate(**inputs, max_new_tokens=max_new_tokens)
+        with torch.inference_mode():
+            output_ids = self.model.generate(
+                input_ids,
+                images=image_tensor,
+                do_sample=True,
+                temperature=0.2,
+                max_new_tokens=max_new_tokens,
+                use_cache=True,
+            )
         
-        # Decode and clean up the response
-        full_response = self.processor.decode(output[0], skip_special_tokens=True)
-        assistant_response = full_response.split("ASSISTANT:")[-1].strip()
+        # Decode and return the response
+        response_text = self.tokenizer.decode(output_ids[0, input_ids.shape[1]:]).strip()
         
-        return {"generated_text": assistant_response}
+        return {"generated_text": response_text}
+
+# This Conversation class is a helper taken from LLaVA's library for correct prompt formatting
+import dataclasses
+from enum import auto, Enum
+from typing import List, Tuple
+
+class SeparatorStyle(Enum):
+    SINGLE = auto()
+    TWO = auto()
+
+@dataclasses.dataclass
+class Conversation:
+    system: str
+    roles: List[str]
+    messages: List[List[str]]
+    offset: int
+    sep_style: SeparatorStyle = SeparatorStyle.SINGLE
+    sep: str = "###"
+    sep2: str = None
+    version: str = "Unknown"
+
+    def get_prompt(self):
+        if self.sep_style == SeparatorStyle.SINGLE:
+            ret = self.system + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ": " + message + self.sep
+                else:
+                    ret += role + ":"
+            return ret
+        elif self.sep_style == SeparatorStyle.TWO:
+            seps = [self.sep, self.sep2]
+            ret = self.system + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+            return ret
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+    def append_message(self, role, message):
+        self.messages.append([role, message])