""" PacksAI V6 - HF Inference Endpoint Handler Dual-Backbone (DINOv2-Large + SigLIP2) with Multi-Scale Fusion, Patch Attention Pooling, Cross-Modal Attention, T2B Calibration, Monte Carlo Dropout, and optional Test-Time Augmentation. """ import os import json import base64 import io import pickle import logging from typing import Any, Dict import torch import torch.nn as nn import torch.nn.functional as F import numpy as np from PIL import Image from torchvision import transforms from transformers import Dinov2Model, AutoModel logger = logging.getLogger(__name__) # ============================================================ # KONFIGURATION V6 # ============================================================ CONFIG = { 'data': { 'fields': ['country', 'category', 'design_range_shape_1st_step', 'gender', 'age_group'], 'targets': ['curiosity', 'uniqueness', 'purchase_intention', 'brand_fit', 'appealing_design', 'premiumness', 'quality'], 'kpi_groups': { 'emotional': ['curiosity', 'uniqueness'], 'commercial': ['purchase_intention', 'brand_fit'], 'aesthetic': ['appealing_design', 'premiumness', 'quality'] }, 'scale_min': 1.0, 'scale_max': 5.0, }, 'model': { 'dinov2_backbone': 'facebook/dinov2-large', 'dinov2_dim': 1024, 'dinov2_image_size': 518, 'siglip_backbone': 'google/siglip2-so400m-patch14-384', 'siglip_dim': 1152, 'siglip_image_size': 384, 'use_dual_backbone': True, 'multi_scale_layers': [8, 16, 24], 'patch_attn_heads': 8, 'patch_attn_queries': 4, 'tab_emb_dim': 64, 'tab_proj_dim': 256, 'hidden_dim': 1024, 'head_hidden_dim': 256, 'dropout': 0.3, }, 'inference': { 'mc_dropout_samples': 30, 'tta_augmentations': 5, 'use_tta': True, } } KPI_INFO = { 'curiosity': {'label': 'Curiosity', 'icon': '🔍', 'group': 'Emotional'}, 'uniqueness': {'label': 'Uniqueness', 'icon': '✨', 'group': 'Emotional'}, 'purchase_intention': {'label': 'Purchase Intention', 'icon': '🛒', 'group': 'Commercial'}, 'brand_fit': {'label': 'Brand Fit', 'icon': '🎯', 'group': 'Commercial'}, 'appealing_design': {'label': 'Appealing Design', 'icon': '🎨', 'group': 'Aesthetic'}, 'premiumness': {'label': 'Premiumness', 'icon': '💎', 'group': 'Aesthetic'}, 'quality': {'label': 'Quality', 'icon': '⭐', 'group': 'Aesthetic'}, } DISPLAY_ORDER = [ 'curiosity', 'uniqueness', 'purchase_intention', 'brand_fit', 'quality', 'appealing_design', 'premiumness', ] # ============================================================ # MODELL-DEFINITIONEN V6 # ============================================================ class CategoricalEmbeddingLayer(nn.Module): def __init__(self, field_cardinalities, embedding_dim=64): super().__init__() self.field_names = list(field_cardinalities.keys()) self.embeddings = nn.ModuleDict({ name: nn.Embedding(card, embedding_dim) for name, card in field_cardinalities.items() }) def forward(self, x): embs = [self.embeddings[name](x[:, i]) for i, name in enumerate(self.field_names)] return torch.cat(embs, dim=1) class PatchAttentionPool(nn.Module): def __init__(self, dim, num_heads=8, num_queries=4, dropout=0.1): super().__init__() self.num_queries = num_queries self.queries = nn.Parameter(torch.randn(1, num_queries, dim) * 0.02) self.attention = nn.MultiheadAttention(dim, num_heads, dropout=dropout, batch_first=True) self.norm = nn.LayerNorm(dim) self.fuse = nn.Linear(num_queries * dim, dim) def forward(self, patch_tokens): B = patch_tokens.size(0) q = self.queries.expand(B, -1, -1) attended, attn_weights = self.attention(q, patch_tokens, patch_tokens, need_weights=True, average_attn_weights=True) attended = self.norm(attended) pooled = self.fuse(attended.reshape(B, -1)) attn_avg = attn_weights.mean(dim=1) if attn_weights.dim() == 3 else attn_weights return pooled, attn_avg class MultiScaleFusion(nn.Module): def __init__(self, dim, num_scales, dropout=0.1): super().__init__() self.num_scales = num_scales self.scale_weights = nn.Parameter(torch.ones(num_scales) / num_scales) self.proj = nn.Sequential( nn.Linear(dim * num_scales, dim), nn.LayerNorm(dim), nn.GELU(), nn.Dropout(dropout), ) def forward(self, scale_features): weights = torch.softmax(self.scale_weights, dim=0) weighted = [f * w for f, w in zip(scale_features, weights)] concat = torch.cat(scale_features, dim=-1) return self.proj(concat), weights class CrossModalAttention(nn.Module): def __init__(self, dim_a, dim_b, hidden_dim, num_heads=8, dropout=0.1): super().__init__() self.proj_a = nn.Linear(dim_a, hidden_dim) self.proj_b = nn.Linear(dim_b, hidden_dim) self.cross_a_to_b = nn.MultiheadAttention(hidden_dim, num_heads, dropout=dropout, batch_first=True) self.cross_b_to_a = nn.MultiheadAttention(hidden_dim, num_heads, dropout=dropout, batch_first=True) self.norm_a = nn.LayerNorm(hidden_dim) self.norm_b = nn.LayerNorm(hidden_dim) self.fuse = nn.Sequential( nn.Linear(hidden_dim * 2, hidden_dim), nn.LayerNorm(hidden_dim), nn.GELU(), nn.Dropout(dropout), ) def forward(self, feat_a, feat_b): a = self.proj_a(feat_a).unsqueeze(1) b = self.proj_b(feat_b).unsqueeze(1) a_attended, _ = self.cross_a_to_b(a, b, b) b_attended, _ = self.cross_b_to_a(b, a, a) a_out = self.norm_a(a + a_attended).squeeze(1) b_out = self.norm_b(b + b_attended).squeeze(1) return self.fuse(torch.cat([a_out, b_out], dim=-1)) class TabularAttentionFusion(nn.Module): def __init__(self, vision_dim, tab_dim, hidden_dim): super().__init__() self.query = nn.Linear(tab_dim, hidden_dim) self.key = nn.Linear(vision_dim, hidden_dim) self.value = nn.Linear(vision_dim, hidden_dim) self.scale = hidden_dim ** -0.5 def forward(self, vision_feat, tab_feat): q = self.query(tab_feat) k = self.key(vision_feat) v = self.value(vision_feat) attn = torch.softmax(q * k * self.scale, dim=-1) return torch.cat([vision_feat, tab_feat, attn * v], dim=1) class KPIHead(nn.Module): def __init__(self, input_dim, hidden_dim, num_outputs, dropout=0.3): super().__init__() self.head = nn.Sequential( nn.Linear(input_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden_dim, hidden_dim // 2), nn.LayerNorm(hidden_dim // 2), nn.GELU(), nn.Dropout(dropout / 2), nn.Linear(hidden_dim // 2, num_outputs) ) def forward(self, x): return self.head(x) class MultimodalKPIModelV6(nn.Module): def __init__(self, config, field_cardinalities): super().__init__() self.config = config self.targets = config['data']['targets'] self.kpi_groups = config['data']['kpi_groups'] self.scale_min = config['data']['scale_min'] self.scale_max = config['data']['scale_max'] self.scale_range = self.scale_max - self.scale_min mcfg = config['model'] self.use_dual = mcfg['use_dual_backbone'] self.multi_scale_layers = mcfg['multi_scale_layers'] # DINOv2 Backbone self.dinov2 = Dinov2Model.from_pretrained(mcfg['dinov2_backbone']) self.dinov2_dim = mcfg['dinov2_dim'] self.dinov2_patch_pools = nn.ModuleList([ PatchAttentionPool(self.dinov2_dim, mcfg['patch_attn_heads'], mcfg['patch_attn_queries'], mcfg['dropout']) for _ in self.multi_scale_layers ]) self.multi_scale_fusion = MultiScaleFusion(self.dinov2_dim, len(self.multi_scale_layers), mcfg['dropout']) # SigLIP2 Backbone if self.use_dual: self.siglip = AutoModel.from_pretrained(mcfg['siglip_backbone']) if hasattr(self.siglip, 'vision_model'): self.siglip_vision = self.siglip.vision_model else: self.siglip_vision = self.siglip self.siglip_dim = mcfg['siglip_dim'] self.siglip_patch_pool = PatchAttentionPool( self.siglip_dim, mcfg['patch_attn_heads'], mcfg['patch_attn_queries'], mcfg['dropout'] ) self.cross_modal = CrossModalAttention( self.dinov2_dim, self.siglip_dim, mcfg['hidden_dim'], num_heads=mcfg['patch_attn_heads'], dropout=mcfg['dropout'] ) fusion_input_dim = mcfg['hidden_dim'] else: fusion_input_dim = self.dinov2_dim # Tabular Encoder self.categorical_encoder = CategoricalEmbeddingLayer( field_cardinalities, mcfg['tab_emb_dim'] ) self.tab_input_dim = len(field_cardinalities) * mcfg['tab_emb_dim'] self.tab_projection = nn.Sequential( nn.Linear(self.tab_input_dim, mcfg['tab_proj_dim']), nn.LayerNorm(mcfg['tab_proj_dim']), nn.GELU(), nn.Dropout(mcfg['dropout']), ) # Vision-Tabular Fusion self.tab_fusion = TabularAttentionFusion( fusion_input_dim, mcfg['tab_proj_dim'], mcfg['hidden_dim'] // 2 ) fused_dim = fusion_input_dim + mcfg['tab_proj_dim'] + mcfg['hidden_dim'] // 2 # Shared Trunk self.shared_fusion = nn.Sequential( nn.Linear(fused_dim, mcfg['hidden_dim']), nn.LayerNorm(mcfg['hidden_dim']), nn.GELU(), nn.Dropout(mcfg['dropout']), ) # Heads head_hidden = mcfg['head_hidden_dim'] shared_dim = mcfg['hidden_dim'] self.mean_heads = nn.ModuleDict({ 'emotional': KPIHead(shared_dim, head_hidden, 2, mcfg['dropout']), 'commercial': KPIHead(shared_dim, head_hidden, 2, mcfg['dropout']), 'aesthetic': KPIHead(shared_dim, head_hidden, 3, mcfg['dropout']), }) self.t2b_heads = nn.ModuleDict({ 'emotional': KPIHead(shared_dim, head_hidden, 2, mcfg['dropout']), 'commercial': KPIHead(shared_dim, head_hidden, 2, mcfg['dropout']), 'aesthetic': KPIHead(shared_dim, head_hidden, 3, mcfg['dropout']), }) # Contrastive projection (not used in inference, but needed for state_dict) self.contrastive_proj = nn.Sequential( nn.Linear(shared_dim, 256), nn.GELU(), nn.Linear(256, 128), ) self.kpi_to_head_idx = {} for group_name, kpis in self.kpi_groups.items(): for i, kpi in enumerate(kpis): self.kpi_to_head_idx[kpi] = (group_name, i) def _extract_dinov2_features(self, images): outputs = self.dinov2(pixel_values=images, output_hidden_states=True) hidden_states = outputs.hidden_states scale_features = [] for i, layer_idx in enumerate(self.multi_scale_layers): hs = hidden_states[layer_idx] patch_tokens = hs[:, 1:] pooled, _ = self.dinov2_patch_pools[i](patch_tokens) scale_features.append(pooled) fused, _ = self.multi_scale_fusion(scale_features) return fused def _extract_siglip_features(self, images): outputs = self.siglip_vision(pixel_values=images, output_hidden_states=False) hidden = outputs.last_hidden_state pooled, _ = self.siglip_patch_pool(hidden) return pooled def forward(self, dinov2_images, categorical, siglip_images=None): dinov2_feat = self._extract_dinov2_features(dinov2_images) if self.use_dual and siglip_images is not None: siglip_feat = self._extract_siglip_features(siglip_images) vision_feat = self.cross_modal(dinov2_feat, siglip_feat) else: vision_feat = dinov2_feat cat_emb = self.categorical_encoder(categorical) tab_feat = self.tab_projection(cat_emb) fused = self.tab_fusion(vision_feat, tab_feat) shared = self.shared_fusion(fused) mean_outputs = {name: head(shared) for name, head in self.mean_heads.items()} t2b_outputs = {name: head(shared) for name, head in self.t2b_heads.items()} mean_preds_list, t2b_preds_list = [], [] for kpi in self.targets: group_name, idx = self.kpi_to_head_idx[kpi] mean_logit = mean_outputs[group_name][:, idx:idx+1] mean_pred = self.scale_min + self.scale_range * torch.sigmoid(mean_logit) mean_preds_list.append(mean_pred) t2b_logit = t2b_outputs[group_name][:, idx:idx+1] t2b_preds_list.append(torch.sigmoid(t2b_logit)) return { 'preds': torch.cat(mean_preds_list, dim=1), 'preds_t2b': torch.cat(t2b_preds_list, dim=1), } # ============================================================ # ENCODER # ============================================================ class CategoricalEncoder: def __init__(self): self.encoders = {} self.cardinalities = {} def encode(self, field, value): if field not in self.encoders: return 0 if value is None or str(value).lower() in ['nan', 'none', '', 'null']: return 1 return self.encoders[field].get(str(value), 1) @classmethod def load(cls, path): enc = cls() with open(path, 'r') as f: data = json.load(f) enc.encoders = data['encoders'] enc.cardinalities = data['cardinalities'] return enc # ============================================================ # HF INFERENCE ENDPOINT HANDLER V6 # ============================================================ class EndpointHandler: def __init__(self, path=""): self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') logger.info(f"Initialisiere V6 EndpointHandler auf {self.device}...") # Encoder laden encoder_path = os.path.join(path, "categorical_encoder.json") self.encoder = CategoricalEncoder.load(encoder_path) logger.info("Encoder geladen") # Modell laden model_path = os.path.join(path, "best_model_v6.pt") self.model = MultimodalKPIModelV6(CONFIG, self.encoder.cardinalities) checkpoint = torch.load(model_path, map_location=self.device, weights_only=False) if 'model' in checkpoint: self.model.load_state_dict(checkpoint['model']) elif 'model_state_dict' in checkpoint: self.model.load_state_dict(checkpoint['model_state_dict']) else: self.model.load_state_dict(checkpoint) self.model.to(self.device) self.model.eval() logger.info(f"V6 Modell geladen auf {self.device}") # Calibratoren laden self.calibrators = {} for cal_name in ['t2b_calibrators_v6.pkl', 't2b_calibrators.pkl']: calibrator_path = os.path.join(path, cal_name) if os.path.exists(calibrator_path): with open(calibrator_path, 'rb') as f: self.calibrators = pickle.load(f) logger.info(f"Calibratoren geladen: {list(self.calibrators.keys())}") break # Image transforms (dual backbone) self.dinov2_transform = transforms.Compose([ transforms.Resize((518, 518)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) self.siglip_transform = transforms.Compose([ transforms.Resize((384, 384)), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ]) # TTA augmentations self.tta_aug = transforms.Compose([ transforms.RandomHorizontalFlip(p=0.5), transforms.RandomRotation(10), transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1), ]) self.mc_samples = CONFIG['inference']['mc_dropout_samples'] self.tta_n = CONFIG['inference']['tta_augmentations'] self.use_tta = CONFIG['inference']['use_tta'] def _calibrate_t2b(self, t2b_values): calibrated = {} for i, kpi in enumerate(CONFIG['data']['targets']): raw_value = float(t2b_values[i]) if kpi in self.calibrators: calibrated[kpi] = float(self.calibrators[kpi].predict([[raw_value]])[0]) else: calibrated[kpi] = raw_value return calibrated def _prepare_images(self, image): dinov2_tensor = self.dinov2_transform(image).unsqueeze(0).to(self.device) siglip_tensor = self.siglip_transform(image).unsqueeze(0).to(self.device) return dinov2_tensor, siglip_tensor def _predict_single(self, image, categorical): """Single prediction with optional TTA.""" self.model.eval() with torch.inference_mode(): if self.use_tta and self.tta_n > 1: all_mean, all_t2b = [], [] # Original d, s = self._prepare_images(image) out = self.model(d, categorical, siglip_images=s) all_mean.append(out['preds'][0].cpu().numpy()) all_t2b.append(out['preds_t2b'][0].cpu().numpy()) # Augmented passes for _ in range(self.tta_n - 1): aug_img = self.tta_aug(image) d, s = self._prepare_images(aug_img) out = self.model(d, categorical, siglip_images=s) all_mean.append(out['preds'][0].cpu().numpy()) all_t2b.append(out['preds_t2b'][0].cpu().numpy()) return np.mean(all_mean, axis=0), np.mean(all_t2b, axis=0) else: d, s = self._prepare_images(image) out = self.model(d, categorical, siglip_images=s) return out['preds'][0].cpu().numpy(), out['preds_t2b'][0].cpu().numpy() def _predict_with_uncertainty(self, image, categorical): """Monte Carlo Dropout for confidence intervals.""" self.model.train() mean_samples, t2b_samples = [], [] with torch.no_grad(): for _ in range(self.mc_samples): d, s = self._prepare_images(image) out = self.model(d, categorical, siglip_images=s) mean_samples.append(out['preds'][0].cpu().numpy()) t2b_samples.append(out['preds_t2b'][0].cpu().numpy()) self.model.eval() mean_samples = np.array(mean_samples) t2b_samples = np.array(t2b_samples) return { 'mean_pred': mean_samples.mean(axis=0), 'mean_std': mean_samples.std(axis=0), 'mean_ci_lower': np.percentile(mean_samples, 2.5, axis=0), 'mean_ci_upper': np.percentile(mean_samples, 97.5, axis=0), 't2b_pred': t2b_samples.mean(axis=0), 't2b_std': t2b_samples.std(axis=0), 't2b_ci_lower': np.percentile(t2b_samples, 2.5, axis=0), 't2b_ci_upper': np.percentile(t2b_samples, 97.5, axis=0), } def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: inputs = data.get("inputs", data) # Decode image image_b64 = inputs.get("image", "") image_bytes = base64.b64decode(image_b64) image = Image.open(io.BytesIO(image_bytes)) if image.mode != 'RGB': image = image.convert('RGB') # Encode categorical features categorical = torch.tensor([[ self.encoder.encode('country', inputs.get('country')), self.encoder.encode('category', inputs.get('category')), self.encoder.encode('design_range_shape_1st_step', inputs.get('design_type')), self.encoder.encode('gender', inputs.get('gender')), self.encoder.encode('age_group', inputs.get('age_group')), ]], dtype=torch.long, device=self.device) with_ci = inputs.get('with_confidence_interval', False) if with_ci: results = self._predict_with_uncertainty(image, categorical) raw_mean = results['mean_pred'] raw_t2b = results['t2b_pred'] calibrated_t2b = self._calibrate_t2b(raw_t2b) calibrated_ci_lower = self._calibrate_t2b(results['t2b_ci_lower']) calibrated_ci_upper = self._calibrate_t2b(results['t2b_ci_upper']) else: raw_mean, raw_t2b = self._predict_single(image, categorical) calibrated_t2b = self._calibrate_t2b(raw_t2b) results = None # Build response kpis = [] for i, name in enumerate(CONFIG['data']['targets']): info = KPI_INFO[name] value = float(np.clip(raw_mean[i], 1.0, 5.0)) percent = (value - 1.0) / 4.0 * 100 t2b = round(calibrated_t2b[name] * 100, 1) kpi_data = { 'name': name, 'label': info['label'], 'icon': info['icon'], 'value': round(value, 2), 'max': 5, 'min': 1, 'percent': round(percent, 1), 't2b': t2b, 'group': info['group'] } if with_ci and results is not None: kpi_data['value_ci_lower'] = round(float(np.clip(results['mean_ci_lower'][i], 1.0, 5.0)), 2) kpi_data['value_ci_upper'] = round(float(np.clip(results['mean_ci_upper'][i], 1.0, 5.0)), 2) kpi_data['t2b_ci_lower'] = round(calibrated_ci_lower[name] * 100, 1) kpi_data['t2b_ci_upper'] = round(calibrated_ci_upper[name] * 100, 1) ci_width = kpi_data['t2b_ci_upper'] - kpi_data['t2b_ci_lower'] if ci_width < 10: kpi_data['confidence'] = 'high' elif ci_width < 20: kpi_data['confidence'] = 'medium' else: kpi_data['confidence'] = 'low' kpis.append(kpi_data) kpi_order = {name: idx for idx, name in enumerate(DISPLAY_ORDER)} kpis.sort(key=lambda k: kpi_order.get(k['name'], 99)) response = { 'kpis': kpis, 'avg_percent': round(float(np.mean([k['percent'] for k in kpis])), 1), 'avg_mean': round(float(np.mean([k['value'] for k in kpis])), 2), 'avg_t2b': round(float(np.mean([k['t2b'] for k in kpis])), 1), 'version': 'v6', 'calibrated': bool(self.calibrators), } if with_ci: response['confidence_interval'] = True response['ci_level'] = 0.95 return response