Update handler.py

handler.py

@@ -1,4 +1,4 @@
-from typing import Dict, Any, Generator
+from typing import Dict, Any, List
 from PIL import Image
 import torch
 from transformers import AutoModelForCausalLM, AutoProcessor
@@ -36,14 +36,16 @@ class EndpointHandler:
         image = to_channel_dimension_format(image, ChannelDimension.FIRST)
         return torch.tensor(image)
 
-    def stream_response(self, data: Dict[str, Any]) -> Generator[Dict[str, Any], None, None]:
+    def generate_responses(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        results = []
         image = data.get("inputs")
+
         if isinstance(image, str):
             try:
                 image = Image.open(image)
             except Exception as e:
-                yield {"error": f"Failed to open image: {e}"}
-                return
+                results.append({"error": f"Failed to open image: {e}"})
+                return results
 
         try:
             inputs = self.processor.tokenizer(
@@ -54,14 +56,16 @@ class EndpointHandler:
             inputs["pixel_values"] = self.processor.image_processor([image], transform=self.custom_transform)
             inputs = {k: v.to(self.device) for k, v in inputs.items()}
 
-            for generated_ids in self.model.generate(**inputs, bad_words_ids=self.bad_words_ids, max_length=2048, early_stopping=True, return_dict_in_generate=True, output_scores=True):
-                generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-                yield {"label": generated_text, "score": 1.0}
+            generated_ids = self.model.generate(**inputs, bad_words_ids=self.bad_words_ids, max_length=2048, early_stopping=True)
+            generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+            results.append({"label": generated_text, "score": 1.0})
         
         except torch.cuda.CudaError as e:
-            yield {"error": f"CUDA error: {e}"}
+            results.append({"error": f"CUDA error: {e}"})
         except Exception as e:
-            yield {"error": f"Unexpected error: {e}"}
+            results.append({"error": f"Unexpected error: {e}"})
+
+        return results
 
-    def __call__(self, data: Dict[str, Any]) -> Generator[Dict[str, Any], None, None]:
-        return self.stream_response(data)
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        return self.generate_responses(data)