Update handler.py

handler.py

@@ -1,66 +1,91 @@
+from typing import Dict, Any
 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
 
+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
+
 class EndpointHandler:
     def __init__(self, path=""):
-        # The 'path' argument will be the path to your LoRA repo on the Hub, e.g., "Abdulmateen/llava-finetuned"
-        
-        # Define the base model that your LoRA was trained on
-        base_model_id = "llava-hf/llava-1.5-7b-hf"
+        # path is the local path to your repository files
         
-        print("Loading processor...")
-        # Pinning to a specific revision for stability
-        self.processor = AutoProcessor.from_pretrained(base_model_id, revision="a272c74")
-
-        print("Loading base model...")
-        # Load the base model in 4-bit for memory efficiency
-        self.model = LlavaForConditionalGeneration.from_pretrained(
-            base_model_id,
+        # Load the base model in 4-bit for efficiency
+        bnb_config = BitsAndBytesConfig(
             load_in_4bit=True,
-            torch_dtype=torch.float16,
-            device_map="auto"
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
         )
+        
+        # This function loads the base model and applies the LoRA adapter from your repo
+        model_name = get_model_name_from_path(path)
+        self.tokenizer, self.model, self.image_processor, self.context_len = load_pretrained_model(
+            model_path=path, # Loads LoRA from your repo
+            model_base="liuhaotian/llava-v1.5-7b", # Loads the original base model
+            model_name=model_name,
+            quantization_config=bnb_config
+        )
+        self.model.eval()
+        print("Model, tokenizer, and processor loaded.")
 
-        print(f"Loading LoRA adapters from repository path: {path}...")
-        # Load and merge your LoRA adapters onto the base model
-        self.model = PeftModel.from_pretrained(self.model, path)
-        print("✅ Model and adapters loaded successfully.")
-
-    def __call__(self, data: dict) -> dict:
-        # Get the prompt and image from the request payload
-        prompt_text = data.pop("prompt", "Describe the image in detail.")
-        image_url = data.pop("image_url", None)
-        image_b64 = data.pop("image_b64", None)
-        max_new_tokens = data.pop("max_new_tokens", 200)
-
-        # Load image from either a URL or a base64 string
-        if image_url:
-            response = requests.get(image_url)
-            image = Image.open(BytesIO(response.content))
-        elif image_b64:
-            image_bytes = base64.b64decode(image_b64)
-            image = Image.open(BytesIO(image_bytes))
-        else:
-            return {"error": "No image provided. Please use 'image_url' or 'image_b64'."}
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Args:
+            data (:obj:`dict`):
+                A dictionary containing the prompt and image.
+                - "image": A base64 encoded string of the image.
+                - "prompt": The text prompt.
+                - "max_new_tokens" (optional): Max new tokens to generate. Default is 512.
+        Returns:
+            :obj:`dict`: A dictionary with the model's generated text.
+        """
+        # Get inputs
+        base64_image = data.pop("image", None)
+        prompt = data.pop("prompt", None)
+        max_new_tokens = data.pop("max_new_tokens", 512)
 
-        # Format the prompt for LLaVA
-        prompt = f"USER: <image>\n{prompt_text} ASSISTANT:"
+        if not base64_image or not prompt:
+            return {"error": "Missing 'image' or 'prompt' in the request payload."}
 
-        # Process inputs
-        inputs = self.processor(text=prompt, images=image, return_tensors="pt").to("cuda")
+        # --- Image Processing ---
+        try:
+            image_bytes = base64.b64decode(base64_image)
+            image = Image.open(BytesIO(image_bytes)).convert("RGB")
+        except Exception as e:
+            return {"error": f"Failed to decode or open image: {e}"}
 
-        # Generate a response
-        with torch.no_grad():
-            output = self.model.generate(**inputs, max_new_tokens=max_new_tokens)
+        image_tensor = process_images([image], self.image_processor, self.model.config)
+        image_tensor = image_tensor.to(self.model.device, dtype=torch.float16)
 
-        # Decode and clean up the response
-        full_response = self.processor.decode(output[0], skip_special_tokens=True)
-        # Extract only the assistant's part of the response
-        assistant_response = full_response.split("ASSISTANT:")[-1].strip()
+        # --- Prompt Formatting ---
+        conv_mode = "llava_v1" # Or the conversation mode you used during training
+        prompt = DEFAULT_IMAGE_TOKEN + '\n' + prompt
+        conv = conv_mode.conv.copy()
+        conv.append_message(conv.roles[0], prompt)
+        conv.append_message(conv.roles[1], None)
+        prompt_text = conv.get_prompt()
+        
+        input_ids = tokenizer_image_token(prompt_text, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()
+        
+        # --- Generation ---
+        with torch.inference_mode():
+            output_ids = self.model.generate(
+                input_ids,
+                images=image_tensor,
+                do_sample=True,
+                temperature=0.2,
+                top_p=None,
+                num_beams=1,
+                max_new_tokens=max_new_tokens,
+                use_cache=True,
+            )
 
-        return {"generated_text": assistant_response}
+        # --- Decode and Return ---
+        decoded_text = self.tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0].strip()
+        
+        return {"generated_text": decoded_text}