Update handler.py

handler.py

@@ -1,40 +1,88 @@
-import requests
+# handler.py
+
 import torch
-from PIL import Image
 from transformers import MllamaForConditionalGeneration, AutoProcessor
+from PIL import Image
+import base64
+import io
+
+class Llama32VisionHandler:
+    def __init__(self):
+        self.model = None
+        self.processor = None
 
-class EndpointHandler:
-    def __init__(self, model_dir):
-        # Initialize the model and processor from the directory
+    def initialize(self):
+        # Cargar el modelo y el procesador
         model_id = "meta-llama/Llama-3.2-90B-Vision-Instruct"
+
         self.model = MllamaForConditionalGeneration.from_pretrained(
             model_id,
-            torch_dtype=torch.bfloat16,
-            device_map="auto"
+            torch_dtype=torch.bfloat16,  # Usar bfloat16 para eficiencia de memoria
+            device_map="auto",  # Mapear automáticamente el modelo a los dispositivos disponibles
         )
         self.processor = AutoProcessor.from_pretrained(model_id)
-    
-    def process(self, inputs):
-        """
-        Process the input data and return the output.
-        Expecting inputs in the form of a dictionary containing 'image_url' and 'prompt'.
-        """
-        image_url = inputs.get("image_url")
-        prompt = inputs.get("prompt", "If I had to write a haiku for this one, it would be:") 
-        
-        # Process the image
-        image = Image.open(requests.get(image_url, stream=True).raw)
-        
-        # Generate response
+        self.model.eval()
+
+    def handle(self, request):
+        # Asegurarse de que el modelo esté cargado
+        if self.model is None:
+            self.initialize()
+
+        # Extraer imagen y texto de la solicitud
+        image_data = request.get('image', None)
+        text_input = request.get('text', '')
+
+        # Procesar la imagen
+        if image_data:
+            # Si los datos de imagen están en formato base64
+            if isinstance(image_data, str):
+                image_bytes = base64.b64decode(image_data)
+                image = Image.open(io.BytesIO(image_bytes))
+            else:
+                # Si los datos de imagen son bytes crudos
+                image = Image.open(io.BytesIO(image_data))
+        else:
+            image = None  # Manejar casos donde no se proporciona imagen
+
+        # Preparar mensajes para el procesador
         messages = [
-            {"role": "user", "content": [
-                {"type": "image"},
-                {"type": "text", "text": prompt}
-            ]}
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image"},
+                    {"type": "text", "text": text_input}
+                ]
+            }
         ]
+
+        # Aplicar la plantilla de chat a los mensajes
         input_text = self.processor.apply_chat_template(messages, add_generation_prompt=True)
-        model_inputs = self.processor(image, input_text, return_tensors="pt").to(self.model.device)
-        output = self.model.generate(**model_inputs, max_new_tokens=30)
-        
-        # Return the output as a string
-        return self.processor.decode(output[0])
+        # Procesar las entradas
+        inputs = self.processor(image, input_text, return_tensors="pt").to(self.model.device)
+
+        # Generar salida
+        with torch.no_grad():
+            outputs = self.model.generate(**inputs, max_new_tokens=50)
+
+        # Decodificar la salida
+        response = self.processor.decode(outputs[0], skip_special_tokens=True)
+        return response
+
+# Ejemplo de uso
+if __name__ == '__main__':
+    handler = Llama32VisionHandler()
+    # Cargar una imagen de ejemplo y codificarla en base64
+    with open('ruta_a_tu_imagen.jpg', 'rb') as f:
+        image_bytes = f.read()
+    image_base64 = base64.b64encode(image_bytes).decode('utf-8')
+
+    # Crear una solicitud de ejemplo
+    request = {
+        'image': image_base64,
+        'text': 'Por favor, describe esta imagen en detalle.'
+    }
+
+    # Obtener la respuesta del handler
+    response = handler.handle(request)
+    print(response)
+#000