Update handler.py

handler.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Any
+from typing import Dict, Any, Generator
 from PIL import Image
 import torch
 from transformers import AutoModelForCausalLM, AutoProcessor
@@ -8,16 +8,8 @@ from transformers.image_transforms import resize, to_channel_dimension_format
 class EndpointHandler:
     def __init__(self, model_path: str):
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.processor = AutoProcessor.from_pretrained(
-            model_path,
-            # token=api_token
-        )
-        self.model = AutoModelForCausalLM.from_pretrained(
-            model_path,
-            # token=api_token,
-            trust_remote_code=True,
-            torch_dtype=torch.bfloat16,
-        ).to(self.device)
+        self.processor = AutoProcessor.from_pretrained(model_path)
+        self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, torch_dtype=torch.bfloat16).to(self.device)
         self.image_seq_len = self.model.config.perceiver_config.resampler_n_latents
         self.bos_token = self.processor.tokenizer.bos_token
         self.bad_words_ids = self.processor.tokenizer(["<image>", "<fake_token_around_image>"], add_special_tokens=False).input_ids
@@ -44,25 +36,32 @@ class EndpointHandler:
         image = to_channel_dimension_format(image, ChannelDimension.FIRST)
         return torch.tensor(image)
 
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+    def stream_response(self, data: Dict[str, Any]) -> Generator[Dict[str, Any], None, None]:
         image = data.get("inputs")
-
         if isinstance(image, str):
-            image = Image.open(image)
+            try:
+                image = Image.open(image)
+            except Exception as e:
+                yield {"error": f"Failed to open image: {e}"}
+                return
 
-        inputs = self.processor.tokenizer(
-            f"{self.bos_token}<fake_token_around_image>{'<image>' * self.image_seq_len}<fake_token_around_image>",
-            return_tensors="pt",
-            add_special_tokens=False,
-        )
-        inputs["pixel_values"] = self.processor.image_processor([image], transform=self.custom_transform)
-        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+        try:
+            inputs = self.processor.tokenizer(
+                f"{self.bos_token}<fake_token_around_image>{'<image>' * self.image_seq_len}<fake_token_around_image>",
+                return_tensors="pt",
+                add_special_tokens=False,
+            )
+            inputs["pixel_values"] = self.processor.image_processor([image], transform=self.custom_transform)
+            inputs = {k: v.to(self.device) for k, v in inputs.items()}
 
-        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]
-        # print(generated_text)
-        # return {"text": generated_text}
-        # Format the output as an array of dictionaries with 'label' and 'score'
-        output = [{"label": generated_text, "score": 1.0}]
+            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}
+        
+        except torch.cuda.CudaError as e:
+            yield {"error": f"CUDA error: {e}"}
+        except Exception as e:
+            yield {"error": f"Unexpected error: {e}"}
 
-        return output
+    def __call__(self, data: Dict[str, Any]) -> Generator[Dict[str, Any], None, None]:
+        return self.stream_response(data)