Update handler.py

handler.py

@@ -1,6 +1,10 @@
 import subprocess
 import sys
 import torch
+import base64
+from io import BytesIO
+from PIL import Image
+import requests
 from transformers import AutoModelForCausalLM, AutoProcessor
 
 def install(package):
@@ -8,8 +12,7 @@ def install(package):
 
 class EndpointHandler:
     def __init__(self, path=""):
-        
-        required_packages = ['timm', 'einops', 'flash-attn']
+        required_packages = ['timm', 'einops', 'flash-attn', 'Pillow']
         for package in required_packages:
             try:
                 install(package)
@@ -17,11 +20,9 @@ class EndpointHandler:
             except Exception as e:
                 print(f"Failed to install {package}: {str(e)}")
         
-        
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         print(f"Using device: {self.device}")
         
-
         self.model_name = "microsoft/Florence-2-base-ft"
         self.model = AutoModelForCausalLM.from_pretrained(
             self.model_name,
@@ -35,28 +36,118 @@ class EndpointHandler:
             revision='refs/pr/6'
         )
         
-
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
-    
+
+    def process_image(self, image_input):
+        if isinstance(image_input, str):
+            # Check if it's a URL
+            if image_input.startswith('http://') or image_input.startswith('https://'):
+                image = Image.open(requests.get(image_input, stream=True).raw)
+            # Check if it's a base64 string
+            elif image_input.startswith('data:image'):
+                image_data = base64.b64decode(image_input.split(',')[1])
+                image = Image.open(BytesIO(image_data))
+            else:
+                raise ValueError("Invalid image input")
+        elif isinstance(image_input, bytes):
+            image = Image.open(BytesIO(image_input))
+        else:
+            raise ValueError("Unsupported image input type")
+        
+        return image
+
     def __call__(self, data):
         try:
+            # Handle different input formats
+            image_input = data.pop("image", None)
+            text_input = data.pop("text", "")
+            
+            # Process image if provided
+            image = self.process_image(image_input) if image_input else None
             
-            inputs = data.pop("inputs", data)
+            # Prepare inputs
+            inputs = self.processor(
+                images=image if image else None,
+                text=text_input,
+                return_tensors="pt"
+            )
+            
+            # Move inputs to device
+            inputs = {k: v.to(self.device) if isinstance(v, torch.Tensor) else v 
+                     for k, v in inputs.items()}
+            
+            # Generate output
+            with torch.no_grad():
+                outputs = self.model.generate(**inputs)
+            
+            # Decode outputs
+            decoded_outputs = self.processor.batch_decode(outputs, skip_special_tokens=True)
+            
+            return {"generated_text": decoded_outputs[0]}
+        
+        except Exception as e:
+            return {"error": str(e)}
+
+# import subprocess
+# import sys
+# import torch
+# from transformers import AutoModelForCausalLM, AutoProcessor
+
+# def install(package):
+#     subprocess.check_call([sys.executable, "-m", "pip", "install", "--no-warn-script-location", package])
+
+# class EndpointHandler:
+#     def __init__(self, path=""):
+        
+#         required_packages = ['timm', 'einops', 'flash-attn']
+#         for package in required_packages:
+#             try:
+#                 install(package)
+#                 print(f"Successfully installed {package}")
+#             except Exception as e:
+#                 print(f"Failed to install {package}: {str(e)}")
+        
+        
+#         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+#         print(f"Using device: {self.device}")
+        
+
+#         self.model_name = "microsoft/Florence-2-base-ft"
+#         self.model = AutoModelForCausalLM.from_pretrained(
+#             self.model_name,
+#             trust_remote_code=True,
+#             revision='refs/pr/6'
+#         ).to(self.device)
+        
+#         self.processor = AutoProcessor.from_pretrained(
+#             self.model_name,
+#             trust_remote_code=True,
+#             revision='refs/pr/6'
+#         )
+        
+
+#         if torch.cuda.is_available():
+#             torch.cuda.empty_cache()
+    
+#     def __call__(self, data):
+#         try:
+            
+#             inputs = data.pop("inputs", data)
             
            
-            processed_inputs = self.processor(inputs, return_tensors="pt")
+#             processed_inputs = self.processor(inputs, return_tensors="pt")
             
 
-            processed_inputs = {k: v.to(self.device) for k, v in processed_inputs.items()}
+#             processed_inputs = {k: v.to(self.device) for k, v in processed_inputs.items()}
             
            
-            with torch.no_grad():
-                outputs = self.model.generate(**processed_inputs)
+#             with torch.no_grad():
+#                 outputs = self.model.generate(**processed_inputs)
             
            
-            decoded_outputs = self.processor.batch_decode(outputs, skip_special_tokens=True)
+#             decoded_outputs = self.processor.batch_decode(outputs, skip_special_tokens=True)
             
-            return {"outputs": decoded_outputs}
-        except Exception as e:
-            return {"error": str(e)}
+#             return {"outputs": decoded_outputs}
+#         except Exception as e:
+#             return {"error": str(e)}