handler.py

"""
Custom Inference Handler for Hugging Face Inference Endpoints
Combines Qwen2.5-VL embedding extraction + MLP classifiers
"""

import torch
import torch.nn as nn
from transformers import AutoProcessor, AutoModelForVision2Seq
from pathlib import Path
import numpy as np
from typing import Dict, Any
import av
import tempfile


class MLPClassifier(nn.Module):
    """MLP classifier matching training architecture"""
    def __init__(self, input_dim, hidden_dim=512, num_classes=4, dropout=0.3):
        super(MLPClassifier, self).__init__()
        
        self.classifier = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, num_classes)
        )
    
    def forward(self, x):
        return self.classifier(x)


class EndpointHandler:
    """
    Custom handler for HF Inference Endpoints
    """
    def __init__(self, path: str):
        """
        Initialize the handler
        Args:
            path: Path to the model directory on HF Hub
        """
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        
        # Load Qwen2.5-VL model for embeddings
        print("Loading Qwen2.5-VL model...")
        self.processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
        self.vision_model = AutoModelForVision2Seq.from_pretrained(
            "Qwen/Qwen2.5-VL-7B-Instruct",
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map="auto" if torch.cuda.is_available() else None
        )
        self.vision_model.eval()
        
        # Load MLP classifiers for each emotion
        self.categories = ["Boredom", "Engagement", "Confusion", "Frustration"]
        self.classifiers = {}
        
        path = Path(path)
        classifiers_dir = path / "classifiers"
        
        print("Loading MLP classifiers...")
        for category in self.categories:
            checkpoint_path = classifiers_dir / f"mlp_{category}_best.pth"
            
            if checkpoint_path.exists():
                # Determine embedding dimension from checkpoint
                checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
                
                # Get input dimension from first layer
                first_layer_weight = checkpoint['model_state_dict']['classifier.0.weight']
                input_dim = first_layer_weight.shape[1]
                
                # Initialize model
                model = MLPClassifier(input_dim=input_dim, num_classes=4)
                model.load_state_dict(checkpoint['model_state_dict'])
                model.to(self.device)
                model.eval()
                
                self.classifiers[category] = model
                print(f"  ✓ Loaded {category} classifier")
            else:
                print(f"  ✗ Missing {category} classifier at {checkpoint_path}")
        
        self.fps = 1  # Frame sampling rate
        
    def extract_image_embeddings(self, image_path: str) -> np.ndarray:
        """Extract embeddings from a single image using Qwen model"""
        from PIL import Image
        
        # Load image
        image = Image.open(image_path).convert('RGB')
        
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image},
                    {"type": "text", "text": "Analyze this image."}
                ]
            }
        ]
        
        with torch.no_grad():
            text = self.processor.apply_chat_template(
                messages,
                add_generation_prompt=True,
                tokenize=False,
            )
            
            inputs = self.processor(
                text=[text],
                images=[image],
                return_tensors="pt",
                padding=True,
            )
            
            inputs = {k: v.to(self.vision_model.device) for k, v in inputs.items()}
            outputs = self.vision_model(**inputs, output_hidden_states=True)
            hidden_states = outputs.hidden_states[-1]
            
            # Average pooling over sequence dimension
            embeddings = hidden_states.mean(dim=1).squeeze(0)
            embeddings = embeddings.cpu().numpy()
        
        return embeddings
    
    def extract_video_embeddings(self, video_path: str) -> np.ndarray:
        """Extract embeddings from video using Qwen model"""
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "video", "video": str(video_path)},
                    {"type": "text", "text": "Analyze this video."}
                ]
            }
        ]
        
        with torch.no_grad():
            inputs = self.processor.apply_chat_template(
                messages,
                fps=self.fps,
                add_generation_prompt=True,
                tokenize=True,
                return_dict=True,
                return_tensors="pt",
            )
            
            inputs = {k: v.to(self.vision_model.device) for k, v in inputs.items()}
            outputs = self.vision_model(**inputs, output_hidden_states=True)
            hidden_states = outputs.hidden_states[-1]
            
            # Average pooling over sequence dimension
            embeddings = hidden_states.mean(dim=1).squeeze(0)
            embeddings = embeddings.cpu().numpy()
        
        return embeddings
    
    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Handle inference request
        
        Args:
            data: Input data containing either:
                - "inputs": base64 encoded image or video file
                - "video_url": URL to video file
        
        Returns:
            Dictionary with predictions for each emotion category
        """
        try:
            import base64
            from PIL import Image
            import io
            
            file_path = None
            is_image = False
            
            # Handle different input formats
            if "inputs" in data:
                # Base64 encoded data
                input_data = data["inputs"]
                
                # Remove data URL prefix if present (e.g., "data:image/png;base64,")
                if ',' in input_data and input_data.startswith('data:'):
                    input_data = input_data.split(',', 1)[1]
                
                file_bytes = base64.b64decode(input_data)
                
                # Try to detect if it's an image or video
                try:
                    # Try to open as image
                    Image.open(io.BytesIO(file_bytes))
                    is_image = True
                    suffix = '.png'
                except:
                    # Assume it's a video
                    is_image = False
                    suffix = '.avi'
                
                # Save to temporary file
                with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
                    tmp.write(file_bytes)
                    file_path = tmp.name
            
            elif "video_url" in data:
                # Download from URL
                import requests
                response = requests.get(data["video_url"])
                
                with tempfile.NamedTemporaryFile(suffix='.avi', delete=False) as tmp:
                    tmp.write(response.content)
                    file_path = tmp.name
                is_image = False
            else:
                return {"error": "No input provided. Use 'inputs' (base64) or 'video_url'"}
            
            # Extract embeddings based on input type
            if is_image:
                embeddings = self.extract_image_embeddings(file_path)
            else:
                embeddings = self.extract_video_embeddings(file_path)
            
            embeddings_tensor = torch.FloatTensor(embeddings).unsqueeze(0).to(self.device)
            
            # Run classifiers
            predictions = {}
            
            with torch.no_grad():
                for category, model in self.classifiers.items():
                    outputs = model(embeddings_tensor)
                    probabilities = torch.softmax(outputs, dim=1)
                    predicted_level = outputs.argmax(dim=1).item()
                    confidence = probabilities[0][predicted_level].item()
                    
                    predictions[category] = {
                        "level": int(predicted_level),
                        "confidence": float(confidence),
                        "probabilities": probabilities[0].cpu().numpy().tolist()
                    }
            
            # Clean up temporary file
            if file_path:
                Path(file_path).unlink(missing_ok=True)
            
            return {
                "success": True,
                "predictions": predictions
            }
            
        except Exception as e:
            return {
                "success": False,
                "error": str(e)
            }