Update handler.py
Browse files- handler.py +8 -6
handler.py
CHANGED
@@ -1,3 +1,4 @@
|
|
|
|
1 |
from typing import Dict, List, Any
|
2 |
from PIL import Image
|
3 |
import torch
|
@@ -16,16 +17,17 @@ class EndpointHandler():
|
|
16 |
self.model = self.model.to(device)
|
17 |
|
18 |
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
|
19 |
-
|
|
|
20 |
# Check if images is None or empty and handle it appropriately
|
21 |
-
if not
|
22 |
-
return {"captions": []}
|
23 |
|
24 |
# Default to "a photography of" if texts not provided
|
25 |
-
texts = data.get("texts", ["a photography of"] * len(
|
26 |
|
27 |
try:
|
28 |
-
raw_images = [Image.open(
|
29 |
processed_inputs = [
|
30 |
self.processor(img, txt, return_tensors="pt") for img, txt in zip(raw_images, texts)
|
31 |
]
|
@@ -43,4 +45,4 @@ class EndpointHandler():
|
|
43 |
except Exception as e:
|
44 |
# Handle or log the exception and optionally return an error message
|
45 |
print(f"Error during processing: {str(e)}")
|
46 |
-
return {"error": str(e)}
|
|
|
1 |
+
import requests
|
2 |
from typing import Dict, List, Any
|
3 |
from PIL import Image
|
4 |
import torch
|
|
|
17 |
self.model = self.model.to(device)
|
18 |
|
19 |
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
|
20 |
+
image_urls = data.get("images")
|
21 |
+
|
22 |
# Check if images is None or empty and handle it appropriately
|
23 |
+
if not image_urls:
|
24 |
+
return {"captions": [], "error": "No images provided"}
|
25 |
|
26 |
# Default to "a photography of" if texts not provided
|
27 |
+
texts = data.get("texts", ["a photography of"] * len(image_urls))
|
28 |
|
29 |
try:
|
30 |
+
raw_images = [Image.open(requests.get(url, stream=True).raw).convert("RGB") for url in image_urls]
|
31 |
processed_inputs = [
|
32 |
self.processor(img, txt, return_tensors="pt") for img, txt in zip(raw_images, texts)
|
33 |
]
|
|
|
45 |
except Exception as e:
|
46 |
# Handle or log the exception and optionally return an error message
|
47 |
print(f"Error during processing: {str(e)}")
|
48 |
+
return {"captions": [], "error": str(e)}
|