handler.py

import os
import io
import torch
import base64
import re
from typing import Dict, List, Any, Union
from PIL import Image
from peft import PeftModel
from transformers import AutoProcessor, Idefics3ForConditionalGeneration, BitsAndBytesConfig

class EndpointHandler:
    def __init__(self, path=""):
        """
        Initialize the model and processor.
        The path parameter contains the path to the model directory.
        """
        # Set device
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # Set model parameters
        self.use_qlora = True
        self.model_path = path  # This will be the directory with your fine-tuned model
        self.base_model_id = "HuggingFaceTB/SmolVLM-Base"  # Base model you used
        
        # Load processor and model
        self.load_model()
        
        # Print some info when the handler is initialized
        print(f"Handler initialized with model from {self.model_path}")
        print(f"Running on device: {self.device}")

    def load_model(self):
        """Load the fine-tuned model"""
        print(f"Loading processor from {self.base_model_id}...")
        self.processor = AutoProcessor.from_pretrained(self.base_model_id)
        
        print(f"Loading model...")
        
        if self.use_qlora:
            # Configure quantization for inference
            bnb_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=torch.bfloat16
            )
            
            # Load base model with quantization
            base_model = Idefics3ForConditionalGeneration.from_pretrained(
                self.base_model_id,
                quantization_config=bnb_config,
                torch_dtype=torch.bfloat16,
                device_map="auto"
            )
            
            # Load the PEFT adapter from local path
            self.model = PeftModel.from_pretrained(base_model, self.model_path)
            print("Loaded model with PEFT adapter")
        else:
            # Load the full fine-tuned model
            self.model = Idefics3ForConditionalGeneration.from_pretrained(
                self.model_path,
                torch_dtype=torch.bfloat16,
                device_map="auto"
            )
            print("Loaded full fine-tuned model")

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        Run inference on the input data.
        
        Args:
            data: A dictionary with the following keys:
                - inputs: The input image in various formats (file, URL, base64)
                - parameters: Optional parameters for inference
        
        Returns:
            A list of dictionaries with the results
        """
        # Extract parameters
        parameters = data.get("parameters", {})
        max_new_tokens = parameters.get("max_new_tokens", 256)
        temperature = parameters.get("temperature", 0.6)
        top_p = parameters.get("top_p", 0.9)
        do_sample = parameters.get("do_sample", True)
        
        # Get custom question if provided, otherwise use default
        question = parameters.get("question", 
                                "Please describe the person in the image's emotions, expression, and body language.")
        
        # Process input image
        inputs = data.get("inputs", None)
        if inputs is None:
            return [{"error": "No inputs provided"}]
        
        # Handle different input formats
        image = self._process_input(inputs)
        if image is None:
            return [{"error": "Failed to process image input"}]
        
        # Create message following the same format used in training
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "Identify the emotions in this photograph. Focus on the person's facial expression."},
                    {"type": "image"},
                    {"type": "text", "text": question}
                ]
            }
        ]
        
        # Apply chat template and generate input
        try:
            text = self.processor.apply_chat_template(messages, add_generation_prompt=True)
            inputs = self.processor(text=text, images=image, return_tensors="pt").to(self.device)
            
            # Generate prediction
            with torch.no_grad():
                output_ids = self.model.generate(
                    **inputs,
                    max_new_tokens=max_new_tokens,
                    do_sample=do_sample,
                    temperature=temperature,
                    top_p=top_p,
                )
            
            # Decode the output
            output = self.processor.tokenizer.decode(output_ids[0], skip_special_tokens=True)
            
            # Clean up the output
            cleaned_output = self._clean_output(output)
            
            return [{"generated_text": cleaned_output}]
        
        except Exception as e:
            import traceback
            return [{"error": f"Error during inference: {str(e)}", 
                     "traceback": traceback.format_exc()}]

    def _clean_output(self, output: str) -> str:
        """
        Clean up the model output to extract only the relevant emotion information.
        
        Args:
            output: Raw output from the model
        
        Returns:
            str: Cleaned output containing only the emotion analysis
        """
        # First, clean up document-style markup
        output = re.sub(r'<row_\d+_col_\d+>|<global-img>|<[^>]+>', '', output)
        
        # Extract the portion after "assistant:" or "Assistant:"
        if re.search(r'assistant:\s*', output, re.IGNORECASE):
            output = re.split(r'assistant:\s*', output, flags=re.IGNORECASE)[-1].strip()
        
        # Remove any "User:" prefix that might remain
        output = re.sub(r'User:.*?(?=Assistant:|$)', '', output, re.IGNORECASE)
        
        # Try to extract just the emotion statement
        emotion_match = re.search(r'This person displays.+?intensity\.', output)
        if emotion_match:
            return emotion_match.group(0)
        
        # If specific emotion statement not found, clean up any remaining noise
        # Remove newlines and excessive spaces
        output = re.sub(r'\s+', ' ', output).strip()
        
        # If the output is still very long, try to extract the most relevant sentence
        if len(output) > 100:
            sentences = re.split(r'[.!?]\s+', output)
            emotion_sentences = [s for s in sentences if any(term in s.lower() for term in 
                                ['emotion', 'feeling', 'express', 'display', 'mood', 'anxiety', 
                                 'happy', 'sad', 'angry', 'fear', 'neutral', 'surprise'])]
            
            if emotion_sentences:
                return emotion_sentences[0] + '.'
        
        return output

    def _process_input(self, inputs: Union[str, bytes, Dict]) -> Union[Image.Image, str, None]:
        """
        Process the input data to get a PIL Image.
        
        Handles different input formats:
        - Base64 encoded image
        - URL
        - Raw bytes
        - Dictionary with file content
        
        Returns:
            PIL.Image, str, or None: The processed image or URL, or None if processing failed
        """
        try:
            # Handle dictionary with binary data (file upload)
            if isinstance(inputs, dict) and "file" in inputs:
                if isinstance(inputs["file"], bytes):
                    return Image.open(io.BytesIO(inputs["file"]))
                return None
            
            # Handle base64 encoded image
            if isinstance(inputs, str) and inputs.startswith("data:image"):
                base64_data = inputs.split(",")[1]
                image_bytes = base64.b64decode(base64_data)
                return Image.open(io.BytesIO(image_bytes))
            
            # Handle URL (will be downloaded by the processor)
            if isinstance(inputs, str) and (inputs.startswith("http://") or inputs.startswith("https://")):
                # For URLs, we can let the processor handle it
                return inputs
            
            # Handle raw bytes
            if isinstance(inputs, bytes):
                return Image.open(io.BytesIO(inputs))
            
            # Handle base64 without data URI prefix
            if isinstance(inputs, str):
                try:
                    image_bytes = base64.b64decode(inputs)
                    return Image.open(io.BytesIO(image_bytes))
                except:
                    # Not valid base64, try other methods
                    pass
            
            return None
        
        except Exception as e:
            print(f"Error processing input: {str(e)}")
            return None