Update handler.py

handler.py

@@ -1,12 +1,12 @@
-from typing import Dict, List, Any
+from typing import Dict, Any
 import torch
 from transformers import LlavaNextProcessor, LlavaNextForConditionalGeneration, BitsAndBytesConfig
 from PIL import Image
 import requests
 from io import BytesIO
-import re
+import base64
 
-class EndpointHandler():
+class EndpointHandler:
     def __init__(self, path=""):
         # Configuració de la quantització
         quantization_config = BitsAndBytesConfig(
@@ -23,29 +23,40 @@ class EndpointHandler():
             device_map="auto"
         )
 
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        logs = []
+        logs.append("Iniciant processament de la petició.")
+        
         inputs = data.get("inputs")
         if not inputs:
-            return {"error": "Invalid input format. 'inputs' key is missing."}
+            logs.append("Format d'entrada invàlid. Manca la clau 'inputs'.")
+            return {"error": "Invalid input format. 'inputs' key is missing.", "logs": logs}
 
         image_url = inputs.get("url")
+        image_data = inputs.get("image_data")
         prompt = inputs.get("prompt")
-        
-        if not image_url or not prompt:
-            return {"error": "Both 'url' and 'prompt' must be provided in 'inputs'."}
-        
-        # Depuració
-        debug_info = {
-            "stage": "processing input",
-            "image_url": image_url,
-            "prompt": prompt
-        }
+
+        if not prompt:
+            logs.append("S'ha de proporcionar 'prompt' en 'inputs'.")
+            return {"error": "The 'prompt' must be provided in 'inputs'.", "logs": logs}
+
+        if not image_url and not image_data:
+            logs.append("S'ha de proporcionar 'url' o 'image_data' en 'inputs'.")
+            return {"error": "Either 'url' or 'image_data' must be provided in 'inputs'.", "logs": logs}
+
+        logs.append(f"Processant entrada: url={image_url}, image_data={'present' if image_data else 'absent'}, prompt={prompt}")
 
         try:
-            response = requests.get(image_url, stream=True)
-            image = Image.open(response.raw)
+            if image_url:
+                logs.append(f"Carregant imatge des de URL: {image_url}")
+                response = requests.get(image_url, stream=True)
+                image = Image.open(response.raw)
+            elif image_data:
+                logs.append("Carregant imatge des de dades d'imatge en brut.")
+                image = Image.open(BytesIO(base64.b64decode(image_data)))
             
             if image.format == 'PNG':
+                logs.append("Convertint imatge PNG a JPG.")
                 image = image.convert('RGB')
                 buffer = BytesIO()
                 image.save(buffer, format="JPEG")
@@ -53,32 +64,18 @@ class EndpointHandler():
                 image = Image.open(buffer)
 
         except Exception as e:
-            debug_info["stage"] = "loading image"
-            debug_info["error"] = str(e)
-            return debug_info
+            logs.append(f"Error carregant imatge: {str(e)}")
+            return {"error": str(e), "logs": logs}
 
         try:
+            logs.append("Processant imatge amb el model.")
             inputs = self.processor(prompt, image, return_tensors="pt").to("cuda")
             output = self.model.generate(**inputs, max_new_tokens=100)
             result = self.processor.decode(output[0], skip_special_tokens=True)
-            
-            scores = self.extract_scores(result)
-            sorted_scores = sorted(scores.items(), key=lambda item: item[1], reverse=True)
-            return sorted_scores
+            logs.append("Processament complet.")
+            return {"input_prompt": prompt, "model_output": result, "logs": logs}
 
         except Exception as e:
-            debug_info["stage"] = "processing model"
-            debug_info["error"] = str(e)
-            return debug_info
-
-    def extract_scores(self, response):
-        scores = {}
-        result_part = response.split("[/INST]")[-1].strip()
-        pattern = re.compile(r'(\d+)\.\s*(.*?):\s*(\d+)')
-        matches = pattern.findall(result_part)
-        for match in matches:
-            category_number = int(match[0])
-            category_name = match[1].strip()
-            score = int(match[2])
-            scores[category_name] = score
-        return scores
+            logs.append(f"Error processant el model: {str(e)}")
+            return {"error": str(e), "logs": logs}
+