Update handler.py

handler.py

@@ -1,18 +1,64 @@
 from typing import Dict, Any
-from PIL import Image, ImageDraw
-import io
+import os
+import requests
+from io import BytesIO
+from PIL import Image
+import torch
+from torchvision import transforms
+from transformers import AutoModelForImageSegmentation
+
+# Configuración
+torch.set_float32_matmul_precision(["high", "highest"][0])
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 class EndpointHandler():
     def __init__(self, path=''):
-        pass
+        # Cargamos el modelo BiRefNet original (Efectivo y rápido)
+        self.model = AutoModelForImageSegmentation.from_pretrained(
+            'zhengpeng7/BiRefNet', 
+            trust_remote_code=True
+        )
+        self.model.to(device)
+        self.model.eval()
+        self.model.half()
 
     def __call__(self, data: Dict[str, Any]):
-        # IGNORAMOS la entrada.
-        # Creamos una imagen VERDE de 800x800 desde cero.
-        img = Image.new('RGB', (800, 800), color=(0, 255, 0)) # Verde Brillante
+        # 1. RECIBIR IMAGEN (Entrada Blindada)
+        image_src = data["inputs"]
+        image = None
+        
+        if isinstance(image_src, Image.Image):
+            image = image_src
+        elif isinstance(image_src, str):
+            if image_src.startswith('http'):
+                image = Image.open(BytesIO(requests.get(image_src).content))
+            else:
+                image = Image.open(image_src)
+        else:
+            image = Image.open(BytesIO(image_src))
+
+        # 2. LIMPIEZA: Aseguramos RGB (Color Real)
+        image = image.convert("RGB")
+        orig_size = image.size
+        
+        # 3. PROCESAMIENTO IA
+        transform = transforms.Compose([
+            transforms.Resize((1024, 1024)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ])
+        
+        input_tensor = transform(image).unsqueeze(0).to(device).half()
+        
+        with torch.no_grad():
+            preds = self.model(input_tensor)[-1].sigmoid().cpu()
+        
+        # 4. MÁSCARA
+        pred = preds[0].squeeze()
+        mask_pil = transforms.ToPILImage()(pred)
+        mask_pil = mask_pil.resize(orig_size, resample=Image.Resampling.LANCZOS)
         
-        # Le escribimos un texto o dibujamos un círculo azul para confirmar
-        draw = ImageDraw.Draw(img)
-        draw.ellipse((200, 200, 600, 600), fill=(0, 0, 255)) # Círculo Azul
+        # 5. APLICACIÓN FINAL (Sin tocar colores)
+        image.putalpha(mask_pil)
         
-        return img
+        return image