feat: defaults to base64

Author: justinmerrell

builder/requirements.txt

@@ -1,8 +1,7 @@
 # Required Python packages get listed here, one per line.
 # Reccomended to lock the version number to avoid unexpected changes.
 
-runpod==1.0.0
-# git+https://github.com/runpod/runpod-python
+runpod==1.2.0
 
 diffusers==0.12.1
 torch==1.13.1 --extra-index-url=https://download.pytorch.org/whl/cu116

src/rp_handler.py

@@ -1,6 +1,7 @@
 ''' infer.py for runpod worker '''
 
 import os
+import base64
 import predict
 import argparse
 
@@ -12,6 +13,18 @@
 from rp_schema import INPUT_SCHEMA
 
 
+def upload_or_base64_encode(file_name, img_path):
+    """
+    Uploads image to S3 bucket if it is available, otherwise returns base64 encoded image.
+    """
+    if os.environ.get('BUCKET_ENDPOINT_URL', False):
+        return upload_file_to_bucket(file_name, img_path)
+
+    with open(img_path, "rb") as image_file:
+        encoded_string = base64.b64encode(image_file.read())
+        return encoded_string.decode("utf-8")
+
+
 def run(job):
     '''
     Run inference on the model.
@@ -58,10 +71,10 @@ def run(job):
     job_output = []
     for index, img_path in enumerate(img_paths):
         file_name = f"{job['id']}_{index}.png"
-        image_url = upload_file_to_bucket(file_name, img_path)
+        image_return = upload_or_base64_encode(file_name, img_path)
 
         job_output.append({
-            "image": image_url,
+            "image": image_return,
             "seed": validated_input['seed'] + index
         })