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'||<[^>]+>', '', 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