Update handler.py

handler.py

@@ -1,88 +1,52 @@
-# handler.py
-
 import torch
 from transformers import MllamaForConditionalGeneration, AutoProcessor
 from PIL import Image
 import base64
 import io
 
-class Llama32VisionHandler:
-    def __init__(self):
-        self.model = None
-        self.processor = None
-
-    def initialize(self):
-        # Cargar el modelo y el procesador
-        model_id = "meta-llama/Llama-3.2-90B-Vision-Instruct"
+# Load model and processor globally
+model_id = "meta-llama/Llama-3.2-90B-Vision-Instruct"
 
-        self.model = MllamaForConditionalGeneration.from_pretrained(
-            model_id,
-            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)
-        self.model.eval()
+model = MllamaForConditionalGeneration.from_pretrained(
+    model_id,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+processor = AutoProcessor.from_pretrained(model_id)
 
-    def handle(self, request):
-        # Asegurarse de que el modelo esté cargado
-        if self.model is None:
-            self.initialize()
+def handler(event, context):
+    try:
+        # Parse inputs
+        inputs = event.get('inputs', {})
+        image_base64 = inputs.get('image')
+        prompt = inputs.get('prompt', '')
 
-        # Extraer imagen y texto de la solicitud
-        image_data = request.get('image', None)
-        text_input = request.get('text', '')
+        if not image_base64 or not prompt:
+            return {'error': 'Both "image" and "prompt" are required in inputs.'}
 
-        # 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
+        # Decode the base64 image
+        image_bytes = base64.b64decode(image_base64)
+        image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
 
-        # Preparar mensajes para el procesador
+        # Prepare the message
         messages = [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image"},
-                    {"type": "text", "text": text_input}
-                ]
-            }
+            {"role": "user", "content": [
+                {"type": "image"},
+                {"type": "text", "text": prompt}
+            ]}
         ]
+        input_text = processor.apply_chat_template(messages, add_generation_prompt=True)
 
-        # Aplicar la plantilla de chat a los mensajes
-        input_text = self.processor.apply_chat_template(messages, add_generation_prompt=True)
-        # 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
+        # Process inputs
+        inputs = processor(image, input_text, return_tensors="pt").to(model.device)
 
-# 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')
+        # Generate output
+        output_ids = model.generate(**inputs, max_new_tokens=50)
+        generated_text = processor.decode(output_ids[0], skip_special_tokens=True)
 
-    # Crear una solicitud de ejemplo
-    request = {
-        'image': image_base64,
-        'text': 'Por favor, describe esta imagen en detalle.'
-    }
+        # Return the result
+        return {'generated_text': generated_text}
 
-    # Obtener la respuesta del handler
-    response = handler.handle(request)
-    print(response)
-#000
+    except Exception as e:
+        return {'error': str(e)}
+#111