Update handler.py

handler.py

@@ -1,122 +1,820 @@
 from typing import Any, Dict
 import torch
 from torchvision import transforms
-from PIL import Image
+from PIL import Image, ImageDraw, ImageOps
 import base64
 import io
 import numpy as np
-from transformers import AutoModelForImageClassification, AutoImageProcessor
+from transformers import AutoModelForImageClassification, AutoImageProcessor, AutoConfig
 import torch.nn.functional as F
 import json
 import re
+import gc
+import sys
+import traceback
 
-class SimpleGradCAM:
+# Gestion des dépendances optionnelles
+HAS_MATPLOTLIB = False
+try:
+    import matplotlib.pyplot as plt
+    import matplotlib.cm as cm
+    HAS_MATPLOTLIB = True
+    print("✅ Matplotlib disponible - Grad-CAM avancé activé")
+except ImportError:
+    print("⚠️  Matplotlib non disponible - Utilisation de PIL pour Grad-CAM")
+
+class OptimizedGradCAM:
+    """Version optimisée de Grad-CAM avec nettoyage automatique"""
+    
     def __init__(self, model, target_layer):
         self.model = model
         self.target_layer = target_layer
         self.gradients = None
         self.activations = None
+        self.hooks = []
         
-        # Hook pour capturer les gradients
-        self.target_layer.register_backward_hook(self.save_gradients)
-        self.target_layer.register_forward_hook(self.save_activations)
+        # Enregistrer les hooks avec nettoyage automatique
+        if target_layer is not None:
+            hook1 = self.target_layer.register_backward_hook(self.save_gradients)
+            hook2 = self.target_layer.register_forward_hook(self.save_activations)
+            self.hooks = [hook1, hook2]
+        else:
+            print("⚠️  Aucune couche cible trouvée - Grad-CAM désactivé")
     
     def save_gradients(self, module, grad_input, grad_output):
-        self.gradients = grad_output[0]
+        if grad_output[0] is not None:
+            self.gradients = grad_output[0].detach()
     
     def save_activations(self, module, input, output):
-        self.activations = output
+        self.activations = output.detach()
     
     def generate_cam(self, input_tensor, class_idx=None):
-        # Forward pass
-        output = self.model(input_tensor)
+        """Génère la carte de saillance Grad-CAM"""
+        if self.target_layer is None:
+            return None
+            
+        try:
+            # Forward pass
+            output = self.model(input_tensor)
+            
+            if class_idx is None:
+                class_idx = output.logits.argmax(dim=1).item()
+            
+            # Backward pass
+            self.model.zero_grad()
+            output.logits[0, class_idx].backward(retain_graph=False)
+            
+            if self.gradients is None or self.activations is None:
+                print("⚠️  Gradients ou activations manquants")
+                return None
+            
+            # Generate CAM
+            gradients = self.gradients[0]  # (C, H, W)
+            activations = self.activations[0]  # (C, H, W)
+            
+            # Moyenne globale des gradients
+            weights = torch.mean(gradients, dim=(1, 2))  # (C,)
+            
+            # CAM = somme pondérée des activations
+            cam = torch.zeros(activations.shape[1:], device=activations.device)  # (H, W)
+            for i, w in enumerate(weights):
+                cam += w * activations[i, :, :]
+            
+            # ReLU et normalisation
+            cam = F.relu(cam)
+            if cam.max() > 0:
+                cam = cam / cam.max()
+            
+            return cam.detach().cpu().numpy()
+            
+        except Exception as e:
+            print(f"⚠️  Erreur lors de la génération CAM: {e}")
+            return None
+        finally:
+            # Nettoyage explicite
+            if self.gradients is not None:
+                self.gradients = None
+            if self.activations is not None:
+                self.activations = None
+    
+    def cleanup(self):
+        """Nettoie les hooks et libère la mémoire"""
+        for hook in self.hooks:
+            try:
+                hook.remove()
+            except:
+                pass
+        self.hooks = []
+        self.gradients = None
+        self.activations = None
+    
+    def __del__(self):
+        """Nettoyage automatique lors de la destruction"""
+        self.cleanup()
+
+def get_last_conv_layer_safe(model):
+    """Trouve la dernière couche de convolution de manière sécurisée"""
+    try:
+        last_conv = None
+        conv_layers = []
         
-        if class_idx is None:
-            class_idx = output.logits.argmax(dim=1).item()
+        for name, module in model.named_modules():
+            if isinstance(module, (torch.nn.Conv2d, torch.nn.AdaptiveAvgPool2d)):
+                conv_layers.append((name, module))
+        
+        if conv_layers:
+            last_conv = conv_layers[-1][1]
+            print(f"✅ Couche cible trouvée: {conv_layers[-1][0]}")
+        else:
+            print("⚠️  Aucune couche de convolution trouvée")
+            
+        return last_conv
+    except Exception as e:
+        print(f"⚠️  Erreur lors de la recherche de couche: {e}")
+        return None
+
+def create_gradcam_overlay_pil(original_image, cam_array):
+    """Crée une superposition Grad-CAM en utilisant PIL (sans matplotlib)"""
+    try:
+        if cam_array is None:
+            return None
+            
+        # Convertir CAM en image
+        cam_normalized = (cam_array * 255).astype(np.uint8)
+        cam_img = Image.fromarray(cam_normalized, mode='L')
         
-        # Backward pass
-        self.model.zero_grad()
-        output.logits[0, class_idx].backward()
+        # Redimensionner au format de l'image originale
+        cam_resized = cam_img.resize(original_image.size, Image.Resampling.LANCZOS)
         
-        # Generate CAM
-        gradients = self.gradients[0]  # (C, H, W)
-        activations = self.activations[0]  # (C, H, W)
+        # Créer une heatmap colorée (rouge pour les zones importantes)
+        # Convertir en RGB et appliquer une colormap simple
+        cam_array_resized = np.array(cam_resized)
         
-        # Moyenne globale des gradients
-        weights = torch.mean(gradients, dim=(1, 2))  # (C,)
+        # Créer une colormap simple (bleu -> rouge)
+        heatmap = np.zeros((cam_array_resized.shape[0], cam_array_resized.shape[1], 3), dtype=np.uint8)
+        heatmap[:, :, 0] = cam_array_resized  # Rouge
+        heatmap[:, :, 2] = 255 - cam_array_resized  # Bleu inversé
         
-        # CAM = somme pondérée des activations
-        cam = torch.zeros(activations.shape[1:])  # (H, W)
-        for i, w in enumerate(weights):
-            cam += w * activations[i, :, :]
+        heatmap_img = Image.fromarray(heatmap, 'RGB')
         
-        # ReLU et normalisation
-        cam = F.relu(cam)
-        cam = cam / cam.max() if cam.max() > 0 else cam
+        # Mélanger avec l'image originale
+        blended = Image.blend(original_image.convert('RGB'), heatmap_img, alpha=0.4)
         
-        return cam.detach().cpu().numpy()
+        # Convertir en base64
+        buffer = io.BytesIO()
+        blended.save(buffer, format='PNG', optimize=True)
+        buffer.seek(0)
+        
+        return base64.b64encode(buffer.getvalue()).decode('utf-8')
+        
+    except Exception as e:
+        print(f"⚠️  Erreur lors de la création de l'overlay PIL: {e}")
+        return None
 
-def get_last_conv_layer(model):
-    """Trouve la dernière couche de convolution du modèle"""
-    last_conv = None
-    for name, module in model.named_modules():
-        if isinstance(module, (torch.nn.Conv2d, torch.nn.AdaptiveAvgPool2d)):
-            last_conv = module
-    return last_conv
+def create_gradcam_overlay_matplotlib(original_image, cam_array):
+    """Crée une superposition Grad-CAM en utilisant matplotlib (si disponible)"""
+    try:
+        if not HAS_MATPLOTLIB or cam_array is None:
+            return None
+            
+        # Redimensionner CAM
+        cam_resized = np.array(Image.fromarray((cam_array * 255).astype(np.uint8)).resize(
+            original_image.size, Image.Resampling.LANCZOS
+        )) / 255.0
+        
+        # Créer la figure
+        fig, ax = plt.subplots(figsize=(8, 8), dpi=100)
+        ax.imshow(original_image)
+        ax.imshow(cam_resized, cmap='jet', alpha=0.5)
+        ax.axis('off')
+        
+        # Sauvegarder en base64
+        buffer = io.BytesIO()
+        plt.savefig(buffer, format='png', bbox_inches='tight', pad_inches=0, dpi=100)
+        plt.close(fig)  # Important: fermer la figure
+        buffer.seek(0)
+        
+        return base64.b64encode(buffer.getvalue()).decode('utf-8')
+        
+    except Exception as e:
+        print(f"⚠️  Erreur lors de la création de l'overlay matplotlib: {e}")
+        if 'fig' in locals():
+            plt.close(fig)
+        return None
 
 class EndpointHandler:
     def __init__(self, path=""):
-        print("🚀 VerifAI Handler V2 - Initialisation")
-        print("📋 Modèle: haywoodsloan/ai-image-detector-deploy (Haute Précision)")
+        print("🚀 VerifAI Handler V2 FIXED - Initialisation")
+        print("📋 Modèle: haywoodsloan/ai-image-detector-deploy (Version Corrigée)")
+        
+        self.model = None
+        self.processor = None
+        self.grad_cam = None
+        self.model_labels = {}
         
         try:
-            # Modèle unique : Haywoodsloan AI Image Detector
-            print("🔄 Chargement du modèle: haywoodsloan/ai-image-detector-deploy")
+            # Vérification de la disponibilité du modèle
             self.model_name = "haywoodsloan/ai-image-detector-deploy"
+            
+            if not self._verify_model_exists():
+                raise Exception(f"Modèle {self.model_name} non accessible")
+            
+            # Chargement du modèle avec gestion d'erreurs
+            print("🔄 Chargement du modèle...")
             self.processor = AutoImageProcessor.from_pretrained(self.model_name)
-            self.model = AutoModelForImageClassification.from_pretrained(self.model_name)
+            self.model = AutoModelForImageClassification.from_pretrained(
+                self.model_name,
+                torch_dtype=torch.float32  # Force float32 pour la compatibilité
+            )
             self.model.eval()
             
-            # Configuration Grad-CAM
-            target_layer = get_last_conv_layer(self.model)
-            self.grad_cam = SimpleGradCAM(self.model, target_layer) if target_layer else None
-            self.model_labels = self.model.config.id2label
+            # Configuration Grad-CAM sécurisée
+            target_layer = get_last_conv_layer_safe(self.model)
+            if target_layer is not None:
+                self.grad_cam = OptimizedGradCAM(self.model, target_layer)
+                print("✅ Grad-CAM activé")
+            else:
+                print("⚠️  Grad-CAM désactivé (aucune couche compatible)")
+            
+            # Récupérer les labels
+            if hasattr(self.model.config, 'id2label'):
+                self.model_labels = self.model.config.id2label
+            else:
+                self.model_labels = {0: "Real", 1: "Fake"}  # Fallback
             
             print("✅ Modèle chargé avec succès")
             print(f"📋 Étiquettes du modèle: {self.model_labels}")
-            print("🎯 VerifAI Handler V2 prêt!")
+            print("🎯 VerifAI Handler V2 FIXED prêt!")
             
         except Exception as e:
-            print(f"❌ Erreur lors du chargement: {e}")
-            raise e
-
+            print(f"❌ Erreur lors de l'initialisation: {e}")
+            print(f"🔍 Traceback: {traceback.format_exc()}")
+            # Ne pas faire échouer l'initialisation, mais signaler l'erreur
+            self.model = None
+            self.processor = None
+    
+    def _verify_model_exists(self):
+        """Vérifie que le modèle existe avant de le charger"""
+        try:
+            config = AutoConfig.from_pretrained(self.model_name)
+            print(f"✅ Modèle {self.model_name} vérifié")
+            return True
+        except Exception as e:
+            print(f"❌ Modèle {self.model_name} non accessible: {e}")
+            return False
+    
     def _normalize_label(self, label: str) -> str:
         """Normalise les étiquettes pour qu'elles soient cohérentes."""
+        if not isinstance(label, str):
+            label = str(label)
+            
         label_lower = label.lower()
         if re.search(r'real|human|authentic', label_lower):
             return "Human"
         if re.search(r'fake|generated|ai|artificial', label_lower):
             return "AI Generated"
         return "Unknown"
+    
+    def _cleanup_memory(self):
+        """Nettoie la mémoire explicitement"""
+        try:
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            gc.collect()
+        except:
+            pass
+    
+    def __call__(self, data):
+        # Vérification de l'état du handler
+        if self.model is None or self.processor is None:
+            return {
+                "status": "error",
+                "error": "Handler non initialisé correctement",
+                "prediction": 0,
+                "predicted_class_name": "Error",
+                "confidence": 0.0,
+                "class_probabilities": {"Human": 0.0, "AI Generated": 0.0},
+                "cam_image": None,
+                "version": "2.0-fixed",
+                "handler_name": "VerifAI Handler V2 FIXED"
+            }
+        
+        try:
+            # Traitement de l'image avec validation
+            image_data = data.get("inputs") or data
+            if not image_data:
+                raise ValueError("Aucune donnée d'image fournie")
+            
+            # Décodage sécurisé de l'image
+            try:
+                image_bytes = base64.b64decode(image_data)
+                image = Image.open(io.BytesIO(image_bytes))
+                
+                # Validation et conversion
+                if image.mode != 'RGB':
+                    image = image.convert('RGB')
+                    
+                # Validation de la taille
+                if image.size[0] * image.size[1] > 4096 * 4096:
+                    image = image.resize((1024, 1024), Image.Resampling.LANCZOS)
+                    print("⚠️  Image redimensionnée pour éviter les problèmes de mémoire")
+                    
+            except Exception as e:
+                raise ValueError(f"Erreur lors du décodage de l'image: {e}")
+
+            # Prédiction avec gestion d'erreurs
+            print("🔄 VerifAI V2 FIXED - Analyse en cours...")
+            
+            try:
+                inputs = self.processor(image, return_tensors="pt")
+                
+                with torch.no_grad():
+                    outputs = self.model(**inputs)
+                    logits = outputs.logits
+                    probabilities = F.softmax(logits, dim=-1)[0]
+                    predicted_class_id = logits.argmax().item()
+                    
+            except Exception as e:
+                raise RuntimeError(f"Erreur lors de l'inférence: {e}")
+            
+            # Traitement des résultats
+            class_probs = {}
+            for class_id, prob in enumerate(probabilities):
+                label_str = self.model_labels.get(class_id, f"Class {class_id}")
+                normalized_label = self._normalize_label(label_str)
+                if normalized_label != "Unknown":
+                    class_probs[normalized_label] = float(prob)
+
+            # S'assurer que les deux classes existent
+            class_probs.setdefault("Human", 0.0)
+            class_probs.setdefault("AI Generated", 0.0)
+
+            prediction_label = self._normalize_label(self.model_labels.get(predicted_class_id, "Unknown"))
+            confidence = class_probs.get(prediction_label, 0.0)
+            
+            # Déterminer l'ID de prédiction pour la compatibilité
+            prediction_id = 1 if prediction_label == "AI Generated" else 0
+            
+            print(f"🔍 VerifAI V2 FIXED Résultat: {prediction_label} (confiance: {confidence:.3f})")
 
+            # Génération du Grad-CAM avec fallback
+            cam_image_b64 = None
+            if self.grad_cam is not None:
+                try:
+                    print("🎨 Génération du Grad-CAM...")
+                    cam = self.grad_cam.generate_cam(inputs['pixel_values'], predicted_class_id)
+                    
+                    if cam is not None:
+                        # Essayer matplotlib d'abord, puis PIL
+                        if HAS_MATPLOTLIB:
+                            cam_image_b64 = create_gradcam_overlay_matplotlib(image, cam)
+                        
+                        if cam_image_b64 is None:
+                            cam_image_b64 = create_gradcam_overlay_pil(image, cam)
+                        
+                        if cam_image_b64:
+                            print("✅ Grad-CAM généré avec succès")
+                        else:
+                            print("⚠️  Échec de la génération Grad-CAM")
+                    
+                except Exception as e:
+                    print(f"⚠️  Erreur Grad-CAM: {e}")
+                    cam_image_b64 = None
+
+            # Nettoyage mémoire
+            self._cleanup_memory()
+
+            # Construction de la réponse compatible
+            return {
+                "status": "success",
+                "prediction": prediction_id,
+                "predicted_class_name": prediction_label,
+                "confidence": confidence,
+                "class_probabilities": class_probs,
+                "cam_image": cam_image_b64,
+                "model_info": {
+                    "model_name": self.model_name,
+                    "handler_version": "verifai-v2-fixed",
+                    "precision_mode": "high",
+                    "raw_prediction_id": predicted_class_id,
+                    "raw_labels": self.model_labels,
+                    "grad_cam_method": "matplotlib" if HAS_MATPLOTLIB else "pil"
+                },
+                "reliability": "TRÈS ÉLEVÉE",
+                "version": "2.0-fixed",
+                "handler_name": "VerifAI Handler V2 FIXED",
+                "deployment_note": "VERIFAI HANDLER V2 FIXED - PRODUCTION READY",
+                "fixes_applied": [
+                    "Gestion d'erreurs robuste",
+                    "Fallback PIL pour Grad-CAM",
+                    "Nettoyage mémoire automatique",
+                    "Validation d'entrée renforcée"
+                ]
+            }
+
+        except Exception as e:
+            print(f"❌ Erreur dans VerifAI Handler V2 FIXED: {e}")
+            print(f"🔍 Traceback: {traceback.format_exc()}")
+            
+            # Nettoyage en cas d'erreur
+            self._cleanup_memory()
+            
+            return {
+                "status": "error",
+                "error": str(e),
+                "prediction": 0,
+                "predicted_class_name": "Error",
+                "confidence": 0.0,
+                "class_probabilities": {"Human": 0.0, "AI Generated": 0.0},
+                "cam_image": None,
+                "version": "2.0-fixed",
+                "handler_name": "VerifAI Handler V2 FIXED",
+                "error_details": {
+                    "error_type": type(e).__name__,
+                    "traceback": traceback.format_exc()[-500:]  # Dernières 500 chars
+                }
+            }
+    
+    def __del__(self):
+        """Nettoyage lors de la destruction de l'instance"""
+        try:
+            if hasattr(self, 'grad_cam') and self.grad_cam is not None:
+                self.grad_cam.cleanup()
+            self._cleanup_memory()
+        except:
+            pass
+
+# Test de fonctionnement si exécuté directement
+if __name__ == "__main__":
+    print("🧪 TEST DU HANDLER VERIFAI V2 FIXED")
+    print("=" * 50)
+    
+    try:
+        # Initialisation
+        handler = EndpointHandler()
+        
+        if handler.model is not None:
+            print("✅ Initialisation réussie")
+            
+            # Test avec une image simple
+            print("🔄 Test avec image de base...")
+            test_img = Image.new('RGB', (224, 224), color='red')
+            buffer = io.BytesIO()
+            test_img.save(buffer, format='JPEG')
+            test_data = base64.b64encode(buffer.getvalue()).decode('utf-8')
+            
+            result = handler({"inputs": test_data})
+            print(f"📊 Résultat: {result['status']}")
+            if result['status'] == 'success':
+                print(f"🎯 Prédiction: {result['predicted_class_name']} ({result['confidence']:.3f})")
+                print("✅ Handler fonctionnel!")
+            else:
+                print(f"❌ Erreur: {result.get('error', 'Inconnue')}")
+        else:
+            print("❌ Échec de l'initialisation")
+            
+    except Exception as e:
+        print(f"❌ Erreur de test: {e}")
+        print(f"🔍 Traceback: {traceback.format_exc()}") from typing import Any, Dict
+import torch
+from torchvision import transforms
+from PIL import Image, ImageDraw, ImageOps
+import base64
+import io
+import numpy as np
+from transformers import AutoModelForImageClassification, AutoImageProcessor, AutoConfig
+import torch.nn.functional as F
+import json
+import re
+import gc
+import sys
+import traceback
+
+# Gestion des dépendances optionnelles
+HAS_MATPLOTLIB = False
+try:
+    import matplotlib.pyplot as plt
+    import matplotlib.cm as cm
+    HAS_MATPLOTLIB = True
+    print("✅ Matplotlib disponible - Grad-CAM avancé activé")
+except ImportError:
+    print("⚠️  Matplotlib non disponible - Utilisation de PIL pour Grad-CAM")
+
+class OptimizedGradCAM:
+    """Version optimisée de Grad-CAM avec nettoyage automatique"""
+    
+    def __init__(self, model, target_layer):
+        self.model = model
+        self.target_layer = target_layer
+        self.gradients = None
+        self.activations = None
+        self.hooks = []
+        
+        # Enregistrer les hooks avec nettoyage automatique
+        if target_layer is not None:
+            hook1 = self.target_layer.register_backward_hook(self.save_gradients)
+            hook2 = self.target_layer.register_forward_hook(self.save_activations)
+            self.hooks = [hook1, hook2]
+        else:
+            print("⚠️  Aucune couche cible trouvée - Grad-CAM désactivé")
+    
+    def save_gradients(self, module, grad_input, grad_output):
+        if grad_output[0] is not None:
+            self.gradients = grad_output[0].detach()
+    
+    def save_activations(self, module, input, output):
+        self.activations = output.detach()
+    
+    def generate_cam(self, input_tensor, class_idx=None):
+        """Génère la carte de saillance Grad-CAM"""
+        if self.target_layer is None:
+            return None
+            
+        try:
+            # Forward pass
+            output = self.model(input_tensor)
+            
+            if class_idx is None:
+                class_idx = output.logits.argmax(dim=1).item()
+            
+            # Backward pass
+            self.model.zero_grad()
+            output.logits[0, class_idx].backward(retain_graph=False)
+            
+            if self.gradients is None or self.activations is None:
+                print("⚠️  Gradients ou activations manquants")
+                return None
+            
+            # Generate CAM
+            gradients = self.gradients[0]  # (C, H, W)
+            activations = self.activations[0]  # (C, H, W)
+            
+            # Moyenne globale des gradients
+            weights = torch.mean(gradients, dim=(1, 2))  # (C,)
+            
+            # CAM = somme pondérée des activations
+            cam = torch.zeros(activations.shape[1:], device=activations.device)  # (H, W)
+            for i, w in enumerate(weights):
+                cam += w * activations[i, :, :]
+            
+            # ReLU et normalisation
+            cam = F.relu(cam)
+            if cam.max() > 0:
+                cam = cam / cam.max()
+            
+            return cam.detach().cpu().numpy()
+            
+        except Exception as e:
+            print(f"⚠️  Erreur lors de la génération CAM: {e}")
+            return None
+        finally:
+            # Nettoyage explicite
+            if self.gradients is not None:
+                self.gradients = None
+            if self.activations is not None:
+                self.activations = None
+    
+    def cleanup(self):
+        """Nettoie les hooks et libère la mémoire"""
+        for hook in self.hooks:
+            try:
+                hook.remove()
+            except:
+                pass
+        self.hooks = []
+        self.gradients = None
+        self.activations = None
+    
+    def __del__(self):
+        """Nettoyage automatique lors de la destruction"""
+        self.cleanup()
+
+def get_last_conv_layer_safe(model):
+    """Trouve la dernière couche de convolution de manière sécurisée"""
+    try:
+        last_conv = None
+        conv_layers = []
+        
+        for name, module in model.named_modules():
+            if isinstance(module, (torch.nn.Conv2d, torch.nn.AdaptiveAvgPool2d)):
+                conv_layers.append((name, module))
+        
+        if conv_layers:
+            last_conv = conv_layers[-1][1]
+            print(f"✅ Couche cible trouvée: {conv_layers[-1][0]}")
+        else:
+            print("⚠️  Aucune couche de convolution trouvée")
+            
+        return last_conv
+    except Exception as e:
+        print(f"⚠️  Erreur lors de la recherche de couche: {e}")
+        return None
+
+def create_gradcam_overlay_pil(original_image, cam_array):
+    """Crée une superposition Grad-CAM en utilisant PIL (sans matplotlib)"""
+    try:
+        if cam_array is None:
+            return None
+            
+        # Convertir CAM en image
+        cam_normalized = (cam_array * 255).astype(np.uint8)
+        cam_img = Image.fromarray(cam_normalized, mode='L')
+        
+        # Redimensionner au format de l'image originale
+        cam_resized = cam_img.resize(original_image.size, Image.Resampling.LANCZOS)
+        
+        # Créer une heatmap colorée (rouge pour les zones importantes)
+        # Convertir en RGB et appliquer une colormap simple
+        cam_array_resized = np.array(cam_resized)
+        
+        # Créer une colormap simple (bleu -> rouge)
+        heatmap = np.zeros((cam_array_resized.shape[0], cam_array_resized.shape[1], 3), dtype=np.uint8)
+        heatmap[:, :, 0] = cam_array_resized  # Rouge
+        heatmap[:, :, 2] = 255 - cam_array_resized  # Bleu inversé
+        
+        heatmap_img = Image.fromarray(heatmap, 'RGB')
+        
+        # Mélanger avec l'image originale
+        blended = Image.blend(original_image.convert('RGB'), heatmap_img, alpha=0.4)
+        
+        # Convertir en base64
+        buffer = io.BytesIO()
+        blended.save(buffer, format='PNG', optimize=True)
+        buffer.seek(0)
+        
+        return base64.b64encode(buffer.getvalue()).decode('utf-8')
+        
+    except Exception as e:
+        print(f"⚠️  Erreur lors de la création de l'overlay PIL: {e}")
+        return None
+
+def create_gradcam_overlay_matplotlib(original_image, cam_array):
+    """Crée une superposition Grad-CAM en utilisant matplotlib (si disponible)"""
+    try:
+        if not HAS_MATPLOTLIB or cam_array is None:
+            return None
+            
+        # Redimensionner CAM
+        cam_resized = np.array(Image.fromarray((cam_array * 255).astype(np.uint8)).resize(
+            original_image.size, Image.Resampling.LANCZOS
+        )) / 255.0
+        
+        # Créer la figure
+        fig, ax = plt.subplots(figsize=(8, 8), dpi=100)
+        ax.imshow(original_image)
+        ax.imshow(cam_resized, cmap='jet', alpha=0.5)
+        ax.axis('off')
+        
+        # Sauvegarder en base64
+        buffer = io.BytesIO()
+        plt.savefig(buffer, format='png', bbox_inches='tight', pad_inches=0, dpi=100)
+        plt.close(fig)  # Important: fermer la figure
+        buffer.seek(0)
+        
+        return base64.b64encode(buffer.getvalue()).decode('utf-8')
+        
+    except Exception as e:
+        print(f"⚠️  Erreur lors de la création de l'overlay matplotlib: {e}")
+        if 'fig' in locals():
+            plt.close(fig)
+        return None
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        print("🚀 VerifAI Handler V2 FIXED - Initialisation")
+        print("📋 Modèle: haywoodsloan/ai-image-detector-deploy (Version Corrigée)")
+        
+        self.model = None
+        self.processor = None
+        self.grad_cam = None
+        self.model_labels = {}
+        
+        try:
+            # Vérification de la disponibilité du modèle
+            self.model_name = "haywoodsloan/ai-image-detector-deploy"
+            
+            if not self._verify_model_exists():
+                raise Exception(f"Modèle {self.model_name} non accessible")
+            
+            # Chargement du modèle avec gestion d'erreurs
+            print("🔄 Chargement du modèle...")
+            self.processor = AutoImageProcessor.from_pretrained(self.model_name)
+            self.model = AutoModelForImageClassification.from_pretrained(
+                self.model_name,
+                torch_dtype=torch.float32  # Force float32 pour la compatibilité
+            )
+            self.model.eval()
+            
+            # Configuration Grad-CAM sécurisée
+            target_layer = get_last_conv_layer_safe(self.model)
+            if target_layer is not None:
+                self.grad_cam = OptimizedGradCAM(self.model, target_layer)
+                print("✅ Grad-CAM activé")
+            else:
+                print("⚠️  Grad-CAM désactivé (aucune couche compatible)")
+            
+            # Récupérer les labels
+            if hasattr(self.model.config, 'id2label'):
+                self.model_labels = self.model.config.id2label
+            else:
+                self.model_labels = {0: "Real", 1: "Fake"}  # Fallback
+            
+            print("✅ Modèle chargé avec succès")
+            print(f"📋 Étiquettes du modèle: {self.model_labels}")
+            print("🎯 VerifAI Handler V2 FIXED prêt!")
+            
+        except Exception as e:
+            print(f"❌ Erreur lors de l'initialisation: {e}")
+            print(f"🔍 Traceback: {traceback.format_exc()}")
+            # Ne pas faire échouer l'initialisation, mais signaler l'erreur
+            self.model = None
+            self.processor = None
+    
+    def _verify_model_exists(self):
+        """Vérifie que le modèle existe avant de le charger"""
+        try:
+            config = AutoConfig.from_pretrained(self.model_name)
+            print(f"✅ Modèle {self.model_name} vérifié")
+            return True
+        except Exception as e:
+            print(f"❌ Modèle {self.model_name} non accessible: {e}")
+            return False
+    
+    def _normalize_label(self, label: str) -> str:
+        """Normalise les étiquettes pour qu'elles soient cohérentes."""
+        if not isinstance(label, str):
+            label = str(label)
+            
+        label_lower = label.lower()
+        if re.search(r'real|human|authentic', label_lower):
+            return "Human"
+        if re.search(r'fake|generated|ai|artificial', label_lower):
+            return "AI Generated"
+        return "Unknown"
+    
+    def _cleanup_memory(self):
+        """Nettoie la mémoire explicitement"""
+        try:
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            gc.collect()
+        except:
+            pass
+    
     def __call__(self, data):
+        # Vérification de l'état du handler
+        if self.model is None or self.processor is None:
+            return {
+                "status": "error",
+                "error": "Handler non initialisé correctement",
+                "prediction": 0,
+                "predicted_class_name": "Error",
+                "confidence": 0.0,
+                "class_probabilities": {"Human": 0.0, "AI Generated": 0.0},
+                "cam_image": None,
+                "version": "2.0-fixed",
+                "handler_name": "VerifAI Handler V2 FIXED"
+            }
+        
         try:
-            # Traitement de l'image
+            # Traitement de l'image avec validation
             image_data = data.get("inputs") or data
-            image_bytes = base64.b64decode(image_data)
-            image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
+            if not image_data:
+                raise ValueError("Aucune donnée d'image fournie")
+            
+            # Décodage sécurisé de l'image
+            try:
+                image_bytes = base64.b64decode(image_data)
+                image = Image.open(io.BytesIO(image_bytes))
+                
+                # Validation et conversion
+                if image.mode != 'RGB':
+                    image = image.convert('RGB')
+                    
+                # Validation de la taille
+                if image.size[0] * image.size[1] > 4096 * 4096:
+                    image = image.resize((1024, 1024), Image.Resampling.LANCZOS)
+                    print("⚠️  Image redimensionnée pour éviter les problèmes de mémoire")
+                    
+            except Exception as e:
+                raise ValueError(f"Erreur lors du décodage de l'image: {e}")
 
-            # Prédiction avec le modèle unique
-            print("🔄 VerifAI V2 - Analyse en cours...")
-            inputs = self.processor(image, return_tensors="pt")
+            # Prédiction avec gestion d'erreurs
+            print("🔄 VerifAI V2 FIXED - Analyse en cours...")
             
-            with torch.no_grad():
-                outputs = self.model(**inputs)
-                logits = outputs.logits
-                probabilities = F.softmax(logits, dim=-1)[0]
-                predicted_class_id = logits.argmax().item()
+            try:
+                inputs = self.processor(image, return_tensors="pt")
+                
+                with torch.no_grad():
+                    outputs = self.model(**inputs)
+                    logits = outputs.logits
+                    probabilities = F.softmax(logits, dim=-1)[0]
+                    predicted_class_id = logits.argmax().item()
+                    
+            except Exception as e:
+                raise RuntimeError(f"Erreur lors de l'inférence: {e}")
             
-            # Normaliser les probabilités
+            # Traitement des résultats
             class_probs = {}
             for class_id, prob in enumerate(probabilities):
                 label_str = self.model_labels.get(class_id, f"Class {class_id}")
@@ -128,40 +826,42 @@ class EndpointHandler:
             class_probs.setdefault("Human", 0.0)
             class_probs.setdefault("AI Generated", 0.0)
 
-            prediction_label = self._normalize_label(self.model_labels.get(predicted_class_id))
+            prediction_label = self._normalize_label(self.model_labels.get(predicted_class_id, "Unknown"))
             confidence = class_probs.get(prediction_label, 0.0)
             
             # Déterminer l'ID de prédiction pour la compatibilité
             prediction_id = 1 if prediction_label == "AI Generated" else 0
             
-            print(f"🔍 VerifAI V2 Résultat: {prediction_label} (confiance: {confidence:.3f})")
+            print(f"🔍 VerifAI V2 FIXED Résultat: {prediction_label} (confiance: {confidence:.3f})")
 
-            # Génération du Grad-CAM
+            # Génération du Grad-CAM avec fallback
             cam_image_b64 = None
-            if self.grad_cam:
+            if self.grad_cam is not None:
                 try:
-                    print("🎨 Génération du Grad-CAM VerifAI...")
-                    cam = self.grad_cam.generate_cam(
-                        inputs['pixel_values'], 
-                        predicted_class_id
-                    )
-                    cam_resized = np.array(Image.fromarray((cam * 255).astype(np.uint8)).resize(image.size))
+                    print("🎨 Génération du Grad-CAM...")
+                    cam = self.grad_cam.generate_cam(inputs['pixel_values'], predicted_class_id)
                     
-                    import matplotlib.pyplot as plt
-                    plt.figure(figsize=(10, 5))
-                    plt.imshow(image)
-                    plt.imshow(cam_resized, cmap='jet', alpha=0.5)
-                    plt.axis('off')
+                    if cam is not None:
+                        # Essayer matplotlib d'abord, puis PIL
+                        if HAS_MATPLOTLIB:
+                            cam_image_b64 = create_gradcam_overlay_matplotlib(image, cam)
+                        
+                        if cam_image_b64 is None:
+                            cam_image_b64 = create_gradcam_overlay_pil(image, cam)
+                        
+                        if cam_image_b64:
+                            print("✅ Grad-CAM généré avec succès")
+                        else:
+                            print("⚠️  Échec de la génération Grad-CAM")
                     
-                    buf = io.BytesIO()
-                    plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
-                    buf.seek(0)
-                    cam_image_b64 = base64.b64encode(buf.read()).decode('utf-8')
-                    plt.close()
-                    print("✅ Grad-CAM VerifAI généré avec succès")
                 except Exception as e:
-                    print(f"⚠️ Erreur Grad-CAM: {e}")
+                    print(f"⚠️  Erreur Grad-CAM: {e}")
+                    cam_image_b64 = None
 
+            # Nettoyage mémoire
+            self._cleanup_memory()
+
+            # Construction de la réponse compatible
             return {
                 "status": "success",
                 "prediction": prediction_id,
@@ -171,27 +871,85 @@ class EndpointHandler:
                 "cam_image": cam_image_b64,
                 "model_info": {
                     "model_name": self.model_name,
-                    "handler_version": "verifai-v2",
+                    "handler_version": "verifai-v2-fixed",
                     "precision_mode": "high",
                     "raw_prediction_id": predicted_class_id,
-                    "raw_labels": self.model_labels
+                    "raw_labels": self.model_labels,
+                    "grad_cam_method": "matplotlib" if HAS_MATPLOTLIB else "pil"
                 },
                 "reliability": "TRÈS ÉLEVÉE",
-                "version": "2.0",
-                "handler_name": "VerifAI Handler V2",
-                "deployment_note": "VERIFAI HANDLER V2 - HAUTE PRÉCISION"
+                "version": "2.0-fixed",
+                "handler_name": "VerifAI Handler V2 FIXED",
+                "deployment_note": "VERIFAI HANDLER V2 FIXED - PRODUCTION READY",
+                "fixes_applied": [
+                    "Gestion d'erreurs robuste",
+                    "Fallback PIL pour Grad-CAM",
+                    "Nettoyage mémoire automatique",
+                    "Validation d'entrée renforcée"
+                ]
             }
 
         except Exception as e:
-            print(f"❌ Erreur dans VerifAI Handler V2: {e}")
+            print(f"❌ Erreur dans VerifAI Handler V2 FIXED: {e}")
+            print(f"🔍 Traceback: {traceback.format_exc()}")
+            
+            # Nettoyage en cas d'erreur
+            self._cleanup_memory()
+            
             return {
                 "status": "error",
                 "error": str(e),
                 "prediction": 0,
                 "predicted_class_name": "Error",
                 "confidence": 0.0,
-                "class_probabilities": {},
+                "class_probabilities": {"Human": 0.0, "AI Generated": 0.0},
                 "cam_image": None,
-                "version": "2.0",
-                "handler_name": "VerifAI Handler V2"
-            } 
+                "version": "2.0-fixed",
+                "handler_name": "VerifAI Handler V2 FIXED",
+                "error_details": {
+                    "error_type": type(e).__name__,
+                    "traceback": traceback.format_exc()[-500:]  # Dernières 500 chars
+                }
+            }
+    
+    def __del__(self):
+        """Nettoyage lors de la destruction de l'instance"""
+        try:
+            if hasattr(self, 'grad_cam') and self.grad_cam is not None:
+                self.grad_cam.cleanup()
+            self._cleanup_memory()
+        except:
+            pass
+
+# Test de fonctionnement si exécuté directement
+if __name__ == "__main__":
+    print("🧪 TEST DU HANDLER VERIFAI V2 FIXED")
+    print("=" * 50)
+    
+    try:
+        # Initialisation
+        handler = EndpointHandler()
+        
+        if handler.model is not None:
+            print("✅ Initialisation réussie")
+            
+            # Test avec une image simple
+            print("🔄 Test avec image de base...")
+            test_img = Image.new('RGB', (224, 224), color='red')
+            buffer = io.BytesIO()
+            test_img.save(buffer, format='JPEG')
+            test_data = base64.b64encode(buffer.getvalue()).decode('utf-8')
+            
+            result = handler({"inputs": test_data})
+            print(f"📊 Résultat: {result['status']}")
+            if result['status'] == 'success':
+                print(f"🎯 Prédiction: {result['predicted_class_name']} ({result['confidence']:.3f})")
+                print("✅ Handler fonctionnel!")
+            else:
+                print(f"❌ Erreur: {result.get('error', 'Inconnue')}")
+        else:
+            print("❌ Échec de l'initialisation")
+            
+    except Exception as e:
+        print(f"❌ Erreur de test: {e}")
+        print(f"🔍 Traceback: {traceback.format_exc()}")