Update code/inference.py

code/inference.py

@@ -3,6 +3,8 @@ from PIL import Image
 from io import BytesIO
 import torch
 from transformers import AutoTokenizer
+import boto3
+import tempfile
 
 from llava.model import LlavaLlamaForCausalLM
 from llava.utils import disable_torch_init
@@ -68,6 +70,14 @@ def predict_fn(data, model_and_tokenizer):
     if image_file.startswith("http") or image_file.startswith("https"):
         response = requests.get(image_file)
         image = Image.open(BytesIO(response.content)).convert("RGB")
+    elif image_file.startswith("s3://"):
+        s3 = boto3.client("s3")
+        s3_path = s3_path[5:]
+        bucket = s3_path.split('/')[0]
+        s3_key = '/'.join(s3_path.split('/')[1:])
+        with tempfile.NamedTemporaryFile() as temp_file:
+            s3.download_file(bucket, s3_key, temp_file.name)
+        image = Image.open(temp_file).convert("RGB")
     else:
         image = Image.open(image_file).convert("RGB")