Update handler.py

handler.py

@@ -1,140 +1,73 @@
 import torch
-from transformers import BitsAndBytesConfig
+from transformers import AutoProcessor, LlavaForConditionalGeneration
+from peft import PeftModel
 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=""):
-        # 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
+        # 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,
             device_map="auto"
         )
-        self.model.eval()
-        print("✅ Base model and LoRA adapter loaded successfully using LLaVA's official loader.")
+        
+        print(f"Loading LoRA adapters from repository path: {path}...")
+        self.model = PeftModel.from_pretrained(self.model, path)
+        print("✅ Model and adapters loaded successfully.")
 
     def __call__(self, data: dict) -> dict:
-        # Get the payload from the 'inputs' key
+        # --- THIS IS THE FIX ---
+        # The Inference Endpoint wraps the payload in an "inputs" key.
+        # We must extract our data from there first.
         payload = data.pop("inputs", data)
-            
+        # --- END OF FIX ---
+
+        # 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()
-                image = Image.open(BytesIO(response.content)).convert("RGB")
+                response.raise_for_status() # Raise an exception for bad status codes
+                image = Image.open(BytesIO(response.content))
             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)).convert("RGB")
+                image = Image.open(BytesIO(image_bytes))
             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'."}
 
-        # 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>",
-        )
+        # Format the prompt for LLaVA
+        prompt = f"USER: <image>\n{prompt_text} ASSISTANT:"
         
-        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()
-        
-        input_ids = tokenizer_image_token(prompt_with_template, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()
+        # Process inputs
+        inputs = self.processor(text=prompt, images=image, return_tensors="pt").to("cuda")
         
         # Generate a response
-        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,
-            )
+        with torch.no_grad():
+            output = self.model.generate(**inputs, max_new_tokens=max_new_tokens)
         
-        # Decode and return the response
-        response_text = self.tokenizer.decode(output_ids[0, input_ids.shape[1]:]).strip()
+        # 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()
         
-        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])
+        return {"generated_text": assistant_response}